初始化提交

backend/app/services/mcp.py

@@ -0,0 +1,472 @@
+"""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,
+    )