mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
feat: implement NewStar class and associated utilities for plugin runtime management
This commit is contained in:
Soulter
@@ -10,7 +10,7 @@ let AstrBot core interact with both types of stars without needing to know the u
|
||||
"""
|
||||
|
||||
from .stars.virtual import VirtualStar
|
||||
from .stars.new import NewStdioStar, NewWebSocketStar
|
||||
from .stars.new_star import NewStdioStar, NewWebSocketStar
|
||||
from ..api.star.context import Context
|
||||
# from .types import StarURI, StarType
|
||||
|
||||
|
||||
@@ -7,22 +7,25 @@ from typing import Any, AsyncGenerator
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel
|
||||
from ...api.event.astr_message_event import AstrMessageEvent
|
||||
from ...api.star.star import StarMetadata
|
||||
from ..stars.registry import EventType, StarHandlerMetadata
|
||||
from .registry import EventType, StarHandlerMetadata
|
||||
from ..rpc.jsonrpc import (
|
||||
JSONRPCErrorData,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
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
|
||||
from .new_star_utils import (
|
||||
ClientHandshakeHandler,
|
||||
PluginRequestHandler,
|
||||
HandlerProxyFactory,
|
||||
)
|
||||
|
||||
|
||||
class NewStar(VirtualStar):
|
||||
@@ -53,6 +56,11 @@ class NewStar(VirtualStar):
|
||||
# Use RPCRequestHelper for managing requests
|
||||
self._rpc_helper = RPCRequestHelper()
|
||||
|
||||
# Initialize specialized handlers
|
||||
self._handshake_handler = ClientHandshakeHandler(self._rpc_helper)
|
||||
self._plugin_request_handler = PluginRequestHandler(context)
|
||||
self._handler_proxy_factory = HandlerProxyFactory(client, self._rpc_helper)
|
||||
|
||||
# Set up message handler
|
||||
self._client.set_message_handler(self._handle_message)
|
||||
|
||||
@@ -78,69 +86,11 @@ class NewStar(VirtualStar):
|
||||
]:
|
||||
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))
|
||||
|
||||
async def _handle_plugin_request(self, request: JSONRPCRequest) -> None:
|
||||
"""Handle a JSON-RPC request from the plugin (plugin calling core methods).
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request from the plugin
|
||||
"""
|
||||
result: Any = None
|
||||
try:
|
||||
# Handle core methods that plugins might call
|
||||
method = request.method
|
||||
params = request.params
|
||||
|
||||
if method == "call_context_function":
|
||||
ctx = self._context
|
||||
func_full_name = params.get("name", "")
|
||||
args = params.get("args", {})
|
||||
logger.debug(
|
||||
f"plugin called call_context_function: {func_full_name} with args: {args}"
|
||||
# Plugin is calling a method on the core - delegate to PluginRequestHandler
|
||||
asyncio.create_task(
|
||||
self._plugin_request_handler.handle_request(message, self._client)
|
||||
)
|
||||
|
||||
# Get the registered function from context
|
||||
func = ctx.get_registered_function(func_full_name)
|
||||
if func is None:
|
||||
raise ValueError(f"Function not found: {func_full_name}")
|
||||
|
||||
# Call the function
|
||||
import inspect
|
||||
|
||||
if inspect.iscoroutinefunction(func):
|
||||
result = await func(**args)
|
||||
else:
|
||||
result = func(**args)
|
||||
|
||||
logger.debug(f"call_context_function result: {result}")
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
|
||||
# Send success response
|
||||
response = JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
result={
|
||||
"data": result,
|
||||
},
|
||||
)
|
||||
await self._client.send_message(response)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling plugin request: {e}")
|
||||
# Send error response
|
||||
error_response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32603,
|
||||
message=str(e),
|
||||
),
|
||||
)
|
||||
await self._client.send_message(error_response)
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Start the plugin process and establish connection."""
|
||||
# Start the client (which may start a subprocess for STDIO)
|
||||
@@ -153,110 +103,16 @@ class NewStar(VirtualStar):
|
||||
Returns:
|
||||
Plugin metadata including name, version, handlers, etc.
|
||||
"""
|
||||
logger.info("Performing handshake with plugin...")
|
||||
# Delegate to ClientHandshakeHandler
|
||||
(
|
||||
self._metadata,
|
||||
self._handlers,
|
||||
) = await self._handshake_handler.perform_handshake(self._client)
|
||||
|
||||
response = await self._rpc_helper.call_rpc(
|
||||
self._client,
|
||||
HandshakeRequest(
|
||||
jsonrpc="2.0",
|
||||
id=self._rpc_helper._generate_request_id(),
|
||||
method="handshake",
|
||||
),
|
||||
)
|
||||
# Set up handler proxies
|
||||
self._handler_proxy_factory.setup_handlers(self._handlers)
|
||||
|
||||
if not isinstance(response, JSONRPCSuccessResponse):
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
result = response.result
|
||||
|
||||
if isinstance(result, dict):
|
||||
# Parse metadata
|
||||
for star_name, star_info in result.items():
|
||||
handlers_data = star_info.pop("handlers", None)
|
||||
metadata = StarMetadata(**star_info)
|
||||
self._metadata[star_name] = metadata
|
||||
|
||||
# Get handlers
|
||||
self._handlers = []
|
||||
|
||||
for handler_data in handlers_data:
|
||||
handler_meta = StarHandlerMetadata(
|
||||
event_type=EventType(handler_data["event_type"]),
|
||||
handler_full_name=handler_data["handler_full_name"],
|
||||
handler_name=handler_data["handler_name"],
|
||||
handler_module_path=handler_data["handler_module_path"],
|
||||
handler=self._create_handler_proxy(
|
||||
handler_data["handler_full_name"]
|
||||
),
|
||||
event_filters=[],
|
||||
desc=handler_data.get("desc", ""),
|
||||
extras_configs=handler_data.get("extras_configs", {}),
|
||||
)
|
||||
self._handlers.append(handler_meta)
|
||||
|
||||
logger.info(
|
||||
f"Handshake complete: {len(self._metadata)} stars loaded, {self._metadata.keys()}, {len(self._handlers)} handlers registered."
|
||||
)
|
||||
logger.info(f"Registered {len(self._handlers)} handlers")
|
||||
|
||||
return self._metadata
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
def _create_handler_proxy(self, handler_full_name: str):
|
||||
"""Create a proxy function that calls the handler via RPC.
|
||||
|
||||
Args:
|
||||
handler_full_name: The full name of the handler
|
||||
|
||||
Returns:
|
||||
An async function that proxies calls to the remote handler.
|
||||
The function may return a direct result or an async generator for streaming.
|
||||
"""
|
||||
|
||||
async def handler_proxy(event: AstrMessageEvent, **kwargs):
|
||||
"""Proxy function for remote handler invocation.
|
||||
|
||||
Returns either a direct result or an async generator for streaming handlers.
|
||||
"""
|
||||
request_id = self._rpc_helper._generate_request_id()
|
||||
request = CallHandlerRequest(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
method="call_handler",
|
||||
params=CallHandlerRequest.Params(
|
||||
handler_full_name=handler_full_name,
|
||||
event=AstrMessageEventModel.from_event(event),
|
||||
args=kwargs,
|
||||
),
|
||||
)
|
||||
queue = await self._rpc_helper.call_rpc_streaming(self._client, request)
|
||||
|
||||
try:
|
||||
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)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
raise RuntimeError(f"RPC call to {handler_full_name} timed out")
|
||||
|
||||
return handler_proxy
|
||||
|
||||
def _deserialize_result(self, result: Any) -> Any:
|
||||
"""Deserialize result from JSON-RPC response.
|
||||
|
||||
Args:
|
||||
result: The result from the plugin
|
||||
|
||||
Returns:
|
||||
Deserialized result object
|
||||
"""
|
||||
# For now, return as-is
|
||||
# In practice, we might want to reconstruct MessageEventResult etc.
|
||||
return result
|
||||
return self._metadata
|
||||
|
||||
def get_triggered_handlers(
|
||||
self, event: AstrMessageEvent
|
||||
265
src/astrbot_sdk/runtime/stars/new_star_utils.py
Normal file
265
src/astrbot_sdk/runtime/stars/new_star_utils.py
Normal file
@@ -0,0 +1,265 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
from typing import Any, AsyncGenerator
|
||||
from loguru import logger
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel
|
||||
from ...api.star.star import StarMetadata
|
||||
from ..rpc.client import JSONRPCClient
|
||||
from ..rpc.request_helper import RPCRequestHelper
|
||||
from ..rpc.jsonrpc import (
|
||||
JSONRPCSuccessResponse,
|
||||
JSONRPCRequest,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCErrorData,
|
||||
)
|
||||
from ..types import CallHandlerRequest, HandshakeRequest
|
||||
from .registry import StarHandlerMetadata, EventType
|
||||
|
||||
|
||||
class HandlerProxyFactory:
|
||||
"""Creates proxy functions for remote handler invocation."""
|
||||
|
||||
def __init__(self, client: JSONRPCClient, rpc_helper: RPCRequestHelper):
|
||||
"""Initialize the handler proxy factory.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
rpc_helper: RPC request helper for making RPC calls
|
||||
"""
|
||||
self._client = client
|
||||
self._rpc_helper = rpc_helper
|
||||
|
||||
def create_handler_proxy(self, handler_full_name: str):
|
||||
"""Create a proxy function that calls the handler via RPC.
|
||||
|
||||
Args:
|
||||
handler_full_name: The full name of the handler
|
||||
|
||||
Returns:
|
||||
An async generator function that proxies calls to the remote handler.
|
||||
"""
|
||||
|
||||
async def handler_proxy(
|
||||
event: AstrMessageEvent, **kwargs
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Proxy function for remote handler invocation.
|
||||
|
||||
Yields results from the remote handler via streaming.
|
||||
"""
|
||||
request_id = self._rpc_helper._generate_request_id()
|
||||
request = CallHandlerRequest(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
method="call_handler",
|
||||
params=CallHandlerRequest.Params(
|
||||
handler_full_name=handler_full_name,
|
||||
event=AstrMessageEventModel.from_event(event),
|
||||
args=kwargs,
|
||||
),
|
||||
)
|
||||
queue = await self._rpc_helper.call_rpc_streaming(self._client, request)
|
||||
|
||||
try:
|
||||
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)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
raise RuntimeError(f"RPC call to {handler_full_name} timed out")
|
||||
|
||||
return handler_proxy
|
||||
|
||||
def setup_handlers(self, handlers: list[StarHandlerMetadata]) -> None:
|
||||
"""Set up handler proxies for all handlers.
|
||||
|
||||
Args:
|
||||
handlers: List of handler metadata to set up
|
||||
"""
|
||||
for handler in handlers:
|
||||
handler.handler = self.create_handler_proxy(handler.handler_full_name)
|
||||
logger.info(f"Set up {len(handlers)} handler proxies")
|
||||
|
||||
def _deserialize_result(self, result: Any) -> Any:
|
||||
"""Deserialize result from JSON-RPC response.
|
||||
|
||||
Args:
|
||||
result: The result from the plugin
|
||||
|
||||
Returns:
|
||||
Deserialized result object
|
||||
"""
|
||||
# For now, return as-is
|
||||
# In practice, we might want to reconstruct MessageEventResult etc.
|
||||
return result
|
||||
|
||||
|
||||
class ClientHandshakeHandler:
|
||||
"""Handles the handshake protocol to retrieve plugin metadata."""
|
||||
|
||||
def __init__(self, rpc_helper: RPCRequestHelper):
|
||||
"""Initialize the handshake handler.
|
||||
|
||||
Args:
|
||||
rpc_helper: RPC request helper for making RPC calls
|
||||
"""
|
||||
self._rpc_helper = rpc_helper
|
||||
|
||||
async def perform_handshake(
|
||||
self, client: JSONRPCClient
|
||||
) -> tuple[dict[str, StarMetadata], list[StarHandlerMetadata]]:
|
||||
"""Perform handshake to retrieve plugin metadata.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
|
||||
Returns:
|
||||
Tuple of (metadata dict, handlers list)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If handshake fails
|
||||
"""
|
||||
logger.info("Performing handshake with plugin...")
|
||||
|
||||
response = await self._rpc_helper.call_rpc(
|
||||
client,
|
||||
HandshakeRequest(
|
||||
jsonrpc="2.0",
|
||||
id=self._rpc_helper._generate_request_id(),
|
||||
method="handshake",
|
||||
),
|
||||
)
|
||||
|
||||
if not isinstance(response, JSONRPCSuccessResponse):
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
result = response.result
|
||||
|
||||
if not isinstance(result, dict):
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
metadata_dict: dict[str, StarMetadata] = {}
|
||||
handlers_list: list[StarHandlerMetadata] = []
|
||||
|
||||
# Placeholder handler that will be replaced later
|
||||
def _placeholder_handler(*args, **kwargs):
|
||||
raise NotImplementedError("Handler proxy not set up yet")
|
||||
|
||||
# Parse metadata
|
||||
for star_name, star_info in result.items():
|
||||
handlers_data = star_info.pop("handlers", None)
|
||||
metadata = StarMetadata(**star_info)
|
||||
metadata_dict[star_name] = metadata
|
||||
|
||||
# Parse handlers
|
||||
if handlers_data:
|
||||
for handler_data in handlers_data:
|
||||
handler_meta = StarHandlerMetadata(
|
||||
event_type=EventType(handler_data["event_type"]),
|
||||
handler_full_name=handler_data["handler_full_name"],
|
||||
handler_name=handler_data["handler_name"],
|
||||
handler_module_path=handler_data["handler_module_path"],
|
||||
handler=_placeholder_handler, # Will be replaced by HandlerProxyFactory
|
||||
event_filters=[],
|
||||
desc=handler_data.get("desc", ""),
|
||||
extras_configs=handler_data.get("extras_configs", {}),
|
||||
)
|
||||
handlers_list.append(handler_meta)
|
||||
|
||||
logger.info(
|
||||
f"Handshake complete: {len(metadata_dict)} stars loaded, "
|
||||
f"{metadata_dict.keys()}, {len(handlers_list)} handlers registered."
|
||||
)
|
||||
|
||||
return metadata_dict, handlers_list
|
||||
|
||||
|
||||
class PluginRequestHandler:
|
||||
"""Handles JSON-RPC requests from plugins calling core methods."""
|
||||
|
||||
def __init__(self, context: Any):
|
||||
"""Initialize the plugin request handler.
|
||||
|
||||
Args:
|
||||
context: Context instance for managing managers and their functions
|
||||
"""
|
||||
self._context = context
|
||||
|
||||
async def handle_request(
|
||||
self, request: JSONRPCRequest, client: JSONRPCClient
|
||||
) -> None:
|
||||
"""Handle a JSON-RPC request from the plugin.
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request from the plugin
|
||||
client: The client to send response back
|
||||
"""
|
||||
result: Any = None
|
||||
try:
|
||||
method = request.method
|
||||
params = request.params or {}
|
||||
|
||||
if method == "call_context_function":
|
||||
result = await self._handle_context_function_call(params)
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
|
||||
# Send success response
|
||||
response = JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
result={
|
||||
"data": result,
|
||||
},
|
||||
)
|
||||
await client.send_message(response)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling plugin request: {e}")
|
||||
# Send error response
|
||||
error_response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32603,
|
||||
message=str(e),
|
||||
),
|
||||
)
|
||||
await client.send_message(error_response)
|
||||
|
||||
async def _handle_context_function_call(self, params: dict) -> Any:
|
||||
"""Handle call_context_function requests.
|
||||
|
||||
Args:
|
||||
params: Request parameters containing function name and args
|
||||
|
||||
Returns:
|
||||
Result from the function call
|
||||
|
||||
Raises:
|
||||
ValueError: If function is not found
|
||||
"""
|
||||
func_full_name = params.get("name", "")
|
||||
args = params.get("args", {})
|
||||
|
||||
logger.debug(
|
||||
f"Plugin called call_context_function: {func_full_name} with args: {args}"
|
||||
)
|
||||
|
||||
# Get the registered function from context
|
||||
func = self._context.get_registered_function(func_full_name)
|
||||
if func is None:
|
||||
raise ValueError(f"Function not found: {func_full_name}")
|
||||
|
||||
# Call the function
|
||||
if inspect.iscoroutinefunction(func):
|
||||
result = await func(**args)
|
||||
else:
|
||||
result = func(**args)
|
||||
|
||||
logger.debug(f"call_context_function result: {result}")
|
||||
return result
|
||||
Reference in New Issue
Block a user