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

refactor(agent): unify request context routing

Browse Source
This commit is contained in:
chengyongru
2026-07-10 17:54:34 +08:00
committed by Xubin Ren
parent bb3b449e09
commit 42d7ad34a4
20 changed files with 363 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/tools/context.py
+11
View File
@@ -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
+12 -19
View File
@@ -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
+4 -15
View File
@@ -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
+25 -27
View File
@@ -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
+7
View File
@@ -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 -11
View File
@@ -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
+14 -21
View File
@@ -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(),
)