fix(runtime): preserve json transport compatibility

This commit is contained in:
letr
2026-03-14 20:21:34 +08:00
parent 98a5eddfcb
commit a979e92568
5 changed files with 41 additions and 17 deletions

View File

@@ -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):

View File

@@ -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:

View File

@@ -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)
),

View File

@@ -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",

View File

@@ -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