diff --git a/src-new/astrbot_sdk/_legacy_api.py b/src-new/astrbot_sdk/_legacy_api.py index 291a3a91d..6fbfafe25 100644 --- a/src-new/astrbot_sdk/_legacy_api.py +++ b/src-new/astrbot_sdk/_legacy_api.py @@ -119,7 +119,9 @@ class LegacyContext: ctx = self.require_runtime_context() # 正确序列化 MessageChain 对象 # 优先使用 get_plain_text() 方法(旧版 MessageChain) - if hasattr(message_chain, "get_plain_text") and callable(message_chain.get_plain_text): + if hasattr(message_chain, "get_plain_text") and callable( + message_chain.get_plain_text + ): text = message_chain.get_plain_text() elif hasattr(message_chain, "to_text") and callable(message_chain.to_text): text = message_chain.to_text() diff --git a/src-new/astrbot_sdk/compat.py b/src-new/astrbot_sdk/compat.py index e6388df80..8632db153 100644 --- a/src-new/astrbot_sdk/compat.py +++ b/src-new/astrbot_sdk/compat.py @@ -1,4 +1,9 @@ -from ._legacy_api import CommandComponent, Context, LegacyContext, LegacyConversationManager +from ._legacy_api import ( + CommandComponent, + Context, + LegacyContext, + LegacyConversationManager, +) __all__ = [ "CommandComponent", diff --git a/src-new/astrbot_sdk/decorators.py b/src-new/astrbot_sdk/decorators.py index 58e59f640..4b783f96f 100644 --- a/src-new/astrbot_sdk/decorators.py +++ b/src-new/astrbot_sdk/decorators.py @@ -17,7 +17,9 @@ HANDLER_META_ATTR = "__astrbot_handler_meta__" @dataclass(slots=True) class HandlerMeta: - trigger: CommandTrigger | MessageTrigger | EventTrigger | ScheduleTrigger | None = None + trigger: CommandTrigger | MessageTrigger | EventTrigger | ScheduleTrigger | None = ( + None + ) priority: int = 0 permissions: Permissions = field(default_factory=Permissions) diff --git a/src-new/astrbot_sdk/events.py b/src-new/astrbot_sdk/events.py index ba58948a5..6c3c77c7e 100644 --- a/src-new/astrbot_sdk/events.py +++ b/src-new/astrbot_sdk/events.py @@ -37,7 +37,9 @@ class MessageEvent: self.raw = raw or {} self._reply_handler = reply_handler if self._reply_handler is None and context is not None: - self._reply_handler = lambda text: context.platform.send(self.session_id, text) + self._reply_handler = lambda text: context.platform.send( + self.session_id, text + ) @classmethod def from_payload( diff --git a/src-new/astrbot_sdk/protocol/legacy_adapter.py b/src-new/astrbot_sdk/protocol/legacy_adapter.py index a2f851162..7a99cfd36 100644 --- a/src-new/astrbot_sdk/protocol/legacy_adapter.py +++ b/src-new/astrbot_sdk/protocol/legacy_adapter.py @@ -55,7 +55,9 @@ class LegacyErrorResponse(_LegacyResponse): LegacyMessage = LegacyRequest | LegacySuccessResponse | LegacyErrorResponse -LegacyToV4Message = InitializeMessage | InvokeMessage | ResultMessage | EventMessage | CancelMessage +LegacyToV4Message = ( + InitializeMessage | InvokeMessage | ResultMessage | EventMessage | CancelMessage +) class LegacyAdapter: @@ -93,7 +95,11 @@ class LegacyAdapter: self, payload: LegacyRequest | dict[str, Any], ) -> InitializeMessage | InvokeMessage | EventMessage | CancelMessage: - message = payload if isinstance(payload, LegacyRequest) else LegacyRequest.model_validate(payload) + message = ( + payload + if isinstance(payload, LegacyRequest) + else LegacyRequest.model_validate(payload) + ) params = message.params or {} method = message.method @@ -164,7 +170,9 @@ class LegacyAdapter: return EventMessage( id=request_id, phase="failed", - error=ErrorPayload.model_validate(self._coerce_error_payload(error)), + error=ErrorPayload.model_validate( + self._coerce_error_payload(error) + ), ) return EventMessage(id=request_id, phase="completed") @@ -185,13 +193,22 @@ class LegacyAdapter: self, payload: LegacySuccessResponse | dict[str, Any], ) -> InitializeMessage | ResultMessage: - message = payload if isinstance(payload, LegacySuccessResponse) else LegacySuccessResponse.model_validate(payload) + message = ( + payload + if isinstance(payload, LegacySuccessResponse) + else LegacySuccessResponse.model_validate(payload) + ) request_id = self._request_id(message.id, "legacy-result") - if request_id in self._pending_handshake_ids or self._looks_like_handshake_payload(message.result): + if ( + request_id in self._pending_handshake_ids + or self._looks_like_handshake_payload(message.result) + ): self._pending_handshake_ids.discard(request_id) payload_dict = self._as_dict(message.result, field_name="data") - peer_name, peer_version = self._legacy_peer_from_handshake_payload(payload_dict) + peer_name, peer_version = self._legacy_peer_from_handshake_payload( + payload_dict + ) return InitializeMessage( id=request_id, protocol_version=self.protocol_version, @@ -217,7 +234,11 @@ class LegacyAdapter: self, payload: LegacyErrorResponse | dict[str, Any], ) -> ResultMessage: - message = payload if isinstance(payload, LegacyErrorResponse) else LegacyErrorResponse.model_validate(payload) + message = ( + payload + if isinstance(payload, LegacyErrorResponse) + else LegacyErrorResponse.model_validate(payload) + ) request_id = self._request_id(message.id, "legacy-error") kind = None if request_id in self._pending_handshake_ids: @@ -227,7 +248,9 @@ class LegacyAdapter: id=request_id, kind=kind, success=False, - error=ErrorPayload.model_validate(self._legacy_error_to_payload(message.error)), + error=ErrorPayload.model_validate( + self._legacy_error_to_payload(message.error) + ), ) def build_legacy_handshake_request(self, request_id: str) -> dict[str, Any]: @@ -263,8 +286,12 @@ class LegacyAdapter: "method": "call_handler", "params": { "handler_full_name": handler_full_name, - "event": self._as_dict(message.input.get("event"), field_name="data"), - "args": self._as_dict(message.input.get("args"), field_name="value"), + "event": self._as_dict( + message.input.get("event"), field_name="data" + ), + "args": self._as_dict( + message.input.get("args"), field_name="value" + ), }, } @@ -275,7 +302,9 @@ class LegacyAdapter: "method": "call_context_function", "params": { "name": str(message.input.get("name", "")), - "args": self._as_dict(message.input.get("args"), field_name="value"), + "args": self._as_dict( + message.input.get("args"), field_name="value" + ), }, } @@ -350,7 +379,9 @@ class LegacyAdapter: def _looks_like_handshake_payload(value: Any) -> bool: if not isinstance(value, dict) or not value: return False - return all(isinstance(item, dict) and "handlers" in item for item in value.values()) + return all( + isinstance(item, dict) and "handlers" in item for item in value.values() + ) @staticmethod def _coerce_error_payload(value: dict[str, Any]) -> dict[str, Any]: @@ -394,7 +425,9 @@ class LegacyAdapter: return handlers @staticmethod - def _legacy_handler_to_descriptor(handler_data: dict[str, Any]) -> HandlerDescriptor: + def _legacy_handler_to_descriptor( + handler_data: dict[str, Any], + ) -> HandlerDescriptor: extras_configs = handler_data.get("extras_configs") extras = extras_configs if isinstance(extras_configs, dict) else {} handler_id = str( @@ -425,7 +458,9 @@ class LegacyAdapter: display_name = str(message.metadata.get("display_name") or plugin_name) module_path = str(message.metadata.get("module_path") or f"{plugin_name}.main") root_dir_name = str(message.metadata.get("root_dir_name") or plugin_name) - handlers = [self._descriptor_to_legacy_handler(item) for item in message.handlers] + handlers = [ + self._descriptor_to_legacy_handler(item) for item in message.handlers + ] return { module_path: { "name": plugin_name, @@ -438,7 +473,9 @@ class LegacyAdapter: "reserved": bool(message.metadata.get("reserved", False)), "activated": bool(message.metadata.get("activated", True)), "config": message.metadata.get("config"), - "star_handler_full_names": [item["handler_full_name"] for item in handlers], + "star_handler_full_names": [ + item["handler_full_name"] for item in handlers + ], "display_name": display_name, "logo_path": message.metadata.get("logo_path"), "handlers": handlers, diff --git a/src-new/astrbot_sdk/runtime/bootstrap.py b/src-new/astrbot_sdk/runtime/bootstrap.py index 921515c56..6c747c0c0 100644 --- a/src-new/astrbot_sdk/runtime/bootstrap.py +++ b/src-new/astrbot_sdk/runtime/bootstrap.py @@ -114,7 +114,9 @@ class WorkerSession: try: await self.peer.start() # 同时监听初始化完成和连接关闭,避免 worker 崩溃时等满超时 - init_task = asyncio.create_task(self.peer.wait_until_remote_initialized(timeout=None)) + init_task = asyncio.create_task( + self.peer.wait_until_remote_initialized(timeout=None) + ) closed_task = asyncio.create_task(self.peer.wait_closed()) done, pending = await asyncio.wait( {init_task, closed_task}, @@ -128,7 +130,9 @@ class WorkerSession: pass if closed_task in done: - raise RuntimeError(f"插件 {self.plugin.name} worker 进程在初始化阶段退出") + raise RuntimeError( + f"插件 {self.plugin.name} worker 进程在初始化阶段退出" + ) self.handlers = list(self.peer.remote_handlers) @@ -147,7 +151,9 @@ class WorkerSession: try: self.on_closed() except Exception: - logger.exception("on_closed callback failed for plugin {}", self.plugin.name) + logger.exception( + "on_closed callback failed for plugin {}", self.plugin.name + ) async def stop(self) -> None: if self.peer is not None: @@ -267,11 +273,17 @@ class SupervisorRuntime: self.handler_to_worker[handler.id] = session aggregated_handlers = list(self.handler_to_worker.keys()) - logger.info("Loaded plugins: {}", ", ".join(sorted(self.loaded_plugins)) or "none") + logger.info( + "Loaded plugins: {}", ", ".join(sorted(self.loaded_plugins)) or "none" + ) await self.peer.start() await self.peer.initialize( - [handler for session in self.worker_sessions.values() for handler in session.handlers], + [ + handler + for session in self.worker_sessions.values() + for handler in session.handlers + ], metadata={ "plugins": sorted(self.loaded_plugins), "skipped_plugins": self.skipped_plugins, @@ -349,7 +361,9 @@ class PluginWorkerRuntime: peer=self.peer, handlers=self.loaded_plugin.handlers, ) - self._lifecycle_context = RuntimeContext(peer=self.peer, plugin_id=self.plugin.name) + self._lifecycle_context = RuntimeContext( + peer=self.peer, plugin_id=self.plugin.name + ) self.peer.set_invoke_handler(self._handle_invoke) self.peer.set_cancel_handler(self.dispatcher.cancel) @@ -391,7 +405,8 @@ class PluginWorkerRuntime: positional_params = [ parameter for parameter in signature.parameters.values() - if parameter.kind in ( + if parameter.kind + in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ) diff --git a/src-new/astrbot_sdk/runtime/capability_router.py b/src-new/astrbot_sdk/runtime/capability_router.py index 1ac3281a5..5470f4c34 100644 --- a/src-new/astrbot_sdk/runtime/capability_router.py +++ b/src-new/astrbot_sdk/runtime/capability_router.py @@ -42,9 +42,7 @@ class CapabilityRouter: def descriptors(self) -> list[CapabilityDescriptor]: return [ - entry.descriptor - for entry in self._registrations.values() - if entry.exposed + entry.descriptor for entry in self._registrations.values() if entry.exposed ] def register( @@ -57,9 +55,13 @@ class CapabilityRouter: exposed: bool = True, ) -> None: if not CAPABILITY_NAME_PATTERN.fullmatch(descriptor.name): - raise ValueError(f"capability 名称必须匹配 {{namespace}}.{{method}}:{descriptor.name}") + raise ValueError( + f"capability 名称必须匹配 {{namespace}}.{{method}}:{descriptor.name}" + ) if exposed and descriptor.name.startswith(RESERVED_CAPABILITY_PREFIXES): - raise ValueError(f"保留 capability 命名空间仅供框架内部使用:{descriptor.name}") + raise ValueError( + f"保留 capability 命名空间仅供框架内部使用:{descriptor.name}" + ) self._registrations[descriptor.name] = _CapabilityRegistration( descriptor=descriptor, call_handler=call_handler, @@ -105,11 +107,15 @@ class CapabilityRouter: "required": required, } - async def llm_chat(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def llm_chat( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: prompt = str(payload.get("prompt", "")) return {"text": f"Echo: {prompt}"} - async def llm_chat_raw(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def llm_chat_raw( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: prompt = str(payload.get("prompt", "")) text = f"Echo: {prompt}" return { @@ -133,7 +139,9 @@ class CapabilityRouter: await asyncio.sleep(0) yield {"text": char} - async def memory_search(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def memory_search( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: query = str(payload.get("query", "")) items = [ {"key": key, "value": value} @@ -142,7 +150,9 @@ class CapabilityRouter: ] return {"items": items} - async def memory_save(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def memory_save( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: key = str(payload.get("key", "")) value = payload.get("value") if not isinstance(value, dict): @@ -150,14 +160,20 @@ class CapabilityRouter: self.memory_store[key] = value return {} - async def memory_delete(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def memory_delete( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: self.memory_store.pop(str(payload.get("key", "")), None) return {} - async def db_get(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def db_get( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: return {"value": self.db_store.get(str(payload.get("key", "")))} - async def db_set(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def db_set( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: key = str(payload.get("key", "")) value = payload.get("value") if not isinstance(value, dict): @@ -165,18 +181,24 @@ class CapabilityRouter: self.db_store[key] = value return {} - async def db_delete(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def db_delete( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: self.db_store.pop(str(payload.get("key", "")), None) return {} - async def db_list(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def db_list( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: prefix = payload.get("prefix") keys = sorted(self.db_store.keys()) if isinstance(prefix, str): keys = [item for item in keys if item.startswith(prefix)] return {"keys": keys} - async def platform_send(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def platform_send( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: session = str(payload.get("session", "")) text = str(payload.get("text", "")) message_id = f"msg_{len(self.sent_messages) + 1}" @@ -189,7 +211,9 @@ class CapabilityRouter: ) return {"message_id": message_id} - async def platform_send_image(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def platform_send_image( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: session = str(payload.get("session", "")) image_url = str(payload.get("image_url", "")) message_id = f"img_{len(self.sent_messages) + 1}" @@ -202,7 +226,9 @@ class CapabilityRouter: ) return {"message_id": message_id} - async def platform_get_members(_request_id: str, payload: dict[str, Any], _token) -> dict[str, Any]: + async def platform_get_members( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: session = str(payload.get("session", "")) return { "members": [ @@ -239,7 +265,9 @@ class CapabilityRouter: cancelable=True, ), stream_handler=llm_stream, - finalize=lambda chunks: {"text": "".join(item.get("text", "") for item in chunks)}, + finalize=lambda chunks: { + "text": "".join(item.get("text", "") for item in chunks) + }, ) self.register( CapabilityDescriptor( @@ -254,7 +282,9 @@ class CapabilityRouter: CapabilityDescriptor( name="memory.save", description="保存记忆", - input_schema=obj_schema(["key", "value"], key={"type": "string"}, value={"type": "object"}), + input_schema=obj_schema( + ["key", "value"], key={"type": "string"}, value={"type": "object"} + ), output_schema=obj_schema([]), ), call_handler=memory_save, @@ -281,7 +311,9 @@ class CapabilityRouter: CapabilityDescriptor( name="db.set", description="写入 KV", - input_schema=obj_schema(["key", "value"], key={"type": "string"}, value={"type": "object"}), + input_schema=obj_schema( + ["key", "value"], key={"type": "string"}, value={"type": "object"} + ), output_schema=obj_schema([]), ), call_handler=db_set, @@ -308,7 +340,11 @@ class CapabilityRouter: CapabilityDescriptor( name="platform.send", description="发送消息", - input_schema=obj_schema(["session", "text"], session={"type": "string"}, text={"type": "string"}), + input_schema=obj_schema( + ["session", "text"], + session={"type": "string"}, + text={"type": "string"}, + ), output_schema=obj_schema(["message_id"], message_id={"type": "string"}), ), call_handler=platform_send, @@ -317,7 +353,11 @@ class CapabilityRouter: CapabilityDescriptor( name="platform.send_image", description="发送图片", - input_schema=obj_schema(["session", "image_url"], session={"type": "string"}, image_url={"type": "string"}), + input_schema=obj_schema( + ["session", "image_url"], + session={"type": "string"}, + image_url={"type": "string"}, + ), output_schema=obj_schema(["message_id"], message_id={"type": "string"}), ), call_handler=platform_send_image, diff --git a/src-new/astrbot_sdk/runtime/handler_dispatcher.py b/src-new/astrbot_sdk/runtime/handler_dispatcher.py index 3c4df8238..df0baae57 100644 --- a/src-new/astrbot_sdk/runtime/handler_dispatcher.py +++ b/src-new/astrbot_sdk/runtime/handler_dispatcher.py @@ -23,7 +23,9 @@ class HandlerDispatcher: if loaded is None: raise LookupError(f"handler not found: {handler_id}") - ctx = Context(peer=self._peer, plugin_id=self._plugin_id, cancel_token=cancel_token) + ctx = Context( + peer=self._peer, plugin_id=self._plugin_id, cancel_token=cancel_token + ) event = MessageEvent.from_payload(message.input.get("event", {})) event.bind_reply_handler(lambda text: ctx.platform.send(event.session_id, text)) if loaded.legacy_context is not None: @@ -56,7 +58,9 @@ class HandlerDispatcher: legacy_args: dict[str, Any] | None = None, ) -> None: try: - result = loaded.callable(*self._build_args(loaded.callable, event, ctx, legacy_args)) + result = loaded.callable( + *self._build_args(loaded.callable, event, ctx, legacy_args) + ) if inspect.isasyncgen(result): async for item in result: await self._consume_legacy_result(item, event) diff --git a/src-new/astrbot_sdk/runtime/loader.py b/src-new/astrbot_sdk/runtime/loader.py index 13025693f..de0737275 100644 --- a/src-new/astrbot_sdk/runtime/loader.py +++ b/src-new/astrbot_sdk/runtime/loader.py @@ -88,7 +88,9 @@ def load_plugin_spec(plugin_dir: Path) -> PluginSpec: requirements_path = plugin_dir / "requirements.txt" manifest_data = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} runtime = manifest_data.get("runtime") or {} - python_version = runtime.get("python") or f"{sys.version_info.major}.{sys.version_info.minor}" + python_version = ( + runtime.get("python") or f"{sys.version_info.major}.{sys.version_info.minor}" + ) return PluginSpec( name=str(manifest_data.get("name") or plugin_dir.name), plugin_dir=plugin_dir, @@ -119,7 +121,9 @@ def discover_plugins(plugins_dir: Path) -> PluginDiscoveryResult: skipped_plugins[entry.name] = "missing requirements.txt" continue try: - manifest_data = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} + manifest_data = ( + yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} + ) except Exception as exc: skipped_plugins[entry.name] = f"failed to parse plugin.yaml: {exc}" continue @@ -306,9 +310,7 @@ def load_plugin(plugin: PluginSpec) -> LoadedPlugin: meta = get_handler_meta(func) if meta is None or meta.trigger is None: continue - handler_id = ( - f"{plugin.name}:{instance.__class__.__module__}.{instance.__class__.__name__}.{name}" - ) + handler_id = f"{plugin.name}:{instance.__class__.__module__}.{instance.__class__.__name__}.{name}" handlers.append( LoadedHandler( descriptor=HandlerDescriptor( diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py index 0f7603afe..5fe85357f 100644 --- a/src-new/astrbot_sdk/runtime/peer.py +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -20,7 +20,9 @@ from ..protocol.messages import ( from .capability_router import StreamExecution InitializeHandler = Callable[[InitializeMessage], Awaitable[InitializeOutput]] -InvokeHandler = Callable[[InvokeMessage, CancelToken], Awaitable[dict[str, Any] | StreamExecution]] +InvokeHandler = Callable[ + [InvokeMessage, CancelToken], Awaitable[dict[str, Any] | StreamExecution] +] CancelHandler = Callable[[str], Awaitable[None]] @@ -102,7 +104,9 @@ class Peer: ) -> InitializeOutput: self._ensure_usable() request_id = self._next_id() - future: asyncio.Future[ResultMessage] = asyncio.get_running_loop().create_future() + future: asyncio.Future[ResultMessage] = ( + asyncio.get_running_loop().create_future() + ) self._pending_results[request_id] = future await self._send( InitializeMessage( @@ -119,7 +123,9 @@ class Peer: if not result.success: self._unusable = True await self.stop() - raise AstrBotError.from_payload(result.error.model_dump() if result.error else {}) + raise AstrBotError.from_payload( + result.error.model_dump() if result.error else {} + ) output = InitializeOutput.model_validate(result.output) self.remote_peer = output.peer self.remote_capabilities = output.capabilities @@ -139,7 +145,9 @@ class Peer: if stream: raise ValueError("stream=True 请使用 invoke_stream()") request_id = request_id or self._next_id() - future: asyncio.Future[ResultMessage] = asyncio.get_running_loop().create_future() + future: asyncio.Future[ResultMessage] = ( + asyncio.get_running_loop().create_future() + ) self._pending_results[request_id] = future await self._send( InvokeMessage( @@ -151,7 +159,9 @@ class Peer: ) result = await future if not result.success: - raise AstrBotError.from_payload(result.error.model_dump() if result.error else {}) + raise AstrBotError.from_payload( + result.error.model_dump() if result.error else {} + ) return result.output async def invoke_stream( @@ -190,7 +200,9 @@ class Peer: if item.phase == "completed": break if item.phase == "failed": - raise AstrBotError.from_payload(item.error.model_dump() if item.error else {}) + raise AstrBotError.from_payload( + item.error.model_dump() if item.error else {} + ) finally: self._pending_streams.pop(request_id, None) @@ -222,7 +234,11 @@ class Peer: task = asyncio.create_task(self._handle_invoke(message)) token = CancelToken() self._inbound_tasks[message.id] = (task, token) - task.add_done_callback(lambda _task, request_id=message.id: self._inbound_tasks.pop(request_id, None)) + task.add_done_callback( + lambda _task, request_id=message.id: self._inbound_tasks.pop( + request_id, None + ) + ) return if isinstance(message, CancelMessage): await self._handle_cancel(message) @@ -268,12 +284,16 @@ class Peer: execution = await self._invoke_handler(message, token) if message.stream: if not isinstance(execution, StreamExecution): - raise AstrBotError.protocol_error("stream=true 必须返回 StreamExecution") + raise AstrBotError.protocol_error( + "stream=true 必须返回 StreamExecution" + ) await self._send(EventMessage(id=message.id, phase="started")) chunks: list[dict[str, Any]] = [] async for chunk in execution.iterator: chunks.append(chunk) - await self._send(EventMessage(id=message.id, phase="delta", data=chunk)) + await self._send( + EventMessage(id=message.id, phase="delta", data=chunk) + ) await self._send( EventMessage( id=message.id, @@ -284,7 +304,9 @@ class Peer: return if isinstance(execution, StreamExecution): raise AstrBotError.protocol_error("stream=false 不能返回流式执行对象") - await self._send(ResultMessage(id=message.id, success=True, output=execution)) + await self._send( + ResultMessage(id=message.id, success=True, output=execution) + ) except asyncio.CancelledError: await self._send_cancelled_termination(message) except LookupError as exc: @@ -293,7 +315,9 @@ class Peer: except AstrBotError as exc: await self._send_error_result(message, exc) except Exception as exc: - await self._send_error_result(message, AstrBotError.internal_error(str(exc))) + await self._send_error_result( + message, AstrBotError.internal_error(str(exc)) + ) async def _handle_cancel(self, message: CancelMessage) -> None: inbound = self._inbound_tasks.get(message.id) @@ -310,7 +334,9 @@ class Peer: if future is None: queue = self._pending_streams.get(message.id) if queue is not None: - await queue.put(AstrBotError.protocol_error("stream=true 调用不应收到 result")) + await queue.put( + AstrBotError.protocol_error("stream=true 调用不应收到 result") + ) return # 检查 future 是否已完成(可能被调用方取消) if not future.done(): @@ -321,11 +347,15 @@ class Peer: if queue is None: future = self._pending_results.get(message.id) if future is not None and not future.done(): - future.set_exception(AstrBotError.protocol_error("stream=false 调用不应收到 event")) + future.set_exception( + AstrBotError.protocol_error("stream=false 调用不应收到 event") + ) return await queue.put(message) - async def _send_error_result(self, message: InvokeMessage, error: AstrBotError) -> None: + async def _send_error_result( + self, message: InvokeMessage, error: AstrBotError + ) -> None: if message.stream: await self._send( EventMessage( @@ -343,7 +373,9 @@ class Peer: ) ) - async def _reject_initialize(self, message: InitializeMessage, error: AstrBotError) -> None: + async def _reject_initialize( + self, message: InitializeMessage, error: AstrBotError + ) -> None: await self._send( ResultMessage( id=message.id, diff --git a/src/astrbot_sdk/api/components/command.py b/src/astrbot_sdk/api/components/command.py index a99a1381f..af63250ac 100644 --- a/src/astrbot_sdk/api/components/command.py +++ b/src/astrbot_sdk/api/components/command.py @@ -1,2 +1,2 @@ class CommandComponent: - pass \ No newline at end of file + pass diff --git a/src/astrbot_sdk/api/event/filter.py b/src/astrbot_sdk/api/event/filter.py index 80d47658a..c6fe7d5d7 100644 --- a/src/astrbot_sdk/api/event/filter.py +++ b/src/astrbot_sdk/api/event/filter.py @@ -8,20 +8,33 @@ from ...runtime.stars.filter.platform_adapter_type import ( PlatformAdapterType, PlatformAdapterTypeFilter, ) -from ...runtime.stars.registry.register import register_after_message_sent as after_message_sent +from ...runtime.stars.registry.register import ( + register_after_message_sent as after_message_sent, +) from ...runtime.stars.registry.register import register_command as command from ...runtime.stars.registry.register import register_command_group as command_group from ...runtime.stars.registry.register import register_custom_filter as custom_filter -from ...runtime.stars.registry.register import register_event_message_type as event_message_type +from ...runtime.stars.registry.register import ( + register_event_message_type as event_message_type, +) + # from ...runtime.stars.registry.register import register_llm_tool as llm_tool -from ...runtime.stars.registry.register import register_on_astrbot_loaded as on_astrbot_loaded +from ...runtime.stars.registry.register import ( + register_on_astrbot_loaded as on_astrbot_loaded, +) from ...runtime.stars.registry.register import ( register_on_decorating_result as on_decorating_result, ) from ...runtime.stars.registry.register import register_on_llm_request as on_llm_request -from ...runtime.stars.registry.register import register_on_llm_response as on_llm_response -from ...runtime.stars.registry.register import register_on_platform_loaded as on_platform_loaded -from ...runtime.stars.registry.register import register_permission_type as permission_type +from ...runtime.stars.registry.register import ( + register_on_llm_response as on_llm_response, +) +from ...runtime.stars.registry.register import ( + register_on_platform_loaded as on_platform_loaded, +) +from ...runtime.stars.registry.register import ( + register_permission_type as permission_type, +) from ...runtime.stars.registry.register import ( register_platform_adapter_type as platform_adapter_type, ) diff --git a/src/astrbot_sdk/cli/__init__.py b/src/astrbot_sdk/cli/__init__.py index ed9c2c342..ef8650159 100644 --- a/src/astrbot_sdk/cli/__init__.py +++ b/src/astrbot_sdk/cli/__init__.py @@ -1,3 +1,3 @@ from .main import cli -__all__ = ['cli'] +__all__ = ["cli"] diff --git a/src/astrbot_sdk/runtime/stars/registry/__init__.py b/src/astrbot_sdk/runtime/stars/registry/__init__.py index 307656104..6639cf560 100644 --- a/src/astrbot_sdk/runtime/stars/registry/__init__.py +++ b/src/astrbot_sdk/runtime/stars/registry/__init__.py @@ -76,6 +76,7 @@ class StarHandlerMetadata: p.pop("event_filters") return p + class StarHandlerRegistry(Generic[T]): def __init__(self): self.star_handlers_map: dict[str, StarHandlerMetadata] = {} diff --git a/src/astrbot_sdk/runtime/stars/registry/register.py b/src/astrbot_sdk/runtime/stars/registry/register.py index 531bee25d..6cd82ba4f 100644 --- a/src/astrbot_sdk/runtime/stars/registry/register.py +++ b/src/astrbot_sdk/runtime/stars/registry/register.py @@ -25,6 +25,7 @@ from ..filter.platform_adapter_type import ( from ..filter.regex import RegexFilter from ..registry import star_handlers_registry, StarHandlerMetadata, EventType + def get_handler_full_name(awaitable: Callable[..., Awaitable[Any]]) -> str: """获取 Handler 的全名""" return f"{awaitable.__module__}:{awaitable.__qualname__}" diff --git a/src/astrbot_sdk/runtime/stars/star_manager.py b/src/astrbot_sdk/runtime/stars/star_manager.py index c3f574883..fdf90a590 100644 --- a/src/astrbot_sdk/runtime/stars/star_manager.py +++ b/src/astrbot_sdk/runtime/stars/star_manager.py @@ -24,18 +24,18 @@ class StarManager: root_dir = Path.cwd() else: root_dir = Path(root_dir).resolve() - + path = root_dir / "plugin.yaml" if not path.exists(): logger.warning("No plugin.yaml found in the current directory.") return [] - + # Add the plugin directory to sys.path so we can import its modules root_dir_str = str(root_dir) if root_dir_str not in sys.path: sys.path.insert(0, root_dir_str) logger.debug(f"Added {root_dir_str} to sys.path") - + with open(path, "r", encoding="utf-8") as f: data = yaml.safe_load(f) diff --git a/src/astrbot_sdk/runtime/supervisor.py b/src/astrbot_sdk/runtime/supervisor.py index 659bb079c..3084946e1 100644 --- a/src/astrbot_sdk/runtime/supervisor.py +++ b/src/astrbot_sdk/runtime/supervisor.py @@ -310,7 +310,9 @@ class WorkerRuntime: result = response.result if not isinstance(result, dict): - raise RuntimeError(f"Invalid handshake payload for plugin {self.plugin.name}") + raise RuntimeError( + f"Invalid handshake payload for plugin {self.plugin.name}" + ) self.raw_handshake = result self.handlers = self._parse_handlers(result) @@ -493,7 +495,9 @@ class SupervisorRuntime: async def _handle_message(self, message: JSONRPCMessage) -> None: if isinstance(message, JSONRPCRequest): if message.method == "handshake": - await self.server.send_message(self._build_handshake_response(message.id)) + await self.server.send_message( + self._build_handshake_response(message.id) + ) return if message.method == "call_handler": await self._route_call_handler(message) @@ -527,7 +531,9 @@ class SupervisorRuntime: JSONRPCErrorResponse( jsonrpc="2.0", id=message.id, - error=JSONRPCErrorData(code=-32602, message=f"Invalid params: {exc}"), + error=JSONRPCErrorData( + code=-32602, message=f"Invalid params: {exc}" + ), ) ) return diff --git a/src/astrbot_sdk/runtime/types.py b/src/astrbot_sdk/runtime/types.py index f8819764d..4ded1b400 100644 --- a/src/astrbot_sdk/runtime/types.py +++ b/src/astrbot_sdk/runtime/types.py @@ -5,6 +5,7 @@ from .rpc.jsonrpc import JSONRPCRequest from typing import Any, Literal, Type from ..api.event.astr_message_event import AstrMessageEventModel + class HandshakeRequest(JSONRPCRequest): class Params(BaseModel): pass @@ -34,33 +35,33 @@ class CallHandlerRequest(JSONRPCRequest): class HandlerStreamStartNotification(JSONRPCRequest): """Notification sent when a handler stream starts.""" - + class Params(BaseModel): id: str | None # The original request ID handler_full_name: str - + method: Literal["handler_stream_start"] = "handler_stream_start" params: Params # type: ignore[assignment] class HandlerStreamUpdateNotification(JSONRPCRequest): """Notification sent when a handler stream has new data.""" - + class Params(BaseModel): id: str | None # The original request ID handler_full_name: str data: Any # The streamed data - + method: Literal["handler_stream_update"] = "handler_stream_update" params: Params # type: ignore[assignment] class HandlerStreamEndNotification(JSONRPCRequest): """Notification sent when a handler stream ends.""" - + class Params(BaseModel): id: str | None # The original request ID handler_full_name: str - + method: Literal["handler_stream_end"] = "handler_stream_end" params: Params # type: ignore[assignment] diff --git a/src/astrbot_sdk/tests/benchmark_8_plugins_resource_usage.py b/src/astrbot_sdk/tests/benchmark_8_plugins_resource_usage.py index 41d6c7c33..df7582e7c 100644 --- a/src/astrbot_sdk/tests/benchmark_8_plugins_resource_usage.py +++ b/src/astrbot_sdk/tests/benchmark_8_plugins_resource_usage.py @@ -255,9 +255,7 @@ async def run_benchmark(plugins_dir: Path, python_executable: str) -> dict[str, handshake_error = f"{exc.__class__.__name__}: {exc}" measured_at = time.perf_counter() - metrics = ( - collect_process_tree_metrics(client_process.pid) if client_process else {} - ) + metrics = collect_process_tree_metrics(client_process.pid) if client_process else {} loaded_plugins = sorted( metadata_item.name for metadata_item in metadata.values() diff --git a/src/astrbot_sdk/tests/test_supervisor.py b/src/astrbot_sdk/tests/test_supervisor.py index 7e9cc082b..27eab71ae 100644 --- a/src/astrbot_sdk/tests/test_supervisor.py +++ b/src/astrbot_sdk/tests/test_supervisor.py @@ -258,9 +258,7 @@ class SupervisorRuntimeTest(unittest.IsolatedAsyncioTestCase): ["plugin_one.main", "plugin_two.main"], ) - handler_full_name = ( - "commands.plugin_two:SampleCommand.handle_plugin_two" - ) + handler_full_name = "commands.plugin_two:SampleCommand.handle_plugin_two" await self.server.handler( CallHandlerRequest( jsonrpc="2.0", @@ -294,7 +292,9 @@ class SupervisorRuntimeTest(unittest.IsolatedAsyncioTestCase): ) ) - self.assertEqual(self.server.sent_messages[-1].result, {"handled_by": "plugin_two"}) + self.assertEqual( + self.server.sent_messages[-1].result, {"handled_by": "plugin_two"} + ) await runtime.stop() async def test_routes_context_response_back_to_matching_worker(self) -> None: diff --git a/tests_v4/test_api_contract.py b/tests_v4/test_api_contract.py index ca6d2223b..db668f4e4 100644 --- a/tests_v4/test_api_contract.py +++ b/tests_v4/test_api_contract.py @@ -4,7 +4,12 @@ import unittest from unittest.mock import patch from astrbot_sdk import MessageEvent, Star, on_command -from astrbot_sdk._legacy_api import CommandComponent, MIGRATION_DOC_URL, LegacyContext, _warned_methods +from astrbot_sdk._legacy_api import ( + CommandComponent, + MIGRATION_DOC_URL, + LegacyContext, + _warned_methods, +) from astrbot_sdk.clients._proxy import CapabilityProxy from astrbot_sdk.clients.memory import MemoryClient diff --git a/tests_v4/test_legacy_adapter.py b/tests_v4/test_legacy_adapter.py index 2e8875418..dfbcaee53 100644 --- a/tests_v4/test_legacy_adapter.py +++ b/tests_v4/test_legacy_adapter.py @@ -8,11 +8,18 @@ from astrbot_sdk.protocol.legacy_adapter import ( LEGACY_HANDSHAKE_METADATA_KEY, LegacyAdapter, ) -from astrbot_sdk.protocol.messages import EventMessage, InitializeMessage, PeerInfo, ResultMessage +from astrbot_sdk.protocol.messages import ( + EventMessage, + InitializeMessage, + PeerInfo, + ResultMessage, +) class LegacyAdapterTest(unittest.TestCase): - def test_call_handler_roundtrip_preserves_handler_name_in_stream_notifications(self) -> None: + def test_call_handler_roundtrip_preserves_handler_name_in_stream_notifications( + self, + ) -> None: adapter = LegacyAdapter() invoke = adapter.legacy_to_v4( @@ -45,7 +52,9 @@ class LegacyAdapterTest(unittest.TestCase): ) self.assertEqual(started["method"], "handler_stream_start") - self.assertEqual(delta["params"]["handler_full_name"], "commands.demo:MyPlugin.handle") + self.assertEqual( + delta["params"]["handler_full_name"], "commands.demo:MyPlugin.handle" + ) self.assertEqual(delta["params"]["data"], {"text": "hi"}) self.assertEqual(completed["method"], "handler_stream_end") self.assertEqual(response["result"], {"handled_by": "demo"}) @@ -139,7 +148,9 @@ class LegacyAdapterTest(unittest.TestCase): "reserved": False, "activated": True, "config": None, - "star_handler_full_names": ["commands.plugin_one:SampleCommand.handle_plugin_one"], + "star_handler_full_names": [ + "commands.plugin_one:SampleCommand.handle_plugin_one" + ], "display_name": "plugin_one", "logo_path": None, "handlers": [ @@ -165,7 +176,10 @@ class LegacyAdapterTest(unittest.TestCase): self.assertIsInstance(initialize, InitializeMessage) self.assertEqual(initialize.peer.name, "plugin_one") - self.assertEqual(initialize.handlers[0].id, "commands.plugin_one:SampleCommand.handle_plugin_one") + self.assertEqual( + initialize.handlers[0].id, + "commands.plugin_one:SampleCommand.handle_plugin_one", + ) self.assertEqual(initialize.handlers[0].priority, 7) self.assertEqual( initialize.metadata[LEGACY_HANDSHAKE_METADATA_KEY], diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py index 465b509b1..a882ff10f 100644 --- a/tests_v4/test_peer.py +++ b/tests_v4/test_peer.py @@ -6,10 +6,18 @@ 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.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 astrbot_sdk.runtime.transport import ( + WebSocketClientTransport, + WebSocketServerTransport, +) from tests_v4.helpers import make_transport_pair @@ -25,11 +33,14 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): 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={}, - )) + 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( @@ -72,7 +83,9 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): 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()) + await self.left.send( + EventMessage(id="req-1", phase="started").model_dump_json() + ) with self.assertRaises(AstrBotError) as raised: await task @@ -88,8 +101,12 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): 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()) + 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: @@ -103,7 +120,11 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): name="slow.stream", description="slow stream", input_schema={"type": "object", "properties": {}, "required": []}, - output_schema={"type": "object", "properties": {"count": {"type": "number"}}, "required": ["count"]}, + output_schema={ + "type": "object", + "properties": {"count": {"type": "number"}}, + "required": ["count"], + }, supports_stream=True, cancelable=True, ) @@ -161,11 +182,14 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): 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={}, - )) + 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( @@ -238,7 +262,9 @@ class CapabilityRouterContractTest(unittest.TestCase): ) self.assertIn(name, str(raised.exception)) - def test_reserved_capability_namespaces_are_rejected_for_exposed_registrations(self) -> None: + 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: @@ -250,7 +276,9 @@ class CapabilityRouterContractTest(unittest.TestCase): ) self.assertIn(name, str(raised.exception)) - def test_reserved_capability_namespaces_remain_available_for_hidden_internal_registrations(self) -> None: + def test_reserved_capability_namespaces_remain_available_for_hidden_internal_registrations( + self, + ) -> None: router = CapabilityRouter() router.register( CapabilityDescriptor( diff --git a/tests_v4/test_runtime.py b/tests_v4/test_runtime.py index 1ac536d62..eadb01281 100644 --- a/tests_v4/test_runtime.py +++ b/tests_v4/test_runtime.py @@ -68,11 +68,14 @@ class RuntimeIntegrationTest(unittest.IsolatedAsyncioTestCase): peer_info=PeerInfo(name="outer-core", role="core", version="v4"), ) self.core.set_initialize_handler( - lambda _message: asyncio.sleep(0, result=InitializeOutput( - peer=PeerInfo(name="outer-core", role="core", version="v4"), - capabilities=[], - metadata={}, - )) + lambda _message: asyncio.sleep( + 0, + result=InitializeOutput( + peer=PeerInfo(name="outer-core", role="core", version="v4"), + capabilities=[], + metadata={}, + ), + ) ) await self.core.start() @@ -108,7 +111,9 @@ class RuntimeIntegrationTest(unittest.IsolatedAsyncioTestCase): }, request_id="call-v4", ) - texts = [item.get("text") for item in runtime.capability_router.sent_messages] + texts = [ + item.get("text") for item in runtime.capability_router.sent_messages + ] self.assertEqual(texts, ["Echo: hello", "Echo: stream"]) finally: await runtime.stop() @@ -142,7 +147,9 @@ class RuntimeIntegrationTest(unittest.IsolatedAsyncioTestCase): }, request_id="call-compat", ) - texts = [item.get("text") for item in runtime.capability_router.sent_messages] + texts = [ + item.get("text") for item in runtime.capability_router.sent_messages + ] self.assertEqual(len(texts), 4) self.assertIn("Created conversation ID", texts[0]) finally: diff --git a/tests_v4/test_script_migrations.py b/tests_v4/test_script_migrations.py index 91f333324..df2f840b1 100644 --- a/tests_v4/test_script_migrations.py +++ b/tests_v4/test_script_migrations.py @@ -13,7 +13,10 @@ from astrbot_sdk.protocol.messages import InitializeOutput, PeerInfo from astrbot_sdk.runtime.bootstrap import PluginWorkerRuntime, SupervisorRuntime from astrbot_sdk.runtime.capability_router import CapabilityRouter from astrbot_sdk.runtime.peer import Peer -from astrbot_sdk.runtime.transport import WebSocketClientTransport, WebSocketServerTransport +from astrbot_sdk.runtime.transport import ( + WebSocketClientTransport, + WebSocketServerTransport, +) from tests_v4.helpers import FakeEnvManager, make_transport_pair @@ -108,13 +111,19 @@ def write_benchmark_plugin(plugins_dir: Path, index: int) -> None: class StartClientMigrationTest(unittest.IsolatedAsyncioTestCase): - async def test_websocket_plugin_worker_supports_handshake_and_handler_invoke(self) -> None: + async def test_websocket_plugin_worker_supports_handshake_and_handler_invoke( + self, + ) -> None: with tempfile.TemporaryDirectory() as temp_dir: plugin_root = Path(temp_dir) / "websocket_plugin" write_websocket_plugin(plugin_root) - server_transport = WebSocketServerTransport(host="127.0.0.1", port=0, path="/ws") - runtime = PluginWorkerRuntime(plugin_dir=plugin_root, transport=server_transport) + server_transport = WebSocketServerTransport( + host="127.0.0.1", port=0, path="/ws" + ) + runtime = PluginWorkerRuntime( + plugin_dir=plugin_root, transport=server_transport + ) runtime_task = asyncio.create_task(runtime.start()) try: @@ -126,13 +135,17 @@ class StartClientMigrationTest(unittest.IsolatedAsyncioTestCase): core_router = CapabilityRouter() core_peer = Peer( transport=WebSocketClientTransport(url=server_transport.url), - peer_info=PeerInfo(name="websocket-core", role="core", version="v4"), + peer_info=PeerInfo( + name="websocket-core", role="core", version="v4" + ), ) core_peer.set_initialize_handler( lambda _message: asyncio.sleep( 0, result=InitializeOutput( - peer=PeerInfo(name="websocket-core", role="core", version="v4"), + peer=PeerInfo( + name="websocket-core", role="core", version="v4" + ), capabilities=core_router.descriptors(), metadata={}, ), @@ -152,8 +165,12 @@ class StartClientMigrationTest(unittest.IsolatedAsyncioTestCase): await asyncio.wait_for(runtime_task, timeout=5) await core_peer.wait_until_remote_initialized() handler_id = core_peer.remote_handlers[0].id - self.assertEqual(core_peer.remote_metadata["plugin_id"], "websocket_plugin") - self.assertEqual(core_peer.remote_handlers[0].trigger.command, "hello") + self.assertEqual( + core_peer.remote_metadata["plugin_id"], "websocket_plugin" + ) + self.assertEqual( + core_peer.remote_handlers[0].trigger.command, "hello" + ) await core_peer.invoke( "handler.invoke", @@ -205,7 +222,9 @@ class BenchmarkMigrationTest(unittest.IsolatedAsyncioTestCase): async def asyncTearDown(self) -> None: await self.core.stop() - async def test_benchmark_style_runtime_report_covers_multi_plugin_workers(self) -> None: + async def test_benchmark_style_runtime_report_covers_multi_plugin_workers( + self, + ) -> None: plugin_count = 8 with tempfile.TemporaryDirectory() as temp_dir: plugins_dir = Path(temp_dir) / "plugins" @@ -226,15 +245,24 @@ class BenchmarkMigrationTest(unittest.IsolatedAsyncioTestCase): worker_pids = sorted( process.pid for session in runtime.worker_sessions.values() - if (process := getattr(getattr(session.peer, "transport", None), "_process", None)) is not None + if ( + process := getattr( + getattr(session.peer, "transport", None), "_process", None + ) + ) + is not None and process.returncode is None ) report = { "plugin_count": plugin_count, "loaded_plugin_count": len(runtime.loaded_plugins), "loaded_plugins": sorted(runtime.loaded_plugins), - "aggregated_handler_ids": list(self.core.remote_metadata["aggregated_handler_ids"]), - "startup_total_duration_ms": round((measured_at - started_at) * 1000, 2), + "aggregated_handler_ids": list( + self.core.remote_metadata["aggregated_handler_ids"] + ), + "startup_total_duration_ms": round( + (measured_at - started_at) * 1000, 2 + ), "worker_pids": worker_pids, } @@ -265,7 +293,10 @@ class BenchmarkMigrationTest(unittest.IsolatedAsyncioTestCase): [f"plugin_{index:03d}" for index in range(plugin_count)], ) self.assertGreaterEqual(report["startup_total_duration_ms"], 0) - self.assertIn("plugin_002:bench_002", runtime.capability_router.sent_messages[-1]["text"]) + self.assertIn( + "plugin_002:bench_002", + runtime.capability_router.sent_messages[-1]["text"], + ) finally: await runtime.stop() diff --git a/tests_v4/test_supervisor_migration.py b/tests_v4/test_supervisor_migration.py index ab1d1ed83..c815138b8 100644 --- a/tests_v4/test_supervisor_migration.py +++ b/tests_v4/test_supervisor_migration.py @@ -44,7 +44,7 @@ def write_plugin( manifest_lines.extend( [ "runtime:", - f" python: \"{python_version}\"", + f' python: "{python_version}"', ] ) manifest_lines.extend( @@ -56,7 +56,9 @@ def write_plugin( " description: hello", ] ) - (plugin_dir / "plugin.yaml").write_text("\n".join(manifest_lines) + "\n", encoding="utf-8") + (plugin_dir / "plugin.yaml").write_text( + "\n".join(manifest_lines) + "\n", encoding="utf-8" + ) if include_requirements: (plugin_dir / "requirements.txt").write_text("", encoding="utf-8") @@ -125,7 +127,9 @@ class SupervisorMigrationTest(unittest.IsolatedAsyncioTestCase): async def asyncTearDown(self) -> None: await self.core.stop() - async def test_supervisor_aggregates_handlers_and_routes_target_plugin(self) -> None: + async def test_supervisor_aggregates_handlers_and_routes_target_plugin( + self, + ) -> None: with tempfile.TemporaryDirectory() as temp_dir: plugins_dir = Path(temp_dir) / "plugins" write_plugin( @@ -150,8 +154,12 @@ class SupervisorMigrationTest(unittest.IsolatedAsyncioTestCase): await runtime.start() await self.core.wait_until_remote_initialized() - self.assertEqual(sorted(runtime.loaded_plugins), ["plugin_one", "plugin_two"]) - self.assertEqual(self.core.remote_metadata["plugins"], ["plugin_one", "plugin_two"]) + self.assertEqual( + sorted(runtime.loaded_plugins), ["plugin_one", "plugin_two"] + ) + self.assertEqual( + self.core.remote_metadata["plugins"], ["plugin_one", "plugin_two"] + ) self.assertEqual(len(self.core.remote_handlers), 2) handler_id = next( @@ -173,7 +181,9 @@ class SupervisorMigrationTest(unittest.IsolatedAsyncioTestCase): request_id="call-route", ) - texts = [item.get("text") for item in runtime.capability_router.sent_messages] + texts = [ + item.get("text") for item in runtime.capability_router.sent_messages + ] self.assertEqual(texts, ["plugin_two handled"]) finally: await runtime.stop()