"""协议对等端模块。 定义 Peer 类，封装双向传输通道上的消息收发、初始化握手、能力调用、 流式事件转发与取消处理。这里的 peer 指"通信对端/本端"这一网络协议概念， 而不是业务上的用户、群聊或会话对象。 核心职责： - 消息序列化/反序列化 - 初始化握手协议 - 能力调用（同步/流式） - 取消处理 - 连接生命周期管理 消息处理： 入站: ResultMessage -> 唤醒等待的 Future EventMessage -> 投递到流式队列 InitializeMessage -> 调用 initialize_handler InvokeMessage -> 创建任务调用 invoke_handler CancelMessage -> 取消对应的任务 出站: initialize() -> InitializeMessage invoke() -> InvokeMessage(stream=False) invoke_stream() -> InvokeMessage(stream=True) cancel() -> CancelMessage 与旧版对比： 旧版 JSON-RPC: - 分离的 JSONRPCClient 和 JSONRPCServer - 通过 method 字段区分操作类型 - 使用 JSONRPCRequest/Response 消息类型 - 流式通过独立的 notification 实现 - 无统一的取消机制 新版 Peer: - 统一的 Peer 抽象，既是客户端也是服务端 - 通过 type 字段区分消息类型 - 使用 InitializeMessage/InvokeMessage/EventMessage 等 - 流式通过 EventMessage(phase=delta) 实现 - 统一的 CancelMessage 取消机制 使用示例： # 作为客户端发起调用 peer = Peer(transport=transport, peer_info=PeerInfo(...)) await peer.start() output = await peer.initialize(handlers) result = await peer.invoke("llm.chat", {"prompt": "hello"}) # 作为服务端处理调用 peer.set_invoke_handler(my_handler) await peer.start() 消息处理流程： 入站消息: ResultMessage -> 唤醒等待的 Future EventMessage -> 投递到流式队列 InitializeMessage -> 调用 _initialize_handler InvokeMessage -> 创建任务调用 _invoke_handler CancelMessage -> 取消对应的任务 出站消息: initialize() -> InitializeMessage invoke() -> InvokeMessage(stream=False) invoke_stream() -> InvokeMessage(stream=True) cancel() -> CancelMessage 取消机制： - CancelToken 用于检查取消状态 - 入站任务在收到 CancelMessage 时被取消 - 早到取消：在任务执行前检查 cancel_token，避免竞态条件 `Peer` 把 `Transport` 和 v4 协议消息模型接起来，负责： - 握手与远端元数据缓存 - 请求 ID 关联 - 非流式 / 流式调用分发 - 取消传播 - 连接异常时的统一收口 它本身不做业务路由，真正的执行逻辑交给 `CapabilityRouter` 或 `HandlerDispatcher`。 """ from __future__ import annotations import asyncio import inspect from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any from ..context import CancelToken from ..errors import AstrBotError, ErrorCodes from ..protocol.messages import ( CancelMessage, ErrorPayload, EventMessage, InitializeMessage, InitializeOutput, InvokeMessage, PeerInfo, ResultMessage, parse_message, ) from .capability_router import StreamExecution InitializeHandler = Callable[[InitializeMessage], Awaitable[InitializeOutput]] InvokeHandler = Callable[ [InvokeMessage, CancelToken], Awaitable[dict[str, Any] | StreamExecution] ] CancelHandler = Callable[[str], Awaitable[None]] class Peer: """表示协议连接中的一个对等端。 `Peer` 封装一条双向传输通道上的消息收发、初始化握手、能力调用、 流式事件转发与取消处理。这里的 `peer` 指“通信对端/本端”这一网络 协议概念，而不是业务上的用户、群聊或会话对象。 """ def __init__( self, *, transport, peer_info: PeerInfo, protocol_version: str = "1.0", ) -> None: """创建一个协议对等端实例。 Args: transport: 底层传输实现，负责发送字符串消息并回调入站消息。 peer_info: 当前端点对外声明的身份信息。 protocol_version: 当前端点支持的协议版本，用于初始化握手校验。 """ self.transport = transport self.peer_info = peer_info self.protocol_version = protocol_version self.remote_peer: PeerInfo | None = None self.remote_handlers = [] self.remote_provided_capabilities = [] self.remote_capabilities = [] self.remote_capability_map: dict[str, Any] = {} self.remote_provided_capability_map: dict[str, Any] = {} self.remote_metadata: dict[str, Any] = {} self._initialize_handler: InitializeHandler | None = None self._invoke_handler: InvokeHandler | None = None self._cancel_handler: CancelHandler | None = None self._counter = 0 self._closed = asyncio.Event() self._unusable = False self._stopping = False self._pending_results: dict[str, asyncio.Future[ResultMessage]] = {} self._pending_streams: dict[str, asyncio.Queue[Any]] = {} self._inbound_tasks: dict[ str, tuple[asyncio.Task[None], CancelToken, asyncio.Event] ] = {} self._remote_initialized = asyncio.Event() self._transport_watch_task: asyncio.Task[None] | None = None def set_initialize_handler(self, handler: InitializeHandler) -> None: """注册处理远端 `initialize` 请求的握手处理器。""" self._initialize_handler = handler def set_invoke_handler(self, handler: InvokeHandler) -> None: """注册处理远端 `invoke` 请求的能力调用处理器。""" self._invoke_handler = handler def set_cancel_handler(self, handler: CancelHandler) -> None: """注册处理远端 `cancel` 请求的取消回调。""" self._cancel_handler = handler async def start(self) -> None: """启动传输层并将原始入站消息绑定到当前 `Peer`。""" self._closed.clear() self._unusable = False self._stopping = False self._remote_initialized.clear() self.transport.set_message_handler(self._handle_raw_message) await self.transport.start() self._transport_watch_task = asyncio.create_task(self._watch_transport_closed()) async def stop(self) -> None: """关闭 `Peer` 并清理所有挂起中的请求、流和入站任务。""" if self._closed.is_set(): return self._stopping = True # 终止所有挂起的 RPC，避免调用方永久挂起 for future in list(self._pending_results.values()): if not future.done(): future.set_exception(AstrBotError.internal_error("连接已关闭")) self._pending_results.clear() for queue in list(self._pending_streams.values()): await queue.put(AstrBotError.internal_error("连接已关闭")) self._pending_streams.clear() # 取消所有入站任务 for task, token, _started in list(self._inbound_tasks.values()): token.cancel() task.cancel() self._inbound_tasks.clear() await self.transport.stop() self._closed.set() async def wait_closed(self) -> None: """等待底层传输彻底关闭。""" await self.transport.wait_closed() async def _watch_transport_closed(self) -> None: """监视底层传输的意外关闭，并主动失败挂起调用。""" try: await self.transport.wait_closed() if self._closed.is_set() or self._stopping: return await self._fail_connection( AstrBotError( code=ErrorCodes.NETWORK_ERROR, message="连接已关闭", hint="请检查对端进程或传输连接", retryable=True, ) ) finally: current_task = asyncio.current_task() if self._transport_watch_task is current_task: self._transport_watch_task = None async def wait_until_remote_initialized(self, timeout: float | None = 30.0) -> None: """等待远端完成初始化握手。 Args: timeout: 等待秒数。传入 `None` 表示无限等待。 """ init_waiter = asyncio.create_task(self._remote_initialized.wait()) closed_waiter = asyncio.create_task(self.wait_closed()) try: done, pending = await asyncio.wait( {init_waiter, closed_waiter}, timeout=timeout, return_when=asyncio.FIRST_COMPLETED, ) if not done: raise TimeoutError() if init_waiter in done: return raise AstrBotError.protocol_error("连接在初始化完成前关闭") finally: for task in (init_waiter, closed_waiter): if not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass async def initialize( self, handlers, *, provided_capabilities=None, metadata: dict[str, Any] | None = None, ) -> InitializeOutput: """向远端发送初始化请求并缓存远端声明的能力信息。 Args: handlers: 当前端点声明可接收的处理器列表。 metadata: 附带给远端的握手元数据。 Returns: 远端返回的初始化结果。 """ self._ensure_usable() request_id = self._next_id() future: asyncio.Future[ResultMessage] = ( asyncio.get_running_loop().create_future() ) self._pending_results[request_id] = future await self._send( InitializeMessage( id=request_id, protocol_version=self.protocol_version, peer=self.peer_info, handlers=list(handlers), provided_capabilities=list(provided_capabilities or []), metadata=metadata or {}, ) ) result = await future if result.kind != "initialize_result": raise AstrBotError.protocol_error("initialize 必须收到 initialize_result") if not result.success: self._unusable = True await self.stop() raise AstrBotError.from_payload( result.error.model_dump() if result.error else {} ) output = InitializeOutput.model_validate(result.output) self.remote_peer = output.peer self.remote_capabilities = output.capabilities self.remote_capability_map = {item.name: item for item in output.capabilities} self.remote_metadata = output.metadata return output async def invoke( self, capability: str, payload: dict[str, Any], *, stream: bool = False, request_id: str | None = None, ) -> dict[str, Any]: """发起一次非流式能力调用并等待最终结果。 Args: capability: 远端能力名。 payload: 调用输入。 stream: 必须为 `False`；流式场景应改用 `invoke_stream()`。 request_id: 可选的请求 ID；未提供时自动生成。 """ self._ensure_usable() if stream: raise ValueError("stream=True 请使用 invoke_stream()") request_id = request_id or self._next_id() future: asyncio.Future[ResultMessage] = ( asyncio.get_running_loop().create_future() ) self._pending_results[request_id] = future await self._send( InvokeMessage( id=request_id, capability=capability, input=payload, stream=False, ) ) result = await future if not result.success: raise AstrBotError.from_payload( result.error.model_dump() if result.error else {} ) return result.output async def invoke_stream( self, capability: str, payload: dict[str, Any], *, request_id: str | None = None, ) -> AsyncIterator[EventMessage]: """发起一次流式能力调用并返回事件迭代器。 调用方会收到 `delta` 事件，`started` 会被内部吞掉， `completed` 用于结束迭代，`failed` 会转换为异常抛出。 Args: capability: 远端能力名。 payload: 调用输入。 request_id: 可选的请求 ID；未提供时自动生成。 """ self._ensure_usable() request_id = request_id or self._next_id() queue: asyncio.Queue[Any] = asyncio.Queue() self._pending_streams[request_id] = queue await self._send( InvokeMessage( id=request_id, capability=capability, input=payload, stream=True, ) ) async def iterator() -> AsyncIterator[EventMessage]: try: while True: item = await queue.get() if isinstance(item, Exception): raise item if not isinstance(item, EventMessage): raise AstrBotError.protocol_error("流式调用收到非法事件") if item.phase == "started": continue if item.phase == "delta": yield item continue if item.phase == "completed": break if item.phase == "failed": raise AstrBotError.from_payload( item.error.model_dump() if item.error else {} ) finally: self._pending_streams.pop(request_id, None) return iterator() async def cancel(self, request_id: str, reason: str = "user_cancelled") -> None: """向远端发送取消请求，尝试中止指定 ID 的在途调用。""" await self._send(CancelMessage(id=request_id, reason=reason)) def _next_id(self) -> str: """生成当前连接内递增的消息 ID。""" self._counter += 1 return f"msg_{self._counter:04d}" def _ensure_usable(self) -> None: """确保连接仍处于可用状态，否则立即抛出协议错误。""" if self._unusable: raise AstrBotError.protocol_error("连接已进入不可用状态") async def _handle_raw_message(self, payload: str) -> None: """解析原始消息并分发到对应的消息处理分支。""" try: message = parse_message(payload) if isinstance(message, ResultMessage): await self._handle_result(message) return if isinstance(message, EventMessage): await self._handle_event(message) return if isinstance(message, InitializeMessage): await self._handle_initialize(message) return if isinstance(message, InvokeMessage): token = CancelToken() started = asyncio.Event() task = asyncio.create_task(self._handle_invoke(message, token, started)) self._inbound_tasks[message.id] = (task, token, started) task.add_done_callback( lambda _task, request_id=message.id: self._inbound_tasks.pop( request_id, None ) ) return if isinstance(message, CancelMessage): await self._handle_cancel(message) return except Exception as exc: if isinstance(exc, AstrBotError): error = exc else: error = AstrBotError.protocol_error(f"无法解析协议消息: {exc}") await self._fail_connection(error) raise error from exc async def _handle_initialize(self, message: InitializeMessage) -> None: """处理远端发起的初始化握手并返回握手结果。""" self.remote_peer = message.peer self.remote_handlers = message.handlers self.remote_provided_capabilities = message.provided_capabilities self.remote_provided_capability_map = { item.name: item for item in message.provided_capabilities } self.remote_metadata = message.metadata if self._initialize_handler is None: await self._reject_initialize( message, AstrBotError.protocol_error("对端不接受 initialize"), ) return if message.protocol_version != self.protocol_version: await self._reject_initialize( message, AstrBotError.protocol_version_mismatch( f"服务端支持协议版本 {self.protocol_version}，客户端请求版本 {message.protocol_version}" ), ) return output = await self._initialize_handler(message) await self._send( ResultMessage( id=message.id, kind="initialize_result", success=True, output=output.model_dump(), ) ) self._remote_initialized.set() async def _handle_invoke( self, message: InvokeMessage, token: CancelToken, started: asyncio.Event, ) -> None: """处理远端发起的能力调用，并按流式或非流式协议返回结果。""" try: started.set() token.raise_if_cancelled() if self._invoke_handler is None: raise AstrBotError.capability_not_found(message.capability) execution = await self._invoke_handler(message, token) if inspect.isawaitable(execution): execution = await execution if message.stream: if not isinstance(execution, StreamExecution): raise AstrBotError.protocol_error( "stream=true 必须返回 StreamExecution" ) await self._send(EventMessage(id=message.id, phase="started")) chunks: list[dict[str, Any]] = [] async for chunk in execution.iterator: chunks.append(chunk) await self._send( EventMessage(id=message.id, phase="delta", data=chunk) ) await self._send( EventMessage( id=message.id, phase="completed", output=execution.finalize(chunks), ) ) return if isinstance(execution, StreamExecution): raise AstrBotError.protocol_error("stream=false 不能返回流式执行对象") await self._send( ResultMessage(id=message.id, success=True, output=execution) ) except asyncio.CancelledError: await self._send_cancelled_termination(message) except LookupError as exc: error = AstrBotError.invalid_input(str(exc)) await self._send_error_result(message, error) except AstrBotError as exc: await self._send_error_result(message, exc) except Exception as exc: await self._send_error_result( message, AstrBotError.internal_error(str(exc)) ) async def _handle_cancel(self, message: CancelMessage) -> None: """处理远端取消请求并终止对应的入站任务。""" inbound = self._inbound_tasks.get(message.id) if inbound is None: return task, token, started = inbound token.cancel() if self._cancel_handler is not None: await self._cancel_handler(message.id) if started.is_set(): task.cancel() async def _handle_result(self, message: ResultMessage) -> None: """处理非流式结果消息并唤醒等待中的调用方。""" future = self._pending_results.pop(message.id, None) if future is None: queue = self._pending_streams.get(message.id) if queue is not None: await queue.put( AstrBotError.protocol_error("stream=true 调用不应收到 result") ) return # 检查 future 是否已完成（可能被调用方取消） if not future.done(): future.set_result(message) async def _handle_event(self, message: EventMessage) -> None: """处理流式事件消息并投递到对应请求的事件队列。""" queue = self._pending_streams.get(message.id) if queue is None: future = self._pending_results.get(message.id) if future is not None and not future.done(): future.set_exception( AstrBotError.protocol_error("stream=false 调用不应收到 event") ) return await queue.put(message) async def _send_error_result( self, message: InvokeMessage, error: AstrBotError ) -> None: """根据调用模式，将错误编码为 `result` 或失败事件发回远端。""" if message.stream: await self._send( EventMessage( id=message.id, phase="failed", error=ErrorPayload.model_validate(error.to_payload()), ) ) return await self._send( ResultMessage( id=message.id, success=False, error=ErrorPayload.model_validate(error.to_payload()), ) ) async def _reject_initialize( self, message: InitializeMessage, error: AstrBotError ) -> None: """拒绝一次初始化握手，并把连接标记为不可继续使用。""" await self._send( ResultMessage( id=message.id, kind="initialize_result", success=False, error=ErrorPayload.model_validate(error.to_payload()), ) ) self._unusable = True self._remote_initialized.set() await self.stop() async def _send_cancelled_termination(self, message: InvokeMessage) -> None: """把本端取消执行转换为标准化的取消错误响应。""" error = AstrBotError.cancelled() await self._send_error_result(message, error) async def _fail_connection(self, error: AstrBotError) -> None: """把连接标记为不可用，并让所有等待中的调用尽快失败。""" if self._unusable: return self._unusable = True self._remote_initialized.set() for future in list(self._pending_results.values()): if not future.done(): future.set_exception(error) self._pending_results.clear() for queue in list(self._pending_streams.values()): await queue.put(error) self._pending_streams.clear() for task, token, _started in list(self._inbound_tasks.values()): token.cancel() task.cancel() self._inbound_tasks.clear() asyncio.create_task(self.stop()) async def _send(self, message) -> None: """序列化协议消息并通过底层传输发送出去。""" await self.transport.send(message.model_dump_json(exclude_none=True))