mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
9c5ecf0a13
- Introduced LoadedCapability class to manage plugin capabilities. - Updated load_plugin function to discover and load capabilities alongside handlers. - Enhanced Peer class to handle remote provided capabilities during initialization. - Added tests for capability registration and invocation in legacy plugins. - Improved MessageChain and message component handling in legacy plugins. - Added comprehensive tests for legacy plugin integration and compatibility. - Updated protocol messages to include provided capabilities in initialization. - Enhanced top-level module imports to include capability-related functions.
376 lines
13 KiB
Python
376 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import unittest
|
|
|
|
from astrbot_sdk.context import CancelToken
|
|
from astrbot_sdk.errors import AstrBotError
|
|
from astrbot_sdk.protocol.descriptors import CapabilityDescriptor
|
|
from astrbot_sdk.protocol.messages import (
|
|
EventMessage,
|
|
InitializeOutput,
|
|
PeerInfo,
|
|
ResultMessage,
|
|
)
|
|
from astrbot_sdk.runtime.capability_router import CapabilityRouter, StreamExecution
|
|
from astrbot_sdk.runtime.peer import Peer
|
|
from astrbot_sdk.runtime.transport import (
|
|
WebSocketClientTransport,
|
|
WebSocketServerTransport,
|
|
)
|
|
|
|
from tests_v4.helpers import make_transport_pair
|
|
|
|
|
|
class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self.left, self.right = make_transport_pair()
|
|
|
|
async def test_initialize_and_call_builtin_capabilities(self) -> None:
|
|
router = CapabilityRouter()
|
|
core = Peer(
|
|
transport=self.left,
|
|
peer_info=PeerInfo(name="core", role="core", version="v4"),
|
|
)
|
|
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"),
|
|
)
|
|
|
|
await core.start()
|
|
await plugin.start()
|
|
await plugin.initialize([])
|
|
|
|
result = await plugin.invoke("llm.chat", {"prompt": "hello"})
|
|
self.assertEqual(result["text"], "Echo: hello")
|
|
|
|
stream = await plugin.invoke_stream("llm.stream_chat", {"prompt": "hi"})
|
|
chunks = [event.data["text"] async for event in stream]
|
|
self.assertEqual("".join(chunks), "Echo: hi")
|
|
|
|
await plugin.stop()
|
|
await core.stop()
|
|
|
|
async def test_initialize_carries_remote_provided_capabilities(self) -> None:
|
|
provided = CapabilityDescriptor(
|
|
name="demo.echo",
|
|
description="Echo text",
|
|
input_schema={"type": "object", "properties": {"text": {"type": "string"}}},
|
|
output_schema={
|
|
"type": "object",
|
|
"properties": {"echo": {"type": "string"}},
|
|
},
|
|
)
|
|
|
|
async def init_handler(message):
|
|
self.assertEqual(
|
|
[item.name for item in message.provided_capabilities], ["demo.echo"]
|
|
)
|
|
return InitializeOutput(
|
|
peer=PeerInfo(name="core", role="core", version="v4"),
|
|
capabilities=[],
|
|
metadata={},
|
|
)
|
|
|
|
core = Peer(
|
|
transport=self.left,
|
|
peer_info=PeerInfo(name="core", role="core", version="v4"),
|
|
)
|
|
core.set_initialize_handler(init_handler)
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
|
|
await core.start()
|
|
await plugin.start()
|
|
await plugin.initialize([], provided_capabilities=[provided])
|
|
|
|
self.assertEqual(
|
|
[item.name for item in core.remote_provided_capabilities],
|
|
["demo.echo"],
|
|
)
|
|
await plugin.stop()
|
|
await core.stop()
|
|
|
|
async def test_stream_false_receiving_event_is_protocol_error(self) -> None:
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
await self.left.start()
|
|
await plugin.start()
|
|
|
|
task = asyncio.create_task(
|
|
plugin.invoke("llm.chat", {"prompt": "bad"}, request_id="req-1")
|
|
)
|
|
await asyncio.sleep(0)
|
|
await self.left.send(
|
|
EventMessage(id="req-1", phase="started").model_dump_json()
|
|
)
|
|
|
|
with self.assertRaises(AstrBotError) as raised:
|
|
await task
|
|
self.assertEqual(raised.exception.code, "protocol_error")
|
|
await plugin.stop()
|
|
await self.left.stop()
|
|
|
|
async def test_stream_true_receiving_result_is_protocol_error(self) -> None:
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
await self.left.start()
|
|
await plugin.start()
|
|
|
|
stream = await plugin.invoke_stream(
|
|
"llm.stream_chat", {"prompt": "bad"}, request_id="stream-1"
|
|
)
|
|
await self.left.send(
|
|
ResultMessage(id="stream-1", success=True, output={}).model_dump_json()
|
|
)
|
|
|
|
with self.assertRaises(AstrBotError) as raised:
|
|
async for _ in stream:
|
|
pass
|
|
self.assertEqual(raised.exception.code, "protocol_error")
|
|
await plugin.stop()
|
|
await self.left.stop()
|
|
|
|
async def test_invalid_inbound_message_fails_pending_calls(self) -> None:
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
await self.left.start()
|
|
await plugin.start()
|
|
|
|
task = asyncio.create_task(
|
|
plugin.invoke("llm.chat", {"prompt": "bad"}, request_id="req-invalid")
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
with self.assertRaises(AstrBotError) as raised_send:
|
|
await self.left.send("[]")
|
|
self.assertEqual(raised_send.exception.code, "protocol_error")
|
|
|
|
with self.assertRaises(AstrBotError) as raised_task:
|
|
await task
|
|
self.assertEqual(raised_task.exception.code, "protocol_error")
|
|
|
|
await asyncio.wait_for(plugin.wait_closed(), timeout=1.0)
|
|
await self.left.stop()
|
|
|
|
async def test_cancel_waits_for_failed_terminal_event(self) -> None:
|
|
descriptor = CapabilityDescriptor(
|
|
name="slow.stream",
|
|
description="slow stream",
|
|
input_schema={"type": "object", "properties": {}, "required": []},
|
|
output_schema={
|
|
"type": "object",
|
|
"properties": {"count": {"type": "number"}},
|
|
"required": ["count"],
|
|
},
|
|
supports_stream=True,
|
|
cancelable=True,
|
|
)
|
|
|
|
async def init_handler(_message):
|
|
return InitializeOutput(
|
|
peer=PeerInfo(name="core", role="core", version="v4"),
|
|
capabilities=[descriptor],
|
|
metadata={},
|
|
)
|
|
|
|
async def invoke_handler(message, token: CancelToken):
|
|
async def iterator():
|
|
while True:
|
|
token.raise_if_cancelled()
|
|
await asyncio.sleep(0.01)
|
|
yield {"text": "x"}
|
|
|
|
return StreamExecution(
|
|
iterator=iterator(),
|
|
finalize=lambda chunks: {"count": len(chunks)},
|
|
)
|
|
|
|
core = Peer(
|
|
transport=self.left,
|
|
peer_info=PeerInfo(name="core", role="core", version="v4"),
|
|
)
|
|
core.set_initialize_handler(init_handler)
|
|
core.set_invoke_handler(invoke_handler)
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
|
|
await core.start()
|
|
await plugin.start()
|
|
await plugin.initialize([])
|
|
|
|
stream = await plugin.invoke_stream("slow.stream", {}, request_id="cancel-1")
|
|
first = await anext(stream)
|
|
self.assertEqual(first.data["text"], "x")
|
|
await plugin.cancel("cancel-1")
|
|
|
|
with self.assertRaises(AstrBotError) as raised:
|
|
await anext(stream)
|
|
self.assertEqual(raised.exception.code, "cancelled")
|
|
await plugin.stop()
|
|
await core.stop()
|
|
|
|
async def test_websocket_transport_smoke(self) -> None:
|
|
router = CapabilityRouter()
|
|
server_transport = WebSocketServerTransport(port=0)
|
|
core = Peer(
|
|
transport=server_transport,
|
|
peer_info=PeerInfo(name="core", role="core", version="v4"),
|
|
)
|
|
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"),
|
|
)
|
|
await plugin.start()
|
|
await plugin.initialize([])
|
|
result = await plugin.invoke("llm.chat", {"prompt": "ws"})
|
|
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)
|
|
await asyncio.wait_for(plugin.wait_closed(), timeout=1.0)
|
|
self.assertTrue(core._closed)
|
|
self.assertTrue(plugin._closed)
|
|
|
|
async def test_wait_until_remote_initialized_raises_if_connection_closes_first(
|
|
self,
|
|
) -> None:
|
|
plugin = Peer(
|
|
transport=self.right,
|
|
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
|
|
)
|
|
|
|
await self.left.start()
|
|
await plugin.start()
|
|
await plugin.stop()
|
|
|
|
with self.assertRaises(AstrBotError) as raised:
|
|
await plugin.wait_until_remote_initialized(timeout=None)
|
|
self.assertEqual(raised.exception.code, "protocol_error")
|
|
|
|
await self.left.stop()
|
|
|
|
|
|
class CapabilityRouterContractTest(unittest.TestCase):
|
|
def test_capability_names_must_match_namespace_method_format(self) -> None:
|
|
router = CapabilityRouter()
|
|
for name in ("llm", "llm.chat.extra", "LLM.chat", "llm.Chat"):
|
|
with self.assertRaises(ValueError) as raised:
|
|
router.register(
|
|
CapabilityDescriptor(
|
|
name=name,
|
|
description="invalid",
|
|
)
|
|
)
|
|
self.assertIn(name, str(raised.exception))
|
|
|
|
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()])
|