diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 5b840a39..a2dc2707 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -63,7 +63,7 @@ from nanobot.session.goal_state import ( sustained_goal_active, ) from nanobot.session.history_visibility import HIDDEN_HISTORY_META -from nanobot.session.keys import UNIFIED_SESSION_KEY, session_key_for_channel +from nanobot.session.keys import UNIFIED_SESSION_KEY from nanobot.session.manager import ( Session, SessionManager, @@ -549,32 +549,6 @@ class AgentLoop: """Connect configured MCP servers.""" await agent_context.connect_mcp(self, self.tools) - def _set_tool_context( - self, channel: str, chat_id: str, - message_id: str | None = None, metadata: dict | None = None, - session_key: str | None = None, - ) -> None: - """Update context for all tools that need routing info.""" - from nanobot.agent.tools.context import ContextAware - - effective_key = session_key or session_key_for_channel( - channel, - chat_id, - unified_session=self._unified_session, - ) - request_ctx = RequestContext( - channel=channel, - chat_id=chat_id, - message_id=message_id, - session_key=effective_key, - metadata=dict(metadata or {}), - ) - - for name in self.tools.tool_names: - tool = self.tools.get(name) - if tool and isinstance(tool, ContextAware): - tool.set_context(request_ctx) - @staticmethod def _runtime_chat_id(msg: InboundMessage) -> str: """Return the chat id shown in runtime metadata for the model.""" @@ -835,26 +809,6 @@ class AgentLoop: 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, @@ -880,6 +834,25 @@ class AgentLoop: session_metadata = session.metadata if session is not None else None try: + 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, + 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, + )) result = await self.runner.run(AgentRunSpec( initial_messages=initial_messages, tools=effective_tools, @@ -1251,10 +1224,6 @@ class AgentLoop: if is_subagent and self._persist_subagent_followup(session, msg): logger.debug("Subagent result persisted for session {}", key) self.sessions.save(session) - self._set_tool_context( - channel, chat_id, msg.metadata.get("message_id"), - msg.metadata, session_key=key, - ) current_role = "assistant" if is_subagent else "user" _hist_kwargs: dict[str, Any] = { "max_messages": self._max_messages, @@ -1535,13 +1504,6 @@ class AgentLoop: ctx.session, replay_max_messages=self._max_messages, ) - self._set_tool_context( - ctx.msg.channel, - ctx.msg.chat_id, - ctx.msg.metadata.get("message_id"), - ctx.msg.metadata, - session_key=ctx.session_key, - ) if message_tool := self.tools.get("message"): if isinstance(message_tool, MessageTool): message_tool.start_turn() diff --git a/nanobot/agent/progress_hook.py b/nanobot/agent/progress_hook.py index a9bf6a1e..7a500127 100644 --- a/nanobot/agent/progress_hook.py +++ b/nanobot/agent/progress_hook.py @@ -28,26 +28,16 @@ class AgentProgressHook(AgentHook): 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, ) -> None: super().__init__(reraise=True) self._on_progress = on_progress self._on_stream = on_stream self._on_stream_end = on_stream_end - self._channel = channel - self._chat_id = chat_id - self._message_id = message_id - self._metadata = metadata or {} self._session_key = session_key self._tool_hint_max_length = tool_hint_max_length - self._set_tool_context = set_tool_context self._on_iteration = on_iteration self._stream_buf = "" self._think_extractor = IncrementalThinkExtractor() @@ -124,15 +114,6 @@ class AgentProgressHook(AgentHook): for tc in context.tool_calls: args_str = json.dumps(tc.arguments, ensure_ascii=False) logger.info("Tool call: {}({})", tc.name, args_str[:200]) - if self._set_tool_context: - self._set_tool_context( - self._channel, - self._chat_id, - self._message_id, - self._metadata, - session_key=self._session_key, - ) - async def emit_reasoning(self, reasoning_content: str | None) -> None: """Publish a reasoning chunk; channel plugins decide whether to render.""" if ( diff --git a/nanobot/agent/tools/context.py b/nanobot/agent/tools/context.py index 054801bd..bf80493b 100644 --- a/nanobot/agent/tools/context.py +++ b/nanobot/agent/tools/context.py @@ -1,6 +1,7 @@ """Runtime context for tool construction.""" from __future__ import annotations +from contextlib import contextmanager from contextvars import ContextVar, Token from dataclasses import dataclass, field from typing import Any, Callable, Protocol, runtime_checkable @@ -36,6 +37,16 @@ def reset_request_context(token: Token[RequestContext | None]) -> None: _CURRENT_REQUEST_CONTEXT.reset(token) +@contextmanager +def request_context(ctx: RequestContext): + """Bind one immutable request snapshot and restore the previous value.""" + token = bind_request_context(ctx) + try: + yield ctx + finally: + reset_request_context(token) + + def current_request_context() -> RequestContext | None: return _CURRENT_REQUEST_CONTEXT.get() diff --git a/nanobot/agent/tools/cron.py b/nanobot/agent/tools/cron.py index b0098c07..db30dfdd 100644 --- a/nanobot/agent/tools/cron.py +++ b/nanobot/agent/tools/cron.py @@ -7,7 +7,7 @@ from datetime import datetime from typing import Any from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters -from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.context import current_request_context from nanobot.agent.tools.schema import ( IntegerSchema, StringSchema, @@ -51,19 +51,12 @@ _CRON_PARAMETERS = tool_parameters_schema( @tool_parameters(_CRON_PARAMETERS) -class CronTool(Tool, ContextAware): +class CronTool(Tool): """Tool to schedule reminders and recurring tasks.""" def __init__(self, cron_service: CronService, default_timezone: str = "UTC"): self._cron = cron_service self._default_timezone = default_timezone - self._session_key: ContextVar[str] = ContextVar("cron_session_key", default="") - self._origin_channel: ContextVar[str] = ContextVar("cron_origin_channel", default="") - self._origin_chat_id: ContextVar[str] = ContextVar("cron_origin_chat_id", default="") - self._origin_metadata: ContextVar[dict[str, Any] | None] = ContextVar( - "cron_origin_metadata", - default=None, - ) self._in_cron_context: ContextVar[bool] = ContextVar("cron_in_context", default=False) @classmethod @@ -74,15 +67,17 @@ class CronTool(Tool, ContextAware): def create(cls, ctx: Any) -> Tool: return cls(cron_service=ctx.cron_service, default_timezone=ctx.timezone) - def set_context(self, ctx: RequestContext) -> None: - """Set the current session context for scheduled cron job ownership.""" + @staticmethod + def _request_route() -> tuple[str, str, str, dict[str, Any]]: + """Return routing from the authoritative request snapshot.""" + ctx = current_request_context() + if ctx is None: + return "", "", "", {} raw_key = f"{ctx.channel}:{ctx.chat_id}" if ctx.channel and ctx.chat_id else "" - self._session_key.set( + session_key = ( raw_key if ctx.session_key == UNIFIED_SESSION_KEY else (ctx.session_key or "") ) - self._origin_channel.set(ctx.channel or "") - self._origin_chat_id.set(ctx.chat_id or "") - self._origin_metadata.set(dict(ctx.metadata or {})) + return session_key, ctx.channel or "", ctx.chat_id or "", dict(ctx.metadata or {}) def set_cron_context(self, active: bool): """Mark whether the tool is executing inside a cron job callback.""" @@ -171,11 +166,9 @@ class CronTool(Tool, ContextAware): "describing what to do when the job triggers " "(e.g. the reminder text). Retry including message=\"...\"." ) - session_key = self._session_key.get() + session_key, origin_channel, origin_chat_id, origin_metadata = self._request_route() if not session_key: return ToolResult.error("Error: scheduled cron jobs must be created from a chat session") - origin_channel = self._origin_channel.get() - origin_chat_id = self._origin_chat_id.get() if not origin_channel or not origin_chat_id: return ToolResult.error("Error: scheduled cron jobs must be created from a chat session") if tz and not cron_expr: @@ -218,7 +211,7 @@ class CronTool(Tool, ContextAware): session_key=session_key, origin_channel=origin_channel, origin_chat_id=origin_chat_id, - origin_metadata=dict(self._origin_metadata.get() or {}), + origin_metadata=origin_metadata, ) return f"Created job '{job.name}' (id: {job.id})" diff --git a/nanobot/agent/tools/long_task.py b/nanobot/agent/tools/long_task.py index 887222f6..66b9bd49 100644 --- a/nanobot/agent/tools/long_task.py +++ b/nanobot/agent/tools/long_task.py @@ -16,12 +16,11 @@ There is **no** sub-agent orchestrator and **no** special WebSocket ``agent_ui`` from __future__ import annotations -from contextvars import ContextVar from datetime import datetime from typing import TYPE_CHECKING, Any from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters -from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.context import current_request_context from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext from nanobot.session.goal_state import ( @@ -39,7 +38,7 @@ def _iso_now() -> str: return datetime.now().isoformat() -class _GoalToolsMixin(ContextAware): +class _GoalToolsMixin: """Shared routing context + Session lookup.""" def __init__( @@ -49,19 +48,9 @@ class _GoalToolsMixin(ContextAware): ) -> None: self._sessions = sessions self._runtime_events = runtime_events - # Each subclass gets its own ContextVar so concurrent tasks across - # different tool types (LongTaskTool vs CompleteGoalTool) do not - # interfere with each other. - self._request_ctx: ContextVar[RequestContext | None] = ContextVar( - f"{self.__class__.__name__}_request_ctx", - default=None, - ) - - def set_context(self, ctx: RequestContext) -> None: - self._request_ctx.set(ctx) def _session(self): - request_ctx = self._request_ctx.get() + request_ctx = current_request_context() if request_ctx is None: return None key = request_ctx.session_key @@ -72,7 +61,7 @@ class _GoalToolsMixin(ContextAware): async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None: """Publish authoritative goal metadata as a runtime event.""" runtime_events = self._runtime_events - rc = self._request_ctx.get() + rc = current_request_context() if runtime_events is None or rc is None: return cid = (rc.chat_id or "").strip() diff --git a/nanobot/agent/tools/message.py b/nanobot/agent/tools/message.py index 2a009b22..dad393d5 100644 --- a/nanobot/agent/tools/message.py +++ b/nanobot/agent/tools/message.py @@ -7,7 +7,7 @@ from typing import Any, Awaitable, Callable from loguru import logger from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters -from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.context import current_request_context from nanobot.agent.tools.path_utils import resolve_workspace_path from nanobot.agent.tools.schema import ArraySchema, StringSchema, tool_parameters_schema from nanobot.bus.events import OutboundMessage @@ -45,7 +45,7 @@ from nanobot.security.workspace_access import current_tool_workspace required=["content"], ) ) -class MessageTool(Tool, ContextAware): +class MessageTool(Tool): """Tool to send messages to users on chat channels.""" def __init__( @@ -62,20 +62,10 @@ class MessageTool(Tool, ContextAware): Path(workspace).expanduser() if workspace is not None else get_workspace_path() ) self._restrict_to_workspace = restrict_to_workspace - self._default_channel: ContextVar[str] = ContextVar( - "message_default_channel", default=default_channel - ) - self._default_chat_id: ContextVar[str] = ContextVar( - "message_default_chat_id", default=default_chat_id - ) - self._default_message_id: ContextVar[str | None] = ContextVar( - "message_default_message_id", - default=default_message_id, - ) - self._default_metadata: ContextVar[dict[str, Any]] = ContextVar( - "message_default_metadata", - default={}, - ) + self._fallback_channel = default_channel + 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._turn_delivered_media_var: ContextVar[tuple[str, ...]] = ContextVar( "message_turn_delivered_media", @@ -99,13 +89,6 @@ class MessageTool(Tool, ContextAware): restrict_to_workspace=ctx.config.restrict_to_workspace, ) - def set_context(self, ctx: RequestContext) -> None: - """Set the current message context.""" - self._default_channel.set(ctx.channel) - self._default_chat_id.set(ctx.chat_id) - self._default_message_id.set(ctx.message_id) - self._default_metadata.set(dict(ctx.metadata or {})) - def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None: """Set the callback for sending messages.""" self._send_callback = callback @@ -199,8 +182,23 @@ class MessageTool(Tool, ContextAware): for row in buttons ): return ToolResult.error("Error: buttons must be a list of list of strings") - default_channel = self._default_channel.get() - default_chat_id = self._default_chat_id.get() + request_ctx = current_request_context() + default_channel = ( + request_ctx.channel if request_ctx is not None else self._fallback_channel + ) + default_chat_id = ( + request_ctx.chat_id if request_ctx is not None else self._fallback_chat_id + ) + default_message_id = ( + request_ctx.message_id + if request_ctx is not None + else self._fallback_message_id + ) + default_metadata = ( + request_ctx.metadata + if request_ctx is not None + else self._fallback_metadata + ) channel = channel or default_channel explicit_chat_id = chat_id if ( @@ -224,7 +222,7 @@ class MessageTool(Tool, ContextAware): # to the wrong chat entirely. same_target = channel == default_channel and chat_id == default_chat_id if same_target: - message_id = message_id or self._default_message_id.get() + message_id = message_id or default_message_id else: message_id = None @@ -240,7 +238,7 @@ class MessageTool(Tool, ContextAware): except (OSError, PermissionError, ValueError) as e: return ToolResult.error(f"Error: media path is not allowed: {str(e)}") - metadata = dict(self._default_metadata.get()) if same_target else {} + metadata = dict(default_metadata) if same_target else {} if message_id: metadata["message_id"] = message_id if self._record_channel_delivery_var.get() or media: diff --git a/nanobot/agent/tools/registry.py b/nanobot/agent/tools/registry.py index 4cb608b4..de3c019e 100644 --- a/nanobot/agent/tools/registry.py +++ b/nanobot/agent/tools/registry.py @@ -4,6 +4,7 @@ import json from typing import Any from nanobot.agent.tools.base import Tool, ToolResult +from nanobot.agent.tools.context import ContextAware, current_request_context def is_tool_error_result(name: str, result: Any) -> bool: @@ -109,6 +110,12 @@ class ToolRegistry: ) ) + # Compatibility for external tools that still implement the legacy + # setter protocol. Built-ins read the authoritative ContextVar + # directly and never copy routing state. + if isinstance(tool, ContextAware) and (ctx := current_request_context()) is not None: + tool.set_context(ctx) + params = self._coerce_params(tool, params) if not isinstance(params, dict): return tool, params, ( diff --git a/nanobot/agent/tools/self.py b/nanobot/agent/tools/self.py index 57cc2d4d..d0310a9b 100644 --- a/nanobot/agent/tools/self.py +++ b/nanobot/agent/tools/self.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any from loguru import logger from nanobot.agent.tools.base import Tool, ToolResult -from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.context import current_request_context from nanobot.agent.tools.runtime_state import RuntimeState from nanobot.config_base import Base @@ -41,7 +41,7 @@ def _is_subagent_status(value: Any) -> bool: return isinstance(value, SubagentStatus) -class MyTool(Tool, ContextAware): +class MyTool(Tool): """Check and set the agent loop's runtime configuration.""" _plugin_discoverable = False # Requires AgentLoop reference; registered manually @@ -111,8 +111,6 @@ class MyTool(Tool, ContextAware): def __init__(self, runtime_state: RuntimeState, modify_allowed: bool = True) -> None: self._runtime_state = runtime_state self._modify_allowed = modify_allowed - self._channel = "" - self._chat_id = "" def __deepcopy__(self, memo: dict[int, Any]) -> MyTool: cls = self.__class__ @@ -120,14 +118,8 @@ class MyTool(Tool, ContextAware): memo[id(self)] = result result._runtime_state = self._runtime_state result._modify_allowed = self._modify_allowed - result._channel = self._channel - result._chat_id = self._chat_id return result - def set_context(self, ctx: RequestContext) -> None: - self._channel = ctx.channel - self._chat_id = ctx.chat_id - @property def name(self) -> str: return "my" @@ -184,7 +176,12 @@ class MyTool(Tool, ContextAware): } def _audit(self, action: str, detail: str) -> None: - session = f"{self._channel}:{self._chat_id}" if self._channel else "unknown" + ctx = current_request_context() + session = ( + ctx.session_key or f"{ctx.channel}:{ctx.chat_id}" + if ctx is not None and ctx.channel + else "unknown" + ) logger.info("self.{} | {} | session:{}", action, detail, session) # ------------------------------------------------------------------ diff --git a/nanobot/agent/tools/spawn.py b/nanobot/agent/tools/spawn.py index 420afc04..eebf042c 100644 --- a/nanobot/agent/tools/spawn.py +++ b/nanobot/agent/tools/spawn.py @@ -2,11 +2,10 @@ from __future__ import annotations -from contextvars import ContextVar from typing import TYPE_CHECKING, Any from nanobot.agent.tools.base import Tool, tool_parameters -from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.context import current_request_context from nanobot.agent.tools.schema import NumberSchema, StringSchema, tool_parameters_schema from nanobot.security.workspace_access import current_workspace_scope @@ -30,30 +29,16 @@ if TYPE_CHECKING: required=["task"], ) ) -class SpawnTool(Tool, ContextAware): +class SpawnTool(Tool): """Tool to spawn a subagent for background task execution.""" def __init__(self, manager: "SubagentManager"): self._manager = manager - self._origin_channel: ContextVar[str] = ContextVar("spawn_origin_channel", default="cli") - self._origin_chat_id: ContextVar[str] = ContextVar("spawn_origin_chat_id", default="direct") - self._session_key: ContextVar[str] = ContextVar("spawn_session_key", default="cli:direct") - self._origin_message_id: ContextVar[str | None] = ContextVar( - "spawn_origin_message_id", - default=None, - ) @classmethod def create(cls, ctx: Any) -> Tool: return cls(manager=ctx.subagent_manager) - def set_context(self, ctx: RequestContext) -> None: - """Set the origin context for subagent announcements.""" - self._origin_channel.set(ctx.channel) - self._origin_chat_id.set(ctx.chat_id) - self._session_key.set(ctx.session_key or f"{ctx.channel}:{ctx.chat_id}") - self._origin_message_id.set(ctx.message_id) - @property def name(self) -> str: return "spawn" @@ -84,13 +69,21 @@ class SpawnTool(Tool, ContextAware): f"({running}/{limit} running). Wait for a running subagent " f"to complete before spawning a new one." ) + request_ctx = current_request_context() + origin_channel = request_ctx.channel if request_ctx is not None else "cli" + origin_chat_id = request_ctx.chat_id if request_ctx is not None else "direct" + session_key = ( + request_ctx.session_key or f"{origin_channel}:{origin_chat_id}" + if request_ctx is not None + else "cli:direct" + ) return await self._manager.spawn( task=task, label=label, - origin_channel=self._origin_channel.get(), - origin_chat_id=self._origin_chat_id.get(), - session_key=self._session_key.get(), - origin_message_id=self._origin_message_id.get(), + origin_channel=origin_channel, + origin_chat_id=origin_chat_id, + session_key=session_key, + origin_message_id=request_ctx.message_id if request_ctx is not None else None, temperature=temperature, workspace_scope=current_workspace_scope(), ) diff --git a/nanobot/agent/turn_hooks.py b/nanobot/agent/turn_hooks.py index f7d6d640..537c4520 100644 --- a/nanobot/agent/turn_hooks.py +++ b/nanobot/agent/turn_hooks.py @@ -32,7 +32,6 @@ class AgentTurnHookSpec: 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) @@ -48,13 +47,8 @@ def build_agent_turn_hook(spec: AgentTurnHookSpec) -> AgentHook: 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, ) if spec.ephemeral and not spec.run_extra_hooks_for_ephemeral: diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 98c51449..7be00b48 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -7,6 +7,7 @@ import pytest from nanobot.agent.context import ContextBuilder from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.bus.events import InboundMessage from nanobot.bus.outbound_events import ( GoalStatusEvent, @@ -1090,20 +1091,24 @@ async def test_process_direct_skip_user_persist_does_not_save_retry_user( ] -def test_set_tool_context_uses_effective_key_for_spawn_tool(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_request_context_uses_effective_key_for_spawn_tool(tmp_path: Path) -> None: loop = _make_full_loop(tmp_path) spawn_tool = loop.tools.get("spawn") assert spawn_tool is not None + spawn_tool._manager.spawn = AsyncMock(return_value="started") # type: ignore[attr-defined] - loop._set_tool_context( - "discord", - "thread-777", + with request_context(RequestContext( + channel="discord", + chat_id="thread-777", session_key="discord:parent-456:thread:thread-777", - ) + )): + await spawn_tool.execute(task="inspect context") - assert spawn_tool._origin_channel.get() == "discord" # type: ignore[attr-defined] - assert spawn_tool._origin_chat_id.get() == "thread-777" # type: ignore[attr-defined] - assert spawn_tool._session_key.get() == "discord:parent-456:thread:thread-777" # type: ignore[attr-defined] + call = spawn_tool._manager.spawn.await_args.kwargs # type: ignore[attr-defined] + assert call["origin_channel"] == "discord" + assert call["origin_chat_id"] == "thread-777" + assert call["session_key"] == "discord:parent-456:thread:thread-777" @pytest.mark.asyncio @@ -1422,21 +1427,25 @@ def test_subagent_followup_skips_empty_content() -> None: assert session.messages == [] -def test_set_tool_context_passes_thread_session_key_to_spawn(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_request_context_passes_thread_session_key_to_spawn(tmp_path: Path) -> None: loop = _make_full_loop(tmp_path) + spawn_tool = loop.tools.get("spawn") + assert spawn_tool is not None + spawn_tool._manager.spawn = AsyncMock(return_value="started") # type: ignore[attr-defined] - loop._set_tool_context( - "slack", - "C123", + with request_context(RequestContext( + channel="slack", + chat_id="C123", message_id="msg-123", metadata={"slack": {"thread_ts": "1700.42", "channel_type": "channel"}}, session_key="slack:C123:1700.42", - ) + )): + await spawn_tool.execute(task="inspect thread") - spawn_tool = loop.tools.get("spawn") - assert spawn_tool is not None - assert spawn_tool._session_key.get() == "slack:C123:1700.42" - assert spawn_tool._origin_message_id.get() == "msg-123" + call = spawn_tool._manager.spawn.await_args.kwargs # type: ignore[attr-defined] + assert call["session_key"] == "slack:C123:1700.42" + assert call["origin_message_id"] == "msg-123" @pytest.mark.asyncio diff --git a/tests/agent/test_loop_tool_context.py b/tests/agent/test_loop_tool_context.py index 44b3b353..101f580a 100644 --- a/tests/agent/test_loop_tool_context.py +++ b/tests/agent/test_loop_tool_context.py @@ -24,15 +24,15 @@ class _ContextRecordingTool: def __init__(self) -> None: self.contexts: list[dict] = [] - def set_context(self, ctx: RequestContext) -> None: + async def execute(self, **_kwargs) -> str: + ctx = current_request_context() + assert ctx is not None self.contexts.append({ "channel": ctx.channel, "chat_id": ctx.chat_id, "metadata": ctx.metadata, "session_key": ctx.session_key, }) - - async def execute(self, **_kwargs) -> str: return "created" @@ -55,7 +55,7 @@ class _Tools: @pytest.mark.asyncio -async def test_loop_hook_preserves_metadata_when_resetting_tool_context(tmp_path: Path) -> None: +async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> None: provider = MagicMock() calls = {"n": 0} diff --git a/tests/agent/tools/test_long_task.py b/tests/agent/tools/test_long_task.py index 7edda8c3..ca46c2f0 100644 --- a/tests/agent/tools/test_long_task.py +++ b/tests/agent/tools/test_long_task.py @@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from nanobot.agent.loop import AgentLoop -from nanobot.agent.tools.context import RequestContext +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.long_task import ( CompleteGoalTool, LongTaskTool, @@ -24,15 +24,16 @@ from nanobot.session.webui_turns import WebuiTurnCoordinator def _tools(sm: SessionManager) -> tuple[LongTaskTool, CompleteGoalTool]: lt = LongTaskTool(sessions=sm) cg = CompleteGoalTool(sessions=sm) - rc = RequestContext( + return lt, cg + + +def _request_context(chat_id: str = "c1") -> RequestContext: + return RequestContext( channel="websocket", - chat_id="c1", - session_key="websocket:c1", + chat_id=chat_id, + session_key=f"websocket:{chat_id}", metadata={}, ) - lt.set_context(rc) - cg.set_context(rc) - return lt, cg @pytest.mark.asyncio @@ -40,7 +41,8 @@ async def test_long_task_records_goal_metadata(tmp_path): sm = SessionManager(tmp_path) lt, _cg = _tools(sm) - out = await lt.execute(goal="Do the thing", ui_summary="thing") + with request_context(_request_context()): + out = await lt.execute(goal="Do the thing", ui_summary="thing") assert "Goal recorded" in out sess = sm.get_or_create("websocket:c1") @@ -56,8 +58,9 @@ async def test_long_task_rejects_second_active_goal(tmp_path): sm = SessionManager(tmp_path) lt, _cg = _tools(sm) - await lt.execute(goal="First") - out = await lt.execute(goal="Second") + with request_context(_request_context()): + await lt.execute(goal="First") + out = await lt.execute(goal="Second") assert "already active" in out @@ -66,8 +69,9 @@ async def test_complete_goal_closes_active_goal(tmp_path): sm = SessionManager(tmp_path) lt, cg = _tools(sm) - await lt.execute(goal="X") - out = await cg.execute(recap="Done.") + with request_context(_request_context()): + await lt.execute(goal="X") + out = await cg.execute(recap="Done.") assert "marked complete" in out sess = sm.get_or_create("websocket:c1") @@ -84,19 +88,23 @@ async def test_goal_tools_keep_request_context_per_task(tmp_path): ctx_a = RequestContext(channel="websocket", chat_id="a", session_key="websocket:a") ctx_b = RequestContext(channel="websocket", chat_id="b", session_key="websocket:b") - lt.set_context(ctx_a) - task_a = asyncio.create_task(lt.execute(goal="Goal A")) - lt.set_context(ctx_b) - task_b = asyncio.create_task(lt.execute(goal="Goal B")) + async def start_goal(ctx: RequestContext, goal: str) -> str: + with request_context(ctx): + return await lt.execute(goal=goal) + + task_a = asyncio.create_task(start_goal(ctx_a, "Goal A")) + task_b = asyncio.create_task(start_goal(ctx_b, "Goal B")) await asyncio.gather(task_a, task_b) assert sm.get_or_create("websocket:a").metadata[GOAL_STATE_KEY]["objective"] == "Goal A" assert sm.get_or_create("websocket:b").metadata[GOAL_STATE_KEY]["objective"] == "Goal B" - cg.set_context(ctx_a) - done_a = asyncio.create_task(cg.execute(recap="Done A")) - cg.set_context(ctx_b) - done_b = asyncio.create_task(cg.execute(recap="Done B")) + async def complete_goal(ctx: RequestContext, recap: str) -> str: + with request_context(ctx): + return await cg.execute(recap=recap) + + done_a = asyncio.create_task(complete_goal(ctx_a, "Done A")) + done_b = asyncio.create_task(complete_goal(ctx_b, "Done B")) await asyncio.gather(done_a, done_b) assert sm.get_or_create("websocket:a").metadata[GOAL_STATE_KEY]["recap"] == "Done A" @@ -104,19 +112,19 @@ async def test_goal_tools_keep_request_context_per_task(tmp_path): @pytest.mark.asyncio -async def test_goal_tools_context_isolated_across_tool_types(tmp_path): - """LongTaskTool and CompleteGoalTool must not share routing context.""" +async def test_goal_tools_share_authoritative_request_context(tmp_path): + """Both goal tools resolve routing from the same request snapshot.""" sm = SessionManager(tmp_path) lt = LongTaskTool(sessions=sm) cg = CompleteGoalTool(sessions=sm) ctx = RequestContext(channel="websocket", chat_id="a", session_key="websocket:a") - lt.set_context(ctx) - assert cg._request_ctx.get() is None + with request_context(ctx): + assert lt._session() is sm.get_or_create("websocket:a") + assert cg._session() is sm.get_or_create("websocket:a") - cg.set_context(ctx) - assert lt._request_ctx.get() is ctx - assert cg._request_ctx.get() is ctx + assert lt._session() is None + assert cg._session() is None @pytest.mark.asyncio @@ -137,9 +145,8 @@ async def test_long_task_publishes_goal_state_ws_after_save(tmp_path): session_key="websocket:chat-99", metadata={}, ) - lt.set_context(rc) - - await lt.execute(goal="Objective alpha", ui_summary="alpha") + with request_context(rc): + await lt.execute(goal="Objective alpha", ui_summary="alpha") bus.publish_outbound.assert_awaited_once() call = bus.publish_outbound.await_args.args[0] @@ -172,12 +179,11 @@ async def test_complete_goal_publishes_inactive_goal_state_ws(tmp_path): session_key="websocket:chat-z", metadata={}, ) - lt.set_context(rc) - await lt.execute(goal="X") + with request_context(rc): + await lt.execute(goal="X") - bus.publish_outbound.reset_mock() - cg.set_context(rc) - await cg.execute(recap="Done.") + bus.publish_outbound.reset_mock() + await cg.execute(recap="Done.") bus.publish_outbound.assert_awaited_once() call = bus.publish_outbound.await_args.args[0] @@ -190,7 +196,8 @@ async def test_complete_goal_without_active_is_noop_message(tmp_path): sm = SessionManager(tmp_path) _lt, cg = _tools(sm) - out = await cg.execute(recap="n/a") + with request_context(_request_context()): + out = await cg.execute(recap="n/a") assert "No active" in out @@ -198,7 +205,8 @@ async def test_complete_goal_without_active_is_noop_message(tmp_path): async def test_long_task_skips_ws_publish_without_bus(tmp_path): sm = SessionManager(tmp_path) lt, _cg = _tools(sm) - out = await lt.execute(goal="Solo", ui_summary="s") + with request_context(_request_context()): + out = await lt.execute(goal="Solo", ui_summary="s") assert "Goal recorded" in out diff --git a/tests/agent/tools/test_self_tool.py b/tests/agent/tools/test_self_tool.py index 59fbe7f3..8e777a09 100644 --- a/tests/agent/tools/test_self_tool.py +++ b/tests/agent/tools/test_self_tool.py @@ -4,11 +4,12 @@ from __future__ import annotations import time from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from pydantic import BaseModel +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.self import MyTool # --------------------------------------------------------------------------- @@ -1121,14 +1122,26 @@ class TestLastUsageInSummary: # --------------------------------------------------------------------------- -# set_context (audit session tracking) +# request context (audit session tracking) # --------------------------------------------------------------------------- -class TestSetContext: +class TestRequestContext: - def test_set_context_stores_channel_and_chat_id(self): - from nanobot.agent.tools.context import RequestContext + def test_audit_reads_bound_session(self): tool = _make_tool() - tool.set_context(RequestContext(channel="feishu", chat_id="oc_abc123")) - assert tool._channel == "feishu" - assert tool._chat_id == "oc_abc123" + ctx = RequestContext( + channel="feishu", + chat_id="oc_abc123", + session_key="feishu:oc_abc123", + ) + + with patch("nanobot.agent.tools.self.logger.info") as info: + with request_context(ctx): + tool._audit("modify", "temperature = 0.2") + + info.assert_called_once_with( + "self.{} | {} | session:{}", + "modify", + "temperature = 0.2", + "feishu:oc_abc123", + ) diff --git a/tests/agent/tools/test_subagent_tools.py b/tests/agent/tools/test_subagent_tools.py index 6a037981..96fdb8c7 100644 --- a/tests/agent/tools/test_subagent_tools.py +++ b/tests/agent/tools/test_subagent_tools.py @@ -159,19 +159,18 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path): mgr.runner.run = AsyncMock(side_effect=fake_run) - from nanobot.agent.tools.context import RequestContext + from nanobot.agent.tools.context import RequestContext, request_context tool = SpawnTool(mgr) - tool.set_context(RequestContext(channel="test", chat_id="c1", session_key="test:c1")) + with request_context(RequestContext(channel="test", chat_id="c1", session_key="test:c1")): + # First spawn succeeds + result = await tool.execute(task="first task") + assert "started" in result - # First spawn succeeds - result = await tool.execute(task="first task") - assert "started" in result - - # Second spawn should be rejected (default limit is 1) - result = await tool.execute(task="second task") - assert "Cannot spawn subagent" in result - assert "concurrency limit reached" in result + # Second spawn should be rejected (default limit is 1) + result = await tool.execute(task="second task") + assert "Cannot spawn subagent" in result + assert "concurrency limit reached" in result # Release the first subagent release.set() diff --git a/tests/cron/test_cron_tool_list.py b/tests/cron/test_cron_tool_list.py index d2099d45..b164532e 100644 --- a/tests/cron/test_cron_tool_list.py +++ b/tests/cron/test_cron_tool_list.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone import pytest -from nanobot.agent.tools.context import RequestContext +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.cron import CronTool from nanobot.cron.service import CronService from nanobot.cron.types import CronJob, CronJobState, CronPayload, CronSchedule @@ -321,11 +321,10 @@ def test_remove_protected_dream_job_returns_clear_feedback(tmp_path) -> None: def test_add_cron_job_defaults_to_tool_timezone(tmp_path) -> None: tool = _make_tool_with_tz(tmp_path, "Asia/Shanghai") - tool.set_context( + with request_context( RequestContext(channel="telegram", chat_id="chat-1", session_key="telegram:chat-1") - ) - - result = tool._add_job(None, "Morning standup", None, "0 8 * * *", None, None) + ): + result = tool._add_job(None, "Morning standup", None, "0 8 * * *", None, None) assert result.startswith("Created job") job = tool._cron.list_jobs()[0] @@ -334,11 +333,12 @@ def test_add_cron_job_defaults_to_tool_timezone(tmp_path) -> None: def test_add_at_job_uses_default_timezone_for_naive_datetime(tmp_path) -> None: tool = _make_tool_with_tz(tmp_path, "Asia/Shanghai") - tool.set_context( + with request_context( RequestContext(channel="telegram", chat_id="chat-1", session_key="telegram:chat-1") - ) - - result = tool._add_job(None, "Morning reminder", None, None, None, "2026-03-25T08:00:00") + ): + result = tool._add_job( + None, "Morning reminder", None, None, None, "2026-03-25T08:00:00" + ) assert result.startswith("Created job") job = tool._cron.list_jobs()[0] @@ -348,11 +348,10 @@ def test_add_at_job_uses_default_timezone_for_naive_datetime(tmp_path) -> None: def test_add_job_binds_current_session_key(tmp_path) -> None: tool = _make_tool(tmp_path) - tool.set_context( + with request_context( RequestContext(channel="telegram", chat_id="chat-1", session_key="telegram:chat-1") - ) - - result = tool._add_job(None, "Morning standup", 60, None, None, None) + ): + result = tool._add_job(None, "Morning standup", 60, None, None, None) assert result.startswith("Created job") job = tool._cron.list_jobs()[0] @@ -366,9 +365,8 @@ def test_add_job_binds_current_session_key(tmp_path) -> None: def test_add_job_requires_session_key(tmp_path) -> None: tool = _make_tool(tmp_path) - tool.set_context(RequestContext(channel="telegram", chat_id="chat-1")) - - result = tool._add_job(None, "Background refresh", 60, None, None, None) + with request_context(RequestContext(channel="telegram", chat_id="chat-1")): + result = tool._add_job(None, "Background refresh", 60, None, None, None) assert result == "Error: scheduled cron jobs must be created from a chat session" assert tool._cron.list_jobs() == [] @@ -403,11 +401,10 @@ def test_validate_params_requires_message_only_for_add(tmp_path) -> None: def test_add_job_empty_message_returns_actionable_error(tmp_path) -> None: tool = _make_tool(tmp_path) - tool.set_context( + with request_context( RequestContext(channel="telegram", chat_id="chat-1", session_key="telegram:chat-1") - ) - - result = tool._add_job(None, "", 60, None, None, None) + ): + result = tool._add_job(None, "", 60, None, None, None) assert "action='add' requires a non-empty 'message'" in result assert "Retry including message=" in result @@ -417,11 +414,15 @@ def test_add_job_captures_owner_and_origin_without_legacy_delivery_fields(tmp_pa """CronTool stores owner/session identity separately from origin delivery context.""" tool = _make_tool(tmp_path) meta = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}} - tool.set_context(RequestContext( - channel="slack", chat_id="C99", metadata=meta, session_key="slack:C99:111.222" - )) - - result = tool._add_job("test", "say hi", 60, None, None, None) + with request_context( + RequestContext( + channel="slack", + chat_id="C99", + metadata=meta, + session_key="slack:C99:111.222", + ) + ): + result = tool._add_job("test", "say hi", 60, None, None, None) assert "Created job" in result jobs = tool._cron.list_jobs() diff --git a/tests/cron/test_cron_tool_schema_contract.py b/tests/cron/test_cron_tool_schema_contract.py index b0454b9f..9fadd1e2 100644 --- a/tests/cron/test_cron_tool_schema_contract.py +++ b/tests/cron/test_cron_tool_schema_contract.py @@ -9,9 +9,11 @@ and tightens the runtime error for ``add`` without ``message``. from __future__ import annotations +from collections.abc import Iterator + import pytest -from nanobot.agent.tools.context import RequestContext +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.cron import CronTool from nanobot.agent.tools.registry import ToolRegistry @@ -39,14 +41,14 @@ class _SvcStub: @pytest.fixture -def registry() -> ToolRegistry: +def registry() -> Iterator[ToolRegistry]: tool = CronTool(_SvcStub(), default_timezone="UTC") - tool.set_context( - RequestContext(channel="channel", chat_id="chat-id", session_key="channel:chat-id") - ) reg = ToolRegistry() reg.register(tool) - return reg + with request_context( + RequestContext(channel="channel", chat_id="chat-id", session_key="channel:chat-id") + ): + yield reg class TestSchemaContract: diff --git a/tests/test_tool_contextvars.py b/tests/test_tool_contextvars.py index 4300ff90..1d6a8d2b 100644 --- a/tests/test_tool_contextvars.py +++ b/tests/test_tool_contextvars.py @@ -4,8 +4,7 @@ import asyncio import pytest -from nanobot.agent.loop import AgentLoop -from nanobot.agent.tools.context import RequestContext +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.cron import CronTool from nanobot.agent.tools.message import MessageTool from nanobot.agent.tools.spawn import SpawnTool @@ -26,16 +25,16 @@ async def test_message_tool_keeps_task_local_context() -> None: tool = MessageTool(send_callback=send_callback) async def task_one() -> str: - tool.set_context(RequestContext(channel="feishu", chat_id="chat-a")) - entered.set() - await release.wait() - return await tool.execute(content="one") + with request_context(RequestContext(channel="feishu", chat_id="chat-a")): + entered.set() + await release.wait() + return await tool.execute(content="one") async def task_two() -> str: await entered.wait() - tool.set_context(RequestContext(channel="email", chat_id="chat-b")) - release.set() - return await tool.execute(content="two") + with request_context(RequestContext(channel="email", chat_id="chat-b")): + release.set() + return await tool.execute(content="two") result_one, result_two = await asyncio.gather(task_one(), task_two()) @@ -75,16 +74,16 @@ async def test_spawn_tool_keeps_task_local_context() -> None: tool = SpawnTool(_Manager()) async def task_one() -> str: - tool.set_context(RequestContext(channel="whatsapp", chat_id="chat-a")) - entered.set() - await release.wait() - return await tool.execute(task="one") + with request_context(RequestContext(channel="whatsapp", chat_id="chat-a")): + entered.set() + await release.wait() + return await tool.execute(task="one") async def task_two() -> str: await entered.wait() - tool.set_context(RequestContext(channel="telegram", chat_id="chat-b")) - release.set() - return await tool.execute(task="two") + with request_context(RequestContext(channel="telegram", chat_id="chat-b")): + release.set() + return await tool.execute(task="two") result_one, result_two = await asyncio.gather(task_one(), task_two()) @@ -101,20 +100,20 @@ async def test_cron_tool_keeps_task_local_context(tmp_path) -> None: release = asyncio.Event() async def task_one() -> str: - tool.set_context( + with request_context( RequestContext(channel="feishu", chat_id="chat-a", session_key="feishu:chat-a") - ) - entered.set() - await release.wait() - return await tool.execute(action="add", message="first", every_seconds=60) + ): + entered.set() + await release.wait() + return await tool.execute(action="add", message="first", every_seconds=60) async def task_two() -> str: await entered.wait() - tool.set_context( + with request_context( RequestContext(channel="email", chat_id="chat-b", session_key="email:chat-b") - ) - release.set() - return await tool.execute(action="add", message="second", every_seconds=60) + ): + release.set() + return await tool.execute(action="add", message="second", every_seconds=60) result_one, result_two = await asyncio.gather(task_one(), task_two()) @@ -133,24 +132,25 @@ async def test_cron_tool_keeps_task_local_context(tmp_path) -> None: @pytest.mark.asyncio -async def test_message_tool_basic_set_context_and_execute() -> None: - """Single task: set_context then execute should route correctly.""" +async def test_message_tool_basic_request_context_and_execute() -> None: + """A bound request context should route a single execution correctly.""" seen: list[tuple[str, str, str]] = [] async def send_callback(msg): seen.append((msg.channel, msg.chat_id, msg.content)) tool = MessageTool(send_callback=send_callback) - tool.set_context(RequestContext(channel="telegram", chat_id="chat-123", message_id="msg-456")) - - result = await tool.execute(content="hello") + with request_context( + RequestContext(channel="telegram", chat_id="chat-123", message_id="msg-456") + ): + result = await tool.execute(content="hello") assert result == "Message sent to telegram:chat-123" assert seen == [("telegram", "chat-123", "hello")] @pytest.mark.asyncio -async def test_message_tool_default_values_without_set_context() -> None: - """Without set_context, constructor defaults should be used.""" +async def test_message_tool_default_values_without_request_context() -> None: + """Without a request context, constructor defaults should be used.""" seen: list[tuple[str, str, str]] = [] async def send_callback(msg): @@ -168,8 +168,8 @@ async def test_message_tool_default_values_without_set_context() -> None: @pytest.mark.asyncio -async def test_spawn_tool_basic_set_context_and_execute() -> None: - """Single task: set_context then execute should pass correct origin.""" +async def test_spawn_tool_basic_request_context_and_execute() -> None: + """A bound request context should provide the correct origin.""" seen: list[tuple[str, str, str]] = [] class _Manager: @@ -194,16 +194,15 @@ async def test_spawn_tool_basic_set_context_and_execute() -> None: return f"ok: {task}" tool = SpawnTool(_Manager()) - tool.set_context(RequestContext(channel="feishu", chat_id="chat-abc")) - - result = await tool.execute(task="do something") + with request_context(RequestContext(channel="feishu", chat_id="chat-abc")): + result = await tool.execute(task="do something") assert result == "ok: do something" assert seen == [("feishu", "chat-abc", "feishu:chat-abc")] @pytest.mark.asyncio -async def test_spawn_tool_default_values_without_set_context() -> None: - """Without set_context, default cli:direct should be used.""" +async def test_spawn_tool_default_values_without_request_context() -> None: + """Without a request context, default cli:direct should be used.""" seen: list[tuple[str, str, str]] = [] class _Manager: @@ -234,14 +233,13 @@ async def test_spawn_tool_default_values_without_set_context() -> None: @pytest.mark.asyncio -async def test_cron_tool_basic_set_context_and_execute(tmp_path) -> None: - """Single task: set_context then add job should use correct target.""" +async def test_cron_tool_basic_request_context_and_execute(tmp_path) -> None: + """A bound request context should provide the correct cron owner.""" tool = CronTool(CronService(tmp_path / "jobs.json")) - tool.set_context( + with request_context( RequestContext(channel="wechat", chat_id="user-789", session_key="wechat:user-789") - ) - - result = await tool.execute(action="add", message="standup", every_seconds=300) + ): + result = await tool.execute(action="add", message="standup", every_seconds=300) assert result.startswith("Created job") jobs = tool._cron.list_jobs() @@ -256,23 +254,15 @@ async def test_webui_cron_tool_uses_origin_session_when_unified_enabled(tmp_path """WebUI-created cron jobs stay attached to the creating chat.""" tool = CronTool(CronService(tmp_path / "jobs.json")) - class _Tools: - tool_names = ["cron"] - - def get(self, name: str): - return tool if name == "cron" else None - - loop = object.__new__(AgentLoop) - loop._unified_session = True - loop.tools = _Tools() - loop._set_tool_context( - "websocket", - "chat-123", - metadata={"webui": True}, - session_key=UNIFIED_SESSION_KEY, - ) - - result = await tool.execute(action="add", message="standup", every_seconds=300) + with request_context( + RequestContext( + channel="websocket", + chat_id="chat-123", + metadata={"webui": True}, + session_key=UNIFIED_SESSION_KEY, + ) + ): + result = await tool.execute(action="add", message="standup", every_seconds=300) assert result.startswith("Created job") jobs = tool._cron.list_jobs() @@ -287,16 +277,15 @@ async def test_webui_cron_tool_uses_origin_session_when_unified_enabled(tmp_path async def test_cron_tool_preserves_thread_scoped_session_key(tmp_path) -> None: """Channel-provided thread session keys should remain the cron owner.""" tool = CronTool(CronService(tmp_path / "jobs.json")) - tool.set_context( + with request_context( RequestContext( channel="slack", chat_id="C123", metadata={"slack": {"thread_ts": "1700.42"}}, session_key="slack:C123:1700.42", ) - ) - - result = await tool.execute(action="add", message="check thread", every_seconds=300) + ): + result = await tool.execute(action="add", message="check thread", every_seconds=300) assert result.startswith("Created job") jobs = tool._cron.list_jobs() @@ -309,7 +298,7 @@ async def test_cron_tool_preserves_thread_scoped_session_key(tmp_path) -> None: @pytest.mark.asyncio async def test_cron_tool_no_context_returns_error(tmp_path) -> None: - """Without set_context, add should fail with a clear error.""" + """Without a request context, add should fail with a clear error.""" tool = CronTool(CronService(tmp_path / "jobs.json")) result = await tool.execute(action="add", message="test", every_seconds=60) diff --git a/tests/tools/test_message_tool.py b/tests/tools/test_message_tool.py index 3da2f628..73f58dba 100644 --- a/tests/tools/test_message_tool.py +++ b/tests/tools/test_message_tool.py @@ -2,6 +2,7 @@ import os import pytest +from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.message import MessageTool from nanobot.bus.events import OutboundMessage from nanobot.config.paths import get_workspace_path @@ -108,11 +109,9 @@ async def test_message_tool_inherits_metadata_for_same_target() -> None: tool = MessageTool(send_callback=_send) slack_meta = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}} - from nanobot.agent.tools.context import RequestContext - tool.set_context(RequestContext(channel="slack", chat_id="C123", metadata=slack_meta)) - - await tool.execute(content="thread reply") + with request_context(RequestContext(channel="slack", chat_id="C123", metadata=slack_meta)): + await tool.execute(content="thread reply") assert sent[0].metadata == slack_meta @@ -125,18 +124,23 @@ async def test_message_tool_clears_metadata_when_context_has_none() -> None: sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext + rich_context = RequestContext( + channel="slack", + chat_id="C123", + metadata={"slack": {"thread_ts": "111.222", "channel_type": "channel"}}, + ) + with request_context(rich_context): + await tool.execute(content="thread reply") + sent.clear() - tool.set_context( + with request_context( RequestContext( channel="slack", chat_id="C123", - metadata={"slack": {"thread_ts": "111.222", "channel_type": "channel"}}, + metadata={}, ), - ) - tool.set_context(RequestContext(channel="slack", chat_id="C123", metadata={})) - - await tool.execute(content="plain reply") + ): + await tool.execute(content="plain reply") assert sent[0].metadata == {} @@ -149,17 +153,14 @@ async def test_message_tool_does_not_inherit_metadata_for_cross_target() -> None sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - - tool.set_context( + with request_context( RequestContext( channel="slack", chat_id="C123", metadata={"slack": {"thread_ts": "111.222", "channel_type": "channel"}}, ), - ) - - await tool.execute(content="channel reply", channel="slack", chat_id="C999") + ): + await tool.execute(content="channel reply", channel="slack", chat_id="C999") assert sent[0].metadata == {} @@ -337,15 +338,17 @@ async def test_message_tool_tracks_turn_media_for_same_target(tmp_path) -> None: sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - - tool.set_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})) - tool.start_turn() f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - await tool.execute(content="see file", channel="websocket", chat_id="chat-1", media=[str(f)]) - - assert tool.turn_delivered_media_paths() == [str(f.resolve())] + with request_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})): + tool.start_turn() + await tool.execute( + content="see file", + channel="websocket", + chat_id="chat-1", + media=[str(f)], + ) + assert tool.turn_delivered_media_paths() == [str(f.resolve())] @pytest.mark.asyncio @@ -354,15 +357,13 @@ async def test_message_tool_start_turn_clears_tracked_media(tmp_path) -> None: pass tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - - tool.set_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})) - tool.start_turn() f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - await tool.execute(content="see file", media=[str(f)]) - tool.start_turn() - assert tool.turn_delivered_media_paths() == [] + with request_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})): + tool.start_turn() + await tool.execute(content="see file", media=[str(f)]) + tool.start_turn() + assert tool.turn_delivered_media_paths() == [] @pytest.mark.asyncio @@ -371,18 +372,16 @@ async def test_message_tool_cross_target_does_not_track_turn_media(tmp_path) -> pass tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - - tool.set_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})) f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - await tool.execute( - content="see file", - channel="telegram", - chat_id="tg-other", - media=[str(f)], - ) - assert tool.turn_delivered_media_paths() == [] + with request_context(RequestContext(channel="websocket", chat_id="chat-1", metadata={})): + await tool.execute( + content="see file", + channel="telegram", + chat_id="tg-other", + media=[str(f)], + ) + assert tool.turn_delivered_media_paths() == [] @pytest.mark.asyncio @@ -393,18 +392,16 @@ async def test_message_tool_rejects_wrong_explicit_ws_chat_id(tmp_path) -> None: sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - conv = "550e8400-e29b-41d4-a716-446655440000" - tool.set_context(RequestContext(channel="websocket", chat_id=conv, metadata={})) f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - result = await tool.execute( - content="see file", - channel="websocket", - chat_id="anon-deadbeefcafe", - media=[str(f)], - ) + with request_context(RequestContext(channel="websocket", chat_id=conv, metadata={})): + result = await tool.execute( + content="see file", + channel="websocket", + chat_id="anon-deadbeefcafe", + media=[str(f)], + ) assert result.startswith("Error: chat_id does not match") assert sent == [] @@ -417,18 +414,16 @@ async def test_message_tool_allows_ws_explicit_when_matches_context(tmp_path) -> sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - conv = "550e8400-e29b-41d4-a716-446655440000" - tool.set_context(RequestContext(channel="websocket", chat_id=conv, metadata={})) f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - result = await tool.execute( - content="see file", - channel="websocket", - chat_id=conv, - media=[str(f)], - ) + with request_context(RequestContext(channel="websocket", chat_id=conv, metadata={})): + result = await tool.execute( + content="see file", + channel="websocket", + chat_id=conv, + media=[str(f)], + ) assert result.startswith("Message sent") assert sent[0].chat_id == conv @@ -442,18 +437,16 @@ async def test_message_tool_cli_context_may_target_other_ws_chat(tmp_path) -> No sent.append(msg) tool = MessageTool(send_callback=_send) - from nanobot.agent.tools.context import RequestContext - target = "550e8400-e29b-41d4-a716-446655440000" - tool.set_context(RequestContext(channel="cli", chat_id="direct", metadata={})) f = tmp_path / "doc.md" f.write_text("hello", encoding="utf-8") - result = await tool.execute( - content="ping", - channel="websocket", - chat_id=target, - media=[str(f)], - ) + with request_context(RequestContext(channel="cli", chat_id="direct", metadata={})): + result = await tool.execute( + content="ping", + channel="websocket", + chat_id=target, + media=[str(f)], + ) assert result.startswith("Message sent") assert sent[0].channel == "websocket" assert sent[0].chat_id == target diff --git a/tests/tools/test_message_tool_suppress.py b/tests/tools/test_message_tool_suppress.py index 1a08311e..b06f5413 100644 --- a/tests/tools/test_message_tool_suppress.py +++ b/tests/tools/test_message_tool_suppress.py @@ -156,11 +156,12 @@ class TestMessageToolTurnTracking: def test_sent_in_turn_tracks_same_target(self) -> None: tool = MessageTool() - from nanobot.agent.tools.context import RequestContext - tool.set_context(RequestContext(channel="feishu", chat_id="chat1")) - assert not tool._sent_in_turn - tool._sent_in_turn = True - assert tool._sent_in_turn + 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()