fix(sdk): use contextvars for per-call hooks to prevent concurrent run() race

Nanobot.run() previously mutated the shared self._loop._extra_hooks
attribute under a try/finally. When two run() calls with different
session_keys execute concurrently, they overwrite each other's hook
lists — the second call saves the first call's hooks as 'prev', and
the finally block restores stale state.

Use a contextvars.ContextVar instead, which is per-task in asyncio.
This gives each concurrent run() call its own hook list without
changing any function signatures. Falls back to self._extra_hooks
when no per-call hooks are set (non-SDK usage).
This commit is contained in:
michaelxer
2026-06-21 15:59:28 +08:00
committed by Xubin Ren
parent fede2ecfc2
commit 0bb1b0b3c2
2 changed files with 15 additions and 7 deletions
+11 -2
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import contextvars
import dataclasses
import os
import time
@@ -75,6 +76,12 @@ if TYPE_CHECKING:
)
from nanobot.cron.service import CronService
# Per-call hooks for SDK Nanobot.run() — avoids mutating shared _extra_hooks
# when multiple run() calls execute concurrently on the same AgentLoop.
_per_call_hooks: contextvars.ContextVar[list[AgentHook] | None] = contextvars.ContextVar(
"_per_call_hooks", default=None,
)
class TurnState(Enum):
RESTORE = auto()
@@ -720,8 +727,10 @@ class AgentLoop:
on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration),
)
hook: AgentHook = loop_hook
if not ephemeral and self._extra_hooks:
hook = CompositeHook([loop_hook] + self._extra_hooks)
per_call = _per_call_hooks.get()
extra_hooks = per_call if per_call is not None else (self._extra_hooks if not ephemeral else None)
if extra_hooks:
hook = CompositeHook([loop_hook] + extra_hooks)
async def _checkpoint(payload: dict[str, Any]) -> None:
if session is None:
+4 -5
View File
@@ -7,7 +7,7 @@ from pathlib import Path
from typing import Any
from nanobot.agent.hook import AgentHook, SDKCaptureHook
from nanobot.agent.loop import AgentLoop
from nanobot.agent.loop import AgentLoop, _per_call_hooks
from nanobot.providers.image_generation import image_gen_provider_configs
@@ -84,15 +84,14 @@ class Nanobot:
hooks: Optional lifecycle hooks for this run.
"""
capture = SDKCaptureHook()
prev = self._loop._extra_hooks
base_hooks = list(hooks) if hooks is not None else list(prev or [])
self._loop._extra_hooks = [capture, *base_hooks]
base_hooks = list(hooks) if hooks is not None else list(self._loop._extra_hooks or [])
token = _per_call_hooks.set([capture, *base_hooks])
try:
response = await self._loop.process_direct(
message, session_key=session_key,
)
finally:
self._loop._extra_hooks = prev
_per_call_hooks.reset(token)
content = (response.content if response else None) or ""
return RunResult(