diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 4adc2084..4ba68894 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -21,9 +21,8 @@ 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, CompositeHook +from nanobot.agent.hook import AgentHook from nanobot.agent.memory import Consolidator -from nanobot.agent.progress_hook import AgentProgressHook from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec from nanobot.agent.subagent import SubagentManager from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context @@ -31,6 +30,7 @@ from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, res from nanobot.agent.tools.message import MessageTool from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.self import MyTool +from nanobot.agent.turn_hooks import AgentTurnHookSpec, build_agent_turn_hook from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.outbound_events import ( RetryWaitEvent, @@ -753,7 +753,7 @@ class AgentLoop: """ self._sync_subagent_runtime_limits() - loop_hook = AgentProgressHook( + hook = build_agent_turn_hook(AgentTurnHookSpec( on_progress=on_progress, on_stream=on_stream, on_stream_end=on_stream_end, @@ -765,11 +765,11 @@ class AgentLoop: 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), - ) - run_hooks = [*self._extra_hooks, *(hooks or [])] - hook: AgentHook = loop_hook - if run_hooks and (not ephemeral or run_extra_hooks_for_ephemeral): - hook = CompositeHook([loop_hook, *run_hooks]) + 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: diff --git a/nanobot/agent/turn_hooks.py b/nanobot/agent/turn_hooks.py new file mode 100644 index 00000000..f94e381a --- /dev/null +++ b/nanobot/agent/turn_hooks.py @@ -0,0 +1,52 @@ +"""Turn-scoped hook assembly for agent runs.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from nanobot.agent.hook import AgentHook, CompositeHook +from nanobot.agent.progress_hook import AgentProgressHook + + +@dataclass(slots=True) +class AgentTurnHookSpec: + """Inputs needed to build the hook chain for one agent turn.""" + + on_progress: Callable[..., Awaitable[None]] | None = None + on_stream: Callable[[str], Awaitable[None]] | None = None + on_stream_end: Callable[..., Awaitable[None]] | None = None + channel: str = "cli" + chat_id: str = "direct" + message_id: str | None = None + metadata: dict[str, Any] | None = None + session_key: str | None = None + tool_hint_max_length: int = 40 + set_tool_context: Callable[..., None] | None = None + on_iteration: Callable[[int], None] | None = None + registered_hooks: list[AgentHook] = field(default_factory=list) + turn_hooks: list[AgentHook] = field(default_factory=list) + ephemeral: bool = False + run_extra_hooks_for_ephemeral: bool = False + + +def build_agent_turn_hook(spec: AgentTurnHookSpec) -> AgentHook: + """Build the hook chain used by ``AgentRunner`` for one turn.""" + progress_hook = AgentProgressHook( + on_progress=spec.on_progress, + on_stream=spec.on_stream, + on_stream_end=spec.on_stream_end, + channel=spec.channel, + chat_id=spec.chat_id, + message_id=spec.message_id, + metadata=spec.metadata, + session_key=spec.session_key, + tool_hint_max_length=spec.tool_hint_max_length, + 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 diff --git a/tests/agent/test_turn_hooks.py b/tests/agent/test_turn_hooks.py new file mode 100644 index 00000000..aaf4dd4a --- /dev/null +++ b/tests/agent/test_turn_hooks.py @@ -0,0 +1,56 @@ +import pytest + +from nanobot.agent.hook import AgentHook, AgentHookContext +from nanobot.agent.turn_hooks import AgentTurnHookSpec, build_agent_turn_hook + + +class RecordingHook(AgentHook): + def __init__(self, events: list[str]) -> None: + super().__init__() + self._events = events + + async def before_iteration(self, context: AgentHookContext) -> None: + self._events.append(f"hook:{context.iteration}") + + +@pytest.mark.asyncio +async def test_turn_hook_builder_runs_progress_hook_before_extra_hooks() -> None: + events: list[str] = [] + + hook = build_agent_turn_hook(AgentTurnHookSpec( + on_iteration=lambda iteration: events.append(f"progress:{iteration}"), + registered_hooks=[RecordingHook(events)], + )) + + await hook.before_iteration(AgentHookContext(iteration=2, messages=[])) + + assert events == ["progress:2", "hook:2"] + + +@pytest.mark.asyncio +async def test_turn_hook_builder_skips_extra_hooks_for_ephemeral_turns_by_default() -> None: + events: list[str] = [] + + hook = build_agent_turn_hook(AgentTurnHookSpec( + registered_hooks=[RecordingHook(events)], + ephemeral=True, + )) + + await hook.before_iteration(AgentHookContext(iteration=1, messages=[])) + + assert events == [] + + +@pytest.mark.asyncio +async def test_turn_hook_builder_can_include_extra_hooks_for_ephemeral_turns() -> None: + events: list[str] = [] + + hook = build_agent_turn_hook(AgentTurnHookSpec( + registered_hooks=[RecordingHook(events)], + ephemeral=True, + run_extra_hooks_for_ephemeral=True, + )) + + await hook.before_iteration(AgentHookContext(iteration=1, messages=[])) + + assert events == ["hook:1"]