refactor(agent): unify request context routing

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()

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 (

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()
 

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})"
 

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()

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:

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, (

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)
 
     # ------------------------------------------------------------------

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(),
         )

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: