nanobot 工具上下文注入机制深度分析

def _register_default_tools(self) -> None:
    # 工具只注册一次
    message_tool = MessageTool(send_callback=self.bus.publish_outbound)
    self.tools.register(message_tool)

    spawn_tool = SpawnTool(manager=self.subagents)
    self.tools.register(spawn_tool)

class MessageTool(Tool):
    def __init__(
        self,
        send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
        default_channel: str = "",
        default_chat_id: str = ""
    ):
        self._send_callback = send_callback
        self._default_channel = default_channel  # 动态上下文
        self._default_chat_id = default_chat_id  # 动态上下文

def set_context(self, channel: str, chat_id: str) -> None:
    """Set the current message context."""
    self._default_channel = channel
    self._default_chat_id = chat_id

async def execute(
    self,
    content: str,
    channel: str | None = None,
    chat_id: str | None = None,
    **kwargs: Any
) -> str:
    # 优先级：显式参数 > 上下文注入值
    channel = channel or self._default_channel
    chat_id = chat_id or self._default_chat_id

    if not channel or not chat_id:
        return "Error: No target channel/chat specified"

    if not self._send_callback:
        return "Error: Message sending not configured"

    msg = OutboundMessage(
        channel=channel,
        chat_id=chat_id,
        content=content
    )

    try:
        await self._send_callback(msg)
        return f"Message sent to {channel}:{chat_id}"
    except Exception as e:
        return f"Error sending message: {str(e)}"

# 获取或创建会话
session = self.sessions.get_or_create(msg.session_key)

# 更新工具上下文
message_tool = self.tools.get("message")
if isinstance(message_tool, MessageTool):
    message_tool.set_context(msg.channel, msg.chat_id)

# 解析原始上下文
if ":" in msg.chat_id:
    parts = msg.chat_id.split(":", 1)
    origin_channel = parts[0]
    origin_chat_id = parts[1]

# 更新工具上下文（使用原始会话）
message_tool = self.tools.get("message")
if isinstance(message_tool, MessageTool):
    message_tool.set_context(origin_channel, origin_chat_id)

MessageTool(default_channel="", default_chat_id="")

# 如果没有调用 set_context()
channel = channel or self._default_channel  # -> ""
if not channel or not chat_id:
    return "Error: No target channel/chat specified"

class SpawnTool(Tool):
    def __init__(self, manager: "SubagentManager"):
        self._manager = manager
        self._origin_channel = "cli"  # 默认值
        self._origin_chat_id = "direct"  # 默认值

def set_context(self, channel: str, chat_id: str) -> None:
    """Set the origin context for subagent announcements."""
    self._origin_channel = channel
    self._origin_chat_id = chat_id

async def execute(self, task: str, label: str | None = None, **kwargs: Any) -> str:
    """Spawn a subagent to execute the given task."""
    return await self._manager.spawn(
        task=task,
        label=label,
        origin_channel=self._origin_channel,  # 传递上下文
        origin_chat_id=self._origin_chat_id,  # 传递上下文
    )

async def spawn(
    self,
    task: str,
    label: str | None = None,
    origin_channel: str = "cli",
    origin_chat_id: str = "direct",
) -> str:
    task_id = str(uuid.uuid4())[:8]
    display_label = label or task[:30] + ("..." if len(task) > 30 else "")

    # 保存原始会话信息
    origin = {
        "channel": origin_channel,
        "chat_id": origin_chat_id,
    }

    # 创建后台任务
    bg_task = asyncio.create_task(
        self._run_subagent(task_id, task, display_label, origin)
    )

    return f"Subagent [{display_label}] started (id: {task_id}). I'll notify you when it completes."

