refactor(sdk): pass run hooks explicitly

nanobot/agent/loop.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextvars
 import dataclasses
 import os
 import time
@@ -76,13 +75,6 @@ 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()
     COMPACT = auto()
@@ -136,6 +128,7 @@ class TurnContext:
 
     ephemeral: bool = False
     tools: ToolRegistry | None = None
+    extra_hooks: list[AgentHook] | None = None
 
     turn_wall_started_at: float = field(default_factory=time.time)
     visible_run_started_at: float | None = None
@@ -701,6 +694,7 @@ class AgentLoop:
         pending_queue: asyncio.Queue | None = None,
         ephemeral: bool = False,
         tools: ToolRegistry | None = None,
+        extra_hooks: list[AgentHook] | None = None,
     ) -> tuple[str | None, list[str], list[dict], str, bool]:
         """Run the agent iteration loop.
 
@@ -727,16 +721,9 @@ class AgentLoop:
             on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration),
         )
         hook: AgentHook = loop_hook
-        per_call = _per_call_hooks.get()
+        turn_hooks = extra_hooks if extra_hooks is not None else self._extra_hooks
-        extra_hooks = None
+        if not ephemeral and turn_hooks:
-        # NOTE: SDK callers always go through run() → process_direct(ephemeral=False),
+            hook = CompositeHook([loop_hook] + turn_hooks)
-        # so per_call hooks are intentionally excluded from ephemeral turns. If a future
-        # caller sets the contextvar and then calls process_direct(ephemeral=True), the
-        # hooks will correctly NOT fire (ephemeral guard takes precedence).
-        if not ephemeral:
-            extra_hooks = per_call if per_call is not None else self._extra_hooks
-        if extra_hooks:
-            hook = CompositeHook([loop_hook] + extra_hooks)
 
         async def _checkpoint(payload: dict[str, Any]) -> None:
             if session is None:
@@ -1255,6 +1242,7 @@ class AgentLoop:
         pending_queue: asyncio.Queue | None = None,
         ephemeral: bool = False,
         tools: ToolRegistry | None = None,
+        extra_hooks: list[AgentHook] | None = None,
     ) -> OutboundMessage | None:
         """Process a single inbound message and return the response."""
         self._refresh_provider_snapshot()
@@ -1287,6 +1275,7 @@ class AgentLoop:
             pending_queue=pending_queue,
             ephemeral=ephemeral,
             tools=tools,
+            extra_hooks=extra_hooks,
         )
 
         while ctx.state is not TurnState.DONE:
@@ -1513,6 +1502,7 @@ class AgentLoop:
             pending_queue=ctx.pending_queue,
             ephemeral=ctx.ephemeral,
             tools=ctx.tools,
+            extra_hooks=ctx.extra_hooks,
         )
         final_content, tools_used, all_msgs, stop_reason, had_injections = result
         ctx.final_content = final_content
@@ -1828,6 +1818,7 @@ class AgentLoop:
         ephemeral: bool = False,
         tools: ToolRegistry | None = None,
         persist_user_message: bool = True,
+        extra_hooks: list[AgentHook] | None = None,
     ) -> OutboundMessage | None:
         """Process a message directly and return the outbound payload."""
         await self._connect_mcp()
@@ -1851,6 +1842,8 @@ class AgentLoop:
                 }
                 if tools is not None:
                     kwargs["tools"] = tools
+                if extra_hooks is not None:
+                    kwargs["extra_hooks"] = extra_hooks
                 return await self._process_message(
                     msg,
                     **kwargs,

nanobot/nanobot.py

@@ -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, _per_call_hooks
+from nanobot.agent.loop import AgentLoop
 from nanobot.providers.image_generation import image_gen_provider_configs
 
 
@@ -85,13 +85,11 @@ class Nanobot:
         """
         capture = SDKCaptureHook()
         base_hooks = list(hooks) if hooks is not None else list(self._loop._extra_hooks or [])
