diff --git a/src/astrbot_sdk/__main__.py b/src/astrbot_sdk/__main__.py new file mode 100644 index 000000000..b0847a229 --- /dev/null +++ b/src/astrbot_sdk/__main__.py @@ -0,0 +1,5 @@ +from .cli.main import cli + + +if __name__ == "__main__": + cli() diff --git a/src/astrbot_sdk/cli/main.py b/src/astrbot_sdk/cli/main.py index ba9de33d9..e070a820d 100644 --- a/src/astrbot_sdk/cli/main.py +++ b/src/astrbot_sdk/cli/main.py @@ -1,8 +1,10 @@ import asyncio import sys + import click from loguru import logger -from ..runtime.serve import run_websocket_server + +from ..runtime.serve import run_plugin_worker, run_supervisor, run_websocket_server def setup_logger(verbose: bool = False): @@ -38,10 +40,34 @@ def cli(ctx, verbose): @cli.command() -@click.option("--port", default=8765, help="WebSocket server port", type=int) +@click.option( + "--plugins-dir", + default="plugins", + type=click.Path(file_okay=False, dir_okay=True, path_type=str), + help="Directory containing plugin folders", +) @click.pass_context -def run(ctx, port: int): - """Start the WebSocket server""" +def run(ctx, plugins_dir: str): + """Start the plugin supervisor over stdio.""" + logger.info(f"Starting plugin supervisor with plugins dir: {plugins_dir}") + asyncio.run(run_supervisor(plugins_dir=plugins_dir)) + + +@cli.command(hidden=True) +@click.option( + "--plugin-dir", + required=True, + type=click.Path(file_okay=False, dir_okay=True, path_type=str), +) +def worker(plugin_dir: str): + """Internal command used by the supervisor to start a worker.""" + asyncio.run(run_plugin_worker(plugin_dir=plugin_dir)) + + +@cli.command(hidden=True) +@click.option("--port", default=8765, help="WebSocket server port", type=int) +def websocket(port: int): + """Legacy websocket runtime entrypoint.""" logger.info(f"Starting WebSocket server on port {port}...") asyncio.run(run_websocket_server(port=port)) diff --git a/src/astrbot_sdk/runtime/rpc/client/stdio.py b/src/astrbot_sdk/runtime/rpc/client/stdio.py index 06acd4548..f42daa5a2 100644 --- a/src/astrbot_sdk/runtime/rpc/client/stdio.py +++ b/src/astrbot_sdk/runtime/rpc/client/stdio.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import os import subprocess from typing import IO, Any @@ -23,6 +24,7 @@ class StdioClient(JSONRPCClient): self, command: list[str], cwd: str | None = None, + env: dict[str, str] | None = None, ) -> None: """Initialize the STDIO client. @@ -33,6 +35,7 @@ class StdioClient(JSONRPCClient): super().__init__() self._command = command self._cwd = cwd + self._env = env or os.environ.copy() self._process: subprocess.Popen | None = None self._stdin: IO[Any] | None = None self._stdout: IO[Any] | None = None @@ -64,6 +67,7 @@ class StdioClient(JSONRPCClient): stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=self._cwd, + env=self._env, text=True, bufsize=1, # Line buffered ) @@ -116,6 +120,11 @@ class StdioClient(JSONRPCClient): # Terminate subprocess if running if self._process: + if self._stdout: + try: + self._stdout.close() + except Exception: + logger.debug("Failed to close subprocess stdin cleanly") logger.info("Terminating subprocess...") self._process.terminate() try: diff --git a/src/astrbot_sdk/runtime/rpc/server/stdio.py b/src/astrbot_sdk/runtime/rpc/server/stdio.py index 004fd7e48..aa8da3325 100644 --- a/src/astrbot_sdk/runtime/rpc/server/stdio.py +++ b/src/astrbot_sdk/runtime/rpc/server/stdio.py @@ -38,6 +38,7 @@ class StdioServer(JSONRPCServer): self._stdout = stdout or sys.stdout self._read_task: asyncio.Task | None = None self._write_lock = asyncio.Lock() + self._closed_event = asyncio.Event() async def start(self) -> None: """Start the server and begin reading from stdin.""" @@ -55,6 +56,7 @@ class StdioServer(JSONRPCServer): return self._running = False + self._closed_event.set() # Cancel read task if self._read_task: @@ -101,6 +103,8 @@ class StdioServer(JSONRPCServer): if not line: # EOF reached logger.info("EOF reached on stdin") + self._running = False + self._closed_event.set() break line = line.strip() @@ -120,8 +124,12 @@ class StdioServer(JSONRPCServer): except Exception as e: logger.error(f"Error in read loop: {e}") finally: + self._closed_event.set() logger.debug("Stopped reading from stdin") + async def wait_closed(self) -> None: + await self._closed_event.wait() + def _parse_message(self, line: str) -> JSONRPCMessage: """Parse a JSON-RPC message from a string. diff --git a/src/astrbot_sdk/runtime/serve.py b/src/astrbot_sdk/runtime/serve.py index 7460733d5..38cb19afe 100644 --- a/src/astrbot_sdk/runtime/serve.py +++ b/src/astrbot_sdk/runtime/serve.py @@ -1,11 +1,56 @@ import asyncio import signal -from .rpc.server import WebSocketServer, StdioServer +import sys +from pathlib import Path +from typing import IO, Any + +from loguru import logger + +from .api.context import Context +from .rpc.server import StdioServer, WebSocketServer from .star_runner import StarRunner from .stars.star_manager import StarManager -from .api.context import Context -from loguru import logger -from typing import IO, Any +from .supervisor import SupervisorRuntime + + +def _install_signal_handlers(stop_event: asyncio.Event) -> None: + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, stop_event.set) + except NotImplementedError: + logger.debug(f"Signal handlers are not supported for {sig}") + + +def _prepare_stdio_transport( + stdin: IO[Any] | None, + stdout: IO[Any] | None, +) -> tuple[IO[Any], IO[Any], IO[Any] | None]: + if stdin is not None and stdout is not None: + return stdin, stdout, None + + transport_stdin = stdin or sys.stdin + transport_stdout = stdout or sys.stdout + original_stdout = sys.stdout + sys.stdout = sys.stderr + return transport_stdin, transport_stdout, original_stdout + + +async def _wait_for_stdio_shutdown( + server: StdioServer, stop_event: asyncio.Event +) -> None: + stop_waiter = asyncio.create_task(stop_event.wait()) + stdio_waiter = asyncio.create_task(server.wait_closed()) + done, pending = await asyncio.wait( + {stop_waiter, stdio_waiter}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + for task in done: + if task.cancelled(): + continue + task.result() async def run_websocket_server( @@ -13,6 +58,7 @@ async def run_websocket_server( port: int = 8765, path: str = "/", heartbeat_interval: int = 30, + plugin_dir: str | Path | None = None, ): server = WebSocketServer( port=port, host=host, path=path, heartbeat=heartbeat_interval @@ -20,13 +66,11 @@ async def run_websocket_server( runner = StarRunner(server) context = Context.default_context(runner=runner) star_manager = StarManager(context=context) - star_manager.discover_star() + star_manager.discover_star(Path(plugin_dir) if plugin_dir else None) await runner.run() stop_event = asyncio.Event() - loop = asyncio.get_running_loop() - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, stop_event.set) + _install_signal_handlers(stop_event) logger.info("Server is running. Press Ctrl+C to stop.") @@ -37,26 +81,55 @@ async def run_websocket_server( await server.stop() -async def start_stdio_server( - stdin: IO[Any] | None = None, stdout: IO[Any] | None = None -): - """Start a JSON-RPC server over stdio.""" - server = StdioServer(stdin=stdin, stdout=stdout) +async def run_supervisor( + plugins_dir: str | Path = "plugins", + stdin: IO[Any] | None = None, + stdout: IO[Any] | None = None, +) -> None: + transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( + stdin, stdout + ) + server = StdioServer(stdin=transport_stdin, stdout=transport_stdout) + supervisor = SupervisorRuntime( + server=server, + plugins_dir=Path(plugins_dir), + ) + + try: + await supervisor.start() + stop_event = asyncio.Event() + _install_signal_handlers(stop_event) + logger.info(f"Plugin supervisor is running with plugins dir: {plugins_dir}") + await _wait_for_stdio_shutdown(server, stop_event) + finally: + logger.info("Shutting down plugin supervisor...") + await supervisor.stop() + if original_stdout is not None: + sys.stdout = original_stdout + + +async def run_plugin_worker( + plugin_dir: str | Path, + stdin: IO[Any] | None = None, + stdout: IO[Any] | None = None, +) -> None: + transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( + stdin, stdout + ) + server = StdioServer(stdin=transport_stdin, stdout=transport_stdout) runner = StarRunner(server) context = Context.default_context(runner=runner) star_manager = StarManager(context=context) - star_manager.discover_star() - await runner.run() - - stop_event = asyncio.Event() - loop = asyncio.get_running_loop() - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, stop_event.set) - - logger.info("Stdio server is running. Press Ctrl+C to stop.") + star_manager.discover_star(Path(plugin_dir)) try: - await stop_event.wait() + await runner.run() + stop_event = asyncio.Event() + _install_signal_handlers(stop_event) + logger.info(f"Plugin worker is running for: {plugin_dir}") + await _wait_for_stdio_shutdown(server, stop_event) finally: - logger.info("Shutting down...") - await server.stop() + logger.info("Shutting down plugin worker...") + await runner.stop() + if original_stdout is not None: + sys.stdout = original_stdout diff --git a/src/astrbot_sdk/runtime/stars/new_star.py b/src/astrbot_sdk/runtime/stars/new_star.py index 520ddcaec..da9fa184d 100644 --- a/src/astrbot_sdk/runtime/stars/new_star.py +++ b/src/astrbot_sdk/runtime/stars/new_star.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import os import inspect +from pathlib import Path from typing import Any, AsyncGenerator from loguru import logger @@ -177,7 +178,7 @@ class NewStdioStar(NewStar): def __init__( self, - plugin_dir: str, + plugins_dir: str, python_executable: str = "python", context: Any = None, **kwargs: Any, @@ -185,18 +186,33 @@ class NewStdioStar(NewStar): """Initialize a STDIO-based NewStar. Args: - plugin_dir: Path to the plugin directory + plugins_dir: Path to the plugins directory python_executable: Python executable to use (defaults to 'python') context: Context instance for managing managers and their functions """ - # Construct the command to start the plugin - if not os.path.exists(plugin_dir): - raise FileNotFoundError(f"Plugin directory not found: {plugin_dir}") + if not os.path.exists(plugins_dir): + raise FileNotFoundError(f"Plugins directory not found: {plugins_dir}") - command = [python_executable, "-m", "astrbot_sdk", "run", "--stdio"] + repo_src_dir = str(Path(__file__).resolve().parents[3]) + env = os.environ.copy() + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + f"{repo_src_dir}{os.pathsep}{existing_pythonpath}" + if existing_pythonpath + else repo_src_dir + ) + + command = [ + python_executable, + "-m", + "astrbot_sdk", + "run", + "--plugins-dir", + plugins_dir, + ] # Create StdioClient with subprocess management - client = StdioClient(command=command, cwd=plugin_dir) + client = StdioClient(command=command, cwd=plugins_dir, env=env) super().__init__(client, context=context) diff --git a/src/astrbot_sdk/runtime/supervisor.py b/src/astrbot_sdk/runtime/supervisor.py new file mode 100644 index 000000000..9dbe0d8b8 --- /dev/null +++ b/src/astrbot_sdk/runtime/supervisor.py @@ -0,0 +1,558 @@ +from __future__ import annotations + +import asyncio +import json +import os +import re +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import yaml +from loguru import logger +from .rpc.client.stdio import StdioClient +from .rpc.jsonrpc import ( + JSONRPCErrorData, + JSONRPCErrorResponse, + JSONRPCMessage, + JSONRPCRequest, + JSONRPCSuccessResponse, +) +from .rpc.request_helper import RPCRequestHelper +from .rpc.server.base import JSONRPCServer +from .stars.registry import EventType, StarHandlerMetadata +from .types import CallHandlerRequest, HandshakeRequest + +STATE_FILE_NAME = ".astrbot-worker-state.json" + + +def _venv_python_path(venv_dir: Path) -> Path: + if os.name == "nt": + return venv_dir / "Scripts" / "python.exe" + return venv_dir / "bin" / "python" + + +@dataclass(slots=True) +class PluginSpec: + name: str + plugin_dir: Path + manifest_path: Path + requirements_path: Path + python_version: str + manifest_data: dict[str, Any] + + +@dataclass(slots=True) +class PluginDiscoveryResult: + plugins: list[PluginSpec] + skipped_plugins: dict[str, str] + + +def discover_plugins(plugins_dir: Path) -> PluginDiscoveryResult: + plugins_root = plugins_dir.resolve() + skipped_plugins: dict[str, str] = {} + plugins: list[PluginSpec] = [] + seen_names: set[str] = set() + + if not plugins_root.exists(): + logger.warning(f"Plugins directory does not exist: {plugins_root}") + return PluginDiscoveryResult([], {}) + + for entry in sorted(plugins_root.iterdir()): + if not entry.is_dir() or entry.name.startswith("."): + continue + + manifest_path = entry / "plugin.yaml" + requirements_path = entry / "requirements.txt" + if not manifest_path.exists(): + logger.warning(f"Skipping {entry}: missing plugin.yaml") + continue + if not requirements_path.exists(): + logger.warning(f"Skipping {entry}: missing requirements.txt") + skipped_plugins[entry.name] = "missing requirements.txt" + continue + + try: + 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 + + plugin_name = manifest_data.get("name") + components = manifest_data.get("components") + runtime = manifest_data.get("runtime") or {} + python_version = runtime.get("python") + + if not isinstance(plugin_name, str) or not plugin_name: + skipped_plugins[entry.name] = "plugin name is required" + continue + if plugin_name in seen_names: + skipped_plugins[plugin_name] = "duplicate plugin name" + continue + if not isinstance(components, list) or not components: + skipped_plugins[plugin_name] = "components must be a non-empty list" + continue + if not isinstance(python_version, str) or not python_version: + skipped_plugins[plugin_name] = "runtime.python is required" + continue + + seen_names.add(plugin_name) + plugins.append( + PluginSpec( + name=plugin_name, + plugin_dir=entry.resolve(), + manifest_path=manifest_path.resolve(), + requirements_path=requirements_path.resolve(), + python_version=python_version, + manifest_data=manifest_data, + ) + ) + + return PluginDiscoveryResult(plugins=plugins, skipped_plugins=skipped_plugins) + + +class PluginEnvironmentManager: + def __init__(self, repo_root: Path, uv_binary: str | None = None) -> None: + self.repo_root = repo_root.resolve() + self.uv_binary = uv_binary or shutil.which("uv") + self.cache_dir = self.repo_root / ".uv-cache" + + def prepare_environment(self, plugin: PluginSpec) -> Path: + if not self.uv_binary: + raise RuntimeError("uv executable not found") + + state_path = plugin.plugin_dir / STATE_FILE_NAME + venv_dir = plugin.plugin_dir / ".venv" + python_path = _venv_python_path(venv_dir) + fingerprint = self._fingerprint(plugin) + state = self._load_state(state_path) + + if ( + not python_path.exists() + or not self._matches_python_version(venv_dir, plugin.python_version) + or state.get("fingerprint") != fingerprint + ): + self._rebuild(plugin, venv_dir, python_path) + self._write_state(state_path, plugin, fingerprint) + + return python_path + + def _rebuild(self, plugin: PluginSpec, venv_dir: Path, python_path: Path) -> None: + if venv_dir.exists(): + shutil.rmtree(venv_dir) + + venv_dir.parent.mkdir(parents=True, exist_ok=True) + self._run_command( + [ + self.uv_binary, + "venv", + "--python", + plugin.python_version, + "--system-site-packages", + "--no-python-downloads", + "--no-managed-python", + str(venv_dir), + ], + cwd=self.repo_root, + command_name=f"create venv for {plugin.name}", + ) + + requirements_text = plugin.requirements_path.read_text(encoding="utf-8").strip() + if not requirements_text: + return + + self._run_command( + [ + self.uv_binary, + "pip", + "install", + "--python", + str(python_path), + "-r", + str(plugin.requirements_path), + ], + cwd=plugin.plugin_dir, + command_name=f"install requirements for {plugin.name}", + ) + + def _run_command( + self, + command: list[str], + *, + cwd: Path, + command_name: str, + ) -> None: + logger.info(f"{command_name}: {' '.join(command)}") + process = subprocess.run( + command, + cwd=str(cwd), + env={**os.environ, "UV_CACHE_DIR": str(self.cache_dir)}, + capture_output=True, + text=True, + check=False, + ) + if process.returncode != 0: + raise RuntimeError( + f"{command_name} failed with exit code {process.returncode}: " + f"{process.stderr.strip() or process.stdout.strip()}" + ) + + @staticmethod + def _fingerprint(plugin: PluginSpec) -> str: + requirements = plugin.requirements_path.read_text(encoding="utf-8") + payload = { + "python_version": plugin.python_version, + "requirements": requirements, + } + return json.dumps(payload, ensure_ascii=True, sort_keys=True) + + @staticmethod + def _load_state(state_path: Path) -> dict[str, Any]: + if not state_path.exists(): + return {} + try: + data = json.loads(state_path.read_text(encoding="utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + @staticmethod + def _write_state(state_path: Path, plugin: PluginSpec, fingerprint: str) -> None: + state_path.write_text( + json.dumps( + { + "plugin": plugin.name, + "python_version": plugin.python_version, + "fingerprint": fingerprint, + }, + ensure_ascii=True, + indent=2, + sort_keys=True, + ), + encoding="utf-8", + ) + + @staticmethod + def _matches_python_version(venv_dir: Path, version: str) -> bool: + pyvenv_cfg = venv_dir / "pyvenv.cfg" + if not pyvenv_cfg.exists(): + return False + try: + content = pyvenv_cfg.read_text(encoding="utf-8") + except OSError: + return False + match = re.search(r"version\s*=\s*(\d+\.\d+)\.\d+", content, re.IGNORECASE) + return match is not None and match.group(1) == version + + +class WorkerRuntime: + def __init__( + self, + plugin: PluginSpec, + server: JSONRPCServer, + repo_root: Path, + env_manager: PluginEnvironmentManager, + ) -> None: + self.plugin = plugin + self.server = server + self.repo_root = repo_root.resolve() + self.env_manager = env_manager + self.rpc_helper = RPCRequestHelper() + self.client: StdioClient | None = None + self.raw_handshake: dict[str, Any] = {} + self.handlers: list[StarHandlerMetadata] = [] + self._context_requests: dict[str, str] = {} + self._forwarded_call_ids: set[str] = set() + + async def start(self) -> None: + python_path = self.env_manager.prepare_environment(self.plugin) + repo_src_dir = str(self.repo_root / "src") + env = os.environ.copy() + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + f"{repo_src_dir}{os.pathsep}{existing_pythonpath}" + if existing_pythonpath + else repo_src_dir + ) + + self.client = StdioClient( + command=[ + str(python_path), + "-m", + "astrbot_sdk", + "worker", + "--plugin-dir", + str(self.plugin.plugin_dir), + ], + cwd=str(self.plugin.plugin_dir), + env=env, + ) + self.client.set_message_handler(self._handle_message) + await self.client.start() + + response = await asyncio.wait_for( + self.rpc_helper.call_rpc( + self.client, + HandshakeRequest( + jsonrpc="2.0", + id=self.rpc_helper._generate_request_id(), + method="handshake", + ), + ), + timeout=15.0, + ) + if not isinstance(response, JSONRPCSuccessResponse): + raise RuntimeError(f"Handshake failed for plugin {self.plugin.name}") + + result = response.result + if not isinstance(result, dict): + raise RuntimeError(f"Invalid handshake payload for plugin {self.plugin.name}") + + self.raw_handshake = result + self.handlers = self._parse_handlers(result) + + async def stop(self) -> None: + if self.client is not None: + await self.client.stop() + + async def forward_call_handler(self, request: JSONRPCRequest) -> None: + if self.client is None: + raise RuntimeError(f"Worker for {self.plugin.name} is not running") + if request.id is not None: + self._forwarded_call_ids.add(str(request.id)) + await self.client.send_message(request) + + async def handle_context_response( + self, message: JSONRPCSuccessResponse | JSONRPCErrorResponse + ) -> bool: + message_id = str(message.id) + worker_request_id = self._context_requests.pop(message_id, None) + if worker_request_id is None: + return False + if self.client is None: + return True + + if isinstance(message, JSONRPCSuccessResponse): + await self.client.send_message( + JSONRPCSuccessResponse( + jsonrpc="2.0", + id=worker_request_id, + result=message.result, + ) + ) + else: + await self.client.send_message( + JSONRPCErrorResponse( + jsonrpc="2.0", + id=worker_request_id, + error=message.error, + ) + ) + return True + + async def _handle_message(self, message: JSONRPCMessage) -> None: + if isinstance(message, (JSONRPCSuccessResponse, JSONRPCErrorResponse)): + if message.id in self.rpc_helper.pending_requests: + self.rpc_helper.resolve_pending_request(message) + return + + if message.id is not None and str(message.id) in self._forwarded_call_ids: + self._forwarded_call_ids.discard(str(message.id)) + await self.server.send_message(message) + return + + if not isinstance(message, JSONRPCRequest): + return + + if message.method in [ + "handler_stream_start", + "handler_stream_update", + "handler_stream_end", + ]: + await self.server.send_message(message) + return + + if message.method != "call_context_function": + logger.warning( + f"Worker {self.plugin.name} sent unknown request: {message.method}" + ) + return + + supervisor_request_id = ( + f"ctx:{self.plugin.name}:{message.id}" + if message.id is not None + else f"ctx:{self.plugin.name}:none" + ) + self._context_requests[supervisor_request_id] = str(message.id) + await self.server.send_message( + JSONRPCRequest( + jsonrpc="2.0", + id=supervisor_request_id, + method=message.method, + params=message.params, + ) + ) + + @staticmethod + def _parse_handlers(handshake_payload: dict[str, Any]) -> list[StarHandlerMetadata]: + handlers: list[StarHandlerMetadata] = [] + + def _placeholder_handler(*args, **kwargs): + raise NotImplementedError("Worker supervisor does not execute handlers") + + for star_info in handshake_payload.values(): + handlers_data = star_info.get("handlers") or [] + for handler_data in handlers_data: + handlers.append( + 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, + event_filters=[], + desc=handler_data.get("desc", ""), + extras_configs=handler_data.get("extras_configs", {}), + ) + ) + return handlers + + +class SupervisorRuntime: + def __init__( + self, + server: JSONRPCServer, + plugins_dir: Path, + *, + env_manager: PluginEnvironmentManager | None = None, + worker_factory: Callable[ + [PluginSpec, JSONRPCServer, Path, PluginEnvironmentManager], WorkerRuntime + ] + | None = None, + ) -> None: + self.server = server + self.plugins_dir = plugins_dir.resolve() + self.repo_root = Path(__file__).resolve().parents[3] + self.env_manager = env_manager or PluginEnvironmentManager(self.repo_root) + self.worker_factory = worker_factory or WorkerRuntime + self.loaded_plugins: list[str] = [] + self.skipped_plugins: dict[str, str] = {} + self._workers_by_name: dict[str, WorkerRuntime] = {} + self._handler_to_worker: dict[str, WorkerRuntime] = {} + + async def start(self) -> None: + discovery = discover_plugins(self.plugins_dir) + self.skipped_plugins = dict(discovery.skipped_plugins) + + for plugin in discovery.plugins: + worker = self.worker_factory( + plugin, + self.server, + self.repo_root, + self.env_manager, + ) + try: + await worker.start() + except Exception as exc: + self.skipped_plugins[plugin.name] = str(exc) + logger.error(f"Failed to start worker for {plugin.name}: {exc}") + await worker.stop() + continue + + duplicate_handlers = [ + handler.handler_full_name + for handler in worker.handlers + if handler.handler_full_name in self._handler_to_worker + ] + if duplicate_handlers: + self.skipped_plugins[plugin.name] = ( + f"duplicate handlers: {', '.join(sorted(duplicate_handlers))}" + ) + await worker.stop() + continue + + self._workers_by_name[plugin.name] = worker + self.loaded_plugins.append(plugin.name) + for handler in worker.handlers: + self._handler_to_worker[handler.handler_full_name] = worker + + self.loaded_plugins.sort() + self.server.set_message_handler(self._handle_message) + await self.server.start() + self._log_startup_summary() + + async def stop(self) -> None: + for worker in list(self._workers_by_name.values()): + await worker.stop() + await self.server.stop() + + 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)) + return + if message.method == "call_handler": + await self._route_call_handler(message) + return + logger.warning(f"Unknown method from core: {message.method}") + return + + for worker in self._workers_by_name.values(): + if await worker.handle_context_response(message): + return + + logger.warning(f"Received response for unknown request id: {message.id}") + + def _build_handshake_response( + self, request_id: str | None + ) -> JSONRPCSuccessResponse: + payload: dict[str, Any] = {} + for worker in self._workers_by_name.values(): + payload.update(worker.raw_handshake) + return JSONRPCSuccessResponse( + jsonrpc="2.0", + id=request_id, + result=payload, + ) + + async def _route_call_handler(self, message: JSONRPCRequest) -> None: + try: + params = CallHandlerRequest.Params.model_validate(message.params) + except Exception as exc: + await self.server.send_message( + JSONRPCErrorResponse( + jsonrpc="2.0", + id=message.id, + error=JSONRPCErrorData(code=-32602, message=f"Invalid params: {exc}"), + ) + ) + return + + worker = self._handler_to_worker.get(params.handler_full_name) + if worker is None: + await self.server.send_message( + JSONRPCErrorResponse( + jsonrpc="2.0", + id=message.id, + error=JSONRPCErrorData( + code=-32601, + message=f"Handler not found: {params.handler_full_name}", + ), + ) + ) + return + + await worker.forward_call_handler(message) + + def _log_startup_summary(self) -> None: + loaded = ", ".join(self.loaded_plugins) if self.loaded_plugins else "none" + logger.info(f"Loaded plugins: {loaded}") + if not self.skipped_plugins: + logger.info("Skipped plugins: none") + return + for plugin_name, reason in sorted(self.skipped_plugins.items()): + logger.warning(f"Skipped plugin {plugin_name}: {reason}") diff --git a/src/astrbot_sdk/tests/test_supervisor.py b/src/astrbot_sdk/tests/test_supervisor.py new file mode 100644 index 000000000..7e9cc082b --- /dev/null +++ b/src/astrbot_sdk/tests/test_supervisor.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import asyncio +import tempfile +import unittest +from pathlib import Path +from typing import Any + +import yaml + +from astrbot_sdk.runtime.rpc.jsonrpc import ( + JSONRPCRequest, + JSONRPCSuccessResponse, +) +from astrbot_sdk.runtime.stars.registry import EventType, StarHandlerMetadata +from astrbot_sdk.runtime.supervisor import ( + PluginEnvironmentManager, + PluginSpec, + SupervisorRuntime, + WorkerRuntime, + discover_plugins, +) +from astrbot_sdk.runtime.types import CallHandlerRequest + + +def write_plugin( + root: Path, + folder_name: str, + *, + plugin_name: str | None = None, + python_version: str | None = "3.12", + include_requirements: bool = True, +) -> Path: + plugin_dir = root / folder_name + commands_dir = plugin_dir / "commands" + commands_dir.mkdir(parents=True, exist_ok=True) + (commands_dir / "__init__.py").write_text("", encoding="utf-8") + + manifest: dict[str, Any] = { + "_schema_version": 2, + "name": plugin_name or folder_name, + "display_name": folder_name, + "desc": "test plugin", + "author": "tester", + "version": "0.1.0", + "components": [ + { + "class": "commands.sample:SampleCommand", + "type": "command", + "name": "hello", + "description": "hello", + } + ], + } + if python_version is not None: + manifest["runtime"] = {"python": python_version} + + (plugin_dir / "plugin.yaml").write_text( + yaml.safe_dump(manifest, sort_keys=False), + encoding="utf-8", + ) + if include_requirements: + (plugin_dir / "requirements.txt").write_text("", encoding="utf-8") + return plugin_dir + + +class FakeServer: + def __init__(self) -> None: + self.handler = None + self.sent_messages: list[Any] = [] + + def set_message_handler(self, handler) -> None: + self.handler = handler + + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + async def send_message(self, message) -> None: + self.sent_messages.append(message) + + +class FakeEnvManager(PluginEnvironmentManager): + def __init__(self) -> None: + self.prepared: list[str] = [] + + def prepare_environment(self, plugin: PluginSpec) -> Path: + self.prepared.append(plugin.name) + return Path("/tmp/fake-python") + + +class FakeWorkerRuntime(WorkerRuntime): + def __init__( + self, + plugin: PluginSpec, + server, + repo_root: Path, + env_manager: PluginEnvironmentManager, + ) -> None: + self.plugin = plugin + self.server = server + self.repo_root = repo_root + self.env_manager = env_manager + self.raw_handshake: dict[str, Any] = {} + self.handlers: list[StarHandlerMetadata] = [] + self.forwarded_requests: list[JSONRPCRequest] = [] + self.received_context_responses: list[Any] = [] + self.stopped = False + + async def start(self) -> None: + handler_full_name = ( + f"commands.{self.plugin.name}:SampleCommand.handle_{self.plugin.name}" + ) + self.raw_handshake = { + f"{self.plugin.name}.main": { + "name": self.plugin.name, + "author": "tester", + "desc": "test plugin", + "version": "0.1.0", + "repo": None, + "module_path": f"{self.plugin.name}.main", + "root_dir_name": self.plugin.plugin_dir.name, + "reserved": False, + "activated": True, + "config": None, + "star_handler_full_names": [handler_full_name], + "display_name": self.plugin.name, + "logo_path": None, + "handlers": [ + { + "event_type": EventType.AdapterMessageEvent.value, + "handler_full_name": handler_full_name, + "handler_name": f"handle_{self.plugin.name}", + "handler_module_path": f"commands.{self.plugin.name}", + "desc": "", + "extras_configs": {}, + } + ], + } + } + self.handlers = [ + StarHandlerMetadata( + event_type=EventType.AdapterMessageEvent, + handler_full_name=handler_full_name, + handler_name=f"handle_{self.plugin.name}", + handler_module_path=f"commands.{self.plugin.name}", + handler=lambda *args, **kwargs: None, + event_filters=[], + ) + ] + + async def stop(self) -> None: + self.stopped = True + + async def forward_call_handler(self, request: JSONRPCRequest) -> None: + self.forwarded_requests.append(request) + handler_full_name = self.handlers[0].handler_full_name + await self.server.send_message( + JSONRPCRequest( + jsonrpc="2.0", + method="handler_stream_start", + params={ + "id": request.id, + "handler_full_name": handler_full_name, + }, + ) + ) + await self.server.send_message( + JSONRPCRequest( + jsonrpc="2.0", + method="handler_stream_update", + params={ + "id": request.id, + "handler_full_name": handler_full_name, + "data": {"plugin": self.plugin.name}, + }, + ) + ) + await self.server.send_message( + JSONRPCRequest( + jsonrpc="2.0", + method="handler_stream_end", + params={ + "id": request.id, + "handler_full_name": handler_full_name, + }, + ) + ) + await self.server.send_message( + JSONRPCSuccessResponse( + jsonrpc="2.0", + id=request.id, + result={"handled_by": self.plugin.name}, + ) + ) + + async def handle_context_response(self, message) -> bool: + if message.id != f"ctx:{self.plugin.name}:1": + return False + self.received_context_responses.append(message) + return True + + +class DiscoverPluginsTest(unittest.TestCase): + def test_discover_plugins_requires_runtime_python(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + write_plugin(root, "plugin_one", plugin_name="plugin_one") + write_plugin( + root, + "plugin_two", + plugin_name="plugin_two", + python_version=None, + ) + write_plugin( + root, + "plugin_three", + plugin_name="plugin_three", + include_requirements=False, + ) + + discovery = discover_plugins(root) + + self.assertEqual([plugin.name for plugin in discovery.plugins], ["plugin_one"]) + self.assertIn("plugin_two", discovery.skipped_plugins) + self.assertIn("plugin_three", discovery.skipped_plugins) + + +class SupervisorRuntimeTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.temp_dir = tempfile.TemporaryDirectory() + self.plugins_dir = Path(self.temp_dir.name) + write_plugin(self.plugins_dir, "plugin_one", plugin_name="plugin_one") + write_plugin(self.plugins_dir, "plugin_two", plugin_name="plugin_two") + self.server = FakeServer() + + async def asyncTearDown(self) -> None: + self.temp_dir.cleanup() + + async def test_handshake_aggregates_workers_and_routes_call_handler(self) -> None: + runtime = SupervisorRuntime( + server=self.server, + plugins_dir=self.plugins_dir, + env_manager=FakeEnvManager(), + worker_factory=FakeWorkerRuntime, + ) + await runtime.start() + + await self.server.handler( + JSONRPCRequest(jsonrpc="2.0", id="handshake-1", method="handshake") + ) + handshake_response = self.server.sent_messages[-1] + self.assertIsInstance(handshake_response, JSONRPCSuccessResponse) + self.assertEqual( + sorted(handshake_response.result.keys()), + ["plugin_one.main", "plugin_two.main"], + ) + + handler_full_name = ( + "commands.plugin_two:SampleCommand.handle_plugin_two" + ) + await self.server.handler( + CallHandlerRequest( + jsonrpc="2.0", + id="call-1", + method="call_handler", + params=CallHandlerRequest.Params( + handler_full_name=handler_full_name, + event={ + "message_str": "hello", + "message_obj": { + "type": "FriendMessage", + "self_id": "bot", + "session_id": "session", + "message_id": "message-id", + "sender": {"user_id": "user-1", "nickname": "User 1"}, + "message": [], + "message_str": "hello", + "raw_message": {}, + "timestamp": 0, + }, + "platform_meta": { + "name": "fake", + "description": "fake", + "id": "fake-1", + }, + "session_id": "session", + "is_at_or_wake_command": True, + }, + args={}, + ), + ) + ) + + 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: + runtime = SupervisorRuntime( + server=self.server, + plugins_dir=self.plugins_dir, + env_manager=FakeEnvManager(), + worker_factory=FakeWorkerRuntime, + ) + await runtime.start() + + await self.server.handler( + JSONRPCSuccessResponse( + jsonrpc="2.0", + id="ctx:plugin_one:1", + result={"data": "ok"}, + ) + ) + + worker = runtime._workers_by_name["plugin_one"] + self.assertEqual(len(worker.received_context_responses), 1) + await runtime.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/test_plugin/commands/hello.py b/test_plugin/commands/hello.py index 67c25f517..183f57331 100644 --- a/test_plugin/commands/hello.py +++ b/test_plugin/commands/hello.py @@ -1,6 +1,7 @@ from astrbot_sdk.api.components.command import CommandComponent from astrbot_sdk.api.event import AstrMessageEvent, filter from astrbot_sdk.api.star.context import Context +from loguru import logger class HelloCommand(CommandComponent): @@ -10,7 +11,7 @@ class HelloCommand(CommandComponent): @filter.command("hello") async def hello(self, event: AstrMessageEvent): ret = await self.context.conversation_manager.new_conversation("hello") - print(f"New conversation created: {ret}") + logger.info(f"New conversation created: {ret}") yield event.plain_result(f"Hello, Astrbot! Created conversation ID: {ret}") yield event.plain_result("Hello, Astrbot!") yield event.plain_result("Hello again, Astrbot!") diff --git a/test_plugin/plugin.yaml b/test_plugin/plugin.yaml index 9cbbad136..25f4f533e 100644 --- a/test_plugin/plugin.yaml +++ b/test_plugin/plugin.yaml @@ -4,6 +4,8 @@ display_name: HelloWorld 插件 desc: 一个简单的问候插件示例 author: Soulter version: 0.1.0 +runtime: + python: "3.12" components: # 组件列表，将支持自动生成 - class: commands.hello:HelloCommand type: command diff --git a/test_plugin/requirements.txt b/test_plugin/requirements.txt new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/test_plugin/requirements.txt @@ -0,0 +1 @@ +