mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-31 16:21:50 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): decouple loop from message tool state (#5559)

Browse Source 
* refactor(agent): decouple loop from message tool state

* refactor(agent): scope message delivery tracking per run

* refactor(agent): clarify message delivery scope name
This commit is contained in:
chengyongru
2026-08-27 17:23:16 +08:00
committed by GitHub
parent d6c112ab74
commit 39de4594d7
3 changed files with 84 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+26 -32
View File
@@ -38,7 +38,7 @@ from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context
from nanobot.agent.tools.exec_session import ExecSessionManager
from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, reset_file_states
from nanobot.agent.tools.message import MessageTool
from nanobot.agent.tools.message import capture_message_deliveries
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.tools.runtime_control import AgentRuntimeControl
from nanobot.agent.turn_delivery import (
@@ -144,7 +144,6 @@ class TurnContext:
final_content: str | None = None
all_messages: list[dict[str, Any]] = field(default_factory=list)
stop_reason: str = ""
had_injections: bool = False
streamed_content: bool = False
input_persisted_early: bool = False
@@ -1723,18 +1722,12 @@ class AgentLoop:
msg: InboundMessage,
final_content: str,
stop_reason: str,
had_injections: bool,
streamed_content: bool,
*,
log_content: bool = True,
turn_latency_ms: int | None = None,
) -> OutboundMessage | None:
"""Assemble the final outbound message from turn results."""
# MessageTool suppression
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
if not had_injections or stop_reason == "empty_final_response":
return None
if log_content:
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview)
@@ -1890,10 +1883,6 @@ class AgentLoop:
)
is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent"
if ctx.kind is TurnKind.USER and (message_tool := self.tools.get("message")):
if isinstance(message_tool, MessageTool):
message_tool.start_turn()
_hist_kwargs: dict[str, Any] = {
"max_tokens": self._replay_token_budget(runtime),
"extend_to_user": is_subagent,
@@ -1994,28 +1983,34 @@ class AgentLoop:
if ctx.visible_run_started_at is None:
ctx.visible_run_started_at = time.time()
await ctx.delivery.running(started_at=ctx.visible_run_started_at)
result = await self._run_agent_loop(
ctx.initial_messages,
runtime=runtime,
on_progress=ctx.on_progress,
on_stream=ctx.on_stream,
on_stream_end=ctx.on_stream_end,
on_retry_wait=ctx.on_retry_wait,
session=ctx.session,
pending_queue=ctx.pending_queue,
ephemeral=ctx.ephemeral,
run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral,
hooks=ctx.hooks,
hook_factories=ctx.hook_factories,
turn_scopes=ctx.turn_scopes,
tools=ctx.tools,
request_context=ctx.request_context,
provider_state=ctx.provider_state,
)
with capture_message_deliveries() as message_sends:
result = await self._run_agent_loop(
ctx.initial_messages,
runtime=runtime,
on_progress=ctx.on_progress,
on_stream=ctx.on_stream,
on_stream_end=ctx.on_stream_end,
on_retry_wait=ctx.on_retry_wait,
session=ctx.session,
pending_queue=ctx.pending_queue,
ephemeral=ctx.ephemeral,
run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral,
hooks=ctx.hooks,
hook_factories=ctx.hook_factories,
turn_scopes=ctx.turn_scopes,
tools=ctx.tools,
request_context=ctx.request_context,
provider_state=ctx.provider_state,
)
ctx.final_content = result.final_content
ctx.all_messages = result.messages
ctx.stop_reason = result.stop_reason
ctx.had_injections = result.had_injections
if (
ctx.kind is TurnKind.USER
and (ctx.delivery.route.channel, ctx.delivery.route.chat_id) in message_sends
and (not result.had_injections or result.stop_reason == "empty_final_response")
):
ctx.suppress_response = True
ctx.usage = result.usage
ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER:
@@ -2084,7 +2079,6 @@ class AgentLoop:
ctx.delivery.delivery_message,
cast(str, ctx.final_content),
ctx.stop_reason,
ctx.had_injections,
ctx.streamed_content,
log_content=ctx.require_session().policy.log_content,
turn_latency_ms=ctx.turn_latency_ms,
nanobot/agent/tools/message.py
+22 -16
View File
@@ -2,9 +2,11 @@
# pyright: reportIncompatibleMethodOverride=false
from collections.abc import Awaitable, Callable, Generator
from contextlib import contextmanager
from contextvars import ContextVar, Token
from pathlib import Path
from typing import Any, Awaitable, Callable, cast
from typing import Any, cast
from loguru import logger
@@ -16,6 +18,22 @@ from nanobot.bus.events import OutboundMessage
from nanobot.config.paths import get_workspace_path
from nanobot.security.workspace_access import current_tool_workspace
_CURRENT_MESSAGE_SENDS: ContextVar[set[tuple[str, str]] | None] = ContextVar(
"message_sends",
default=None,
)
@contextmanager
def capture_message_deliveries() -> Generator[set[tuple[str, str]], None, None]:
"""Record successful MessageTool targets within one agent run."""
sends: set[tuple[str, str]] = set()
token = _CURRENT_MESSAGE_SENDS.set(sends)
try:
yield sends
finally:
_CURRENT_MESSAGE_SENDS.reset(token)
@tool_parameters(
tool_parameters_schema(
@@ -68,7 +86,6 @@ class MessageTool(Tool):
self._fallback_chat_id = default_chat_id
self._fallback_message_id = default_message_id
self._fallback_metadata: dict[str, Any] = {}
self._sent_in_turn_var: ContextVar[bool] = ContextVar("message_sent_in_turn", default=False)
self._suppress_delivery_var: ContextVar[bool] = ContextVar(
"message_suppress_delivery",
default=False,
@@ -87,10 +104,6 @@ class MessageTool(Tool):
"""Set the callback for sending messages."""
self._send_callback = callback
def start_turn(self) -> None:
"""Reset per-turn send tracking."""
self._sent_in_turn = False
def set_suppress_delivery(self, active: bool) -> Token[bool]:
"""Acknowledge but don't deliver tool sends (heartbeat internal check)."""
return self._suppress_delivery_var.set(active)
@@ -99,14 +112,6 @@ class MessageTool(Tool):
"""Restore previous delivery-suppression state."""
self._suppress_delivery_var.reset(token)
@property
def _sent_in_turn(self) -> bool:
return self._sent_in_turn_var.get()
@_sent_in_turn.setter
def _sent_in_turn(self, value: bool) -> None:
self._sent_in_turn_var.set(value)
@property
def name(self) -> str:
return "message"
@@ -244,8 +249,9 @@ class MessageTool(Tool):
try:
await self._send_callback(msg)
if channel == default_channel and chat_id == default_chat_id:
self._sent_in_turn = True
sends = _CURRENT_MESSAGE_SENDS.get()
if sends is not None:
sends.add((channel, chat_id))
media_info = f" with {len(media)} attachments" if media else ""
button_info = (
f" with {sum(len(row) for row in button_rows)} button(s)"
tests/tools/test_message_tool_suppress.py
+36 -18
View File
@@ -24,7 +24,12 @@ class TestMessageToolSuppressLogic:
"""Final reply suppressed only when message tool sends to the same target."""
@pytest.mark.asyncio
async def test_suppress_when_sent_to_same_target(self, tmp_path: Path) -> None:
@pytest.mark.parametrize("ephemeral", [False, True])
async def test_suppress_when_sent_to_same_target(
self,
tmp_path: Path,
ephemeral: bool,
) -> None:
loop = _make_loop(tmp_path)
tool_call = ToolCallRequest(
id="call1", name="message",
@@ -43,7 +48,7 @@ class TestMessageToolSuppressLogic:
mt.set_send_callback(AsyncMock(side_effect=lambda m: sent.append(m)))
msg = InboundMessage(channel="feishu", sender_id="user1", chat_id="chat123", content="Send")
result = await loop._process_message(msg)
result = await loop._process_message(msg, ephemeral=ephemeral)
assert len(sent) == 1
assert result is None # suppressed
@@ -87,6 +92,34 @@ class TestMessageToolSuppressLogic:
assert result is not None
assert "Hello" in result.content
@pytest.mark.asyncio
async def test_internal_message_check_keeps_final_response(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path)
tool_call = ToolCallRequest(
id="call1", name="message",
arguments={"content": "all clear", "channel": "feishu", "chat_id": "chat123"},
)
calls = iter([
LLMResponse(content="", tool_calls=[tool_call]),
LLMResponse(content="Heartbeat summary", tool_calls=[]),
])
loop.provider.chat_with_retry = AsyncMock(side_effect=lambda *a, **kw: next(calls))
loop.tools.get_definitions = MagicMock(return_value=[])
mt = loop.tools.get("message")
assert isinstance(mt, MessageTool)
token = mt.set_suppress_delivery(True)
try:
msg = InboundMessage(
channel="feishu", sender_id="user1", chat_id="chat123", content="Check",
)
result = await loop._process_message(msg)
finally:
mt.reset_suppress_delivery(token)
assert result is not None
assert result.content == "Heartbeat summary"
@pytest.mark.asyncio
async def test_injected_followup_with_message_tool_does_not_emit_empty_fallback(
self, tmp_path: Path
@@ -154,22 +187,7 @@ class TestMessageToolSuppressLogic:
('read foo.txt', True),
]
class TestMessageToolTurnTracking:
def test_sent_in_turn_tracks_same_target(self) -> None:
tool = MessageTool()
from nanobot.agent.tools.context import RequestContext, request_context
with request_context(RequestContext(channel="feishu", chat_id="chat1")):
assert not tool._sent_in_turn
tool._sent_in_turn = True
assert tool._sent_in_turn
def test_start_turn_resets(self) -> None:
tool = MessageTool()
tool._sent_in_turn = True
tool.start_turn()
assert not tool._sent_in_turn
class TestMessageToolSchema:
def test_schema_discourages_current_chat_replies(self) -> None:
tool = MessageTool()