feat: enhance context management with component registration and RPC function handling

src/astrbot_sdk/api/star/context.py

@@ -1,6 +1,52 @@
 from abc import ABC
+from typing import Any, Callable
 from ..basic.conversation_mgr import BaseConversationManager
 
 
 class Context(ABC):
     conversation_manager: BaseConversationManager
+
+    def __init__(self):
+        self._registered_managers: dict[str, Any] = {}
+        self._registered_functions: dict[str, Callable] = {}
+
+    def register_component(self, *components: Any) -> None:
+        """Register a components instance and its public methods.
+
+        This allows the components's methods to be called via RPC using the pattern:
+        ComponentClassName.method_name
+
+        Args:
+            components: The components instance to register
+        """
+        for component in components:
+            class_name = component.__class__.__name__
+            self._registered_managers[class_name] = component
+
+            # Register all public methods (not starting with _)
+            for attr_name in dir(component):
+                if not attr_name.startswith("_"):
+                    attr = getattr(component, attr_name)
+                    if callable(attr):
+                        full_name = f"{class_name}.{attr_name}"
+                        self._registered_functions[full_name] = attr
+
+    def get_registered_function(self, full_name: str) -> Callable | None:
+        """Get a registered function by its full name.
+
+        Args:
+            full_name: Full name in format "ComponentClassName.method_name"
+
+        Returns:
+            The callable function or None if not found
+        """
+
+        return self._registered_functions.get(full_name)
+
+    def list_registered_functions(self) -> list[str]:
+        """List all registered function names.
+
+        Returns:
+            List of full function names in format "ComponentClassName.method_name"
+        """
+        return list(self._registered_functions.keys())

src/astrbot_sdk/runtime/api/context.py

@@ -4,7 +4,10 @@ from .conversation_mgr import ConversationManager
 
 class Context(BaseContext):
     def __init__(self, conversation_manager: ConversationManager):
+        super().__init__()
         self.conversation_manager = conversation_manager
+        # Auto-register the conversation manager
+        self.register_component(self.conversation_manager)
 
     @classmethod
     def default_context(cls, runner=None):

src/astrbot_sdk/runtime/api/conversation_mgr.py

