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:
+11
-2
@@ -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
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user