diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 3d51b376e..a2b967a42 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -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, diff --git a/nanobot/agent/tools/message.py b/nanobot/agent/tools/message.py index 12e008f10..82a5245c5 100644 --- a/nanobot/agent/tools/message.py +++ b/nanobot/agent/tools/message.py @@ -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)" diff --git a/tests/tools/test_message_tool_suppress.py b/tests/tools/test_message_tool_suppress.py index 37f7acd49..80198eff0 100644 --- a/tests/tools/test_message_tool_suppress.py +++ b/tests/tools/test_message_tool_suppress.py @@ -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()