"""MCP server 注册服务(支持 streamable-http / stdio 两种传输协议)。 注册流程: - streamable-http:POST 时实时拉 tools/list → 存 tools_json + register ToolSpec - stdio:POST 时只存 DB,tools_json 留空;前端调用 /start 后才启进程 → 拉 tools/list → register 启动加载(lifespan): - streamable-http 按 tools_json 缓存 register ToolSpec - stdio 跳过(tools_json 空时无 ToolSpec 可注册) """ from __future__ import annotations import json from app.application.tool_runtime.mcp import make_mcp_tool_specs from app.application.tool_runtime.runtime import ToolRuntime from app.core.errors import AppError from app.core.ids import new_mcp_server_id from app.db.session import AsyncSessionLocal from app.integrations.mcp.process_manager import ProcessManager from app.integrations.mcp.transports import get_transport from app.models.mcp_servers import McpServer from app.repositories import mcp_servers as mcp_repo from app.repositories.json_utils import dumps from app.schemas.mcp import McpServerResponse, McpServerStatusResponse from app.services.audit import record_audit_event async def register_mcp_server( *, name: str, transport: str = "streamable-http", endpoint_url: str | None = None, auth_token: str | None = None, command: str | None = None, args_json: str | None = None, env_json: str | None = None, cwd: str | None = None, status: str = "enabled", runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerResponse: async with AsyncSessionLocal() as session: existing = await mcp_repo.get_server_by_name(session, name) if existing is not None: raise AppError("conflict", f"MCP server「{name}」已存在。") server = McpServer( id=new_mcp_server_id(), name=name, transport=transport, endpoint_url=endpoint_url, auth_token=auth_token, command=command, args_json=args_json, env_json=env_json, cwd=cwd, status=status, tools_json="[]", ) # streamable-http:注册时同步拉取工具清单 if transport == "streamable-http" and endpoint_url: transport_instance = get_transport(server) tools = await transport_instance.list_tools() server.tools_json = dumps(tools) if runtime and runtime.http_client is not None: for spec in make_mcp_tool_specs(server, tools, runtime.http_client): runtime.register(spec) if process_manager is not None: process_manager.register_transport(server.id, transport_instance) # stdio:注册时启动子进程并拉取工具清单 elif transport == "stdio": if process_manager is not None: transport_instance = get_transport(server) process_manager.register_transport(server.id, transport_instance) await transport_instance.start() tools = await transport_instance.list_tools() server.tools_json = dumps(tools) if runtime is not None and runtime.http_client is not None: for spec in make_mcp_tool_specs( server, tools, runtime.http_client, transport=transport_instance ): runtime.register(spec) await mcp_repo.add_server(session, server) await record_audit_event( session, action="mcp.server.registered", actor="user", input_summary=name[:200], resource_type="mcp_server", resource_id=server.id, ) await session.commit() await session.refresh(server) return _to_response(server) async def list_mcp_servers(*, status: str | None = None) -> list[McpServerResponse]: async with AsyncSessionLocal() as session: servers = await mcp_repo.list_servers(session, status=status) return [_to_response(s) for s in servers] async def delete_mcp_server( server_id: str, *, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> None: async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server is None: raise AppError("not_found", "MCP server 不存在。") # 先 stop(如果是 stdio 且正在运行) if process_manager is not None: transport = process_manager.get(server_id) if transport and transport.is_running: await transport.stop() process_manager.unregister(server_id) # unregister ToolSpec if runtime is not None and server.tools_json: tools = json.loads(server.tools_json) if server.tools_json else [] for tool in tools: tool_name = tool.get("name") if isinstance(tool, dict) else None if tool_name: runtime.unregister(f"mcp_{server.name}_{tool_name}") name = server.name await session.delete(server) await record_audit_event( session, action="mcp.server.deleted", actor="user", input_summary=name[:200], resource_type="mcp_server", resource_id=server_id, ) await session.commit() async def start_mcp_server( server_id: str, *, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerStatusResponse: """启动 stdio MCP server 子进程,拉取工具清单并注册。""" if process_manager is None: raise AppError("invalid_request", "进程管理器未就绪") transport = process_manager.get(server_id) if transport is None: raise AppError("not_found", f"MCP server {server_id} 未找到") if isinstance(transport, StreamableHttpTransport): raise AppError("invalid_request", "streamable-http 无需启动") await transport.start() # 拉取工具清单 tools = await transport.list_tools() async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server is None: await transport.stop() raise AppError("not_found", "MCP server 不存在") server.tools_json = dumps(tools) await session.commit() # 注册 ToolSpec if runtime is not None and runtime.http_client is not None: from app.application.tool_runtime.mcp import make_mcp_tool_specs async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server: for spec in make_mcp_tool_specs(server, tools, runtime.http_client): runtime.register(spec) status = await process_manager.status(server_id) return McpServerStatusResponse(**status) async def stop_mcp_server( server_id: str, *, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerStatusResponse: """停止 stdio MCP server 子进程并注销 ToolSpec。""" if process_manager is None: raise AppError("invalid_request", "进程管理器未就绪") transport = process_manager.get(server_id) if transport is None: raise AppError("not_found", f"MCP server {server_id} 未找到") if isinstance(transport, StreamableHttpTransport): raise AppError("invalid_request", "streamable-http 无需停止") await transport.stop() # unregister ToolSpec if runtime is not None: async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server and server.tools_json: tools = json.loads(server.tools_json) for tool in tools: name = tool.get("name") if isinstance(tool, dict) else None if name: runtime.unregister(f"mcp_{server.name}_{name}") status = await process_manager.status(server_id) return McpServerStatusResponse(**status) async def restart_mcp_server( server_id: str, *, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerStatusResponse: """重启 stdio MCP server。""" await stop_mcp_server(server_id, runtime=runtime, process_manager=process_manager) return await start_mcp_server(server_id, runtime=runtime, process_manager=process_manager) async def refresh_mcp_server_tools( server_id: str, *, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerResponse: """重新拉取 MCP server 工具清单,更新缓存并注册到运行时。""" async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server is None: raise AppError("not_found", "MCP server 不存在。") # 注销旧工具 old_tools = json.loads(server.tools_json) if server.tools_json else [] if runtime is not None: for tool in old_tools: tool_name = tool.get("name") if isinstance(tool, dict) else None if tool_name: runtime.unregister(f"mcp_{server.name}_{tool_name}") # 重新拉取工具清单(优先用 process_manager 中的传输实例,确保 stdio 进程已启动) transport_instance = None if process_manager is not None: transport_instance = process_manager.get(server_id) if transport_instance is None: transport_instance = get_transport(server) tools = await transport_instance.list_tools() server.tools_json = dumps(tools) # 注册新工具 if runtime is not None and runtime.http_client is not None: for spec in make_mcp_tool_specs(server, tools, runtime.http_client): runtime.register(spec) await record_audit_event( session, action="mcp.server.refreshed", actor="user", input_summary=server.name[:200], resource_type="mcp_server", resource_id=server.id, ) await session.commit() await session.refresh(server) return _to_response(server) def _tool_names_from_json(tools_json: str) -> list[str]: """从 tools_json JSON 字符串提取工具名列表。""" import json try: tools = json.loads(tools_json) if tools_json else [] except (ValueError, TypeError): return [] return [t.get("name") for t in tools if isinstance(t, dict) and t.get("name")] async def update_mcp_server( server_id: str, *, status: str | None = None, endpoint_url: str | None = None, auth_token: str | None = None, command: str | None = None, args_json: str | None = None, env_json: str | None = None, cwd: str | None = None, runtime: ToolRuntime | None = None, process_manager: ProcessManager | None = None, ) -> McpServerResponse: """更新 MCP server 字段(部分更新,不传的字段保持不变)。 更新 endpoint_url 时自动重新拉取工具清单(streamable-http 类型)。 """ async with AsyncSessionLocal() as session: server = await mcp_repo.get_server(session, server_id) if server is None: raise AppError("not_found", "MCP server 不存在。") # 记录旧的工具清单,以便后续注销 old_tools_json = server.tools_json if status is not None: old_status = server.status server.status = status # 启用→禁用:注销运行时中的工具 if ( status == "disabled" and old_status == "enabled" and runtime is not None and old_tools_json ): for tool_name in _tool_names_from_json(old_tools_json): runtime.unregister(f"mcp_{server.name}_{tool_name}") # 禁用→启用(HTTP):重新拉取工具并注册 if ( status == "enabled" and old_status == "disabled" and server.transport == "streamable-http" and server.endpoint_url ): from app.application.tool_runtime.mcp import make_mcp_tool_specs from app.integrations.mcp.transports import get_transport transport_instance = get_transport(server) new_tools = await transport_instance.list_tools() server.tools_json = dumps(new_tools) if runtime is not None and runtime.http_client is not None: for spec in make_mcp_tool_specs(server, new_tools, runtime.http_client): runtime.register(spec) endpoint_changed = endpoint_url is not None and endpoint_url != server.endpoint_url if endpoint_url is not None: server.endpoint_url = endpoint_url if auth_token is not None: server.auth_token = auth_token if command is not None: server.command = command if args_json is not None: server.args_json = args_json if env_json is not None: server.env_json = env_json if cwd is not None: server.cwd = cwd # streamable-http 且 endpoint 变化时:重新拉取工具清单并注册 if endpoint_changed and server.transport == "streamable-http" and server.endpoint_url: from app.application.tool_runtime.mcp import make_mcp_tool_specs from app.integrations.mcp.transports import get_transport # 注销旧工具 if runtime is not None and old_tools_json: old_tools = json.loads(old_tools_json) if old_tools_json else [] for tool in old_tools: tool_name = tool.get("name") if isinstance(tool, dict) else None if tool_name: runtime.unregister(f"mcp_{server.name}_{tool_name}") # 拉取新工具(优先用 process_manager 中的传输实例) transport_instance = None if process_manager is not None: transport_instance = process_manager.get(server_id) if transport_instance is None: transport_instance = get_transport(server) new_tools = await transport_instance.list_tools() server.tools_json = dumps(new_tools) # 注册新工具 if runtime is not None and runtime.http_client is not None: for spec in make_mcp_tool_specs(server, new_tools, runtime.http_client): runtime.register(spec) await record_audit_event( session, action="mcp.server.updated", actor="user", input_summary=server.name[:200], resource_type="mcp_server", resource_id=server.id, ) await session.commit() await session.refresh(server) return _to_response(server) async def get_mcp_server_status( server_id: str, *, process_manager: ProcessManager | None = None, ) -> McpServerStatusResponse: """查询 MCP server 进程状态。""" if process_manager is None: return McpServerStatusResponse( server_id=server_id, name="unknown", transport="unknown", status="unknown", ) status = await process_manager.status(server_id) return McpServerStatusResponse(**status) async def load_registered_mcp_servers( runtime: ToolRuntime, process_manager: ProcessManager | None = None ) -> None: """启动加载:按 transport 类型注册 MCP server。""" if runtime.http_client is None: return async with AsyncSessionLocal() as session: servers = await mcp_repo.list_servers(session, status="enabled") for server in servers: try: if server.transport == "streamable-http": tools = json.loads(server.tools_json) if server.tools_json else [] for spec in make_mcp_tool_specs(server, tools, runtime.http_client): runtime.register(spec) # 创建传输层实例但 no-op start if process_manager is not None: transport_instance = get_transport(server, runtime.http_client) process_manager.register_transport(server.id, transport_instance) elif server.transport == "stdio": # stdio 不自动启动,只注册传输层实例到管理器 if process_manager is not None: transport_instance = get_transport(server) process_manager.register_transport(server.id, transport_instance) except Exception: # noqa: BLE001 pass # 导入用于类型判断 from app.integrations.mcp.transports.http import StreamableHttpTransport # noqa: E402, F811 def _to_response(server: McpServer) -> McpServerResponse: try: tools = json.loads(server.tools_json) if server.tools_json else [] except (ValueError, TypeError): tools = [] return McpServerResponse( id=server.id, name=server.name, transport=server.transport, endpoint_url=server.endpoint_url, has_auth_token=bool(server.auth_token), command=server.command, args_json=server.args_json, env_json=server.env_json, cwd=server.cwd, status=server.status, tools=tools, created_at=server.created_at, updated_at=server.updated_at, )