473 lines
18 KiB
Python
473 lines
18 KiB
Python
"""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,
|
||
)
|