97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
"""MCP stdio 子进程生命周期管理。
|
||
|
||
全局单例，维护 server_id → TransportProtocol 的映射，
|
||
前端通过 API 控制 start/stop/restart/status。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from typing import Any
|
||
|
||
from app.integrations.mcp.transports.base import TransportProtocol
|
||
from app.integrations.mcp.transports.http import StreamableHttpTransport
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 模块级全局实例（避免依赖 request.app.state）
|
||
_default_instance: ProcessManager | None = None
|
||
|
||
|
||
def get_process_manager() -> ProcessManager | None:
|
||
"""返回默认的全局 ProcessManager 实例。"""
|
||
return _default_instance
|
||
|
||
|
||
def set_process_manager(manager: ProcessManager) -> None:
|
||
"""设置全局 ProcessManager 实例（由 lifespan 调用）。"""
|
||
global _default_instance
|
||
_default_instance = manager
|
||
|
||
|
||
class ProcessManager:
|
||
"""MCP 进程管理器（全局单例）。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._transports: dict[str, TransportProtocol] = {}
|
||
|
||
def register_transport(self, server_id: str, transport: TransportProtocol) -> None:
|
||
"""注册传输层实例（不自动 start）。"""
|
||
self._transports[server_id] = transport
|
||
|
||
def unregister(self, server_id: str) -> None:
|
||
"""注销传输层实例（不自动 stop，调用方需先 stop）。"""
|
||
self._transports.pop(server_id, None)
|
||
|
||
def get(self, server_id: str) -> TransportProtocol | None:
|
||
return self._transports.get(server_id)
|
||
|
||
async def start(self, server_id: str) -> None:
|
||
transport = self._transports.get(server_id)
|
||
if transport is None:
|
||
raise ValueError(f"server {server_id} 未注册")
|
||
if isinstance(transport, StreamableHttpTransport):
|
||
raise ValueError("streamable-http 无需启动")
|
||
await transport.start()
|
||
|
||
async def stop(self, server_id: str) -> None:
|
||
transport = self._transports.get(server_id)
|
||
if transport is None:
|
||
raise ValueError(f"server {server_id} 未注册")
|
||
if isinstance(transport, StreamableHttpTransport):
|
||
raise ValueError("streamable-http 无需停止")
|
||
await transport.stop()
|
||
|
||
async def restart(self, server_id: str) -> None:
|
||
transport = self._transports.get(server_id)
|
||
if transport is None:
|
||
raise ValueError(f"server {server_id} 未注册")
|
||
if isinstance(transport, StreamableHttpTransport):
|
||
raise ValueError("streamable-http 无需重启")
|
||
await transport.restart()
|
||
|
||
async def status(self, server_id: str) -> dict[str, Any]:
|
||
transport = self._transports.get(server_id)
|
||
if transport is None:
|
||
return {
|
||
"server_id": server_id,
|
||
"name": "unknown",
|
||
"transport": "unknown",
|
||
"status": "unregistered",
|
||
"pid": None,
|
||
"uptime_seconds": None,
|
||
}
|
||
return await transport.get_status()
|
||
|
||
async def shutdown_all(self) -> None:
|
||
"""停止所有 stdio 子进程（lifespan 关闭时调用）。"""
|
||
for sid, transport in list(self._transports.items()):
|
||
try:
|
||
await transport.stop()
|
||
except Exception: # noqa: BLE001
|
||
logger.warning(
|
||
"停止 MCP 进程失败",
|
||
extra={"server_id": sid},
|
||
)
|
||
self._transports.clear()
|