feat: 接入 MCP 外部工具服务,AI 可调用外部工具获取书外信息

- 后端新增 MCP 客户端(Streamable HTTP / 老式 SSE)与流式 agent 工具循环
- 新增 /api/mcp/servers 管理接口与 mcp_servers 表
- 对话中 OpenAI 兼容模型可调用启用服务器的工具,SSE 推送 tool 事件
- 前端新增 MCP 服务管理弹窗与侧栏入口,流式回答展示工具调用状态
- 修复用户消息气泡缺少删除按钮,并级联删除其名下回答版本

Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
2026-10-24 04:29:34 +08:00
co-authored by Claude Code
parent c7b0cfb92b
commit 8cd390b3ed
13 changed files with 1472 additions and 26 deletions
+42 -5
View File
@@ -14,6 +14,7 @@ from app.rag.engine import (
search_segments, build_chat_prompt, parse_citations,
estimate_tokens, match_chapters, _trim_messages,
)
from app.mcp.agent import agent_stream, build_mcp_context, execute_mcp_tool
logger = logging.getLogger(__name__)
@@ -164,10 +165,24 @@ async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession =
chapter_segments=chapter_segments,
)
total_tokens = estimate_tokens("".join(m["content"] for m in messages))
provider = await resolve_provider(db, model_name)
if provider is None:
raise HTTPException(400, f"模型 {model_name} 未配置,请先在页面「管理模型」中添加或选择可用模型")
# MCP 外部工具:仅当模型原生支持 function calling 时接入工具循环;
# 无 MCP 配置 / 单服务器失败时自动跳过,行为与改造前完全一致
mcp_ctx = await build_mcp_context(db) if provider.supports_tools else None
if mcp_ctx is not None:
system = messages[0]["content"]
system += (
"\n\n## 外部信息工具\n"
"当问题需要书中没有的外部信息(最新动态、时事、网络资料、行业数据等)时,"
f"可调用以下外部 MCP 工具获取:{', '.join(mcp_ctx.names)}\n"
"- 外部信息要说明来源与时效;回答主体仍以书中内容为准,并保留 [[ref:N]] 引用标注。\n"
)
messages[0]["content"] = system
total_tokens = estimate_tokens("".join(m["content"] for m in messages))
if total_tokens > provider.max_context * 0.8:
# 内容感知裁剪:保留 system(含检索原文 + 整章注入),从最旧历史开始丢
messages = _trim_messages(messages, int(provider.max_context * 0.8))
@@ -175,9 +190,25 @@ async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession =
async def stream_response():
full_content = ""
try:
async for token in provider.chat_stream(messages, temperature=0.7):
full_content += token
yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
if mcp_ctx is not None:
async for kind, payload in agent_stream(
messages,
provider,
mcp_ctx.ai_tools,
execute=lambda name, args: execute_mcp_tool(mcp_ctx, name, args),
describe=mcp_ctx.display_name,
):
if kind == "token":
full_content += payload
yield f"data: {json.dumps({'type': 'token', 'content': payload}, ensure_ascii=False)}\n\n"
elif kind == "tool":
yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'start'}, ensure_ascii=False)}\n\n"
elif kind == "tool_done":
yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'done'}, ensure_ascii=False)}\n\n"
else:
async for token in provider.chat_stream(messages, temperature=0.7):
full_content += token
yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
except Exception as e:
logger.exception("流式生成失败(conversation=%s)", conv_id)
yield f"data: {json.dumps({'type': 'error', 'content': str(e)}, ensure_ascii=False)}\n\n"
@@ -266,9 +297,15 @@ async def delete_message(msg_id: str, db: AsyncSession = Depends(get_db)):
msg = await db.get(Message, msg_id)
if not msg:
raise HTTPException(404, "消息不存在")
# 级联删除:该消息名下的所有回答版本(parent 指向它),避免残留孤儿版本
versions = (await db.execute(
select(Message).where(Message.parent_message_id == msg_id)
)).scalars().all()
for v in versions:
await db.delete(v)
await db.delete(msg)
await db.commit()
return {"deleted": msg_id}
return {"deleted": msg_id, "cascade": len(versions)}
@router.delete("/conversations/{conv_id}")
+146
View File
@@ -0,0 +1,146 @@
import json
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.database import get_db
from app.models.mcp_server import McpServer
from app.mcp.client import McpServerConfig, McpClientError, list_mcp_tools, invalidate_mcp_cache
router = APIRouter(prefix="/mcp", tags=["mcp"])
class McpServerIn(BaseModel):
name: str
transport: str = "streamable-http" # streamable-http | sse
url: str
headers: dict[str, str] = {}
enabled: bool = True
class McpServerUpdate(BaseModel):
name: str | None = None
transport: str | None = None
url: str | None = None
headers: dict[str, str] | None = None
enabled: bool | None = None
class McpTestRequest(BaseModel):
"""「测试连接」用临时配置(不落库)。"""
name: str = "test"
transport: str = "streamable-http"
url: str
headers: dict[str, str] = {}
def _to_out(m: McpServer) -> dict:
return {
"id": m.id,
"name": m.name,
"transport": m.transport,
"url": m.url,
"headers": m.headers_dict(),
"enabled": m.enabled,
"created_at": m.created_at.isoformat() if m.created_at else "",
}
@router.get("/servers")
async def list_servers(db: AsyncSession = Depends(get_db)):
result = await db.execute(select(McpServer).order_by(McpServer.created_at))
return [_to_out(m) for m in result.scalars().all()]
@router.post("/servers")
async def add_server(req: McpServerIn, db: AsyncSession = Depends(get_db)):
name = (req.name or "").strip()
url = (req.url or "").strip()
if not name:
raise HTTPException(400, "服务器名称不能为空")
if not url:
raise HTTPException(400, "服务器 URL 不能为空")
if req.transport not in ("streamable-http", "sse"):
raise HTTPException(400, "不支持的传输协议")
row = McpServer(
name=name,
transport=req.transport,
url=url,
headers=json.dumps(req.headers or {}, ensure_ascii=False),
enabled=req.enabled,
)
db.add(row)
await db.commit()
await db.refresh(row)
return _to_out(row)
@router.put("/servers/{server_id}")
async def update_server(server_id: str, req: McpServerUpdate, db: AsyncSession = Depends(get_db)):
row = await db.get(McpServer, server_id)
if not row:
raise HTTPException(404, "服务器不存在")
if req.name is not None:
name = req.name.strip()
if not name:
raise HTTPException(400, "服务器名称不能为空")
row.name = name
if req.transport is not None:
if req.transport not in ("streamable-http", "sse"):
raise HTTPException(400, "不支持的传输协议")
row.transport = req.transport
if req.url is not None:
url = req.url.strip()
if not url:
raise HTTPException(400, "服务器 URL 不能为空")
row.url = url
if req.headers is not None:
row.headers = json.dumps(req.headers, ensure_ascii=False)
if req.enabled is not None:
row.enabled = req.enabled
await db.commit()
await db.refresh(row)
await invalidate_mcp_cache(server_id) # 配置变更 → 失效旧会话 / 工具缓存
return _to_out(row)
@router.delete("/servers/{server_id}")
async def delete_server(server_id: str, db: AsyncSession = Depends(get_db)):
row = await db.get(McpServer, server_id)
if not row:
raise HTTPException(404, "服务器不存在")
await db.delete(row)
await db.commit()
await invalidate_mcp_cache(server_id)
return {"deleted": server_id}
@router.post("/servers/test")
async def test_server(req: McpTestRequest, db: AsyncSession = Depends(get_db)):
"""用临时配置连一次 MCP 服务,拉取工具列表验证连通性(不落库)。"""
url = (req.url or "").strip()
if not url:
raise HTTPException(400, "服务器 URL 不能为空")
if req.transport not in ("streamable-http", "sse"):
raise HTTPException(400, "不支持的传输协议")
cfg = McpServerConfig(
id=f"test-{hash(url)}",
name=(req.name or "test").strip() or "test",
transport=req.transport,
url=url,
headers=req.headers or {},
enabled=True,
)
try:
tools = await list_mcp_tools(cfg, fresh=True)
except McpClientError as e:
return {"ok": False, "error": str(e)}
except Exception as e:
return {"ok": False, "error": f"{type(e).__name__}: {e}"}
return {
"ok": True,
"tools": [{"name": t.name, "description": t.description or "", "input_schema": t.input_schema} for t in tools],
}
+39 -4
View File
@@ -18,24 +18,30 @@ class ChatResponse:
class LLMProvider(ABC):
name: str = "base"
max_context: int = 128_000
# 是否原生支持 function calling(OpenAI 兼容协议);Anthropic 暂未接入工具循环
supports_tools: bool = False
@abstractmethod
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
...
@abstractmethod
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
...
class OpenAIProvider(LLMProvider):
name = "openai"
max_context = 128_000
supports_tools = True
def __init__(self, api_key: str | None = None, base_url: str | None = None, model: str | None = None):
self.api_key = api_key if api_key is not None else os.getenv("OPENAI_API_KEY", "")
self.base_url = base_url if base_url is not None else os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
self.model = model if model is not None else os.getenv("OPENAI_MODEL", "gpt-4o")
# 流式工具循环的工作态:chat_stream 结束时填充 tool_calls / last_content
self.tool_calls: list[dict] = []
self.last_content: str = ""
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
import httpx
@@ -56,15 +62,27 @@ class OpenAIProvider(LLMProvider):
completion_tokens=data.get("usage", {}).get("completion_tokens", 0),
)
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
import httpx
payload = {
"model": self.model,
"messages": messages,
"temperature": temperature,
"stream": True,
}
if tools:
payload["tools"] = tools
# 工作态复位:上一轮(若有)的 tool_calls 不残留
self.tool_calls = []
self.last_content = ""
async with httpx.AsyncClient(timeout=300) as client:
async with client.stream(
"POST",
f"{self.base_url}/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
json={"model": self.model, "messages": messages, "temperature": temperature, "stream": True},
json=payload,
) as resp:
tool_calls_acc: dict[int, dict] = {}
async for line in resp.aiter_lines():
if line.startswith("data: "):
data_str = line[6:]
@@ -79,9 +97,25 @@ class OpenAIProvider(LLMProvider):
delta = choices[0].get("delta", {})
content = delta.get("content", "")
if content:
self.last_content += content
yield content
# 流式 tool_calls:每个 index 首个 delta 带 id/name,后续增量拼 arguments
for tc in delta.get("tool_calls") or []:
idx = tc.get("index", 0)
entry = tool_calls_acc.setdefault(idx, {
"id": "", "type": "function", "function": {"name": "", "arguments": ""},
})
if tc.get("id"):
entry["id"] = tc["id"]
fn = tc.get("function") or {}
if fn.get("name"):
entry["function"]["name"] = fn["name"]
if fn.get("arguments"):
entry["function"]["arguments"] += fn["arguments"]
except json.JSONDecodeError:
continue
if tool_calls_acc:
self.tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)]
class AnthropicProvider(LLMProvider):
@@ -135,8 +169,9 @@ class AnthropicProvider(LLMProvider):
completion_tokens=data.get("usage", {}).get("output_tokens", 0),
)
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
import httpx
# tools 参数为兼容 ABC 而存在:Anthropic 暂未接入工具循环(supports_tools=False)
system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
body = {
"model": self.model,
+3
View File
@@ -10,10 +10,12 @@ from app.core.database import init_db, Base
from app.api.books import router as books_router
from app.api.chat import router as chat_router
from app.api.models import router as models_router
from app.api.mcp import router as mcp_router
from app.models.book import Book, Chapter, Segment
from app.models.chat import Conversation, Message
from app.models.ai_model import AiModel
from app.models.mcp_server import McpServer
@asynccontextmanager
@@ -35,6 +37,7 @@ app.add_middleware(
app.include_router(books_router, prefix="/api")
app.include_router(chat_router, prefix="/api")
app.include_router(models_router, prefix="/api")
app.include_router(mcp_router, prefix="/api")
@app.get("/api/health")
View File
+197
View File
@@ -0,0 +1,197 @@
"""agent — MCP 工具与对话 agent 的桥接层
把启用中的 MCP 服务器暴露的工具转换成 OpenAI 兼容 function calling 格式,
并在对话时执行「流式工具循环」:模型文本即时流式输出,出现工具调用则执行并续一轮。
单服务器工具获取失败只跳过该服务器,绝不阻断对话。
"""
from __future__ import annotations
import json
import logging
import re
from typing import AsyncGenerator, Awaitable, Callable
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.llm.hub import LLMProvider
from app.models.mcp_server import McpServer
from app.mcp.client import McpClientError, McpServerConfig, McpTool, call_mcp_tool, list_mcp_tools
logger = logging.getLogger(__name__)
# MCP 工具名统一加前缀防冲突:mcp_<server短id>_<toolName>
MCP_PREFIX = "mcp_"
# OpenAI function name 限制:^[a-zA-Z0-9_-]{1,64}$
_INVALID_NAME_CHARS = re.compile(r"[^a-zA-Z0-9_-]")
MAX_TOOL_ROUNDS = 3
TOOL_RESULT_CHARS = 6000 # 单条工具结果回传模型的上限(保护上下文预算)
class McpContext:
"""一次对话请求内启用中的 MCP 服务器快照:工具列表 + 短 id → 服务器映射。"""
def __init__(self) -> None:
self.servers: list[McpServerConfig] = []
self.ai_tools: list[dict] = []
self.names: list[str] = [] # 完整 AI 工具名(用于 system 提示词)
self._exec_map: dict[str, McpServerConfig] = {} # 短 id → 服务器
@property
def has_tools(self) -> bool:
return bool(self.ai_tools)
def resolve(self, ai_tool_name: str) -> tuple[McpServerConfig, str]:
"""把 mcp_ 前缀的 AI 工具名还原为 (服务器, 工具名)。"""
rest = ai_tool_name[len(MCP_PREFIX):]
sep = rest.find("_")
if sep <= 0:
raise McpClientError(f"工具名格式异常:{ai_tool_name}")
short_id, tool_name = rest[:sep], rest[sep + 1:]
server = self._exec_map.get(short_id)
if server is None:
raise McpClientError(f"MCP 服务器不存在或已停用({short_id})")
return server, tool_name
def display_name(self, ai_tool_name: str) -> str:
"""前端展示用名称:<服务器名>/<工具名>。"""
try:
server, tool_name = self.resolve(ai_tool_name)
except McpClientError:
return ai_tool_name
return f"{server.name}/{tool_name}"
def _sanitize_tool_name(name: str) -> str:
return _INVALID_NAME_CHARS.sub("_", name)
def _short_server_id(server_id: str, taken: set[str]) -> str:
"""短 id:取 uuid 前缀,必要时加长保证本次请求内唯一(避免截断冲突)。"""
base = server_id.replace("-", "")
n = 8
short = base[:n]
while short in taken:
n += 4
short = base[:n]
taken.add(short)
return short
async def build_mcp_context(db: AsyncSession) -> McpContext | None:
"""收集启用中的 MCP 服务器及其工具;无服务器或全部失败返回 None。"""
result = await db.execute(
select(McpServer).where(McpServer.enabled.is_(True)).order_by(McpServer.created_at)
)
rows = list(result.scalars().all())
if not rows:
return None
ctx = McpContext()
taken: set[str] = set()
for row in rows:
cfg = McpServerConfig(
id=row.id, name=row.name, transport=row.transport,
url=row.url, headers=row.headers_dict(), enabled=row.enabled,
)
try:
tools = await list_mcp_tools(cfg)
except Exception as e:
logger.warning("[mcp] 服务器「%s」工具获取失败,已跳过:%s", row.name, e)
continue
if not tools:
continue
short = _short_server_id(row.id, taken)
ctx.servers.append(cfg)
ctx._exec_map[short] = cfg
for tool in tools:
ai_name = f"{MCP_PREFIX}{short}_{_sanitize_tool_name(tool.name)}"
if len(ai_name) > 64: # OpenAI function name 长度上限
ai_name = ai_name[:64]
ctx.ai_tools.append({
"type": "function",
"function": {
"name": ai_name,
"description": tool.description or f"外部 MCP 工具({row.name} / {tool.name})",
"parameters": tool.input_schema or {"type": "object", "properties": {}},
},
})
ctx.names.append(ai_name)
if not ctx.ai_tools:
return None
return ctx
async def execute_mcp_tool(ctx: McpContext, ai_tool_name: str, args: object) -> str:
"""执行一个 mcp_ 前缀工具;失败返回错误串,绝不断流。"""
try:
server, tool_name = ctx.resolve(ai_tool_name)
except McpClientError as e:
return f"[mcp] {e}"
if not server.enabled:
return f"[mcp] 服务器已停用:{server.name}"
try:
return await call_mcp_tool(server, tool_name, args)
except Exception as e:
msg = str(e) if not isinstance(e, McpClientError) else str(e)
return f"[mcp:{tool_name}] 调用失败:{msg}"
async def agent_stream(
messages: list[dict],
provider: LLMProvider,
tools: list[dict],
execute: Callable[[str, object], Awaitable[str]],
describe: Callable[[str], str] | None = None,
) -> AsyncGenerator[tuple[str, str], None]:
"""流式 agent 循环:每轮携带 tools 请求模型,出现工具调用则执行并续一轮。
产出 (kind, payload) 事件:
- ("token", 文本)
- ("tool", "<服务器>/<工具>") —— 工具调用开始
- ("tool_done", "<服务器>/<工具>") —— 工具执行完成
"""
for _round in range(MAX_TOOL_ROUNDS):
provider.tool_calls = []
provider.last_content = ""
async for token in provider.chat_stream(messages, temperature=0.7, tools=tools):
yield ("token", token)
calls: list[dict] = getattr(provider, "tool_calls", []) or []
if not calls:
return
assistant_msg: dict = {"role": "assistant", "content": getattr(provider, "last_content", "") or None}
for call in calls:
name = (call.get("function") or {}).get("name") or ""
call_id = call.get("id") or f"call_{name}_{_round}"
assistant_msg.setdefault("tool_calls", []).append({
"id": call_id,
"type": "function",
"function": {"name": name, "arguments": (call.get("function") or {}).get("arguments") or "{}"},
})
messages.append(assistant_msg)
for call in assistant_msg["tool_calls"]:
name = call["function"]["name"]
display = describe(name) if describe else name
try:
args = json.loads(call["function"]["arguments"] or "{}")
if not isinstance(args, dict):
args = {"value": args}
except (json.JSONDecodeError, TypeError):
args = {}
yield ("tool", display)
out = await execute(name, args)
yield ("tool_done", display)
messages.append({
"role": "tool",
"tool_call_id": call["id"],
"content": (out or "[工具无返回内容]")[:TOOL_RESULT_CHARS],
})
# 达到最大轮数仍未见收尾:用已收集上下文走无工具流式生成,保证必出最终答案
async for token in provider.chat_stream(messages, temperature=0.7):
yield ("token", token)
+523
View File
@@ -0,0 +1,523 @@
"""client — 最小 MCP 客户端(httpx 手写协议实现,零 MCP 第三方依赖)
支持两种远程传输:
- Streamable HTTP:POST JSON-RPC 到 cfg.url,会话经 `Mcp-Session-Id` 响应头下发后回带。
- 老式 SSE:GET /sse 端点,等 `event: endpoint` 事件拿到 messages URL,再 POST JSON-RPC,
响应经原 SSE 流按请求 id 关联回来。
会话 / 工具列表按服务器 id 内存缓存(TTL 5 分钟),配置变更时用 invalidate_mcp_cache 失效。
任何失败抛 McpClientError(消息为中文、可直接透传给用户),由上层捕获转错误串,绝不阻断对话。
"""
from __future__ import annotations
import asyncio
import itertools
import json
import time
from dataclasses import dataclass, field
from typing import AsyncGenerator, NamedTuple
import httpx
MCP_PROTOCOL_VERSION = "2025-06-18"
CLIENT_NAME = "bookmate"
CLIENT_VERSION = "1.0.0"
REQUEST_TIMEOUT_MS = 15_000
CACHE_TTL_S = 5 * 60
_rpc_seq = itertools.count(1)
class McpClientError(RuntimeError):
"""MCP 调用失败(消息可直接展示给用户)。"""
@dataclass
class McpTool:
"""服务器暴露的一个 MCP 工具(input_schema 本就是 JSON Schema)。"""
name: str
description: str = ""
input_schema: dict | None = None
@dataclass
class McpServerConfig:
"""一个已配置的 MCP 服务器(与 DB 行一一对应)。"""
id: str
name: str
transport: str = "streamable-http"
url: str = ""
headers: dict[str, str] = field(default_factory=dict)
enabled: bool = True
@dataclass
class McpCallResult:
text: str
is_error: bool = False
def next_rpc_id() -> int:
return next(_rpc_seq)
def _resolve_url(url: str, base: str) -> str:
"""把(可能相对)的端点 URL 基于服务器 URL 解析为绝对 URL。"""
from urllib.parse import urljoin
try:
return urljoin(base, url)
except Exception:
return url
def build_headers(cfg: McpServerConfig, extra: dict[str, str]) -> dict[str, str]:
"""合并自定义请求头 + 公共头;空值头会被过滤。"""
headers: dict[str, str] = {}
for k, v in cfg.headers.items():
if k and v:
headers[k] = v
for k, v in extra.items():
if v:
headers[k] = v
headers.setdefault("Accept", "application/json, text/event-stream")
return headers
def expect_result(cfg: McpServerConfig, method: str, res: dict | None) -> dict:
if not res:
raise McpClientError(f"[{cfg.name}] {method} 无响应")
if res.get("error"):
err = res["error"]
raise McpClientError(f"[{cfg.name}] {method} 失败:{err.get('message')}(code {err.get('code')})")
result = res.get("result")
if not isinstance(result, dict):
raise McpClientError(f"[{cfg.name}] {method} 返回异常 result")
return result
def _friendly_http_error(cfg: McpServerConfig, exc: httpx.HTTPError) -> McpClientError:
detail = getattr(exc, "request", None)
if isinstance(exc, httpx.TimeoutException):
return McpClientError(f"[{cfg.name}] 请求超时:{detail.url if detail else cfg.url}")
if isinstance(exc, httpx.ConnectError):
return McpClientError(f"[{cfg.name}] 连接失败:{detail.url if detail else cfg.url}")
return McpClientError(f"[{cfg.name}] 网络错误:{exc}")
# ------------------------------------------------------------------
# SSE 帧解析
# ------------------------------------------------------------------
class SseEvent(NamedTuple):
event: str
data: str
async def read_sse_events(resp: httpx.Response) -> AsyncGenerator[SseEvent, None]:
"""把响应体流解析为 SSE 事件序列(逐行、空行分帧)。"""
event_name = ""
data_lines: list[str] = []
async for line in resp.aiter_lines():
if line == "":
if data_lines:
yield SseEvent(event=event_name or "message", data="\n".join(data_lines))
event_name = ""
data_lines = []
elif line.startswith(":"):
continue
elif line.startswith("event:"):
event_name = line[6:].strip()
elif line.startswith("data:"):
data_lines.append(line[5:].lstrip())
if data_lines:
yield SseEvent(event=event_name or "message", data="\n".join(data_lines))
def parse_jsonrpc(data: str) -> dict | None:
try:
obj = json.loads(data)
if isinstance(obj, dict) and "jsonrpc" in obj:
return obj
except (json.JSONDecodeError, TypeError):
pass
return None
# ------------------------------------------------------------------
# Streamable HTTP 传输
# ------------------------------------------------------------------
class StreamableHttpTransport:
"""每个请求独立 POST,会话经 sessionId 头维持。"""
def __init__(self, cfg: McpServerConfig):
self.cfg = cfg
self.session_id: str | None = None
self.client = httpx.AsyncClient(timeout=REQUEST_TIMEOUT_MS / 1000, follow_redirects=True)
async def close(self) -> None:
await self.client.aclose()
async def post(self, body: dict) -> dict | None:
headers = build_headers(self.cfg, {
"Content-Type": "application/json",
"MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
**(self.session_id and {"Mcp-Session-Id": self.session_id} or {}),
})
try:
resp = await self.client.post(self.cfg.url, headers=headers, json=body)
except httpx.HTTPError as e:
raise _friendly_http_error(self.cfg, e) from e
if resp.status_code >= 400:
raise McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
sid = resp.headers.get("mcp-session-id")
if sid:
self.session_id = sid
ctype = (resp.headers.get("content-type") or "").lower()
if "application/json" in ctype:
if not resp.content: # 通知等无响应体的 200/202
return None
data = resp.json()
return data if isinstance(data, dict) else None
# SSE 响应体:读完整条响应,按请求 id 取对应消息
async for ev in read_sse_events(resp):
parsed = parse_jsonrpc(ev.data)
if parsed and parsed.get("id") == body.get("id"):
return parsed
return None
async def init(self) -> None:
res = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "initialize",
"params": {
"protocolVersion": MCP_PROTOCOL_VERSION,
"capabilities": {},
"clientInfo": {"name": CLIENT_NAME, "version": CLIENT_VERSION},
},
})
expect_result(self.cfg, "initialize", res)
try: # 通知(无 id,忽略错误)
await self.post({"jsonrpc": "2.0", "method": "notifications/initialized"})
except McpClientError:
pass
async def list_tools(self) -> list[McpTool]:
if not self.session_id:
await self.init()
body = {"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list", "params": {}}
res = await self.post(body)
result = expect_result(self.cfg, "tools/list", res)
tools = [McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
for t in result.get("tools") or [] if t.get("name")]
cursor = result.get("nextCursor")
while cursor: # 分页续拉
page = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list",
"params": {"cursor": cursor},
})
page_result = expect_result(self.cfg, "tools/list", page)
tools.extend(McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
for t in page_result.get("tools") or [] if t.get("name"))
cursor = page_result.get("nextCursor")
return tools
async def call_tool(self, name: str, args: object) -> McpCallResult:
if not self.session_id:
await self.init()
res = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/call",
"params": {"name": name, "arguments": args if isinstance(args, dict) else {}},
})
result = expect_result(self.cfg, f"tools/call {name}", res)
return McpCallResult(text=format_tool_content(result.get("content")), is_error=result.get("isError") is True)
# ------------------------------------------------------------------
# 老式 SSE 传输:GET /sse 建连,POST messages URL,响应经 SSE 流按 id 关联
# ------------------------------------------------------------------
class LegacySseTransport:
"""后台 reader 任务持续消费 SSE 流,把 id 匹配的 JSON-RPC 响应派发给挂起的请求。"""
def __init__(self, cfg: McpServerConfig):
self.cfg = cfg
self._pending: dict[int, asyncio.Future] = {}
self._lock = asyncio.Lock()
self._conn_ready = asyncio.Event()
self._client: httpx.AsyncClient | None = None
self._reader_task: asyncio.Task | None = None
self._endpoint_url: str | None = None
self._session_id: str | None = None
self._error: McpClientError | None = None
async def close(self) -> None:
if self._reader_task:
self._reader_task.cancel()
try:
await self._reader_task
except (asyncio.CancelledError, Exception):
pass
self._reader_task = None
if self._client:
await self._client.aclose()
self._client = None
for fut in self._pending.values():
if not fut.done():
fut.cancel()
self._pending.clear()
async def _wait_ready(self) -> None:
try:
await asyncio.wait_for(self._conn_ready.wait(), timeout=REQUEST_TIMEOUT_MS / 1000)
except asyncio.TimeoutError:
raise McpClientError(f"[{self.cfg.name}] SSE 连接超时({REQUEST_TIMEOUT_MS}ms)") from None
if self._error:
raise self._error
async def _ensure_connected(self) -> None:
if self._conn_ready.is_set():
return
if self._reader_task is not None:
await self._wait_ready()
return
async with self._lock:
if self._reader_task is not None:
await self._wait_ready()
return
self._client = httpx.AsyncClient(timeout=None, follow_redirects=True)
try:
# 必须流式取响应:SSE 连接长期不关闭,非流式 get() 会一直等到 EOF
resp = await self._client.send(
httpx.Request("GET", self.cfg.url, headers=build_headers(self.cfg, {})),
stream=True,
)
except httpx.HTTPError as e:
self._error = _friendly_http_error(self.cfg, e)
self._conn_ready.set()
raise self._error from e
if resp.status_code >= 400:
self._error = McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
self._conn_ready.set()
raise self._error
self._reader_task = asyncio.create_task(self._read_loop(resp))
await self._wait_ready() # 等 endpoint 事件就绪后再继续
async def _read_loop(self, resp: httpx.Response) -> None:
"""解析 SSE 流:先等 endpoint 事件,之后把 JSON-RPC 响应分发给挂起 Future。"""
try:
async for ev in read_sse_events(resp):
if self._endpoint_url is None:
if ev.event == "endpoint" and ev.data.strip():
endpoint = ev.data.strip()
self._endpoint_url = _resolve_url(endpoint, self.cfg.url)
try:
u = httpx.URL(self._endpoint_url)
self._session_id = u.params.get("sessionId") or u.params.get("session_id")
except Exception:
pass
self._conn_ready.set()
continue
parsed = parse_jsonrpc(ev.data)
if parsed and isinstance(parsed.get("id"), int) and parsed["id"] in self._pending:
fut = self._pending.pop(parsed["id"])
if not fut.done():
fut.set_result(parsed)
except asyncio.CancelledError:
pass
except Exception as e: # 连接中断 → 所有挂起请求报错
self._error = McpClientError(f"[{self.cfg.name}] SSE 连接中断:{e}")
finally:
if self._endpoint_url is None:
if self._error is None:
self._error = McpClientError(f"[{self.cfg.name}] SSE 端点未返回 endpoint 事件")
self._conn_ready.set()
for fut in self._pending.values():
if not fut.done():
fut.set_exception(self._error or McpClientError(f"[{self.cfg.name}] SSE 连接已关闭"))
self._pending.clear()
async def post(self, body: dict) -> dict | None:
await self._ensure_connected()
if not self._endpoint_url:
raise self._error or McpClientError(f"[{self.cfg.name}] SSE 端点未就绪")
headers = build_headers(self.cfg, {
"Content-Type": "application/json",
"MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
**(self._session_id and {"Mcp-Session-Id": self._session_id} or {}),
})
rid: int = body.get("id")
if rid is not None:
fut: asyncio.Future = asyncio.get_running_loop().create_future()
self._pending[rid] = fut
try:
try:
resp = await self._client.post(self._endpoint_url, headers=headers, json=body)
except httpx.HTTPError as e:
raise _friendly_http_error(self.cfg, e) from e
# 部分服务器直接在 POST 响应体返回 JSON
if resp.status_code >= 400:
raise McpClientError(f"[{self.cfg.name}] HTTP {resp.status_code} {resp.reason_phrase}")
if "application/json" in (resp.headers.get("content-type") or "").lower():
data = resp.json() if resp.content else None
if rid is not None and self._pending.get(rid) is fut:
self._pending.pop(rid, None)
return data if isinstance(data, dict) else None
except Exception:
if rid is not None and self._pending.get(rid) is fut:
self._pending.pop(rid, None)
raise
# 通知(无 id)没有响应:直接返回
if rid is None:
return None
try:
return await asyncio.wait_for(asyncio.shield(fut), timeout=REQUEST_TIMEOUT_MS / 1000)
except asyncio.TimeoutError:
self._pending.pop(rid, None)
raise McpClientError(f"[{self.cfg.name}] 请求超时({REQUEST_TIMEOUT_MS}ms)") from None
async def init(self) -> None:
res = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "initialize",
"params": {
"protocolVersion": MCP_PROTOCOL_VERSION,
"capabilities": {},
"clientInfo": {"name": CLIENT_NAME, "version": CLIENT_VERSION},
},
})
expect_result(self.cfg, "initialize", res)
try:
await self.post({"jsonrpc": "2.0", "method": "notifications/initialized"})
except McpClientError:
pass
async def list_tools(self) -> list[McpTool]:
if not self._endpoint_url:
await self.init()
body = {"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list", "params": {}}
res = await self.post(body)
result = expect_result(self.cfg, "tools/list", res)
tools = [McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
for t in result.get("tools") or [] if t.get("name")]
cursor = result.get("nextCursor")
while cursor:
page = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/list",
"params": {"cursor": cursor},
})
page_result = expect_result(self.cfg, "tools/list", page)
tools.extend(McpTool(name=t.get("name", ""), description=t.get("description") or "", input_schema=t.get("inputSchema"))
for t in page_result.get("tools") or [] if t.get("name"))
cursor = page_result.get("nextCursor")
return tools
async def call_tool(self, name: str, args: object) -> McpCallResult:
if not self._endpoint_url:
await self.init()
res = await self.post({
"jsonrpc": "2.0", "id": next_rpc_id(), "method": "tools/call",
"params": {"name": name, "arguments": args if isinstance(args, dict) else {}},
})
result = expect_result(self.cfg, f"tools/call {name}", res)
return McpCallResult(text=format_tool_content(result.get("content")), is_error=result.get("isError") is True)
# ------------------------------------------------------------------
# 结果格式化
# ------------------------------------------------------------------
def format_tool_content(content: object) -> str:
if not isinstance(content, list):
return "[工具无返回内容]"
parts: list[str] = []
for item in content:
if not isinstance(item, dict):
continue
t = item.get("type")
if t == "text" and item.get("text"):
parts.append(str(item["text"]))
elif t == "image":
data = item.get("data") or ""
parts.append(f"[图片 {item.get('mimeType') or '未知类型'},base64 {len(data)} 字符]")
elif t == "audio":
parts.append(f"[音频 {item.get('mimeType') or '未知类型'}]")
elif t == "resource":
res = item.get("resource") or {}
parts.append(res.get("text") or (f"[资源 {res.get('uri')}]" if res.get("uri") else "[资源]"))
elif item.get("text"):
parts.append(str(item["text"]))
return "\n".join(p for p in parts if p) or "[工具无返回内容]"
# ------------------------------------------------------------------
# 缓存与高层 API
# ------------------------------------------------------------------
_transports: dict[str, StreamableHttpTransport | LegacySseTransport] = {}
_tool_cache: dict[str, tuple[float, list[McpTool]]] = {}
_list_locks: dict[str, asyncio.Lock] = {}
def _get_transport(cfg: McpServerConfig) -> StreamableHttpTransport | LegacySseTransport:
transport = _transports.get(cfg.id)
if transport is None:
transport = LegacySseTransport(cfg) if cfg.transport == "sse" else StreamableHttpTransport(cfg)
_transports[cfg.id] = transport
return transport
async def invalidate_mcp_cache(server_id: str | None = None) -> None:
"""失效某个(或全部)服务器的缓存;配置变更时调用。"""
ids = [server_id] if server_id else list(_transports.keys())
async def _evict(sid: str) -> None:
transport = _transports.pop(sid, None)
if transport is not None:
try:
await transport.close()
except Exception:
pass
_tool_cache.pop(sid, None)
_list_locks.pop(sid, None)
for sid in list(ids):
await _evict(sid)
async def list_mcp_tools(cfg: McpServerConfig, fresh: bool = False) -> list[McpTool]:
"""获取服务器的工具列表(带缓存)。fresh 为 true 时先失效缓存重新拉取(用于「测试连接」)。"""
if fresh:
await invalidate_mcp_cache(cfg.id)
cached = _tool_cache.get(cfg.id)
if not fresh and cached and cached[0] > time.time():
return cached[1]
lock = _list_locks.setdefault(cfg.id, asyncio.Lock())
async with lock:
cached = _tool_cache.get(cfg.id)
if not fresh and cached and cached[0] > time.time():
return cached[1]
transport = _get_transport(cfg)
try:
tools = await transport.list_tools()
except Exception:
await invalidate_mcp_cache(cfg.id) # 失败时释放传输,避免坏会话复用于下一次
raise
_tool_cache[cfg.id] = (time.time() + CACHE_TTL_S, tools)
return tools
async def call_mcp_tool(cfg: McpServerConfig, tool_name: str, args: object) -> str:
"""调用一个 MCP 工具;返回可读文本。失败抛 McpClientError 由上层捕获。"""
transport = _get_transport(cfg)
try:
result = await transport.call_tool(tool_name, args)
except Exception:
await invalidate_mcp_cache(cfg.id) # 会话可能已失效 → 清缓存,允许下次重建
raise
return f"[MCP 工具 {tool_name} 返回错误]\n{result.text}" if result.is_error else result.text
+36
View File
@@ -0,0 +1,36 @@
import json
import uuid
from datetime import datetime
from sqlalchemy import String, Text, Boolean, DateTime, func
from sqlalchemy.orm import Mapped, mapped_column
from app.core.database import Base
def gen_uuid() -> str:
return str(uuid.uuid4())
class McpServer(Base):
"""用户从页面配置的 MCP 服务器(供对话 agent 调用外部工具)。
存放远程 MCP 服务的连接信息:传输协议、URL、附加请求头(含鉴权)。
启用中的服务器会在对话时向模型暴露其工具(mcp_ 前缀)。
"""
__tablename__ = "mcp_servers"
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=gen_uuid)
name: Mapped[str] = mapped_column(String(200)) # 显示名,如 firecrawl
transport: Mapped[str] = mapped_column(String(30), default="streamable-http") # streamable-http | sse
url: Mapped[str] = mapped_column(String(2000)) # 远程端点
headers: Mapped[str] = mapped_column(Text, default="{}") # JSON dict,附加请求头(明文本地存储,与 ai_models 同级信任)
enabled: Mapped[bool] = mapped_column(Boolean, default=True)
created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
def headers_dict(self) -> dict:
try:
data = json.loads(self.headers or "{}")
return data if isinstance(data, dict) else {}
except (json.JSONDecodeError, TypeError):
return {}