async def _announce_result(
    self,
    task_id: str,
    label: str,
    task: str,
    result: str,
    origin: dict[str, str],  # 原始会话信息
    status: str,
) -> None:
    """Announce the subagent result to the main agent via the message bus."""

    announce_content = f"""[Subagent '{label}' {status_text}]

Task: {task}

Result:
{result}

Summarize this naturally for the user. Keep it brief (1-2 sentences). Do not mention technical details like "subagent" or task IDs."""

    # 通过 chat_id 传递原始会话信息
    msg = InboundMessage(
        channel="system",  # 系统消息
        sender_id="subagent",
        chat_id=f"{origin['channel']}:{origin['chat_id']}",  # 编码原始信息
        content=announce_content,
    )

    await self.bus.publish_inbound(msg)

async def _process_system_message(self, msg: InboundMessage) -> OutboundMessage | None:
    # 解析原始会话信息
    if ":" in msg.chat_id:
        parts = msg.chat_id.split(":", 1)
        origin_channel = parts[0]
        origin_chat_id = parts[1]
    else:
        origin_channel = "cli"
        origin_chat_id = msg.chat_id

    # 使用原始会话
    session_key = f"{origin_channel}:{origin_chat_id}"
    session = self.sessions.get_or_create(session_key)

    # 更新工具上下文（恢复原始会话）
    message_tool = self.tools.get("message")
    if isinstance(message_tool, MessageTool):
        message_tool.set_context(origin_channel, origin_chat_id)

    spawn_tool = self.tools.get("spawn")
    if isinstance(spawn_tool, SpawnTool):
        spawn_tool.set_context(origin_channel, origin_chat_id)

    # ... 处理消息并返回响应
    return OutboundMessage(
        channel=origin_channel,
        chat_id=origin_chat_id,
        content=final_content
    )

async def run(self) -> None:
    while self._running:
        # 阻塞等待下一条消息
        msg = await asyncio.wait_for(
            self.bus.consume_inbound(),
            timeout=1.0
        )

        # 处理消息（串行）
        response = await self._process_message(msg)

# 子代理没有 MessageTool 和 SpawnTool
tools = ToolRegistry()
tools.register(ReadFileTool())
tools.register(WriteFileTool())
tools.register(ListDirTool())
tools.register(ExecTool(working_dir=str(self.workspace)))
tools.register(WebSearchTool(api_key=self.brave_api_key))
tools.register(WebFetchTool())
# 注意：没有 MessageTool 和 SpawnTool

# 解析原始会话信息
if ":" in msg.chat_id:
    parts = msg.chat_id.split(":", 1)
    origin_channel = parts[0]
    origin_chat_id = parts[1]
else:
    # 回退到默认值
    origin_channel = "cli"
    origin_chat_id = msg.chat_id

class MessageTool(Tool):
    async def execute(
        self,
        content: str,
        channel: str,  # 必填参数
        chat_id: str,  # 必填参数
    ) -> str:
        msg = OutboundMessage(channel=channel, chat_id=chat_id, content=content)
        await self._send_callback(msg)
        return f"Message sent to {channel}:{chat_id}"

# LLM 需要在每次调用时生成这些参数
tool_calls = [
    {
        "name": "message",
        "arguments": {
            "content": "Hello",
            "channel": "telegram",  # 冗余
            "chat_id": "123"        # 冗余
        }
    }
]

class AgentLoop:
    def __init__(self, ...):
        self._session_tools = {}  # {session_key: ToolRegistry}

    async def _process_message(self, msg: InboundMessage):
        session_key = msg.session_key

        # 为每个会话创建独立工具
        if session_key not in self._session_tools:
            tools = ToolRegistry()
            tools.register(MessageTool(..., default_channel=msg.channel, default_chat_id=msg.chat_id))
            tools.register(SpawnTool(...))
            self._session_tools[session_key] = tools

        tools = self._session_tools[session_key]
        # 使用 tools 执行...

@dataclass
class ToolContext:
    channel: str
    chat_id: str
    session_key: str

class Tool(ABC):
    @abstractmethod
    async def execute(self, ctx: ToolContext, **kwargs: Any) -> str:
        pass

class MessageTool(Tool):
    async def execute(self, ctx: ToolContext, content: str, **kwargs) -> str:
        msg = OutboundMessage(channel=ctx.channel, chat_id=ctx.chat_id, content=content)
        await self._send_callback(msg)
        return "Sent"

import contextvars

# 定义上下文变量
current_channel = contextvars.ContextVar("current_channel", default="")
current_chat_id = contextvars.ContextVar("current_chat_id", default="")

class MessageTool(Tool):
    async def execute(self, content: str, **kwargs) -> str:
        channel = current_channel.get()
        chat_id = current_chat_id.get()
        # ...

def set_context(self, channel: str, chat_id: str) -> None:
    self._default_channel = channel
    self._default_chat_id = chat_id

class MessageTool(Tool):
    def __init__(
        self,
        send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
        default_channel: str = "",
        default_chat_id: str = "",
        default_user_id: str = "",  # 新增
        default_thread_id: str = "",  # 新增
    ):
        self._send_callback = send_callback
        self._default_channel = default_channel
        self._default_chat_id = default_chat_id
        self._default_user_id = default_user_id  # 新增
        self._default_thread_id = default_thread_id  # 新增

    def set_context(
        self,
        channel: str,
        chat_id: str,
        user_id: str | None = None,  # 新增
        thread_id: str | None = None,  # 新增
    ) -> None:
        self._default_channel = channel
        self._default_chat_id = chat_id
        if user_id is not None:
            self._default_user_id = user_id
        if thread_id is not None:
            self._default_thread_id = thread_id

@dataclass
class ToolContext:
    channel: str
    chat_id: str
    user_id: str = ""
    thread_id: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)

class MessageTool(Tool):
    def __init__(
        self,
        send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
        default_context: ToolContext | None = None,
    ):
        self._send_callback = send_callback
        self._context = default_context or ToolContext(channel="", chat_id="")

    def set_context(self, context: ToolContext) -> None:
        self._context = context

    async def execute(
        self,
        content: str,
        channel: str | None = None,
        chat_id: str | None = None,
        **kwargs
    ) -> str:
        # 优先级：显式参数 > 上下文对象
        channel = channel or self._context.channel
        chat_id = chat_id or self._context.chat_id
        user_id = kwargs.get("user_id") or self._context.user_id
        # ...

class MessageTool(Tool):
    def __init__(self, ...):
        self._send_callback = send_callback
        self._context = {}  # 字典存储

    def set_context(self, **kwargs) -> None:
        self._context.update(kwargs)

    def get_context(self, key: str, default: Any = None) -> Any:
        return self._context.get(key, default)

@dataclass
class ToolContext:
    channel: str
    chat_id: str
    user_id: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)

    def inherit(self, **kwargs) -> "ToolContext":
        """创建继承当前上下文的新上下文对象"""
        new_context = ToolContext(
            channel=self.channel,
            chat_id=self.chat_id,
            user_id=self.user_id,
            metadata=self.metadata.copy(),  # 深拷贝
        )
        # 覆盖指定的字段
        for key, value in kwargs.items():
            if hasattr(new_context, key):
                setattr(new_context, key, value)
        return new_context

# 主代理上下文
main_context = ToolContext(
    channel="telegram",
    chat_id="123",
    user_id="alice",
    metadata={"locale": "en", "timezone": "UTC"}
)

# 子代理继承上下文（但修改 chat_id）
sub_context = main_context.inherit(chat_id="subagent_456")
# sub_context: {channel="telegram", chat_id="subagent_456", user_id="alice", ...}

class SubagentManager:
    async def spawn(
        self,
        task: str,
        parent_context: ToolContext | None = None,  # 新增参数
        **kwargs
    ) -> str:
        # 子代理上下文继承（但标记为子代理）
        if parent_context:
            sub_context = parent_context.inherit(
                chat_id=f"subagent:{uuid.uuid4()[:8]}",
                metadata={"parent_channel": parent_context.channel}
            )
        else:
            sub_context = ToolContext(channel="system", chat_id="")

        # 传递给子代理
        bg_task = asyncio.create_task(
            self._run_subagent(task_id, task, display_label, sub_context)
        )

