diff --git a/nanobot/agent/__init__.py b/nanobot/agent/__init__.py index 47d453a8..8b4a0548 100644 --- a/nanobot/agent/__init__.py +++ b/nanobot/agent/__init__.py @@ -1,7 +1,14 @@ """Agent core module.""" from nanobot.agent.context import ContextBuilder -from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext, CompositeHook +from nanobot.agent.hook import ( + AgentHook, + AgentHookContext, + AgentRunHookContext, + AgentTurnHookContext, + AgentTurnHookFactory, + CompositeHook, +) from nanobot.agent.loop import AgentLoop from nanobot.agent.memory import MemoryStore from nanobot.agent.skills import SkillsLoader @@ -11,6 +18,8 @@ __all__ = [ "AgentHook", "AgentHookContext", "AgentRunHookContext", + "AgentTurnHookContext", + "AgentTurnHookFactory", "AgentLoop", "CompositeHook", "ContextBuilder", diff --git a/nanobot/agent/hook.py b/nanobot/agent/hook.py index f9ff37f3..69a5cf4f 100644 --- a/nanobot/agent/hook.py +++ b/nanobot/agent/hook.py @@ -2,7 +2,9 @@ from __future__ import annotations +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from pathlib import Path from typing import Any from loguru import logger @@ -44,6 +46,20 @@ class AgentRunHookContext: exception: BaseException | None = None +@dataclass(slots=True) +class AgentTurnHookContext: + """Turn-local inputs available when constructing per-turn hooks.""" + + on_progress: Callable[..., Awaitable[None]] | None = None + workspace: Path | None = None + channel: str = "cli" + chat_id: str = "direct" + message_id: str | None = None + session_key: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + ephemeral: bool = False + + class AgentHook: """Minimal lifecycle surface for shared runner customization.""" @@ -95,6 +111,9 @@ class AgentHook: return content +AgentTurnHookFactory = Callable[[AgentTurnHookContext], AgentHook | None] + + class CompositeHook(AgentHook): """Fan-out hook that delegates to an ordered list of hooks. diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 4ba68894..7d9d9dac 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -21,7 +21,7 @@ from nanobot.agent.autocompact import AutoCompact from nanobot.agent.automation_turns import publish_next_deferred_turn from nanobot.agent.context import ContextBuilder from nanobot.agent.cron_turns import CronTurnCoordinator -from nanobot.agent.hook import AgentHook +from nanobot.agent.hook import AgentHook, AgentTurnHookFactory from nanobot.agent.memory import Consolidator from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec from nanobot.agent.subagent import SubagentManager @@ -141,6 +141,7 @@ class TurnContext: ephemeral: bool = False run_extra_hooks_for_ephemeral: bool = False hooks: list[AgentHook] = field(default_factory=list) + hook_factories: list[AgentTurnHookFactory] = field(default_factory=list) tools: ToolRegistry | None = None turn_wall_started_at: float = field(default_factory=time.time) @@ -214,6 +215,7 @@ class AgentLoop: session_ttl_minutes: int = 0, consolidation_ratio: float = 0.5, hooks: list[AgentHook] | None = None, + hook_factories: list[AgentTurnHookFactory] | None = None, unified_session: bool = False, disabled_skills: list[str] | None = None, tools_config: ToolsConfig | None = None, @@ -284,6 +286,7 @@ class AgentLoop: self._start_time = time.time() self._last_usage: dict[str, int] = {} self._extra_hooks: list[AgentHook] = hooks or [] + self._hook_factories: list[AgentTurnHookFactory] = hook_factories or [] self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills) self.sessions = session_manager or SessionManager(workspace) @@ -740,6 +743,7 @@ class AgentLoop: ephemeral: bool = False, run_extra_hooks_for_ephemeral: bool = False, hooks: list[AgentHook] | None = None, + hook_factories: list[AgentTurnHookFactory] | None = None, tools: ToolRegistry | None = None, ) -> tuple[str | None, list[str], list[dict], str, bool]: """Run the agent iteration loop. @@ -753,24 +757,6 @@ class AgentLoop: """ self._sync_subagent_runtime_limits() - hook = build_agent_turn_hook(AgentTurnHookSpec( - on_progress=on_progress, - on_stream=on_stream, - on_stream_end=on_stream_end, - channel=channel, - chat_id=chat_id, - message_id=message_id, - metadata=metadata, - session_key=session_key, - tool_hint_max_length=self.tool_hint_max_length, - set_tool_context=self._set_tool_context, - on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration), - registered_hooks=self._extra_hooks, - turn_hooks=list(hooks or []), - ephemeral=ephemeral, - run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral, - )) - async def _checkpoint(payload: dict[str, Any]) -> None: if session is None: return @@ -846,6 +832,27 @@ class AgentLoop: message_metadata=metadata, session_metadata=session.metadata if session is not None else None, ) + effective_tools = tools or self.tools + hook = build_agent_turn_hook(AgentTurnHookSpec( + on_progress=on_progress, + on_stream=on_stream, + on_stream_end=on_stream_end, + channel=channel, + chat_id=chat_id, + message_id=message_id, + metadata=metadata, + session_key=active_session_key, + workspace=effective_scope.project_path, + tool_hint_max_length=self.tool_hint_max_length, + set_tool_context=self._set_tool_context, + on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration), + registered_hook_factories=self._hook_factories, + turn_hook_factories=list(hook_factories or []), + registered_hooks=self._extra_hooks, + turn_hooks=list(hooks or []), + ephemeral=ephemeral, + run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral, + )) request_ctx = RequestContext( channel=channel, chat_id=chat_id, @@ -872,7 +879,7 @@ class AgentLoop: try: result = await self.runner.run(AgentRunSpec( initial_messages=initial_messages, - tools=tools or self.tools, + tools=effective_tools, model=self.model, max_iterations=self.max_iterations, max_tool_result_chars=self.max_tool_result_chars, @@ -1215,6 +1222,7 @@ class AgentLoop: on_stream: Callable[[str], Awaitable[None]] | None = None, on_stream_end: Callable[..., Awaitable[None]] | None = None, pending_queue: asyncio.Queue | None = None, + hook_factories: list[AgentTurnHookFactory] | None = None, ) -> OutboundMessage | None: """Process a system inbound message (e.g. subagent announce).""" channel, chat_id = ( @@ -1276,6 +1284,7 @@ class AgentLoop: metadata=msg.metadata, session_key=key, pending_queue=pending_queue, + hook_factories=hook_factories, ) wall_done = time.time() latency_ms = max(0, int((wall_done - t_wall) * 1000)) @@ -1316,6 +1325,7 @@ class AgentLoop: ephemeral: bool = False, run_extra_hooks_for_ephemeral: bool = False, hooks: list[AgentHook] | None = None, + hook_factories: list[AgentTurnHookFactory] | None = None, tools: ToolRegistry | None = None, ) -> OutboundMessage | None: """Process a single inbound message and return the response.""" @@ -1329,6 +1339,7 @@ class AgentLoop: on_stream=on_stream, on_stream_end=on_stream_end, pending_queue=pending_queue, + hook_factories=hook_factories, ) key = session_key or msg.session_key @@ -1350,6 +1361,7 @@ class AgentLoop: ephemeral=ephemeral, run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral, hooks=list(hooks or []), + hook_factories=list(hook_factories or []), tools=tools, ) @@ -1579,6 +1591,7 @@ class AgentLoop: ephemeral=ctx.ephemeral, run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral, hooks=ctx.hooks, + hook_factories=ctx.hook_factories, tools=ctx.tools, ) final_content, tools_used, all_msgs, stop_reason, had_injections = result @@ -1896,6 +1909,7 @@ class AgentLoop: ephemeral: bool = False, _run_extra_hooks_for_ephemeral: bool = False, hooks: list[AgentHook] | None = None, + hook_factories: list[AgentTurnHookFactory] | None = None, tools: ToolRegistry | None = None, persist_user_message: bool = True, ) -> OutboundMessage | None: @@ -1923,6 +1937,8 @@ class AgentLoop: kwargs["run_extra_hooks_for_ephemeral"] = True if hooks is not None: kwargs["hooks"] = hooks + if hook_factories is not None: + kwargs["hook_factories"] = hook_factories if tools is not None: kwargs["tools"] = tools return await self._process_message( diff --git a/nanobot/agent/tools/self.py b/nanobot/agent/tools/self.py index 48954272..57cc2d4d 100644 --- a/nanobot/agent/tools/self.py +++ b/nanobot/agent/tools/self.py @@ -68,7 +68,7 @@ class MyTool(Tool, ContextAware): "_session_locks", "_active_tasks", "_background_tasks", # Security boundaries (inspect + modify both blocked) "restrict_to_workspace", "channels_config", - "_concurrency_gate", "_unified_session", "_extra_hooks", + "_concurrency_gate", "_unified_session", "_extra_hooks", "_hook_factories", }) READ_ONLY = frozenset({ diff --git a/nanobot/agent/turn_hooks.py b/nanobot/agent/turn_hooks.py index f94e381a..f7d6d640 100644 --- a/nanobot/agent/turn_hooks.py +++ b/nanobot/agent/turn_hooks.py @@ -4,9 +4,17 @@ from __future__ import annotations from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from pathlib import Path from typing import Any -from nanobot.agent.hook import AgentHook, CompositeHook +from loguru import logger + +from nanobot.agent.hook import ( + AgentHook, + AgentTurnHookContext, + AgentTurnHookFactory, + CompositeHook, +) from nanobot.agent.progress_hook import AgentProgressHook @@ -22,9 +30,12 @@ class AgentTurnHookSpec: message_id: str | None = None metadata: dict[str, Any] | None = None session_key: str | None = None + workspace: Path | None = None tool_hint_max_length: int = 40 set_tool_context: Callable[..., None] | None = None on_iteration: Callable[[int], None] | None = None + registered_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list) + turn_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list) registered_hooks: list[AgentHook] = field(default_factory=list) turn_hooks: list[AgentHook] = field(default_factory=list) ephemeral: bool = False @@ -46,7 +57,40 @@ def build_agent_turn_hook(spec: AgentTurnHookSpec) -> AgentHook: set_tool_context=spec.set_tool_context, on_iteration=spec.on_iteration, ) - extra_hooks = [*spec.registered_hooks, *spec.turn_hooks] - if extra_hooks and (not spec.ephemeral or spec.run_extra_hooks_for_ephemeral): - return CompositeHook([progress_hook, *extra_hooks]) - return progress_hook + if spec.ephemeral and not spec.run_extra_hooks_for_ephemeral: + return progress_hook + + turn_context = AgentTurnHookContext( + on_progress=spec.on_progress, + workspace=spec.workspace, + channel=spec.channel, + chat_id=spec.chat_id, + message_id=spec.message_id, + session_key=spec.session_key, + metadata=dict(spec.metadata or {}), + ephemeral=spec.ephemeral, + ) + hook_chain: list[AgentHook] = [progress_hook] + + for factory in spec.registered_hook_factories: + try: + created_hook = factory(turn_context) + except Exception: + logger.exception("Agent turn hook factory failed: {}", factory) + continue + if created_hook is not None: + hook_chain.append(created_hook) + + hook_chain.extend(spec.registered_hooks) + + for factory in spec.turn_hook_factories: + try: + created_hook = factory(turn_context) + except Exception: + logger.exception("Agent turn hook factory failed: {}", factory) + continue + if created_hook is not None: + hook_chain.append(created_hook) + + hook_chain.extend(spec.turn_hooks) + return CompositeHook(hook_chain) if len(hook_chain) > 1 else progress_hook diff --git a/tests/agent/test_hook_composite.py b/tests/agent/test_hook_composite.py index 58fdec6b..be24505a 100644 --- a/tests/agent/test_hook_composite.py +++ b/tests/agent/test_hook_composite.py @@ -6,7 +6,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext, CompositeHook +from nanobot.agent.hook import ( + AgentHook, + AgentHookContext, + AgentRunHookContext, + AgentTurnHookContext, + CompositeHook, +) def _ctx() -> AgentHookContext: @@ -348,7 +354,7 @@ async def test_composite_can_wrap_another_composite(): # --------------------------------------------------------------------------- -def _make_loop(tmp_path, hooks=None): +def _make_loop(tmp_path, hooks=None, hook_factories=None): from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus @@ -363,7 +369,11 @@ def _make_loop(tmp_path, hooks=None): patch("nanobot.agent.loop.Consolidator"): mock_sub_mgr.return_value.cancel_by_session = AsyncMock(return_value=0) loop = AgentLoop( - bus=bus, provider=provider, workspace=tmp_path, hooks=hooks, + bus=bus, + provider=provider, + workspace=tmp_path, + hooks=hooks, + hook_factories=hook_factories, ) return loop @@ -405,6 +415,66 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path): assert "after_run:completed" in events +@pytest.mark.asyncio +async def test_agent_loop_turn_hook_factories_receive_context(tmp_path): + """Turn-scoped hooks can be supplied externally and see turn-local context.""" + from nanobot.providers.base import LLMResponse + + captured: list[tuple[str, AgentTurnHookContext]] = [] + events: list[str] = [] + + class TrackingHook(AgentHook): + def __init__(self, label: str) -> None: + super().__init__() + self._label = label + + async def before_iteration(self, context): + events.append(f"{self._label}:{context.iteration}") + + def factory(label: str): + def _create(context: AgentTurnHookContext) -> AgentHook: + captured.append((label, context)) + return TrackingHook(label) + + return _create + + loop = _make_loop(tmp_path, hook_factories=[factory("registered")]) + loop.provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="done", tool_calls=[], usage={}) + ) + loop.tools.get_definitions = MagicMock(return_value=[]) + + async def on_progress(*args, **kwargs): + pass + + await loop._run_agent_loop( + [{"role": "user", "content": "hi"}], + on_progress=on_progress, + channel="websocket", + chat_id="chat-1", + message_id="msg-1", + metadata={"source": "test"}, + session_key="websocket:chat-1", + hook_factories=[factory("turn")], + ) + + assert events == ["registered:0", "turn:0"] + assert [label for label, _ in captured] == ["registered", "turn"] + assert [context.on_progress for _, context in captured] == [on_progress, on_progress] + assert [context.workspace for _, context in captured] == [tmp_path, tmp_path] + assert [context.channel for _, context in captured] == ["websocket", "websocket"] + assert [context.chat_id for _, context in captured] == ["chat-1", "chat-1"] + assert [context.message_id for _, context in captured] == ["msg-1", "msg-1"] + assert [context.session_key for _, context in captured] == [ + "websocket:chat-1", + "websocket:chat-1", + ] + assert [context.metadata for _, context in captured] == [ + {"source": "test"}, + {"source": "test"}, + ] + + @pytest.mark.asyncio async def test_agent_loop_extra_hook_error_isolation(tmp_path): """A faulty extra hook does not crash the agent loop.""" diff --git a/tests/agent/test_turn_hooks.py b/tests/agent/test_turn_hooks.py index 8ea458bf..46e8b5c7 100644 --- a/tests/agent/test_turn_hooks.py +++ b/tests/agent/test_turn_hooks.py @@ -1,6 +1,6 @@ import pytest -from nanobot.agent.hook import AgentHook, AgentHookContext +from nanobot.agent.hook import AgentHook, AgentHookContext, AgentTurnHookContext from nanobot.agent.turn_hooks import AgentTurnHookSpec, build_agent_turn_hook @@ -44,10 +44,67 @@ async def test_turn_hook_builder_runs_registered_hooks_before_turn_hooks() -> No @pytest.mark.asyncio -async def test_turn_hook_builder_skips_extra_hooks_for_ephemeral_turns_by_default() -> None: +async def test_turn_hook_builder_runs_factories_with_matching_registration_order( + tmp_path, +) -> None: events: list[str] = [] + captured: list[AgentTurnHookContext] = [] + + def factory(label: str): + def _create(context: AgentTurnHookContext) -> AgentHook: + captured.append(context) + return RecordingHook(events, label) + + return _create hook = build_agent_turn_hook(AgentTurnHookSpec( + on_iteration=lambda iteration: events.append(f"progress:{iteration}"), + channel="websocket", + chat_id="chat-1", + message_id="msg-1", + session_key="websocket:chat-1", + workspace=tmp_path, + metadata={"source": "test"}, + registered_hook_factories=[factory("registered_factory")], + registered_hooks=[RecordingHook(events, "registered")], + turn_hook_factories=[factory("turn_factory")], + turn_hooks=[RecordingHook(events, "turn")], + )) + + await hook.before_iteration(AgentHookContext(iteration=2, messages=[])) + + assert events == [ + "progress:2", + "registered_factory:2", + "registered:2", + "turn_factory:2", + "turn:2", + ] + assert [context.workspace for context in captured] == [tmp_path, tmp_path] + assert [context.channel for context in captured] == ["websocket", "websocket"] + assert [context.chat_id for context in captured] == ["chat-1", "chat-1"] + assert [context.message_id for context in captured] == ["msg-1", "msg-1"] + assert [context.session_key for context in captured] == [ + "websocket:chat-1", + "websocket:chat-1", + ] + assert [context.metadata for context in captured] == [ + {"source": "test"}, + {"source": "test"}, + ] + + +@pytest.mark.asyncio +async def test_turn_hook_builder_skips_extra_hooks_for_ephemeral_turns_by_default() -> None: + events: list[str] = [] + factory_calls: list[str] = [] + + def factory(context: AgentTurnHookContext) -> AgentHook: + factory_calls.append(context.channel) + return RecordingHook(events, "factory") + + hook = build_agent_turn_hook(AgentTurnHookSpec( + registered_hook_factories=[factory], registered_hooks=[RecordingHook(events)], ephemeral=True, )) @@ -55,6 +112,7 @@ async def test_turn_hook_builder_skips_extra_hooks_for_ephemeral_turns_by_defaul await hook.before_iteration(AgentHookContext(iteration=1, messages=[])) assert events == [] + assert factory_calls == [] @pytest.mark.asyncio