-        token = _per_call_hooks.set([capture, *base_hooks])
+        response = await self._loop.process_direct(
-        try:
+            message,
-            response = await self._loop.process_direct(
+            session_key=session_key,
-                message, session_key=session_key,
+            extra_hooks=[capture, *base_hooks],
-            )
+        )
-        finally:
-            _per_call_hooks.reset(token)
 
         content = (response.content if response else None) or ""
         return RunResult(
@@ -109,4 +107,3 @@ class Nanobot:
 
     async def __aexit__(self, *exc: object) -> None:
         await self.aclose()
-

tests/test_nanobot_facade.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import asyncio
 import json
 from pathlib import Path
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -65,12 +65,14 @@ async def test_run_returns_result(tmp_path):
 
     assert isinstance(result, RunResult)
     assert result.content == "Hello back!"
-    bot._loop.process_direct.assert_awaited_once_with("hi", session_key="sdk:default")
+    bot._loop.process_direct.assert_awaited_once_with(
+        "hi", session_key="sdk:default", extra_hooks=ANY
+    )
 
 
 @pytest.mark.asyncio
 async def test_run_with_hooks(tmp_path):
-    from nanobot.agent.hook import AgentHook, AgentHookContext
+    from nanobot.agent.hook import AgentHook, AgentHookContext, SDKCaptureHook
     from nanobot.bus.events import OutboundMessage
 
     config_path = _write_config(tmp_path)
@@ -89,6 +91,10 @@ async def test_run_with_hooks(tmp_path):
 
     assert result.content == "done"
     assert bot._loop._extra_hooks == []
+    extra_hooks = bot._loop.process_direct.await_args.kwargs["extra_hooks"]
+    assert len(extra_hooks) == 2
+    assert isinstance(extra_hooks[0], SDKCaptureHook)
+    assert isinstance(extra_hooks[1], TestHook)
 
 
 @pytest.mark.asyncio
@@ -160,7 +166,9 @@ async def test_run_custom_session_key(tmp_path):
     bot._loop.process_direct = AsyncMock(return_value=mock_response)
 
     await bot.run("hi", session_key="user-alice")
-    bot._loop.process_direct.assert_awaited_once_with("hi", session_key="user-alice")
+    bot._loop.process_direct.assert_awaited_once_with(
+        "hi", session_key="user-alice", extra_hooks=ANY
+    )
 
 
 def test_import_from_top_level():
@@ -184,11 +192,8 @@ async def test_run_populates_tools_used_across_iterations(tmp_path):
     config_path = _write_config(tmp_path)
     bot = Nanobot.from_config(config_path, workspace=tmp_path)
 
-    async def fake_process_direct(message, *, session_key):
+    async def fake_process_direct(message, *, session_key, extra_hooks=None):
-        # Whatever hooks the SDK installed are now in the contextvar.
+        extras = extra_hooks or []
-        from nanobot.agent.loop import _per_call_hooks
-
-        extras = _per_call_hooks.get() or []
         messages = [{"role": "user", "content": message}]
         ctx1 = AgentHookContext(iteration=0, messages=messages)
         ctx1.tool_calls = [
@@ -219,10 +224,8 @@ async def test_run_populates_final_messages(tmp_path):
     config_path = _write_config(tmp_path)
     bot = Nanobot.from_config(config_path, workspace=tmp_path)
 
-    async def fake_process_direct(message, *, session_key):
+    async def fake_process_direct(message, *, session_key, extra_hooks=None):
-        from nanobot.agent.loop import _per_call_hooks
+        extras = extra_hooks or []
-
-        extras = _per_call_hooks.get() or []
         messages = [
             {"role": "user", "content": message},
             {"role": "assistant", "content": "hi there"},
@@ -270,10 +273,8 @@ async def test_run_user_hooks_still_fire_alongside_capture(tmp_path):
         async def after_iteration(self, context: AgentHookContext) -> None:
             seen_iterations.append(context.iteration)
 
-    async def fake_process_direct(message, *, session_key):
+    async def fake_process_direct(message, *, session_key, extra_hooks=None):
-        from nanobot.agent.loop import _per_call_hooks
+        extras = extra_hooks or []
-
-        extras = _per_call_hooks.get() or []
         assert len(extras) == 2, f"expected capture + user hook, got {len(extras)}"
         ctx = AgentHookContext(iteration=7, messages=[])
         for h in extras:
@@ -306,16 +307,14 @@ async def test_concurrent_run_hooks_are_isolated_per_call(tmp_path):
     started = 0
     both_started = asyncio.Event()
 
-    async def fake_process_direct(message, *, session_key):
+    async def fake_process_direct(message, *, session_key, extra_hooks=None):
         nonlocal started
-        from nanobot.agent.loop import _per_call_hooks
-
         started += 1
         if started == 2:
             both_started.set()
         await both_started.wait()
 
-        extras = _per_call_hooks.get() or []
+        extras = extra_hooks or []
         messages = [{"role": "user", "content": message}]
         ctx = AgentHookContext(iteration=0, messages=messages)
         ctx.tool_calls = [
@@ -351,9 +350,9 @@ async def test_run_restores_extra_hooks_even_on_populated_iterations(tmp_path):
     sentinel_hook = AgentHook()
     bot._loop._extra_hooks = [sentinel_hook]
 
-    async def fake_process_direct(message, *, session_key):
+    async def fake_process_direct(message, *, session_key, extra_hooks=None):
         ctx = AgentHookContext(iteration=0, messages=[])
-        for h in bot._loop._extra_hooks:
+        for h in extra_hooks or []:
             await h.after_iteration(ctx)
         return OutboundMessage(channel="cli", chat_id="direct", content="done")