From 98a5eddfcbf528eba5d568fec19530afd77abf9a Mon Sep 17 00:00:00 2001
From: letr <letr007@foxmail.com>
Date: Sat, 14 Mar 2026 20:04:34 +0800
Subject: [PATCH] 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,