Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(agent): add turn hook factories

Browse Source
This commit is contained in:
chengyongru
2026-07-08 21:02:12 +08:00
committed by Xubin Ren
parent 04dbf17426
commit 8559458258
7 changed files with 248 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/__init__.py
+10 -1
View File
@@ -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",
nanobot/agent/hook.py
+19
View File
@@ -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.
nanobot/agent/loop.py
+36 -20
View File
@@ -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(
nanobot/agent/tools/self.py
+1 -1
View File
@@ -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({
nanobot/agent/turn_hooks.py
+48 -4
View File
@@ -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])
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
tests/agent/test_hook_composite.py
+73 -3
View File
@@ -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."""
tests/agent/test_turn_hooks.py
+60 -2
View File
@@ -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