@@ -11,7 +11,7 @@ class ConversationManager(BaseConversationManager):
         """
         self.runner = runner
 
-    @rpc_method
+    @rpc_method(return_model=str)
     async def new_conversation(
         self,
         unified_msg_origin: str,

src/astrbot_sdk/runtime/api/util.py

@@ -1,38 +1,92 @@
 import inspect
 from functools import wraps
-from typing import Callable
+from typing import Callable, Type, Any, get_type_hints, TypeVar, overload
 from ..star_runner import StarRunner
 from ..types import CallContextFunctionRequest
 
+F = TypeVar("F", bound=Callable[..., Any])
 
-def rpc_method(func: Callable) -> Callable:
-    """sign as an RPC method."""
 
-    @wraps(func)
-    async def wrapper(self, *args, **kwargs):
-        if not hasattr(self, "runner") or not isinstance(self.runner, StarRunner):
-            raise RuntimeError(
-                f"Class {self.__class__.__name__} is not configured for RPC calls."
+@overload
+def rpc_method(func: F) -> F: ...
+
+
+@overload
+def rpc_method(*, return_model: Type[Any] | None = None) -> Callable[[F], F]: ...
+
+
+def rpc_method(
+    func: F | None = None, *, return_model: Type[Any] | None = None
+) -> F | Callable[[F], F]:
+    """sign as an RPC method.
+
+    Args:
+        func: The function to decorate
+        return_model: The expected return type/model (e.g., dict, BaseModel, etc.)
+                     If not specified, will try to infer from type hints
+
+    Examples:
+        @rpc_method
+        async def get_config(self) -> dict:
+            ...
+
+        @rpc_method(return_model=dict)
+        async def get_config(self):
+            ...
+    """
+
+    def decorator(f: F) -> F:
+        # Try to get return type from type hints if not explicitly provided
+        _return_model = return_model
+        if _return_model is None:
+            try:
+                hints = get_type_hints(f)
+                if "return" in hints:
+                    _return_model = hints["return"]
+            except Exception:
+                pass
+
+        # Store return model as function attribute for potential inspection
+        setattr(f, "__return_model__", _return_model)
+
+        @wraps(f)
+        async def wrapper(self, *args, **kwargs):
+            if not hasattr(self, "runner") or not isinstance(self.runner, StarRunner):
+                raise RuntimeError(
+                    f"Class {self.__class__.__name__} is not configured for RPC calls."
+                )
+            method_name = f"{self.__class__.__name__}.{f.__name__}"
+            sig = inspect.signature(f)
+            bound_args = sig.bind(self, *args, **kwargs)
+            bound_args.apply_defaults()
+            params = dict(bound_args.arguments)
+            params.pop("self")
+
+            runner: StarRunner = getattr(self, "runner")
+
+            result = await runner._call_rpc(
+                CallContextFunctionRequest(
+                    jsonrpc="2.0",
+                    id=runner._generate_request_id(),
+                    method="call_context_function",
+                    params=CallContextFunctionRequest.Params(
+                        name=method_name,
+                        args=params,
+                    ),
+                )
             )
-        method_name = f"{self.__class__.__name__}.{func.__name__}"
-        sig = inspect.signature(func)
-        bound_args = sig.bind(self, *args, **kwargs)
-        bound_args.apply_defaults()
-        params = dict(bound_args.arguments)
-        params.pop("self")
 
-        runner: StarRunner = getattr(self, "runner")
+            # TODO: Process result based on _return_model if needed
+            return result
 
-        return await runner._call_rpc(
-            CallContextFunctionRequest(
-                jsonrpc="2.0",
-                id=runner._generate_request_id(),
-                method="call_context_function",
-                params=CallContextFunctionRequest.Params(
-                    name=method_name,
-                    args=params,
-                ),
-            )
-        )
+        # Also store on wrapper for easy access
+        setattr(wrapper, "__return_model__", _return_model)
+        return wrapper  # type: ignore
 
-    return wrapper
+    # Support both @rpc_method and @rpc_method(return_model=...)
+    if func is None:
+        # Called with arguments: @rpc_method(return_model=...)
+        return decorator
+    else:
+        # Called without arguments: @rpc_method
+        return decorator(func)

src/astrbot_sdk/runtime/galaxy.py

@@ -11,6 +11,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 ..api.star.context import Context
 # from .types import StarURI, StarType
 
 
@@ -19,18 +20,20 @@ class Galaxy:
 
     vs_map: dict[str, VirtualStar] = {}
 
-    async def connect_to_stdio_star(self, star_name: str, config: dict) -> NewStdioStar:
+    async def connect_to_stdio_star(
+        self, context: Context, star_name: str, config: dict
+    ) -> NewStdioStar:
         """Connect to a new-style stdio star given its name."""
-        star = NewStdioStar(**config)
+        star = NewStdioStar(context=context, **config)
         await star.initialize()
         self.vs_map[star_name] = star
         return star
 
     async def connect_to_websocket_star(
-        self, star_name: str, config: dict
+        self, context: Context, star_name: str, config: dict
     ) -> NewWebSocketStar:
         """Connect to a new-style websocket star given its name."""
-        star = NewWebSocketStar(**config)
+        star = NewWebSocketStar(context=context, **config)
         await star.initialize()
         self.vs_map[star_name] = star
         return star

src/astrbot_sdk/runtime/star_runner.py

@@ -17,13 +17,13 @@ class StarRunner:
     def __init__(self, server: JSONRPCServer):
         self.server = server
         self._request_id_counter = 0
-        self.pending_requests: dict[str, asyncio.Future] = {}
+        self.pending_requests: dict[str, asyncio.Future[JSONRPCMessage]] = {}
 
     def _generate_request_id(self) -> str:
         self._request_id_counter += 1
         return str(self._request_id_counter)
 
-    async def _call_rpc(self, message: JSONRPCMessage):
+    async def _call_rpc(self, message: JSONRPCMessage) -> JSONRPCMessage | None:
         if message.id is not None:
             self.pending_requests[message.id] = asyncio.get_event_loop().create_future()
         await self.server.send_message(message)

src/astrbot_sdk/runtime/stars/new.py

@@ -33,12 +33,23 @@ class NewStar(VirtualStar):
     def __init__(
         self,
         client: JSONRPCClient,
+        context: Any = None,
     ) -> None:
         """Initialize a NewStar instance.
 
         Args:
             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] = []
