From 0d7c4994aafd8c79aac7254c7e8fb159ac1ca2e8 Mon Sep 17 00:00:00 2001 From: letr Date: Sat, 14 Mar 2026 19:54:33 +0800 Subject: [PATCH 1/4] feat(runtime): add configurable msgpack wire codec support --- pyproject.toml | 1 + src-new/astrbot_sdk/cli.py | 64 ++++++-- src-new/astrbot_sdk/protocol/wire_codecs.py | 76 ++++++++++ src-new/astrbot_sdk/runtime/bootstrap.py | 46 +++++- src-new/astrbot_sdk/runtime/peer.py | 14 +- src-new/astrbot_sdk/runtime/supervisor.py | 47 ++++-- src-new/astrbot_sdk/runtime/transport.py | 160 +++++++++++++++----- src-new/astrbot_sdk/runtime/worker.py | 23 ++- tests_v4/helpers.py | 9 +- tests_v4/test_peer.py | 45 ++++++ tests_v4/test_transport.py | 83 +++++++--- tests_v4/test_wire_codecs.py | 55 +++++++ 12 files changed, 515 insertions(+), 108 deletions(-) create mode 100644 src-new/astrbot_sdk/protocol/wire_codecs.py create mode 100644 tests_v4/test_wire_codecs.py diff --git a/pyproject.toml b/pyproject.toml index 61a817475..20cdf9083 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "docstring-parser>=0.17.0", "google-genai>=1.50.0", "loguru>=0.7.3", + "msgpack>=1.1.1", "openai>=2.7.2", "pydantic>=2.12.3", "pyyaml>=6.0.3", diff --git a/src-new/astrbot_sdk/cli.py b/src-new/astrbot_sdk/cli.py index 0410585a7..8ca894d23 100644 --- a/src-new/astrbot_sdk/cli.py +++ b/src-new/astrbot_sdk/cli.py @@ -544,6 +544,7 @@ def _handle_dev_meta_command(command: str, state: dict[str, Any]) -> bool: return True return False + def _class_name_for_plugin(value: str) -> str: parts = [part for part in re.split(r"[^a-zA-Z0-9]+", value) if part] if not parts: @@ -866,12 +867,24 @@ def cli(ctx, verbose: bool) -> None: type=click.Path(file_okay=False, dir_okay=True, path_type=Path), help="Directory containing plugin folders", ) -def run(plugins_dir: Path) -> None: +@click.option( + "--wire-codec", + default="json", + show_default=True, + type=click.Choice(["json", "msgpack"]), + help="Wire codec for supervisor/worker transport", +) +def run(plugins_dir: Path, wire_codec: str) -> None: """Start the plugin supervisor over stdio.""" + entrypoint = ( + run_supervisor(plugins_dir=plugins_dir) + if wire_codec == "json" + else run_supervisor(plugins_dir=plugins_dir, wire_codec=wire_codec) + ) _run_async_entrypoint( - run_supervisor(plugins_dir=plugins_dir), + entrypoint, log_message=f"启动插件主管进程，插件目录：{plugins_dir}", - context={"plugins_dir": plugins_dir}, + context={"plugins_dir": plugins_dir, "wire_codec": wire_codec}, ) @@ -1009,7 +1022,15 @@ def dev( required=False, type=click.Path(file_okay=True, dir_okay=False, path_type=Path), ) -def worker(plugin_dir: Path | None, group_metadata: Path | None) -> None: +@click.option( + "--wire-codec", + default="json", + show_default=True, + type=click.Choice(["json", "msgpack"]), +) +def worker( + plugin_dir: Path | None, group_metadata: Path | None, wire_codec: str +) -> None: """Internal command used by the supervisor to start a worker.""" if plugin_dir is None and group_metadata is None: raise click.UsageError("Either --plugin-dir or --group-metadata is required") @@ -1020,23 +1041,42 @@ def worker(plugin_dir: Path | None, group_metadata: Path | None) -> None: target = str(group_metadata or plugin_dir) if group_metadata is not None: - entrypoint = run_plugin_worker(group_metadata=group_metadata) + entrypoint = ( + run_plugin_worker(group_metadata=group_metadata) + if wire_codec == "json" + else run_plugin_worker(group_metadata=group_metadata, wire_codec=wire_codec) + ) else: - entrypoint = run_plugin_worker(plugin_dir=plugin_dir) + entrypoint = ( + run_plugin_worker(plugin_dir=plugin_dir) + if wire_codec == "json" + else run_plugin_worker(plugin_dir=plugin_dir, wire_codec=wire_codec) + ) _run_async_entrypoint( entrypoint, log_message=f"启动插件工作进程：{target}", log_level="debug", - context={"plugin_dir": plugin_dir}, + context={"plugin_dir": plugin_dir, "wire_codec": wire_codec}, ) @cli.command(hidden=True) @click.option("--port", default=8765, type=int, help="WebSocket server port") -def websocket(port: int) -> None: +@click.option( + "--wire-codec", + default="json", + show_default=True, + type=click.Choice(["json", "msgpack"]), +) +def websocket(port: int, wire_codec: str) -> None: """WebSocket runtime entrypoint kept for standalone bridge scenarios.""" - _run_async_entrypoint( - run_websocket_server(port=port), - log_message=f"启动 WebSocket 服务器，端口：{port}", - context={"port": port}, + entrypoint = ( + run_websocket_server(port=port) + if wire_codec == "json" + else run_websocket_server(port=port, wire_codec=wire_codec) + ) + _run_async_entrypoint( + entrypoint, + log_message=f"启动 WebSocket 服务器，端口：{port}", + context={"port": port, "wire_codec": wire_codec}, ) diff --git a/src-new/astrbot_sdk/protocol/wire_codecs.py b/src-new/astrbot_sdk/protocol/wire_codecs.py new file mode 100644 index 000000000..db326553b --- /dev/null +++ b/src-new/astrbot_sdk/protocol/wire_codecs.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Literal + +import msgpack + +from .messages import ProtocolMessage, parse_message + +StdioFraming = Literal["line", "length_prefixed"] +WebSocketFrameType = Literal["text", "binary"] +WireCodecName = Literal["json", "msgpack"] + + +class ProtocolCodec(ABC): + name: WireCodecName + stdio_framing: StdioFraming + websocket_frame_type: WebSocketFrameType + + @abstractmethod + def encode_message(self, message: ProtocolMessage) -> bytes: + raise NotImplementedError + + @abstractmethod + def decode_message(self, payload: bytes | str) -> ProtocolMessage: + raise NotImplementedError + + +class JsonProtocolCodec(ProtocolCodec): + name: WireCodecName = "json" + stdio_framing: StdioFraming = "line" + websocket_frame_type: WebSocketFrameType = "text" + + def encode_message(self, message: ProtocolMessage) -> bytes: + return message.model_dump_json(exclude_none=True).encode("utf-8") + + def decode_message(self, payload: bytes | str) -> ProtocolMessage: + if isinstance(payload, bytes): + return parse_message(payload.decode("utf-8")) + return parse_message(payload) + + +class MsgpackProtocolCodec(ProtocolCodec): + name: WireCodecName = "msgpack" + stdio_framing: StdioFraming = "length_prefixed" + websocket_frame_type: WebSocketFrameType = "binary" + + def encode_message(self, message: ProtocolMessage) -> bytes: + return msgpack.packb( + message.model_dump(exclude_none=True), + use_bin_type=True, + ) + + def decode_message(self, payload: bytes | str) -> ProtocolMessage: + if isinstance(payload, str): + return parse_message(payload) + return parse_message(msgpack.unpackb(payload, raw=False)) + + +def make_protocol_codec(name: WireCodecName | str) -> ProtocolCodec: + if name == "json": + return JsonProtocolCodec() + if name == "msgpack": + return MsgpackProtocolCodec() + raise ValueError(f"未知 wire codec: {name}") + + +__all__ = [ + "JsonProtocolCodec", + "MsgpackProtocolCodec", + "ProtocolCodec", + "StdioFraming", + "WebSocketFrameType", + "WireCodecName", + "make_protocol_codec", +] diff --git a/src-new/astrbot_sdk/runtime/bootstrap.py b/src-new/astrbot_sdk/runtime/bootstrap.py index 7a8706965..e0f26b709 100644 --- a/src-new/astrbot_sdk/runtime/bootstrap.py +++ b/src-new/astrbot_sdk/runtime/bootstrap.py @@ -19,6 +19,7 @@ import sys from pathlib import Path from typing import IO +from ..protocol.wire_codecs import make_protocol_codec from .loader import PluginEnvironmentManager from .supervisor import ( SupervisorRuntime, @@ -49,19 +50,27 @@ __all__ = [ async def run_supervisor( *, plugins_dir: Path = Path("plugins"), - stdin: IO[str] | None = None, - stdout: IO[str] | None = None, + stdin: IO[str] | IO[bytes] | None = None, + stdout: IO[str] | IO[bytes] | None = None, env_manager: PluginEnvironmentManager | None = None, + wire_codec: str = "json", ) -> None: + codec = make_protocol_codec(wire_codec) transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( stdin, stdout, + binary=codec.stdio_framing == "length_prefixed", + ) + transport = StdioTransport( + stdin=transport_stdin, + stdout=transport_stdout, + framing=codec.stdio_framing, ) - transport = StdioTransport(stdin=transport_stdin, stdout=transport_stdout) runtime = SupervisorRuntime( transport=transport, plugins_dir=plugins_dir, env_manager=env_manager, + codec=codec, ) try: @@ -79,26 +88,39 @@ async def run_plugin_worker( *, plugin_dir: Path | None = None, group_metadata: Path | None = None, - stdin: IO[str] | None = None, - stdout: IO[str] | None = None, + stdin: IO[str] | IO[bytes] | None = None, + stdout: IO[str] | IO[bytes] | None = None, + wire_codec: str = "json", ) -> None: if plugin_dir is None and group_metadata is None: raise ValueError("plugin_dir or group_metadata is required") if plugin_dir is not None and group_metadata is not None: raise ValueError("plugin_dir and group_metadata are mutually exclusive") + codec = make_protocol_codec(wire_codec) transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( stdin, stdout, + binary=codec.stdio_framing == "length_prefixed", + ) + transport = StdioTransport( + stdin=transport_stdin, + stdout=transport_stdout, + framing=codec.stdio_framing, ) - transport = StdioTransport(stdin=transport_stdin, stdout=transport_stdout) if group_metadata is not None: runtime = GroupWorkerRuntime( group_metadata_path=group_metadata, transport=transport, + codec=codec, ) else: - runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport) + assert plugin_dir is not None + runtime = PluginWorkerRuntime( + plugin_dir=plugin_dir, + transport=transport, + codec=codec, + ) try: await runtime.start() stop_event = asyncio.Event() @@ -116,10 +138,18 @@ async def run_websocket_server( port: int = 8765, path: str = "/", plugin_dir: Path | None = None, + wire_codec: str = "json", ) -> None: + codec = make_protocol_codec(wire_codec) runtime = PluginWorkerRuntime( plugin_dir=plugin_dir or Path.cwd(), - transport=WebSocketServerTransport(host=host, port=port, path=path), + transport=WebSocketServerTransport( + host=host, + port=port, + path=path, + frame_type=codec.websocket_frame_type, + ), + codec=codec, ) try: await runtime.start() diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py index cd418e719..4c4880143 100644 --- a/src-new/astrbot_sdk/runtime/peer.py +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -5,7 +5,7 @@ 而不是业务上的用户、群聊或会话对象。 核心职责： - - 消息序列化/反序列化 + - 通过可插拔 codec 做消息编解码 - 初始化握手协议 - 能力调用（同步/流式） - 取消处理 @@ -69,7 +69,7 @@ - 入站任务在收到 CancelMessage 时被取消 - 早到取消：在任务执行前检查 cancel_token，避免竞态条件 -`Peer` 把 `Transport` 和 v4 协议消息模型接起来，负责： + `Peer` 把 `Transport`、wire codec 和 v4 协议消息模型接起来，负责： - 握手与远端元数据缓存 - 请求 ID 关联 @@ -100,8 +100,8 @@ from ..protocol.messages import ( InvokeMessage, PeerInfo, ResultMessage, - parse_message, ) +from ..protocol.wire_codecs import JsonProtocolCodec, ProtocolCodec from .capability_router import StreamExecution InitializeHandler = Callable[[InitializeMessage], Awaitable[InitializeOutput]] @@ -181,6 +181,7 @@ class Peer: peer_info: PeerInfo, protocol_version: str = "1.0", supported_protocol_versions: Sequence[str] | None = None, + codec: ProtocolCodec | None = None, ) -> None: """创建一个协议对等端实例。 @@ -192,6 +193,7 @@ class Peer: """ self.transport = transport self.peer_info = peer_info + self.codec = codec or JsonProtocolCodec() self.protocol_version = protocol_version self.supported_protocol_versions = _dedupe_protocol_versions( supported_protocol_versions, @@ -502,10 +504,10 @@ class Peer: if self._unusable: raise AstrBotError.protocol_error("连接已进入不可用状态") - async def _handle_raw_message(self, payload: str) -> None: + async def _handle_raw_message(self, payload: bytes) -> None: """解析原始消息并分发到对应的消息处理分支。""" try: - message = parse_message(payload) + message = self.codec.decode_message(payload) if isinstance(message, ResultMessage): await self._handle_result(message) return @@ -754,4 +756,4 @@ class Peer: async def _send(self, message) -> None: """序列化协议消息并通过底层传输发送出去。""" - await self.transport.send(message.model_dump_json(exclude_none=True)) + await self.transport.send(self.codec.encode_message(message)) diff --git a/src-new/astrbot_sdk/runtime/supervisor.py b/src-new/astrbot_sdk/runtime/supervisor.py index 4df140569..14ad9b109 100644 --- a/src-new/astrbot_sdk/runtime/supervisor.py +++ b/src-new/astrbot_sdk/runtime/supervisor.py @@ -41,13 +41,14 @@ import signal import sys from collections.abc import Callable from pathlib import Path -from typing import IO, Any +from typing import IO, Any, cast from loguru import logger from ..errors import AstrBotError from ..protocol.descriptors import CapabilityDescriptor from ..protocol.messages import EventMessage, InitializeOutput, PeerInfo +from ..protocol.wire_codecs import ProtocolCodec, make_protocol_codec from .capability_router import CapabilityRouter, StreamExecution from .environment_groups import EnvironmentGroup from .loader import ( @@ -79,13 +80,15 @@ def _install_signal_handlers(stop_event: asyncio.Event) -> None: def _prepare_stdio_transport( - stdin: IO[str] | None, - stdout: IO[str] | None, -) -> tuple[IO[str], IO[str], IO[str] | None]: + stdin: IO[str] | IO[bytes] | None, + stdout: IO[str] | IO[bytes] | None, + *, + binary: bool = False, +) -> tuple[IO[str] | IO[bytes], IO[str] | IO[bytes], IO[str] | None]: if stdin is not None and stdout is not None: return stdin, stdout, None - transport_stdin = stdin or sys.stdin - transport_stdout = stdout or sys.stdout + transport_stdin = stdin or (sys.stdin.buffer if binary else sys.stdin) + transport_stdout = stdout or (sys.stdout.buffer if binary else sys.stdout) original_stdout = sys.stdout sys.stdout = sys.stderr return transport_stdin, transport_stdout, original_stdout @@ -128,17 +131,25 @@ class WorkerSession: env_manager: PluginEnvironmentManager, capability_router: CapabilityRouter, on_closed: Callable[[], None] | None = None, + codec: ProtocolCodec | None = None, + wire_codec_name: str = "json", ) -> None: if plugin is None and group is None: raise ValueError("WorkerSession requires either plugin or group") + if group is None and plugin is None: + raise ValueError("WorkerSession requires a plugin when group is absent") self.group = group - self.plugins = list(group.plugins) if group is not None else [plugin] + self.plugins = ( + list(group.plugins) if group is not None else [cast(PluginSpec, plugin)] + ) self.plugin = plugin or self.plugins[0] self.group_id = group.id if group is not None else self.plugin.name self.repo_root = repo_root.resolve() self.env_manager = env_manager self.capability_router = capability_router self.on_closed = on_closed + self.codec = codec or make_protocol_codec(wire_codec_name) + self.wire_codec_name = self.codec.name self.peer: Peer | None = None self.handlers = [] self.provided_capabilities: list[CapabilityDescriptor] = [] @@ -164,10 +175,12 @@ class WorkerSession: command=command, cwd=cwd, env=env, + framing=self.codec.stdio_framing, ) self.peer = Peer( transport=transport, peer_info=PeerInfo(name="astrbot-core", role="core", version="v4"), + codec=self.codec, ) self.peer.set_initialize_handler(self._handle_initialize) self.peer.set_invoke_handler(self._handle_capability_invoke) @@ -225,7 +238,7 @@ class WorkerSession: if self.group is not None: prepare_group = getattr(self.env_manager, "prepare_group_environment", None) if callable(prepare_group): - python_path = prepare_group(self.group) + python_path = cast(Path, prepare_group(self.group)) else: python_path = self.env_manager.prepare_environment(self.plugins[0]) return ( @@ -235,13 +248,16 @@ class WorkerSession: "-m", "astrbot_sdk", "worker", + "--wire-codec", + self.wire_codec_name, "--group-metadata", str(self.group.metadata_path), ], str(self.repo_root), ) - python_path = self.env_manager.prepare_environment(self.plugin) + plugin = self.plugin + python_path = self.env_manager.prepare_environment(plugin) return ( python_path, [ @@ -249,10 +265,12 @@ class WorkerSession: "-m", "astrbot_sdk", "worker", + "--wire-codec", + self.wire_codec_name, "--plugin-dir", - str(self.plugin.plugin_dir), + str(plugin.plugin_dir), ], - str(self.plugin.plugin_dir), + str(plugin.plugin_dir), ) def start_close_watch(self) -> None: @@ -376,15 +394,20 @@ class SupervisorRuntime: transport, plugins_dir: Path, env_manager: PluginEnvironmentManager | None = None, + codec: ProtocolCodec | None = None, + wire_codec_name: str = "json", ) -> None: self.transport = transport self.plugins_dir = plugins_dir.resolve() self.repo_root = Path(__file__).resolve().parents[3] self.env_manager = env_manager or PluginEnvironmentManager(self.repo_root) + self.codec = codec or make_protocol_codec(wire_codec_name) + self.wire_codec_name = self.codec.name self.capability_router = CapabilityRouter() self.peer = Peer( transport=self.transport, peer_info=PeerInfo(name="astrbot-supervisor", role="plugin", version="v4"), + codec=self.codec, ) self.peer.set_invoke_handler(self._handle_upstream_invoke) self.peer.set_cancel_handler(self._handle_upstream_cancel) @@ -623,6 +646,7 @@ class SupervisorRuntime: repo_root=self.repo_root, env_manager=self.env_manager, capability_router=self.capability_router, + wire_codec_name=self.wire_codec_name, on_closed=lambda group_id=group.id: ( self._handle_worker_closed(group_id) ), @@ -636,6 +660,7 @@ class SupervisorRuntime: repo_root=self.repo_root, env_manager=self.env_manager, capability_router=self.capability_router, + wire_codec_name=self.wire_codec_name, on_closed=lambda plugin_name=plugin.name: ( self._handle_worker_closed(plugin_name) ), diff --git a/src-new/astrbot_sdk/runtime/transport.py b/src-new/astrbot_sdk/runtime/transport.py index 3049606f9..811d271b3 100644 --- a/src-new/astrbot_sdk/runtime/transport.py +++ b/src-new/astrbot_sdk/runtime/transport.py @@ -1,7 +1,7 @@ """传输层抽象模块。 -定义 Transport 抽象基类及其实现，负责底层的消息传输。 -传输层只关心"发送字符串"和"接收字符串"，不处理协议细节。 +定义 Transport 抽象基类及其实现，负责底层原始载荷的传输。 +传输层只关心分帧后的 bytes 或 text frame，不处理协议细节。 传输实现： Transport: 抽象基类，定义 start/stop/send/wait_closed 接口 StdioTransport: 标准输入输出传输 @@ -37,7 +37,7 @@ - 支持心跳配置 - WebSocketClientTransport: - 自动重连需要外部实现 - - 传输层只处理字符串，协议由 Peer 层处理 + - 传输层只处理 framed payload，协议由 Peer 层处理 使用示例： # 子进程模式 @@ -58,15 +58,15 @@ # 统一接口 transport.set_message_handler(my_handler) await transport.start() - await transport.send(json_string) + await transport.send(encoded_payload) await transport.stop() -`Transport` 只处理“字符串发出去 / 字符串收进来”这件事，不做协议解析，也不关心 -能力、handler 或迁移适配策略。当前实现包括： +`Transport` 只处理 framed payload，不做协议解析，也不关心能力、handler 或 +legacy 兼容。当前实现包括： -- `StdioTransport`: 子进程或文件对象上的按行文本传输 -- `WebSocketServerTransport`: 单连接 WebSocket 服务端 -- `WebSocketClientTransport`: WebSocket 客户端 +- `StdioTransport`: 子进程或文件对象上的按行或 length-prefixed 传输 +- `WebSocketServerTransport`: 单连接 WebSocket 服务端，支持 text/binary frame +- `WebSocketClientTransport`: WebSocket 客户端，支持 text/binary frame 自动重连、消息重放等策略不在这里实现，统一留给更上层编排。 """ @@ -74,27 +74,57 @@ from __future__ import annotations import asyncio +import io +import struct import sys from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable, Sequence -from typing import IO +from typing import IO, cast import aiohttp from aiohttp import web from loguru import logger -MessageHandler = Callable[[str], Awaitable[None]] +from ..protocol.wire_codecs import StdioFraming, WebSocketFrameType + +MessageHandler = Callable[[bytes], Awaitable[None]] +RawPayload = bytes | str -def _frame_stdio_payload(payload: str) -> str: +def _ensure_bytes(payload: RawPayload) -> bytes: + if isinstance(payload, bytes): + return payload + return payload.encode("utf-8") + + +def _frame_stdio_line_payload(payload: bytes) -> bytes: body = payload - if body.endswith("\r\n"): + if body.endswith(b"\r\n"): body = body[:-2] - elif body.endswith(("\n", "\r")): + elif body.endswith((b"\n", b"\r")): body = body[:-1] - if "\n" in body or "\r" in body: + if b"\n" in body or b"\r" in body: raise ValueError("STDIO payload 不允许包含原始换行符") - return f"{body}\n" + return body + b"\n" + + +def _frame_stdio_length_prefixed_payload(payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + payload + + +def _write_stdio_payload(stream: IO[str] | IO[bytes], payload: bytes) -> None: + if hasattr(stream, "buffer"): + stream.buffer.write(payload) # type: ignore[attr-defined] + stream.flush() # type: ignore[call-arg] + return + if isinstance(stream, io.TextIOBase): + text_stream = cast(IO[str], stream) + text_stream.write(payload.decode("utf-8")) + text_stream.flush() + return + binary_stream = cast(IO[bytes], stream) + binary_stream.write(payload) + binary_stream.flush() class Transport(ABC): @@ -115,14 +145,14 @@ class Transport(ABC): raise NotImplementedError @abstractmethod - async def send(self, payload: str) -> None: + async def send(self, payload: RawPayload) -> None: raise NotImplementedError async def wait_closed(self) -> None: """等待传输层进入关闭状态。""" await self._closed.wait() - async def _dispatch(self, payload: str) -> None: + async def _dispatch(self, payload: bytes) -> None: """把收到的原始载荷转交给上层处理器。""" if self._handler is not None: await self._handler(payload) @@ -132,11 +162,12 @@ class StdioTransport(Transport): def __init__( self, *, - stdin: IO[str] | None = None, - stdout: IO[str] | None = None, + stdin: IO[str] | IO[bytes] | None = None, + stdout: IO[str] | IO[bytes] | None = None, command: Sequence[str] | None = None, cwd: str | None = None, env: dict[str, str] | None = None, + framing: StdioFraming = "line", ) -> None: super().__init__() self._stdin = stdin @@ -144,6 +175,7 @@ class StdioTransport(Transport): self._command = list(command) if command is not None else None self._cwd = cwd self._env = env + self._framing = framing self._process: asyncio.subprocess.Process | None = None self._reader_task: asyncio.Task[None] | None = None @@ -185,12 +217,16 @@ class StdioTransport(Transport): self._process = None self._closed.set() - async def send(self, payload: str) -> None: - line = _frame_stdio_payload(payload) + async def send(self, payload: RawPayload) -> None: + encoded = _ensure_bytes(payload) + if self._framing == "line": + framed = _frame_stdio_line_payload(encoded) + else: + framed = _frame_stdio_length_prefixed_payload(encoded) if self._process is not None: if self._process.stdin is None: raise RuntimeError("STDIO subprocess stdin 不可用") - self._process.stdin.write(line.encode("utf-8")) + self._process.stdin.write(framed) await self._process.stdin.drain() return @@ -199,8 +235,7 @@ class StdioTransport(Transport): def _write() -> None: assert self._stdout is not None - self._stdout.write(line) - self._stdout.flush() + _write_stdio_payload(self._stdout, framed) await asyncio.to_thread(_write) @@ -209,10 +244,18 @@ class StdioTransport(Transport): assert self._process.stdout is not None try: while True: - raw = await self._process.stdout.readline() - if not raw: - break - await self._dispatch(raw.decode("utf-8").rstrip("\r\n")) + if self._framing == "line": + raw = await self._process.stdout.readline() + if not raw: + break + await self._dispatch(raw.rstrip(b"\r\n")) + continue + header = await self._process.stdout.readexactly(4) + length = struct.unpack(">I", header)[0] + payload = await self._process.stdout.readexactly(length) + await self._dispatch(payload) + except asyncio.IncompleteReadError: + pass finally: self._closed.set() @@ -220,10 +263,29 @@ class StdioTransport(Transport): assert self._stdin is not None try: while True: - raw = await asyncio.to_thread(self._stdin.readline) - if not raw: + if self._framing == "line": + raw = await asyncio.to_thread(self._stdin.readline) + if not raw: + break + if isinstance(raw, bytes): + await self._dispatch(raw.rstrip(b"\r\n")) + else: + await self._dispatch(raw.rstrip("\r\n").encode("utf-8")) + continue + header = await asyncio.to_thread(self._stdin.read, 4) + if not header: break - await self._dispatch(raw.rstrip("\r\n")) + if isinstance(header, str): + raise RuntimeError("length_prefixed STDIO 需要二进制 stdin") + if len(header) < 4: + break + length = struct.unpack(">I", header)[0] + payload = await asyncio.to_thread(self._stdin.read, length) + if isinstance(payload, str): + raise RuntimeError("length_prefixed STDIO 需要二进制 stdin") + if len(payload) < length: + break + await self._dispatch(payload) finally: self._closed.set() @@ -236,6 +298,7 @@ class WebSocketServerTransport(Transport): port: int = 8765, path: str = "/", heartbeat: float = 30.0, + frame_type: WebSocketFrameType = "text", ) -> None: super().__init__() self._host = host @@ -243,6 +306,7 @@ class WebSocketServerTransport(Transport): self._actual_port: int | None = None self._path = path self._heartbeat = heartbeat + self._frame_type = frame_type self._app: web.Application | None = None self._runner: web.AppRunner | None = None self._site: web.TCPSite | None = None @@ -259,8 +323,10 @@ class WebSocketServerTransport(Transport): await self._runner.setup() self._site = web.TCPSite(self._runner, self._host, self._port) await self._site.start() - if self._site._server and getattr(self._site._server, "sockets", None): - socket = self._site._server.sockets[0] + server = getattr(self._site, "_server", None) + sockets = getattr(server, "sockets", None) + if sockets: + socket = sockets[0] self._actual_port = socket.getsockname()[1] async def stop(self) -> None: @@ -275,13 +341,17 @@ class WebSocketServerTransport(Transport): self._runner = None self._closed.set() - async def send(self, payload: str) -> None: + async def send(self, payload: RawPayload) -> None: if self._ws is None or self._ws.closed: await asyncio.wait_for(self._connected.wait(), timeout=30.0) if self._ws is None or self._ws.closed: raise RuntimeError("WebSocket 尚未连接") async with self._write_lock: - await self._ws.send_str(payload) + encoded = _ensure_bytes(payload) + if self._frame_type == "text": + await self._ws.send_str(encoded.decode("utf-8")) + else: + await self._ws.send_bytes(encoded) async def _handle_socket(self, request: web.Request) -> web.WebSocketResponse: if self._ws is not None and not self._ws.closed: @@ -299,9 +369,9 @@ class WebSocketServerTransport(Transport): try: async for msg in ws: if msg.type == aiohttp.WSMsgType.TEXT: - await self._dispatch(msg.data) + await self._dispatch(msg.data.encode("utf-8")) elif msg.type == aiohttp.WSMsgType.BINARY: - await self._dispatch(msg.data.decode("utf-8")) + await self._dispatch(msg.data) elif msg.type == aiohttp.WSMsgType.ERROR: logger.error("websocket server error: {}", ws.exception()) break @@ -326,10 +396,12 @@ class WebSocketClientTransport(Transport): *, url: str, heartbeat: float = 30.0, + frame_type: WebSocketFrameType = "text", ) -> None: super().__init__() self._url = url self._heartbeat = heartbeat + self._frame_type = frame_type self._session: aiohttp.ClientSession | None = None self._ws: aiohttp.ClientWebSocketResponse | None = None self._reader_task: asyncio.Task[None] | None = None @@ -359,19 +431,23 @@ class WebSocketClientTransport(Transport): self._session = None self._closed.set() - async def send(self, payload: str) -> None: + async def send(self, payload: RawPayload) -> None: if self._ws is None or self._ws.closed: raise RuntimeError("WebSocket client 尚未连接") - await self._ws.send_str(payload) + encoded = _ensure_bytes(payload) + if self._frame_type == "text": + await self._ws.send_str(encoded.decode("utf-8")) + else: + await self._ws.send_bytes(encoded) async def _read_loop(self) -> None: assert self._ws is not None try: async for msg in self._ws: if msg.type == aiohttp.WSMsgType.TEXT: - await self._dispatch(msg.data) + await self._dispatch(msg.data.encode("utf-8")) elif msg.type == aiohttp.WSMsgType.BINARY: - await self._dispatch(msg.data.decode("utf-8")) + await self._dispatch(msg.data) elif msg.type == aiohttp.WSMsgType.ERROR: logger.error("websocket client error: {}", self._ws.exception()) break diff --git a/src-new/astrbot_sdk/runtime/worker.py b/src-new/astrbot_sdk/runtime/worker.py index 2d3b4626f..7b1b6cc55 100644 --- a/src-new/astrbot_sdk/runtime/worker.py +++ b/src-new/astrbot_sdk/runtime/worker.py @@ -37,6 +37,7 @@ from .._invocation_context import caller_plugin_scope from ..context import Context as RuntimeContext from ..errors import AstrBotError from ..protocol.messages import PeerInfo +from ..protocol.wire_codecs import ProtocolCodec, make_protocol_codec from .handler_dispatcher import CapabilityDispatcher, HandlerDispatcher from .loader import ( LoadedPlugin, @@ -114,13 +115,22 @@ async def run_plugin_lifecycle( class GroupWorkerRuntime: - def __init__(self, *, group_metadata_path: Path, transport) -> None: + def __init__( + self, + *, + group_metadata_path: Path, + transport, + codec: ProtocolCodec | None = None, + wire_codec_name: str = "json", + ) -> None: self.group_metadata_path = group_metadata_path.resolve() self.group_id, self.plugins = _load_group_plugin_specs(self.group_metadata_path) self.transport = transport + self.codec = codec or make_protocol_codec(wire_codec_name) self.peer = Peer( transport=self.transport, peer_info=PeerInfo(name=self.group_id, role="plugin", version="v4"), + codec=self.codec, ) self.skipped_plugins: dict[str, str] = {} self._plugin_states: list[GroupPluginRuntimeState] = [] @@ -284,13 +294,22 @@ class GroupWorkerRuntime: class PluginWorkerRuntime: - def __init__(self, *, plugin_dir: Path, transport) -> None: + def __init__( + self, + *, + plugin_dir: Path, + transport, + codec: ProtocolCodec | None = None, + wire_codec_name: str = "json", + ) -> None: self.plugin = load_plugin_spec(plugin_dir) self.transport = transport + self.codec = codec or make_protocol_codec(wire_codec_name) self.loaded_plugin = load_plugin(self.plugin) self.peer = Peer( transport=self.transport, peer_info=PeerInfo(name=self.plugin.name, role="plugin", version="v4"), + codec=self.codec, ) self.dispatcher = HandlerDispatcher( plugin_id=self.plugin.name, diff --git a/tests_v4/helpers.py b/tests_v4/helpers.py index 57a40cfef..05bf8f0df 100644 --- a/tests_v4/helpers.py +++ b/tests_v4/helpers.py @@ -4,8 +4,9 @@ import asyncio import shutil from types import SimpleNamespace from pathlib import Path +from typing import cast -from astrbot_sdk.runtime.transport import Transport +from astrbot_sdk.runtime.transport import RawPayload, Transport class MemoryTransport(Transport): @@ -36,7 +37,7 @@ class MemoryTransport(Transport): """ self._closed.set() # 设置关闭事件 - async def send(self, payload: str) -> None: + async def send(self, payload: RawPayload) -> None: """发送消息给伙伴。 Args: @@ -48,7 +49,9 @@ class MemoryTransport(Transport): if self.partner is None: raise RuntimeError("MemoryTransport 未连接 partner") # 将消息转发给伙伴的_dispatch方法进行处理 - await self.partner._dispatch(payload) + if isinstance(payload, str): + payload = payload.encode("utf-8") + await self.partner._dispatch(cast(bytes, payload)) def make_transport_pair() -> tuple[MemoryTransport, MemoryTransport]: diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py index 46a77e394..7db9c9340 100644 --- a/tests_v4/test_peer.py +++ b/tests_v4/test_peer.py @@ -13,6 +13,7 @@ from astrbot_sdk.protocol.messages import ( PeerInfo, ResultMessage, ) +from astrbot_sdk.protocol.wire_codecs import MsgpackProtocolCodec from astrbot_sdk.runtime.capability_router import CapabilityRouter, StreamExecution from astrbot_sdk.runtime.peer import Peer from astrbot_sdk.runtime.transport import ( @@ -131,6 +132,50 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): await plugin.stop() await core.stop() + async def test_msgpack_codec_roundtrip(self) -> None: + codec = MsgpackProtocolCodec() + router = CapabilityRouter() + core = Peer( + transport=self.left, + peer_info=PeerInfo(name="core", role="core", version="v4"), + codec=codec, + ) + core.set_initialize_handler( + lambda _message: asyncio.sleep( + 0, + result=InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=router.descriptors(), + metadata={}, + ), + ) + ) + core.set_invoke_handler( + lambda message, token: router.execute( + message.capability, + message.input, + stream=message.stream, + cancel_token=token, + request_id=message.id, + ) + ) + + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + codec=codec, + ) + + await core.start() + await plugin.start() + await plugin.initialize([]) + + result = await plugin.invoke("llm.chat", {"prompt": "hello-msgpack"}) + self.assertEqual(result["text"], "Echo: hello-msgpack") + + await plugin.stop() + await core.stop() + async def test_wait_until_remote_initialized_after_initialize_returns(self) -> None: core = Peer( transport=self.left, diff --git a/tests_v4/test_transport.py b/tests_v4/test_transport.py index 10880fcf5..0251e2630 100644 --- a/tests_v4/test_transport.py +++ b/tests_v4/test_transport.py @@ -5,7 +5,7 @@ Tests for runtime/transport.py - Transport implementations. from __future__ import annotations import asyncio -from io import StringIO +from io import BytesIO, StringIO from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -32,7 +32,7 @@ class TestTransportBase: async def stop(self): pass - async def send(self, message: str): + async def send(self, payload): pass transport = ConcreteTransport() @@ -48,7 +48,7 @@ class TestTransportBase: async def stop(self): pass - async def send(self, message: str): + async def send(self, payload): pass transport = ConcreteTransport() @@ -85,7 +85,7 @@ class TestTransportBase: async def stop(self): pass - async def send(self, message: str): + async def send(self, payload): pass transport = ConcreteTransport() @@ -104,14 +104,14 @@ class TestTransportBase: async def stop(self): pass - async def send(self, message: str): + async def send(self, payload): pass transport = ConcreteTransport() handler = AsyncMock() transport.set_message_handler(handler) - await transport._dispatch("test payload") - handler.assert_called_once_with("test payload") + await transport._dispatch(b"test payload") + handler.assert_called_once_with(b"test payload") @pytest.mark.asyncio async def test_dispatch_without_handler(self): @@ -124,12 +124,12 @@ class TestTransportBase: async def stop(self): pass - async def send(self, message: str): + async def send(self, payload): pass transport = ConcreteTransport() # Should not raise when no handler is set - await transport._dispatch("test payload") + await transport._dispatch(b"test payload") class TestStdioTransportInit: @@ -186,14 +186,13 @@ class TestStdioTransportFileMode: await transport.start() task = transport._reader_task await transport.stop() + assert task is not None assert task.cancelled() or task.done() @pytest.mark.asyncio async def test_send_without_process(self): """send() without process should write to stdout.""" - stdout = MagicMock() - stdout.write = MagicMock() - stdout.flush = MagicMock() + stdout = StringIO() transport = StdioTransport(stdout=stdout) with patch("sys.stdin"): @@ -202,40 +201,35 @@ class TestStdioTransportFileMode: await transport.send("test message") # Should have written the message with newline - stdout.write.assert_called_once_with("test message\n") - stdout.flush.assert_called_once() + assert stdout.getvalue() == "test message\n" await transport.stop() @pytest.mark.asyncio async def test_send_adds_newline_if_missing(self): """send() should add newline if not present.""" - stdout = MagicMock() - stdout.write = MagicMock() - stdout.flush = MagicMock() + stdout = StringIO() transport = StdioTransport(stdout=stdout) with patch("sys.stdin"): await transport.start() await transport.send("message") - stdout.write.assert_called_once_with("message\n") + assert stdout.getvalue() == "message\n" await transport.stop() @pytest.mark.asyncio async def test_send_preserves_existing_newline(self): """send() should not add extra newline.""" - stdout = MagicMock() - stdout.write = MagicMock() - stdout.flush = MagicMock() + stdout = StringIO() transport = StdioTransport(stdout=stdout) with patch("sys.stdin"): await transport.start() await transport.send("message\n") - stdout.write.assert_called_once_with("message\n") + assert stdout.getvalue() == "message\n" await transport.stop() @@ -269,6 +263,28 @@ class TestStdioTransportFileMode: with pytest.raises(RuntimeError, match="stdout"): await transport.send("test") + @pytest.mark.asyncio + async def test_length_prefixed_roundtrip(self): + """length-prefixed stdio should preserve binary payloads.""" + stdin = BytesIO() + stdout = BytesIO() + transport = StdioTransport( + stdin=stdin, + stdout=stdout, + framing="length_prefixed", + ) + received: list[bytes] = [] + + async def handle_message(payload: bytes): + received.append(payload) + + transport.set_message_handler(handle_message) + await transport.send(b"\x81\xa4test\x01") + stdin.write(stdout.getvalue()) + stdin.seek(0) + await transport._read_file_loop() + assert received == [b"\x81\xa4test\x01"] + class TestStdioTransportProcessMode: """Tests for StdioTransport in subprocess mode.""" @@ -423,6 +439,7 @@ class TestWebSocketServerTransportLifecycle: transport._ws = MagicMock() transport._ws.closed = False transport._ws.send_str = AsyncMock() + transport._ws.send_bytes = AsyncMock() # close 也需要是异步的 transport._ws.close = AsyncMock() @@ -443,6 +460,24 @@ class TestWebSocketServerTransportLifecycle: await transport.stop() + @pytest.mark.asyncio + async def test_send_binary_frame_when_configured(self): + transport = WebSocketServerTransport(port=0, heartbeat=0, frame_type="binary") + await transport.start() + + transport._connected.set() + transport._ws = MagicMock() + transport._ws.closed = False + transport._ws.send_str = AsyncMock() + transport._ws.send_bytes = AsyncMock() + transport._ws.close = AsyncMock() + + await transport.send(b"\x81\xa3hi") + transport._ws.send_bytes.assert_called_once_with(b"\x81\xa3hi") + transport._ws.send_str.assert_not_called() + + await transport.stop() + class TestWebSocketClientTransportInit: """Tests for WebSocketClientTransport initialization.""" @@ -527,7 +562,7 @@ class TestTransportIntegration: received_messages = [] - async def handle_message(payload: str): + async def handle_message(payload: bytes): received_messages.append(payload) server.set_message_handler(handle_message) @@ -546,7 +581,7 @@ class TestTransportIntegration: # Wait for message to be received await asyncio.sleep(0.1) - assert "hello from client" in received_messages + assert b"hello from client" in received_messages await client.stop() await server.stop() diff --git a/tests_v4/test_wire_codecs.py b/tests_v4/test_wire_codecs.py new file mode 100644 index 000000000..92f48e41f --- /dev/null +++ b/tests_v4/test_wire_codecs.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import pytest + +from astrbot_sdk.protocol.descriptors import CommandTrigger, HandlerDescriptor +from astrbot_sdk.protocol.messages import InitializeMessage, PeerInfo +from astrbot_sdk.protocol.wire_codecs import ( + JsonProtocolCodec, + MsgpackProtocolCodec, + make_protocol_codec, +) + + +def _sample_initialize_message() -> InitializeMessage: + return InitializeMessage( + id="msg-1", + protocol_version="1.0", + peer=PeerInfo(name="plugin", role="plugin", version="v4"), + handlers=[ + HandlerDescriptor( + id="plugin:hello", + trigger=CommandTrigger(command="hello"), + ) + ], + metadata={"plugin_id": "plugin", "loaded_plugins": ["plugin"]}, + ) + + +class TestJsonProtocolCodec: + def test_roundtrip(self): + codec = JsonProtocolCodec() + message = _sample_initialize_message() + + encoded = codec.encode_message(message) + decoded = codec.decode_message(encoded) + + assert isinstance(encoded, bytes) + assert decoded == message + + +class TestMsgpackProtocolCodec: + def test_roundtrip(self): + codec = MsgpackProtocolCodec() + message = _sample_initialize_message() + + encoded = codec.encode_message(message) + decoded = codec.decode_message(encoded) + + assert isinstance(encoded, bytes) + assert decoded == message + + +def test_make_protocol_codec_rejects_unknown_name(): + with pytest.raises(ValueError, match="未知 wire codec"): + make_protocol_codec("yaml") From 98a5eddfcbf528eba5d568fec19530afd77abf9a Mon Sep 17 00:00:00 2001 From: letr Date: Sat, 14 Mar 2026 20:04:34 +0800 Subject: [PATCH 2/4] fix(runtime): align msgpack framing with transport defaults --- src-new/astrbot_sdk/runtime/peer.py | 3 ++ src-new/astrbot_sdk/runtime/transport.py | 21 ++++++++++-- tests_v4/test_peer.py | 43 ++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py index 4c4880143..26dc0e198 100644 --- a/src-new/astrbot_sdk/runtime/peer.py +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -194,6 +194,9 @@ class Peer: self.transport = transport self.peer_info = peer_info self.codec = codec or JsonProtocolCodec() + configure_for_codec = getattr(self.transport, "configure_for_codec", None) + if callable(configure_for_codec): + configure_for_codec(self.codec) self.protocol_version = protocol_version self.supported_protocol_versions = _dedupe_protocol_versions( supported_protocol_versions, diff --git a/src-new/astrbot_sdk/runtime/transport.py b/src-new/astrbot_sdk/runtime/transport.py index 811d271b3..a5549f67d 100644 --- a/src-new/astrbot_sdk/runtime/transport.py +++ b/src-new/astrbot_sdk/runtime/transport.py @@ -136,6 +136,10 @@ class Transport(ABC): """注册收到原始字符串消息后的回调。""" self._handler = handler + def configure_for_codec(self, codec) -> None: + """Allow transports to align framing or frame type with the selected codec.""" + return None + @abstractmethod async def start(self) -> None: raise NotImplementedError @@ -193,10 +197,17 @@ class StdioTransport(Transport): self._reader_task = asyncio.create_task(self._read_process_loop()) return - self._stdin = self._stdin or sys.stdin - self._stdout = self._stdout or sys.stdout + if self._framing == "length_prefixed": + self._stdin = self._stdin or sys.stdin.buffer + self._stdout = self._stdout or sys.stdout.buffer + else: + self._stdin = self._stdin or sys.stdin + self._stdout = self._stdout or sys.stdout self._reader_task = asyncio.create_task(self._read_file_loop()) + def configure_for_codec(self, codec) -> None: + self._framing = codec.stdio_framing + async def stop(self) -> None: if self._reader_task is not None: self._reader_task.cancel() @@ -389,6 +400,9 @@ class WebSocketServerTransport(Transport): def url(self) -> str: return f"ws://{self._host}:{self.port}{self._path}" + def configure_for_codec(self, codec) -> None: + self._frame_type = codec.websocket_frame_type + class WebSocketClientTransport(Transport): def __init__( @@ -453,3 +467,6 @@ class WebSocketClientTransport(Transport): break finally: self._closed.set() + + def configure_for_codec(self, codec) -> None: + self._frame_type = codec.websocket_frame_type diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py index 7db9c9340..0801317aa 100644 --- a/tests_v4/test_peer.py +++ b/tests_v4/test_peer.py @@ -411,6 +411,49 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): await plugin.stop() await core.stop() + async def test_websocket_transport_smoke_with_msgpack_codec(self) -> None: + codec = MsgpackProtocolCodec() + router = CapabilityRouter() + server_transport = WebSocketServerTransport(port=0) + core = Peer( + transport=server_transport, + peer_info=PeerInfo(name="core", role="core", version="v4"), + codec=codec, + ) + core.set_initialize_handler( + lambda _message: asyncio.sleep( + 0, + result=InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=router.descriptors(), + metadata={}, + ), + ) + ) + core.set_invoke_handler( + lambda message, token: router.execute( + message.capability, + message.input, + stream=message.stream, + cancel_token=token, + request_id=message.id, + ) + ) + + await core.start() + client_transport = WebSocketClientTransport(url=server_transport.url) + plugin = Peer( + transport=client_transport, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + codec=codec, + ) + await plugin.start() + await plugin.initialize([]) + result = await plugin.invoke("llm.chat", {"prompt": "ws-msgpack"}) + self.assertEqual(result["text"], "Echo: ws-msgpack") + await plugin.stop() + await core.stop() + async def test_initialize_failure_closes_receiver_connection(self) -> None: core = Peer( transport=self.left, From a979e9256825ea3b2678185b873b97ed018d600e Mon Sep 17 00:00:00 2001 From: letr Date: Sat, 14 Mar 2026 20:21:34 +0800 Subject: [PATCH 3/4] fix(runtime): preserve json transport compatibility --- src-new/astrbot_sdk/protocol/wire_codecs.py | 6 ++--- src-new/astrbot_sdk/runtime/bootstrap.py | 10 ++----- src-new/astrbot_sdk/runtime/supervisor.py | 10 +++---- tests_v4/test_peer.py | 30 +++++++++++++++++++++ tests_v4/test_wire_codecs.py | 2 +- 5 files changed, 41 insertions(+), 17 deletions(-) diff --git a/src-new/astrbot_sdk/protocol/wire_codecs.py b/src-new/astrbot_sdk/protocol/wire_codecs.py index db326553b..494df4435 100644 --- a/src-new/astrbot_sdk/protocol/wire_codecs.py +++ b/src-new/astrbot_sdk/protocol/wire_codecs.py @@ -18,7 +18,7 @@ class ProtocolCodec(ABC): websocket_frame_type: WebSocketFrameType @abstractmethod - def encode_message(self, message: ProtocolMessage) -> bytes: + def encode_message(self, message: ProtocolMessage) -> bytes | str: raise NotImplementedError @abstractmethod @@ -31,8 +31,8 @@ class JsonProtocolCodec(ProtocolCodec): stdio_framing: StdioFraming = "line" websocket_frame_type: WebSocketFrameType = "text" - def encode_message(self, message: ProtocolMessage) -> bytes: - return message.model_dump_json(exclude_none=True).encode("utf-8") + def encode_message(self, message: ProtocolMessage) -> str: + return message.model_dump_json(exclude_none=True) def decode_message(self, payload: bytes | str) -> ProtocolMessage: if isinstance(payload, bytes): diff --git a/src-new/astrbot_sdk/runtime/bootstrap.py b/src-new/astrbot_sdk/runtime/bootstrap.py index e0f26b709..14fd0138b 100644 --- a/src-new/astrbot_sdk/runtime/bootstrap.py +++ b/src-new/astrbot_sdk/runtime/bootstrap.py @@ -55,22 +55,16 @@ async def run_supervisor( env_manager: PluginEnvironmentManager | None = None, wire_codec: str = "json", ) -> None: - codec = make_protocol_codec(wire_codec) transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( stdin, stdout, - binary=codec.stdio_framing == "length_prefixed", - ) - transport = StdioTransport( - stdin=transport_stdin, - stdout=transport_stdout, - framing=codec.stdio_framing, ) + transport = StdioTransport(stdin=transport_stdin, stdout=transport_stdout) runtime = SupervisorRuntime( transport=transport, plugins_dir=plugins_dir, env_manager=env_manager, - codec=codec, + worker_wire_codec_name=wire_codec, ) try: diff --git a/src-new/astrbot_sdk/runtime/supervisor.py b/src-new/astrbot_sdk/runtime/supervisor.py index 14ad9b109..80211fc88 100644 --- a/src-new/astrbot_sdk/runtime/supervisor.py +++ b/src-new/astrbot_sdk/runtime/supervisor.py @@ -395,14 +395,14 @@ class SupervisorRuntime: plugins_dir: Path, env_manager: PluginEnvironmentManager | None = None, codec: ProtocolCodec | None = None, - wire_codec_name: str = "json", + worker_wire_codec_name: str = "json", ) -> None: self.transport = transport self.plugins_dir = plugins_dir.resolve() self.repo_root = Path(__file__).resolve().parents[3] self.env_manager = env_manager or PluginEnvironmentManager(self.repo_root) - self.codec = codec or make_protocol_codec(wire_codec_name) - self.wire_codec_name = self.codec.name + self.codec = codec or make_protocol_codec("json") + self.worker_wire_codec_name = worker_wire_codec_name self.capability_router = CapabilityRouter() self.peer = Peer( transport=self.transport, @@ -646,7 +646,7 @@ class SupervisorRuntime: repo_root=self.repo_root, env_manager=self.env_manager, capability_router=self.capability_router, - wire_codec_name=self.wire_codec_name, + wire_codec_name=self.worker_wire_codec_name, on_closed=lambda group_id=group.id: ( self._handle_worker_closed(group_id) ), @@ -660,7 +660,7 @@ class SupervisorRuntime: repo_root=self.repo_root, env_manager=self.env_manager, capability_router=self.capability_router, - wire_codec_name=self.wire_codec_name, + wire_codec_name=self.worker_wire_codec_name, on_closed=lambda plugin_name=plugin.name: ( self._handle_worker_closed(plugin_name) ), diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py index 0801317aa..7a8aab17b 100644 --- a/tests_v4/test_peer.py +++ b/tests_v4/test_peer.py @@ -17,6 +17,7 @@ from astrbot_sdk.protocol.wire_codecs import MsgpackProtocolCodec from astrbot_sdk.runtime.capability_router import CapabilityRouter, StreamExecution from astrbot_sdk.runtime.peer import Peer from astrbot_sdk.runtime.transport import ( + Transport, WebSocketClientTransport, WebSocketServerTransport, ) @@ -41,6 +42,21 @@ def make_linked_transport_pair() -> tuple[LinkedMemoryTransport, LinkedMemoryTra return left, right +class RecordingTextTransport(Transport): + def __init__(self) -> None: + super().__init__() + self.sent_payloads: list[bytes | str] = [] + + async def start(self) -> None: + self._closed.clear() + + async def stop(self) -> None: + self._closed.set() + + async def send(self, payload): + self.sent_payloads.append(payload) + + class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self) -> None: self.left, self.right = make_transport_pair() @@ -90,6 +106,20 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): await plugin.stop() await core.stop() + async def test_default_json_codec_preserves_text_transport_payloads(self) -> None: + transport = RecordingTextTransport() + peer = Peer( + transport=transport, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + + await peer.start() + await peer.cancel("request-1") + await peer.stop() + + self.assertEqual(len(transport.sent_payloads), 1) + self.assertIsInstance(transport.sent_payloads[0], str) + async def test_initialize_carries_remote_provided_capabilities(self) -> None: provided = CapabilityDescriptor( name="demo.echo", diff --git a/tests_v4/test_wire_codecs.py b/tests_v4/test_wire_codecs.py index 92f48e41f..cacbac042 100644 --- a/tests_v4/test_wire_codecs.py +++ b/tests_v4/test_wire_codecs.py @@ -34,7 +34,7 @@ class TestJsonProtocolCodec: encoded = codec.encode_message(message) decoded = codec.decode_message(encoded) - assert isinstance(encoded, bytes) + assert isinstance(encoded, str) assert decoded == message From 5bcf259fb60c49b67ff9ee8a942b2e8eb13aa31c Mon Sep 17 00:00:00 2001 From: letr Date: Sat, 14 Mar 2026 20:34:23 +0800 Subject: [PATCH 4/4] fix(cli): scope worker wire codec option --- src-new/astrbot_sdk/cli.py | 18 ++++++++++++------ src-new/astrbot_sdk/runtime/bootstrap.py | 4 ++-- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src-new/astrbot_sdk/cli.py b/src-new/astrbot_sdk/cli.py index 8ca894d23..c00be793f 100644 --- a/src-new/astrbot_sdk/cli.py +++ b/src-new/astrbot_sdk/cli.py @@ -868,23 +868,29 @@ def cli(ctx, verbose: bool) -> None: help="Directory containing plugin folders", ) @click.option( - "--wire-codec", + "--worker-wire-codec", default="json", show_default=True, type=click.Choice(["json", "msgpack"]), - help="Wire codec for supervisor/worker transport", + help="Wire codec for supervisor-to-worker transport", ) -def run(plugins_dir: Path, wire_codec: str) -> None: +def run(plugins_dir: Path, worker_wire_codec: str) -> None: """Start the plugin supervisor over stdio.""" entrypoint = ( run_supervisor(plugins_dir=plugins_dir) - if wire_codec == "json" - else run_supervisor(plugins_dir=plugins_dir, wire_codec=wire_codec) + if worker_wire_codec == "json" + else run_supervisor( + plugins_dir=plugins_dir, + worker_wire_codec=worker_wire_codec, + ) ) _run_async_entrypoint( entrypoint, log_message=f"启动插件主管进程，插件目录：{plugins_dir}", - context={"plugins_dir": plugins_dir, "wire_codec": wire_codec}, + context={ + "plugins_dir": plugins_dir, + "worker_wire_codec": worker_wire_codec, + }, ) diff --git a/src-new/astrbot_sdk/runtime/bootstrap.py b/src-new/astrbot_sdk/runtime/bootstrap.py index 14fd0138b..7d025b27f 100644 --- a/src-new/astrbot_sdk/runtime/bootstrap.py +++ b/src-new/astrbot_sdk/runtime/bootstrap.py @@ -53,7 +53,7 @@ async def run_supervisor( stdin: IO[str] | IO[bytes] | None = None, stdout: IO[str] | IO[bytes] | None = None, env_manager: PluginEnvironmentManager | None = None, - wire_codec: str = "json", + worker_wire_codec: str = "json", ) -> None: transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( stdin, @@ -64,7 +64,7 @@ async def run_supervisor( transport=transport, plugins_dir=plugins_dir, env_manager=env_manager, - worker_wire_codec_name=wire_codec, + worker_wire_codec_name=worker_wire_codec, ) try: