mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-19 10:22:04 +08:00
fix(runtime): preserve json transport compatibility
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user