@@ -120,14 +131,34 @@ class NewStar(VirtualStar):
         Args:
             request: The JSON-RPC request from the plugin
         """
-        result: dict = {}
+        result: Any = None
         try:
             # Handle core methods that plugins might call
             method = request.method
             params = request.params
 
             if method == "call_context_function":
-                logger.debug(f"plugin called call_context_function: {params}")
+                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}"
+                )
+
+                # 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}")
 
@@ -135,7 +166,9 @@ class NewStar(VirtualStar):
             response = JSONRPCSuccessResponse(
                 jsonrpc="2.0",
                 id=request.id,
-                result=result,
+                result={
+                    "data": result,
+                },
             )
             await self._client.send_message(response)
 
@@ -537,6 +570,7 @@ class NewStdioStar(NewStar):
         self,
         plugin_dir: str,
         python_executable: str = "python",
+        context: Any = None,
         **kwargs: Any,
     ) -> None:
         """Initialize a STDIO-based NewStar.
@@ -544,7 +578,7 @@ class NewStdioStar(NewStar):
         Args:
             plugin_dir: Path to the plugin directory
             python_executable: Python executable to use (defaults to 'python')
-            main_script: Main script filename (defaults to 'main.py')
+            context: Context instance for managing managers and their functions
         """
         # Construct the command to start the plugin
         if not os.path.exists(plugin_dir):
@@ -554,7 +588,7 @@ class NewStdioStar(NewStar):
 
         # Create StdioClient with subprocess management
         client = StdioClient(command=command, cwd=plugin_dir)
-        super().__init__(client)
+        super().__init__(client, context=context)
 
 
 class NewWebSocketStar(NewStar):
@@ -569,6 +603,7 @@ class NewWebSocketStar(NewStar):
         url: str,
         heartbeat: float = 30.0,
         reconnect_interval: float = 5.0,
+        context: Any = None,
         **kwargs: Any,
     ) -> None:
         """Initialize a WebSocket-based NewStar.
@@ -577,11 +612,12 @@ class NewWebSocketStar(NewStar):
             url: WebSocket server URL that the plugin will connect to
             heartbeat: Heartbeat interval in seconds
             reconnect_interval: Interval between reconnection attempts in seconds
+            context: Context instance for managing managers and their functions
         """
         client = WebSocketClient(
             url=url, heartbeat=heartbeat, reconnect_interval=reconnect_interval
         )
-        super().__init__(client)
+        super().__init__(client, context=context)
         self._url = url
         self._heartbeat = heartbeat
         self._reconnect_interval = reconnect_interval

src/astrbot_sdk/runtime/stars/virtual.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from ...api.event.astr_message_event import AstrMessageEvent
 from ...api.star.star import StarMetadata
 from .registry import StarHandlerMetadata
+from ...api.star.context import Context
 
 
 class VirtualStar(ABC):
@@ -13,6 +14,8 @@ class VirtualStar(ABC):
     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:

src/astrbot_sdk/tests/start_client.py

@@ -5,12 +5,39 @@ from astrbot_sdk.api.event.astrbot_message import AstrBotMessage, MessageMember
 from astrbot_sdk.api.platform.platform_metadata import PlatformMetadata
 from astrbot_sdk.api.event.message_type import MessageType
 
+from astrbot_sdk.api.star.context import Context
+from astrbot_sdk.api.basic.conversation_mgr import BaseConversationManager
+
+
+class ConversationManager(BaseConversationManager):
+    async def new_conversation(
+        self,
+        unified_msg_origin: str,
+        platform_id: str | None = None,
+        content: list[dict] | None = None,
+        title: str | None = None,
+        persona_id: str | None = None,
+    ) -> str:
+        import uuid
+
+        return str(uuid.uuid4())
+
+
+class TestContext(Context):
+    def __init__(self, conversation_manager: ConversationManager):
+        super().__init__()
+        self.conversation_manager = conversation_manager
+        self.register_component(self.conversation_manager)
+
 
 async def amain():
     galaxy = Galaxy()
+    conversation_manager = ConversationManager()
+    context = TestContext(conversation_manager)
     star = await galaxy.connect_to_websocket_star(
-        "hello",
-        {
+        context=context,
+        star_name="hello",
+        config={
             "url": "ws://127.0.0.1:8765",
         },
     )

test_plugin/commands/hello.py

@@ -11,4 +11,4 @@ class HelloCommand(CommandComponent):
     async def hello(self, event: AstrMessageEvent):
         ret = await self.context.conversation_manager.new_conversation("hello")
         print(f"New conversation created: {ret}")
-        yield event.plain_result("Hello, Astrbot!")
+        yield event.plain_result(f"Hello, Astrbot! Created conversation ID: {ret}")