From 61eecd20242ebf3c8cf1e7cf8c832609b00c4186 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 12 Mar 2026 22:56:27 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E5=AF=B9=E4=BF=9D?= =?UTF-8?q?=E7=95=99=E8=83=BD=E5=8A=9B=E5=91=BD=E5=90=8D=E7=A9=BA=E9=97=B4?= =?UTF-8?q?=E7=9A=84=E6=A3=80=E6=9F=A5=EF=BC=8C=E9=98=B2=E6=AD=A2=E6=9A=B4?= =?UTF-8?q?=E9=9C=B2=E6=B3=A8=E5=86=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../astrbot_sdk/runtime/capability_router.py | 3 + src-new/astrbot_sdk/runtime/peer.py | 43 +++++++------ tests_v4/test_peer.py | 60 +++++++++++++++++++ 3 files changed, 84 insertions(+), 22 deletions(-) diff --git a/src-new/astrbot_sdk/runtime/capability_router.py b/src-new/astrbot_sdk/runtime/capability_router.py index 205befb0d..b197ea3cc 100644 --- a/src-new/astrbot_sdk/runtime/capability_router.py +++ b/src-new/astrbot_sdk/runtime/capability_router.py @@ -12,6 +12,7 @@ from ..protocol.descriptors import CapabilityDescriptor CallHandler = Callable[[str, dict[str, Any], object], Awaitable[dict[str, Any]]] StreamHandler = Callable[[str, dict[str, Any], object], AsyncIterator[dict[str, Any]]] FinalizeHandler = Callable[[list[dict[str, Any]]], dict[str, Any]] +RESERVED_CAPABILITY_PREFIXES = ("handler.", "system.", "internal.") @dataclass(slots=True) @@ -53,6 +54,8 @@ class CapabilityRouter: finalize: FinalizeHandler | None = None, exposed: bool = True, ) -> None: + if exposed and descriptor.name.startswith(RESERVED_CAPABILITY_PREFIXES): + raise ValueError(f"保留 capability 命名空间仅供框架内部使用：{descriptor.name}") self._registrations[descriptor.name] = _CapabilityRegistration( descriptor=descriptor, call_handler=call_handler, diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py index aa6cc2d2d..5af890e69 100644 --- a/src-new/astrbot_sdk/runtime/peer.py +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -216,33 +216,19 @@ class Peer: self.remote_handlers = message.handlers self.remote_metadata = message.metadata if self._initialize_handler is None: - error = AstrBotError.protocol_error("对端不接受 initialize") - await self._send( - ResultMessage( - id=message.id, - kind="initialize_result", - success=False, - error=ErrorPayload.model_validate(error.to_payload()), - ) + await self._reject_initialize( + message, + AstrBotError.protocol_error("对端不接受 initialize"), ) - self._unusable = True - self._remote_initialized.set() return if message.protocol_version != self.protocol_version: - error = AstrBotError.protocol_version_mismatch( - f"服务端支持协议版本 {self.protocol_version}，客户端请求版本 {message.protocol_version}" + await self._reject_initialize( + message, + AstrBotError.protocol_version_mismatch( + f"服务端支持协议版本 {self.protocol_version}，客户端请求版本 {message.protocol_version}" + ), ) - 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() return output = await self._initialize_handler(message) @@ -338,6 +324,19 @@ class Peer: ) ) + 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) diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py index 2b917063e..5ca26d44e 100644 --- a/tests_v4/test_peer.py +++ b/tests_v4/test_peer.py @@ -189,3 +189,63 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): self.assertEqual(result["text"], "Echo: ws") await plugin.stop() await core.stop() + + async def test_initialize_failure_closes_receiver_connection(self) -> None: + core = Peer( + transport=self.left, + peer_info=PeerInfo(name="core", role="core", version="v4"), + protocol_version="1.0", + ) + core.set_initialize_handler( + lambda _message: asyncio.sleep( + 0, + result=InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=[], + metadata={}, + ), + ) + ) + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + protocol_version="2.0", + ) + + await core.start() + await plugin.start() + + with self.assertRaises(AstrBotError) as raised: + await plugin.initialize([]) + self.assertEqual(raised.exception.code, "protocol_version_mismatch") + + await asyncio.wait_for(core.wait_closed(), timeout=1.0) + self.assertTrue(core._closed) + + await plugin.stop() + + +class CapabilityRouterContractTest(unittest.TestCase): + def test_reserved_capability_namespaces_are_rejected_for_exposed_registrations(self) -> None: + router = CapabilityRouter() + for name in ("handler.demo", "system.health", "internal.trace"): + with self.assertRaises(ValueError) as raised: + router.register( + CapabilityDescriptor( + name=name, + description="reserved", + ) + ) + self.assertIn(name, str(raised.exception)) + + def test_reserved_capability_namespaces_remain_available_for_hidden_internal_registrations(self) -> None: + router = CapabilityRouter() + router.register( + CapabilityDescriptor( + name="system.health", + description="internal only", + ), + exposed=False, + ) + + self.assertNotIn("system.health", [item.name for item in router.descriptors()])