diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 43662489..776cb849 100644 --- a/nanobot/agent/loop.py +++ b/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() - extra_hooks = None - # NOTE: SDK callers always go through run() → process_direct(ephemeral=False), - # 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) + turn_hooks = extra_hooks if extra_hooks is not None else self._extra_hooks + if not ephemeral and turn_hooks: + hook = CompositeHook([loop_hook] + turn_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, diff --git a/nanobot/nanobot.py b/nanobot/nanobot.py index 8a9cfc1e..a31338c2 100644 --- a/nanobot/nanobot.py +++ b/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]) - try: - response = await self._loop.process_direct( - message, session_key=session_key, - ) - finally: - _per_call_hooks.reset(token) + response = await self._loop.process_direct( + message, + session_key=session_key, + extra_hooks=[capture, *base_hooks], + ) 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() - diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index 2f8f3f41..9e44afc3 100644 --- a/tests/test_nanobot_facade.py +++ b/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): - # Whatever hooks the SDK installed are now in the contextvar. - from nanobot.agent.loop import _per_call_hooks - - extras = _per_call_hooks.get() or [] + async def fake_process_direct(message, *, session_key, extra_hooks=None): + extras = extra_hooks 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): - from nanobot.agent.loop import _per_call_hooks - - extras = _per_call_hooks.get() or [] + async def fake_process_direct(message, *, session_key, extra_hooks=None): + extras = extra_hooks 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): - from nanobot.agent.loop import _per_call_hooks - - extras = _per_call_hooks.get() or [] + async def fake_process_direct(message, *, session_key, extra_hooks=None): + extras = extra_hooks 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")