@dataclass
class ToolContext:
    channel: str
    chat_id: str

    def validate(self) -> tuple[bool, str]:
        """验证上下文的有效性"""
        if not self.channel:
            return False, "Channel is required"
        if not self.chat_id:
            return False, "Chat ID is required"
        if len(self.channel) > 50:
            return False, "Channel too long"
        return True, ""

class MessageTool(Tool):
    def set_context(self, context: ToolContext) -> None:
        valid, error = context.validate()
        if not valid:
            raise ValueError(f"Invalid context: {error}")
        self._context = context

class ToolRegistry:
    def __init__(self):
        self._tools: dict[str, Tool] = {}
        self._context_expiry: dict[str, datetime] = {}

    def set_context(self, tool_name: str, context: ToolContext, ttl: int = 3600) -> None:
        """设置工具上下文（带过期时间）"""
        tool = self._tools.get(tool_name)
        if tool:
            tool.set_context(context)
            self._context_expiry[tool_name] = datetime.now() + timedelta(seconds=ttl)

    def get_context(self, tool_name: str) -> ToolContext | None:
        """获取工具上下文（检查是否过期）"""
        expiry = self._context_expiry.get(tool_name)
        if expiry and datetime.now() > expiry:
            # 过期则清理
            tool = self._tools.get(tool_name)
            if tool:
                tool.set_context(ToolContext(channel="", chat_id=""))
            del self._context_expiry[tool_name]
            return None
        return tool._context if tool else None

import threading

class ThreadSafeContext:
    def __init__(self):
        self._context = ToolContext(channel="", chat_id="")
        self._lock = threading.Lock()

    def set_context(self, context: ToolContext) -> None:
        with self._lock:
            self._context = context

    def get_context(self) -> ToolContext:
        with self._lock:
            return self._context

def set_context(self, channel: str, chat_id: str) -> None:
    self._default_channel = channel
    self._default_chat_id = chat_id

msg = InboundMessage(
    channel="system",
    sender_id="subagent",
    chat_id=f"{origin['channel']}:{origin['chat_id']}",  # 编码原始信息
    content=announce_content,
)

# 子代理工具（无 MessageTool 和 SpawnTool）
tools.register(ReadFileTool())
tools.register(WriteFileTool())
# ...

async def run(self) -> None:
    while self._running:
        msg = await self.bus.consume_inbound()
        response = await self._process_message(msg)  # 串行处理

def __init__(self, ...):
    self.tools = ToolRegistry()
    self._register_default_tools()  # 只注册一次

def _register_default_tools(self) -> None:
    message_tool = MessageTool(...)
    self.tools.register(message_tool)  # 单例

channel = channel or self._default_channel

def __init__(self, ..., default_channel: str = "", default_chat_id: str = ""):
    self._default_channel = default_channel
    self._default_chat_id = default_chat_id

msg = InboundMessage(
    channel="cli",
    sender_id="user",
    chat_id="direct",
    content=content
)

# 当前：字符串参数
message_tool.set_context("telegram", "123")

# 改进：Context 对象
context = ToolContext(channel="telegram", chat_id="123")
message_tool.set_context(context)

class Tool:
    def on_context_set(self, context: ToolContext) -> None:
        """上下文设置时的回调"""
        pass

    def on_context_clear(self) -> None:
        """上下文清理时的回调"""
        pass

class ContextMiddleware:
    async def before_context_set(self, context: ToolContext) -> ToolContext:
        """上下文设置前的预处理"""
        context.metadata["timestamp"] = datetime.now()
        return context

    async def after_context_set(self, context: ToolContext) -> None:
        """上下文设置后的后处理"""
        logger.debug(f"Context set: {context.channel}:{context.chat_id}")