hunjing/backend/app/services/mcp.py

"""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,
    )