From 1e384d0cd364f40d1225c2465fb99e8d84f1d78b Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Wed, 12 Nov 2025 22:05:04 +0800 Subject: [PATCH] feat: enhance RPC handling with streaming support and notifications --- src/astrbot_sdk/runtime/rpc/request_helper.py | 205 ++++++++++- src/astrbot_sdk/runtime/star_runner.py | 82 +++-- src/astrbot_sdk/runtime/stars/new.py | 325 +++--------------- src/astrbot_sdk/runtime/stars/virtual.py | 39 ++- src/astrbot_sdk/runtime/types.py | 34 ++ src/astrbot_sdk/tests/start_client.py | 4 +- 6 files changed, 339 insertions(+), 350 deletions(-) diff --git a/src/astrbot_sdk/runtime/rpc/request_helper.py b/src/astrbot_sdk/runtime/rpc/request_helper.py index ee39918bb..9e34c5720 100644 --- a/src/astrbot_sdk/runtime/rpc/request_helper.py +++ b/src/astrbot_sdk/runtime/rpc/request_helper.py @@ -1,38 +1,219 @@ import asyncio +from typing import Any +from loguru import logger from .jsonrpc import ( JSONRPCMessage, + JSONRPCRequest, JSONRPCSuccessResponse, JSONRPCErrorResponse, ) from .transport import JSONRPCTransport +from ..types import ( + HandlerStreamStartNotification, + HandlerStreamUpdateNotification, + HandlerStreamEndNotification, +) class RPCRequestHelper: - """Manages RPC communication state and pending requests.""" + """Manages RPC communication state and pending requests. + + Supports both single-response and streaming (multi-response) RPC patterns: + - Single response: Uses asyncio.Future + - Streaming: Uses asyncio.Queue for multiple responses + """ def __init__(self): self._request_id_counter = 0 - self.pending_requests: dict[str, asyncio.Future[JSONRPCMessage]] = {} + self.pending_requests: dict[ + str, asyncio.Future[JSONRPCMessage] | asyncio.Queue[Any] + ] = {} def _generate_request_id(self) -> str: + """Generate a unique request ID.""" self._request_id_counter += 1 return str(self._request_id_counter) async def call_rpc( self, transport_impl: JSONRPCTransport, message: JSONRPCMessage ) -> JSONRPCMessage | None: - """Send RPC request and wait for response.""" - if message.id is not None: - self.pending_requests[message.id] = asyncio.get_event_loop().create_future() + """Send RPC request and wait for a single response. + + Args: + transport_impl: The transport to send the message through + message: The JSON-RPC request message + + Returns: + The JSON-RPC response message, or None if no response expected + """ + if message.id is None: + await transport_impl.send_message(message) + return None + + future: asyncio.Future[JSONRPCMessage] = ( + asyncio.get_event_loop().create_future() + ) + self.pending_requests[message.id] = future await transport_impl.send_message(message) - if message.id is not None: - return await self.pending_requests[message.id] + result = await future + return result + + async def call_rpc_streaming( + self, transport_impl: JSONRPCTransport, message: JSONRPCMessage + ) -> asyncio.Queue[Any]: + """Send RPC request and expect multiple streaming responses. + + The responses will be delivered via notifications with methods: + - handler_stream_start: Stream started + - handler_stream_update: New data available + - handler_stream_end: Stream completed + + Args: + transport_impl: The transport to send the message through + message: The JSON-RPC request message + + Returns: + An asyncio.Queue that will receive streamed results + """ + if message.id is None: + raise ValueError("Streaming RPC calls require a request ID") + + queue: asyncio.Queue[Any] = asyncio.Queue() + self.pending_requests[message.id] = queue + + await transport_impl.send_message(message) + return queue def resolve_pending_request( self, message: JSONRPCSuccessResponse | JSONRPCErrorResponse ): - """Resolve a pending request with the response.""" - if message.id in self.pending_requests: - future = self.pending_requests.pop(message.id) - if not future.done(): - future.set_result(message) + """Resolve a pending request with a response. + + For single-response requests (Future), sets the result/exception. + For streaming requests (Queue), logs completion/error but queue is managed separately. + + Args: + message: The JSON-RPC response message + """ + if message.id not in self.pending_requests: + logger.warning(f"Received response for unknown request ID: {message.id}") + return + + pending = self.pending_requests[message.id] + + if isinstance(pending, asyncio.Future): + # Single response mode + self.pending_requests.pop(message.id) + if not pending.done(): + if isinstance(message, JSONRPCSuccessResponse): + pending.set_result(message) + else: + pending.set_exception( + RuntimeError( + f"RPC Error {message.error.code}: {message.error.message}" + ) + ) + elif isinstance(pending, asyncio.Queue): + # Streaming mode - final response received + if isinstance(message, JSONRPCSuccessResponse): + logger.debug(f"Streaming request {message.id} completed successfully") + else: + logger.error( + f"Streaming request {message.id} failed: {message.error.message}" + ) + # Put error marker in queue + asyncio.create_task( + pending.put({"_error": True, "message": message.error.message}) + ) + + async def handle_stream_notification(self, notification: JSONRPCRequest) -> None: + """Handle incoming streaming notifications. + + Processes handler_stream_start/update/end notifications and updates + the corresponding queue. + + Args: + notification: The streaming notification message + + Raises: + ValueError: If the notification method is not a valid stream notification + """ + # Validate notification method + if notification.method not in [ + "handler_stream_start", + "handler_stream_update", + "handler_stream_end", + ]: + raise ValueError( + f"Invalid stream notification method: {notification.method}" + ) + + # Extract common parameters + params = notification.params + request_id = params.get("id") + + if not request_id or request_id not in self.pending_requests: + logger.warning( + f"Received stream notification for unknown request ID: {request_id}" + ) + return + + pending = self.pending_requests.get(request_id) + if not isinstance(pending, asyncio.Queue): + logger.warning(f"Request {request_id} is not a streaming request") + return + + if notification.method == "handler_stream_start": + try: + typed_notification = HandlerStreamStartNotification.model_validate( + notification.model_dump() + ) + logger.debug( + f"Stream started for handler {typed_notification.params.handler_full_name}" + ) + except Exception as e: + logger.error(f"Invalid handler_stream_start notification: {e}") + # Optionally put a start marker in the queue if needed + # await pending.put({"_stream_start": True}) + + elif notification.method == "handler_stream_update": + try: + typed_notification = HandlerStreamUpdateNotification.model_validate( + notification.model_dump() + ) + # Put the streamed data into the queue + data = typed_notification.params.data + logger.debug(f"Stream update for request {request_id}: {data}") + if data is not None: + await pending.put(data) + except Exception as e: + logger.error(f"Invalid handler_stream_update notification: {e}") + + elif notification.method == "handler_stream_end": + try: + typed_notification = HandlerStreamEndNotification.model_validate( + notification.model_dump() + ) + logger.debug( + f"Stream ended for handler {typed_notification.params.handler_full_name}" + ) + # Put a sentinel value to indicate stream end + await pending.put({"_stream_end": True}) + # Clean up the pending request after a short delay + asyncio.create_task(self._cleanup_stream_request(request_id)) + except Exception as e: + logger.error(f"Invalid handler_stream_end notification: {e}") + + async def _cleanup_stream_request( + self, request_id: str, delay: float = 1.0 + ) -> None: + """Clean up a streaming request after a delay. + + Args: + request_id: The request ID to clean up + delay: Delay before cleanup in seconds + """ + await asyncio.sleep(delay) + if request_id in self.pending_requests: + self.pending_requests.pop(request_id) + logger.debug(f"Cleaned up streaming request {request_id}") diff --git a/src/astrbot_sdk/runtime/star_runner.py b/src/astrbot_sdk/runtime/star_runner.py index f9932bb87..42cd48b23 100644 --- a/src/astrbot_sdk/runtime/star_runner.py +++ b/src/astrbot_sdk/runtime/star_runner.py @@ -11,7 +11,12 @@ from .rpc.jsonrpc import ( JSONRPCErrorData, ) from .rpc.request_helper import RPCRequestHelper -from .types import CallHandlerRequest +from .types import ( + CallHandlerRequest, + HandlerStreamStartNotification, + HandlerStreamUpdateNotification, + HandlerStreamEndNotification, +) class HandshakeHandler: @@ -76,51 +81,62 @@ class HandlerExecutor: ready_to_call, ): """Execute handler and stream results.""" - await self._send_stream_notification( - server, "handler_stream_start", request_id, handler_name + # Send start notification + await server.send_message( + HandlerStreamStartNotification( + jsonrpc="2.0", + method="handler_stream_start", + params=HandlerStreamStartNotification.Params( + id=request_id, + handler_full_name=handler_name, + ), + ) ) try: if inspect.iscoroutine(ready_to_call): result = await ready_to_call - await self._send_stream_notification( - server, "handler_stream_update", request_id, handler_name, result + # Send update notification + await server.send_message( + HandlerStreamUpdateNotification( + jsonrpc="2.0", + method="handler_stream_update", + params=HandlerStreamUpdateNotification.Params( + id=request_id, + handler_full_name=handler_name, + data=result, + ), + ) ) elif inspect.isasyncgen(ready_to_call): async for ret in ready_to_call: - await self._send_stream_notification( - server, "handler_stream_update", request_id, handler_name, ret + # Send update notification for each item + await server.send_message( + HandlerStreamUpdateNotification( + jsonrpc="2.0", + method="handler_stream_update", + params=HandlerStreamUpdateNotification.Params( + id=request_id, + handler_full_name=handler_name, + data=ret, + ), + ) ) except Exception as e: logger.error(f"Error during handler {handler_name}: {e}") finally: - await self._send_stream_notification( - server, "handler_stream_end", request_id, handler_name + # Send end notification + await server.send_message( + HandlerStreamEndNotification( + jsonrpc="2.0", + method="handler_stream_end", + params=HandlerStreamEndNotification.Params( + id=request_id, + handler_full_name=handler_name, + ), + ) ) - async def _send_stream_notification( - self, - server: JSONRPCServer, - method: str, - request_id: str | None, - handler_name: str, - data=None, - ): - """Send a stream notification.""" - params = { - "id": request_id, - "handler_full_name": handler_name, - } - if data is not None: - params["data"] = data - - notification = JSONRPCRequest( - jsonrpc="2.0", - method=method, - params=params, - ) - await server.send_message(notification) - async def _send_error( self, server: JSONRPCServer, request_id: str | None, code: int, message: str ): @@ -173,6 +189,8 @@ class StarRunner: await self.server.send_message(response) elif message.method == "call_handler": await self.handler_executor.execute(message, self.server) + else: + logger.warning(f"Unknown method from client: {message.method}") elif isinstance(message, (JSONRPCSuccessResponse, JSONRPCErrorResponse)): self.rpc_request_helper.resolve_pending_request(message) diff --git a/src/astrbot_sdk/runtime/stars/new.py b/src/astrbot_sdk/runtime/stars/new.py index c32ecd656..3e5c11f65 100644 --- a/src/astrbot_sdk/runtime/stars/new.py +++ b/src/astrbot_sdk/runtime/stars/new.py @@ -2,7 +2,8 @@ from __future__ import annotations import asyncio import os -from typing import Any +import inspect +from typing import Any, AsyncGenerator from loguru import logger @@ -20,6 +21,7 @@ from ..types import CallHandlerRequest, HandshakeRequest from ..rpc.client import JSONRPCClient from ..rpc.client.stdio import StdioClient from ..rpc.client.websocket import WebSocketClient +from ..rpc.request_helper import RPCRequestHelper from .virtual import VirtualStar @@ -33,7 +35,7 @@ class NewStar(VirtualStar): def __init__( self, client: JSONRPCClient, - context: Any = None, + context: Any, ) -> None: """Initialize a NewStar instance. @@ -41,32 +43,19 @@ class NewStar(VirtualStar): client: JSON-RPC client for communication context: Context instance for managing managers and their functions """ - # Import here to avoid circular dependency - from ..api.context import Context - - # Initialize context - if context is None: - context = Context.default_context() - super().__init__(context) self._client = client self._metadata: dict[str, StarMetadata] = {} self._handlers: list[StarHandlerMetadata] = [] - self._request_id_counter = 0 - self._pending_requests: dict[ - str, asyncio.Future[dict] | asyncio.Queue[dict] - ] = {} self._active = False + # Use RPCRequestHelper for managing requests + self._rpc_helper = RPCRequestHelper() + # Set up message handler self._client.set_message_handler(self._handle_message) - def _generate_request_id(self) -> str: - """Generate a unique request ID.""" - self._request_id_counter += 1 - return f"req-{self._request_id_counter}" - async def _handle_message(self, message: JSONRPCMessage) -> None: """Handle incoming JSON-RPC messages from the plugin. @@ -77,41 +66,8 @@ class NewStar(VirtualStar): message, JSONRPCErrorResponse, ): - # This is a response to one of our requests - request_id = message.id - if request_id and request_id in self._pending_requests: - pending = self._pending_requests[request_id] - - # Check if it's a Future or Queue - if isinstance(pending, asyncio.Future): - self._pending_requests.pop(request_id) - if isinstance(message, JSONRPCSuccessResponse): - if not pending.done(): - pending.set_result(message.result) - else: - if not pending.done(): - pending.set_exception( - RuntimeError( - f"RPC Error {message.error.code}: {message.error.message}", - ), - ) - elif isinstance(pending, asyncio.Queue): - if isinstance(message, JSONRPCSuccessResponse): - logger.debug( - f"Streaming handler {request_id} completed successfully" - ) - else: - logger.error( - f"Streaming handler {request_id} failed: {message.error.message}" - ) - # Put error marker in queue - await pending.put( - {"_error": True, "message": message.error.message} - ) - else: - logger.warning( - f"Received response for unknown request ID: {request_id}" - ) + # Delegate to RPCRequestHelper + self._rpc_helper.resolve_pending_request(message) elif isinstance(message, JSONRPCRequest): # Handle notifications from plugin (streaming events or method calls) @@ -120,7 +76,7 @@ class NewStar(VirtualStar): "handler_stream_update", "handler_stream_end", ]: - await self._handle_stream_notification(message) + await self._rpc_helper.handle_stream_notification(message) else: # Plugin is calling a method on the core asyncio.create_task(self._handle_plugin_request(message)) @@ -185,115 +141,6 @@ class NewStar(VirtualStar): ) await self._client.send_message(error_response) - async def _handle_stream_notification(self, notification: JSONRPCRequest) -> None: - """Handle streaming notifications from the plugin. - - Args: - notification: The streaming notification (handler_stream_start/update/end) - """ - params = notification.params - request_id = params.get("id") - - if not request_id or request_id not in self._pending_requests: - logger.warning( - f"Received stream notification for unknown request ID: {request_id}" - ) - return - - pending = self._pending_requests.get(request_id) - if not isinstance(pending, asyncio.Queue): - logger.warning(f"Request {request_id} is not a streaming request") - return - - if notification.method == "handler_stream_start": - logger.debug( - f"Stream started for handler {params.get('handler_full_name')}" - ) - # Optionally put a start marker in the queue - # await pending.put({"_stream_start": True}) - - elif notification.method == "handler_stream_update": - # Put the streamed data into the queue - data = params.get("data") - logger.debug(f"Stream update for request {request_id}: {data}") - if data is not None: - await pending.put(data) - - elif notification.method == "handler_stream_end": - # Mark the end of the stream - logger.debug(f"Stream ended for handler {params.get('handler_full_name')}") - # Put a sentinel value to indicate stream end - await pending.put({"_stream_end": True}) - # Clean up the pending request after a short delay to allow queue to be processed - asyncio.create_task(self._cleanup_stream_request(request_id)) - - async def _cleanup_stream_request( - self, request_id: str, delay: float = 1.0 - ) -> None: - """Clean up a streaming request after a delay. - - Args: - request_id: The request ID to clean up - delay: Delay before cleanup in seconds - """ - await asyncio.sleep(delay) - if request_id in self._pending_requests: - self._pending_requests.pop(request_id) - logger.debug(f"Cleaned up streaming request {request_id}") - - async def _call_rpc(self, request: JSONRPCRequest) -> dict: - """Call a JSON-RPC method on the plugin and wait for response. - - Args: - request: The JSON-RPC request to send - - Returns: - The result from the plugin - - Raises: - RuntimeError: If the RPC call fails - """ - # Create a future to wait for the response - future: asyncio.Future[dict] = asyncio.Future() - - if request.id is not None: - self._pending_requests[request.id] = future - - try: - await self._client.send_message(request) - # Wait for response with timeout - result = await asyncio.wait_for(future, timeout=30.0) - return result - except asyncio.TimeoutError: - if request.id is not None: - self._pending_requests.pop(request.id, None) - raise RuntimeError(f"RPC call to {request.method} timed out") - - async def _call_rpc_streaming( - self, - request: JSONRPCRequest, - ) -> asyncio.Queue[dict]: - """Call a JSON-RPC method on the plugin that returns a stream of results. - - Args: - request: The JSON-RPC request to send - Returns: - An asyncio.Queue that will receive streamed results - """ - # Create a queue to receive streamed results - queue: asyncio.Queue[dict] = asyncio.Queue() - - if request.id is not None: - self._pending_requests[request.id] = queue - - try: - await self._client.send_message(request) - return queue - except Exception as e: - if request.id is not None: - self._pending_requests.pop(request.id, None) - raise RuntimeError(f"RPC streaming call to {request.method} failed: {e}") - async def initialize(self) -> None: """Start the plugin process and establish connection.""" # Start the client (which may start a subprocess for STDIO) @@ -308,13 +155,19 @@ class NewStar(VirtualStar): """ logger.info("Performing handshake with plugin...") - result = await self._call_rpc( + response = await self._rpc_helper.call_rpc( + self._client, HandshakeRequest( - jsonrpc="2.0", id=self._generate_request_id(), method="handshake" - ) + jsonrpc="2.0", + id=self._rpc_helper._generate_request_id(), + method="handshake", + ), ) - print(result, result.__class__) + if not isinstance(response, JSONRPCSuccessResponse): + raise RuntimeError("Handshake failed: Invalid response from plugin") + + result = response.result if isinstance(result, dict): # Parse metadata @@ -365,7 +218,7 @@ class NewStar(VirtualStar): Returns either a direct result or an async generator for streaming handlers. """ - request_id = self._generate_request_id() + request_id = self._rpc_helper._generate_request_id() request = CallHandlerRequest( jsonrpc="2.0", id=request_id, @@ -376,125 +229,22 @@ class NewStar(VirtualStar): args=kwargs, ), ) - - # Create a queue for potential streaming response - queue: asyncio.Queue[dict] = asyncio.Queue() - self._pending_requests[request_id] = queue + queue = await self._rpc_helper.call_rpc_streaming(self._client, request) try: - # Send the request - await self._client.send_message(request) + while True: + item = await asyncio.wait_for(queue.get(), timeout=30.0) + if isinstance(item, dict) and item.get("_stream_end"): + break + if isinstance(item, dict) and item.get("_error"): + raise RuntimeError(item.get("message", "Unknown error")) + yield self._deserialize_result(item) - # Wait for the first response or stream notification - try: - # Set a timeout for the first response - first_response = await asyncio.wait_for(queue.get(), timeout=30.0) - - # Check what type of response we got - if isinstance(first_response, dict): - # Check for stream end (empty stream case) - if first_response.get("_stream_end"): - # Empty stream, return None - self._pending_requests.pop(request_id, None) - return None - - # Check for error - if first_response.get("_error"): - self._pending_requests.pop(request_id, None) - raise RuntimeError( - first_response.get("message", "Unknown error") - ) - - # Check if this is streaming data or a final result - # We peek at the queue to see if more data is coming - # If the queue is empty after a short wait, it's a final result - try: - # Try to get another item with a very short timeout - second_response = await asyncio.wait_for( - queue.get(), timeout=0.1 - ) - # We got a second item, so this is streaming - # Create and return the generator - return self._create_stream_generator( - queue, first_response, second_response - ) - except asyncio.TimeoutError: - # No second item, this might be a final result - # But we should check if stream_end arrives shortly - try: - stream_end = await asyncio.wait_for( - queue.get(), timeout=0.5 - ) - if isinstance(stream_end, dict) and stream_end.get( - "_stream_end" - ): - # This was a single-item stream - self._pending_requests.pop(request_id, None) - return self._deserialize_result(first_response) - else: - # More data arrived, it's streaming - return self._create_stream_generator( - queue, first_response, stream_end - ) - except asyncio.TimeoutError: - # Truly a final result (non-streaming) - self._pending_requests.pop(request_id, None) - return self._deserialize_result(first_response) - else: - # Unexpected response type - self._pending_requests.pop(request_id, None) - return self._deserialize_result(first_response) - - except asyncio.TimeoutError: - # Timeout waiting for response - self._pending_requests.pop(request_id, None) - raise RuntimeError(f"RPC call to {handler_full_name} timed out") - - except Exception: - # Clean up on error - self._pending_requests.pop(request_id, None) - raise + except asyncio.TimeoutError: + raise RuntimeError(f"RPC call to {handler_full_name} timed out") return handler_proxy - async def _create_stream_generator( - self, queue: asyncio.Queue[dict], *initial_items: dict - ): - """Create an async generator that yields items from the stream queue. - - Args: - queue: The queue containing stream items - initial_items: Initial items that were already retrieved from the queue - - Yields: - Items from the stream - """ - # Yield any initial items - for item in initial_items: - if not (isinstance(item, dict) and item.get("_stream_end")): - yield self._deserialize_result(item) - - # Continue yielding items from the queue - while True: - try: - item = await queue.get() - - # Check for end marker - if isinstance(item, dict) and item.get("_stream_end"): - break - - # Check for error marker - if isinstance(item, dict) and item.get("_error"): - raise RuntimeError(item.get("message", "Stream error")) - - # Yield the item - yield self._deserialize_result(item) - - except asyncio.CancelledError: - # Generator was cancelled, stop iteration - logger.debug("Stream generator cancelled") - break - def _deserialize_result(self, result: Any) -> Any: """Deserialize result from JSON-RPC response. @@ -505,7 +255,7 @@ class NewStar(VirtualStar): Deserialized result object """ # For now, return as-is - # In practice, you might want to reconstruct MessageEventResult etc. + # In practice, we might want to reconstruct MessageEventResult etc. return result def get_triggered_handlers( @@ -536,7 +286,7 @@ class NewStar(VirtualStar): event: AstrMessageEvent, *args, **kwargs, - ) -> None: + ) -> AsyncGenerator[Any, None]: """Call a specific handler in the plugin. Args: @@ -546,13 +296,16 @@ class NewStar(VirtualStar): **kwargs: Additional keyword arguments Returns: - Result from the handler + An async generator yielding results from the handler """ logger.debug(f"Calling handler: {handler.handler_name}") # Call the handler proxy - result = await handler.handler(event, *args, **kwargs) # type: ignore - return result + assert inspect.isasyncgenfunction(handler.handler), ( + "Handler proxy must be an async generator function" + ) + async for result in handler.handler(event, **kwargs): + yield result async def stop(self) -> None: """Stop the NewStar and cleanup resources.""" diff --git a/src/astrbot_sdk/runtime/stars/virtual.py b/src/astrbot_sdk/runtime/stars/virtual.py index acdb90959..a8e5e0192 100644 --- a/src/astrbot_sdk/runtime/stars/virtual.py +++ b/src/astrbot_sdk/runtime/stars/virtual.py @@ -9,23 +9,24 @@ from ...api.star.context import Context class VirtualStar(ABC): """Abstract base class for virtual plugin implementations. - + VirtualStar defines the interface for plugins that can run in isolated runtime environments (separate processes). It handles the complete lifecycle of a plugin from initialization to shutdown. """ + def __init__(self, context: Context) -> None: self._context = context @abstractmethod async def initialize(self) -> None: """Establish connection and initialize the plugin. - + This method should: - Start the plugin process (if applicable) - Establish communication channels - Wait for the plugin to be ready - + Raises: RuntimeError: If initialization fails """ @@ -34,15 +35,15 @@ class VirtualStar(ABC): @abstractmethod async def handshake(self) -> StarMetadata: """Perform handshake to retrieve plugin metadata. - + This method should: - Request plugin metadata from the plugin - Cache handler information locally - Validate the plugin's compatibility - + Returns: StarMetadata: Complete plugin metadata including handlers - + Raises: RuntimeError: If handshake fails or times out """ @@ -51,12 +52,12 @@ class VirtualStar(ABC): # @abstractmethod # async def turn_on(self) -> None: # """Attach and prepare resources. Only call when the plugin is not active. - + # This method should: # - Activate the plugin # - Initialize any runtime resources # - Prepare the plugin to handle events - + # Raises: # RuntimeError: If activation fails # """ @@ -65,12 +66,12 @@ class VirtualStar(ABC): # @abstractmethod # async def turn_off(self) -> None: # """Detach and clean up resources. Make the plugin inactive. - + # This method should: # - Deactivate the plugin # - Release runtime resources # - Keep the process running but idle - + # Raises: # RuntimeError: If deactivation fails # """ @@ -82,13 +83,13 @@ class VirtualStar(ABC): event: AstrMessageEvent, ) -> list[StarHandlerMetadata]: """Get the list of handlers that should be triggered for this event. - + This method uses cached handler metadata to determine which handlers should handle the given event. No RPC calls should be made here. - + Args: event: The message event to check - + Returns: List of handler metadata that match the event """ @@ -101,23 +102,23 @@ class VirtualStar(ABC): event: AstrMessageEvent, *args, **kwargs, - ) -> T.Any: + ) -> T.AsyncGenerator[T.Any, None]: """Call a registered handler in the plugin. - + This method should: - Serialize the event and arguments - Call the handler via RPC - Wait for and return the result - + Args: handler: The handler metadata event: The message event *args: Additional positional arguments **kwargs: Additional keyword arguments - + Returns: - The result from the handler - + An async generator yielding results from the handler + Raises: RuntimeError: If the handler call fails or times out """ diff --git a/src/astrbot_sdk/runtime/types.py b/src/astrbot_sdk/runtime/types.py index 5ba10dcd1..f8819764d 100644 --- a/src/astrbot_sdk/runtime/types.py +++ b/src/astrbot_sdk/runtime/types.py @@ -30,3 +30,37 @@ class CallHandlerRequest(JSONRPCRequest): method: Literal["call_handler"] params: Params | dict = Field(default_factory=dict) + + +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/start_client.py b/src/astrbot_sdk/tests/start_client.py index 232bc35bb..d9df59d5c 100644 --- a/src/astrbot_sdk/tests/start_client.py +++ b/src/astrbot_sdk/tests/start_client.py @@ -65,7 +65,9 @@ async def amain(): ), session_id="test_session", ) - await star.call_handler(star._handlers[0], event) + + async for result in star.call_handler(star._handlers[0], event): + print(f"Handler result: {result}") await star.stop()