523 lines
19 KiB
Python
523 lines
19 KiB
Python
"""管理员托管的 Streamable HTTP MCP 外部系统连接器。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import binascii
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from functools import partial
|
||
from typing import Any, cast
|
||
from urllib.parse import urlparse
|
||
|
||
import anyio
|
||
import httpx
|
||
import httpx2
|
||
from mcp import ClientSession
|
||
from mcp.client.streamable_http import streamable_http_client
|
||
from mcp.types import PaginatedRequestParams
|
||
|
||
from .auth import ExternalAuthError, get_auth_strategy
|
||
from .results import MAX_STORED_RESULT_BYTES
|
||
from .runtime_cache import RUNTIME_CACHE
|
||
|
||
_AUTH_CACHE_TTL_SECONDS = 300.0
|
||
_CATALOG_CACHE_TTL_SECONDS = 300.0
|
||
_MAX_CATALOG_TOOLS = 1000
|
||
_SENSITIVE_KEY_RE = re.compile(
|
||
r"(?:password|passwd|secret|token|api[_-]?key|authorization|cookie|credential)",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
class McpConnectorError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def _bool_value(value: Any, default: bool) -> bool:
|
||
if value is None:
|
||
return default
|
||
if isinstance(value, bool):
|
||
return value
|
||
return str(value).strip().lower() not in {"0", "false", "no", "off"}
|
||
|
||
|
||
def _validated_http_url(raw: str, label: str) -> str:
|
||
value = (raw or "").strip().rstrip("/")
|
||
parsed = urlparse(value)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
||
raise McpConnectorError(f"{label} 必须是有效的 http(s) URL")
|
||
if parsed.username or parsed.password:
|
||
raise McpConnectorError(f"{label} 不能内嵌凭据")
|
||
return value
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class McpConfig:
|
||
base_url: str
|
||
mcp_url: str
|
||
login_path: str
|
||
timeout_seconds: float
|
||
max_result_bytes: int
|
||
max_total_result_bytes: int
|
||
verify_tls: bool
|
||
query_guidance: str
|
||
recommended_operation_ids: tuple[str, ...]
|
||
operation_mode: str = "upstream_managed"
|
||
operation_policies: dict[str, str] = field(default_factory=dict)
|
||
auth_type: str = "password_jwt"
|
||
auth_config: dict[str, Any] = field(default_factory=dict)
|
||
expected_server_name: str = ""
|
||
max_response_bytes: int = MAX_STORED_RESULT_BYTES
|
||
|
||
@classmethod
|
||
def from_mapping(cls, data: dict[str, Any]) -> McpConfig:
|
||
mcp_url = _validated_http_url(str(data.get("mcp_url") or ""), "mcp_url")
|
||
parsed = urlparse(mcp_url)
|
||
origin = f"{parsed.scheme}://{parsed.netloc}"
|
||
base_url = _validated_http_url(str(data.get("base_url") or origin), "base_url")
|
||
base = urlparse(base_url)
|
||
if (base.scheme, base.netloc) != (parsed.scheme, parsed.netloc):
|
||
raise McpConnectorError("mcp_url 必须与 base_url 同源")
|
||
login_path = str(data.get("login_path") or "/api/auth/token/").strip()
|
||
if not login_path.startswith("/") or "://" in login_path:
|
||
raise McpConnectorError("login_path 必须是站内绝对路径")
|
||
guidance = str(data.get("query_guidance") or "").strip()
|
||
if len(guidance) > 4000:
|
||
raise McpConnectorError("query_guidance 不能超过 4000 字符")
|
||
raw_recommended = data.get("recommended_operation_ids", [])
|
||
if isinstance(raw_recommended, str):
|
||
raw_recommended = raw_recommended.split(",")
|
||
if not isinstance(raw_recommended, (list, tuple, set)):
|
||
raise McpConnectorError("recommended_operation_ids 必须是字符串数组")
|
||
recommended = tuple(
|
||
dict.fromkeys(
|
||
str(item).strip() for item in raw_recommended if str(item).strip()
|
||
)
|
||
)
|
||
if len(recommended) > 30 or any(len(item) > 200 for item in recommended):
|
||
raise McpConnectorError(
|
||
"recommended_operation_ids 最多 30 项且每项不超过 200 字符"
|
||
)
|
||
expected_name = str(data.get("expected_server_name") or "").strip()
|
||
if len(expected_name) > 200:
|
||
raise McpConnectorError("expected_server_name 不能超过 200 字符")
|
||
max_result = max(4096, min(int(data.get("max_result_bytes", 65536)), 1048576))
|
||
return cls(
|
||
base_url=base_url,
|
||
mcp_url=mcp_url,
|
||
login_path=login_path,
|
||
timeout_seconds=max(1.0, min(float(data.get("timeout_seconds", 15)), 60.0)),
|
||
max_result_bytes=max_result,
|
||
max_total_result_bytes=max(
|
||
max_result,
|
||
min(int(data.get("max_total_result_bytes", 262144)), 4194304),
|
||
),
|
||
verify_tls=_bool_value(data.get("verify_tls"), True),
|
||
query_guidance=guidance,
|
||
recommended_operation_ids=recommended,
|
||
auth_type=str(data.get("auth_type") or "password_jwt").strip(),
|
||
auth_config={
|
||
key: data[key]
|
||
for key in (
|
||
"login_path",
|
||
"username_field",
|
||
"password_field",
|
||
"token_field",
|
||
"auth_header_name",
|
||
"auth_header_template",
|
||
)
|
||
if key in data
|
||
},
|
||
expected_server_name=expected_name,
|
||
max_response_bytes=max(
|
||
65536,
|
||
min(
|
||
int(data.get("max_response_bytes", MAX_STORED_RESULT_BYTES)),
|
||
MAX_STORED_RESULT_BYTES,
|
||
),
|
||
),
|
||
)
|
||
|
||
|
||
class _LimitedAsyncStream(httpx2.AsyncByteStream):
|
||
def __init__(self, stream: httpx2.AsyncByteStream, limit: int):
|
||
self._stream = stream
|
||
self._limit = limit
|
||
|
||
async def __aiter__(self):
|
||
total = 0
|
||
async for chunk in self._stream:
|
||
total += len(chunk)
|
||
if total > self._limit:
|
||
raise McpConnectorError("MCP 响应超过安全下载上限")
|
||
yield chunk
|
||
|
||
async def aclose(self) -> None:
|
||
await self._stream.aclose()
|
||
|
||
|
||
class _LimitedTransport(httpx2.AsyncBaseTransport):
|
||
def __init__(self, *, verify: bool, limit: int):
|
||
self._transport = httpx2.AsyncHTTPTransport(verify=verify, retries=0)
|
||
self._limit = limit
|
||
|
||
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
|
||
response = await self._transport.handle_async_request(request)
|
||
raw_length = response.headers.get("content-length")
|
||
try:
|
||
content_length = int(raw_length) if raw_length else None
|
||
except ValueError:
|
||
content_length = None
|
||
if content_length is not None and content_length > self._limit:
|
||
await response.aclose()
|
||
raise McpConnectorError("MCP 响应超过安全下载上限")
|
||
response.stream = _LimitedAsyncStream(
|
||
cast(httpx2.AsyncByteStream, response.stream),
|
||
self._limit,
|
||
)
|
||
return response
|
||
|
||
async def aclose(self) -> None:
|
||
await self._transport.aclose()
|
||
|
||
|
||
def _auth_cache_ttl(headers: dict[str, str]) -> float:
|
||
authorization = next(
|
||
(value for name, value in headers.items() if name.lower() == "authorization"),
|
||
"",
|
||
)
|
||
token = authorization.split(" ", 1)[-1].strip()
|
||
parts = token.split(".")
|
||
if len(parts) == 3:
|
||
try:
|
||
padding = "=" * (-len(parts[1]) % 4)
|
||
payload = json.loads(
|
||
base64.urlsafe_b64decode(parts[1] + padding).decode("utf-8")
|
||
)
|
||
return max(
|
||
0.0,
|
||
min(
|
||
_AUTH_CACHE_TTL_SECONDS,
|
||
float(payload.get("exp")) - time.time() - 30,
|
||
),
|
||
)
|
||
except (binascii.Error, TypeError, ValueError, UnicodeDecodeError):
|
||
pass
|
||
return _AUTH_CACHE_TTL_SECONDS
|
||
|
||
|
||
def _exception_has_status(exc: BaseException, status_code: int) -> bool:
|
||
response = getattr(exc, "response", None)
|
||
if getattr(response, "status_code", None) == status_code:
|
||
return True
|
||
return any(
|
||
_exception_has_status(child, status_code)
|
||
for child in getattr(exc, "exceptions", ())
|
||
if isinstance(child, BaseException)
|
||
)
|
||
|
||
|
||
def _find_connector_error(exc: BaseException) -> McpConnectorError | None:
|
||
if isinstance(exc, McpConnectorError):
|
||
return exc
|
||
for child in getattr(exc, "exceptions", ()):
|
||
if isinstance(child, BaseException):
|
||
found = _find_connector_error(child)
|
||
if found is not None:
|
||
return found
|
||
return None
|
||
|
||
|
||
def _redacted(value: Any) -> Any:
|
||
if isinstance(value, dict):
|
||
return {
|
||
str(key): "[REDACTED]"
|
||
if _SENSITIVE_KEY_RE.search(str(key))
|
||
else _redacted(item)
|
||
for key, item in value.items()
|
||
}
|
||
if isinstance(value, list):
|
||
return [_redacted(item) for item in value]
|
||
return value
|
||
|
||
|
||
class McpClient:
|
||
def __init__(
|
||
self,
|
||
credentials: dict[str, str],
|
||
cfg: McpConfig,
|
||
*,
|
||
cache_namespace: str = "",
|
||
):
|
||
self.credentials = credentials
|
||
self.cfg = cfg
|
||
identity = json.dumps(
|
||
{
|
||
"namespace": cache_namespace,
|
||
"credentials": credentials,
|
||
"config": cfg.__dict__,
|
||
},
|
||
sort_keys=True,
|
||
ensure_ascii=False,
|
||
default=str,
|
||
)
|
||
self._runtime_identity = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||
|
||
def authenticate(self, *, force: bool = False) -> dict[str, str]:
|
||
if force:
|
||
RUNTIME_CACHE.invalidate_auth(self._runtime_identity)
|
||
else:
|
||
cached = RUNTIME_CACHE.get_auth(self._runtime_identity)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
def load() -> dict[str, str]:
|
||
if not force:
|
||
cached = RUNTIME_CACHE.get_auth(self._runtime_identity)
|
||
if cached is not None:
|
||
return cached
|
||
try:
|
||
with httpx.Client(
|
||
timeout=self.cfg.timeout_seconds,
|
||
verify=self.cfg.verify_tls,
|
||
follow_redirects=False,
|
||
) as client:
|
||
headers = get_auth_strategy(self.cfg.auth_type).headers(
|
||
client=client,
|
||
base_url=self.cfg.base_url,
|
||
credentials=self.credentials,
|
||
config=self.cfg.auth_config,
|
||
)
|
||
except ExternalAuthError as exc:
|
||
raise McpConnectorError(str(exc)) from exc
|
||
RUNTIME_CACHE.set_auth(
|
||
self._runtime_identity,
|
||
headers,
|
||
ttl_seconds=_auth_cache_ttl(headers),
|
||
)
|
||
return dict(headers)
|
||
|
||
return RUNTIME_CACHE.singleflight("auth", self._runtime_identity, load)
|
||
|
||
async def _session_operation(
|
||
self, operation: str, payload: Any, *, force_auth: bool
|
||
):
|
||
headers = self.authenticate(force=force_auth)
|
||
transport = _LimitedTransport(
|
||
verify=self.cfg.verify_tls,
|
||
limit=self.cfg.max_response_bytes,
|
||
)
|
||
async with (
|
||
httpx2.AsyncClient(
|
||
headers=headers,
|
||
timeout=self.cfg.timeout_seconds,
|
||
follow_redirects=False,
|
||
transport=transport,
|
||
) as http_client,
|
||
streamable_http_client(
|
||
self.cfg.mcp_url,
|
||
http_client=http_client,
|
||
) as streams,
|
||
ClientSession(*streams) as session,
|
||
):
|
||
initialized = await session.initialize()
|
||
server_info = initialized.server_info
|
||
if (
|
||
self.cfg.expected_server_name
|
||
and server_info.name != self.cfg.expected_server_name
|
||
):
|
||
raise McpConnectorError(
|
||
"MCP Server 身份不匹配："
|
||
f"期望 {self.cfg.expected_server_name}，实际 {server_info.name}"
|
||
)
|
||
if operation == "list":
|
||
tools: list[dict[str, Any]] = []
|
||
cursor = None
|
||
while True:
|
||
params = (
|
||
PaginatedRequestParams(cursor=cursor)
|
||
if cursor is not None
|
||
else None
|
||
)
|
||
page = await session.list_tools(params=params)
|
||
tools.extend(
|
||
tool.model_dump(by_alias=True, exclude_none=True)
|
||
for tool in page.tools
|
||
)
|
||
if len(tools) > _MAX_CATALOG_TOOLS:
|
||
raise McpConnectorError("MCP 工具目录超过 1000 项安全上限")
|
||
cursor = page.next_cursor
|
||
if not cursor:
|
||
break
|
||
return {
|
||
"server": server_info.model_dump(by_alias=True, exclude_none=True),
|
||
"tools": tools,
|
||
}
|
||
if operation == "call":
|
||
name, arguments = payload
|
||
available: set[str] = set()
|
||
cursor = None
|
||
while True:
|
||
params = (
|
||
PaginatedRequestParams(cursor=cursor)
|
||
if cursor is not None
|
||
else None
|
||
)
|
||
page = await session.list_tools(params=params)
|
||
available.update(tool.name for tool in page.tools)
|
||
if len(available) > _MAX_CATALOG_TOOLS:
|
||
raise McpConnectorError("MCP 工具目录超过 1000 项安全上限")
|
||
cursor = page.next_cursor
|
||
if not cursor:
|
||
break
|
||
if name not in available:
|
||
raise McpConnectorError("MCP 工具已不存在，请重新搜索工具目录")
|
||
result = await session.call_tool(
|
||
name,
|
||
arguments=arguments,
|
||
read_timeout_seconds=self.cfg.timeout_seconds,
|
||
)
|
||
return result.model_dump(by_alias=True, exclude_none=True)
|
||
raise AssertionError(f"unknown MCP operation: {operation}")
|
||
|
||
def _run(self, operation: str, payload: Any = None) -> Any:
|
||
for attempt in range(2):
|
||
try:
|
||
return anyio.run(
|
||
partial(
|
||
self._session_operation,
|
||
operation,
|
||
payload,
|
||
force_auth=attempt == 1,
|
||
)
|
||
)
|
||
except BaseException as exc:
|
||
if attempt == 0 and _exception_has_status(exc, 401):
|
||
RUNTIME_CACHE.invalidate_auth(self._runtime_identity)
|
||
continue
|
||
connector_error = _find_connector_error(exc)
|
||
if connector_error is not None:
|
||
raise connector_error
|
||
raise McpConnectorError(f"MCP 调用失败: {type(exc).__name__}") from exc
|
||
raise McpConnectorError("MCP 认证失败")
|
||
|
||
def _catalog(self, *, force: bool = False) -> dict[str, Any]:
|
||
if force:
|
||
RUNTIME_CACHE.invalidate_mcp_catalog(self._runtime_identity)
|
||
else:
|
||
cached = RUNTIME_CACHE.get_mcp_catalog(self._runtime_identity)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
def load() -> dict[str, Any]:
|
||
if not force:
|
||
cached = RUNTIME_CACHE.get_mcp_catalog(self._runtime_identity)
|
||
if cached is not None:
|
||
return cached
|
||
catalog = self._run("list")
|
||
RUNTIME_CACHE.set_mcp_catalog(
|
||
self._runtime_identity,
|
||
catalog,
|
||
ttl_seconds=_CATALOG_CACHE_TTL_SECONDS,
|
||
)
|
||
return catalog
|
||
|
||
return RUNTIME_CACHE.singleflight("mcp-catalog", self._runtime_identity, load)
|
||
|
||
@staticmethod
|
||
def _operation_id(name: str) -> str:
|
||
return f"mcp/{name}"
|
||
|
||
@staticmethod
|
||
def _remote_name(operation_id: str) -> str:
|
||
value = str(operation_id or "").strip()
|
||
return value.removeprefix("mcp/")
|
||
|
||
def test_connection(self) -> dict[str, Any]:
|
||
catalog = self._catalog(force=True)
|
||
return {
|
||
"server": catalog["server"],
|
||
"tool_count": len(catalog["tools"]),
|
||
}
|
||
|
||
def search(self, query: str, limit: int = 12) -> list[dict[str, Any]]:
|
||
safe_limit = max(1, min(int(limit), 30))
|
||
terms = [term.lower() for term in str(query or "").split() if term]
|
||
recommended = {
|
||
self._remote_name(item): index
|
||
for index, item in enumerate(self.cfg.recommended_operation_ids)
|
||
}
|
||
scored: list[tuple[int, int, str, dict[str, Any]]] = []
|
||
for tool in self._catalog()["tools"]:
|
||
name = str(tool.get("name") or "")
|
||
title = str(tool.get("title") or "")
|
||
description = str(tool.get("description") or "")
|
||
haystack = f"{name} {title} {description}".lower()
|
||
score = sum(
|
||
4 if term in name.lower() else 1 for term in terms if term in haystack
|
||
)
|
||
is_recommended = name in recommended
|
||
if terms and score == 0 and not is_recommended:
|
||
continue
|
||
item = {
|
||
"operation_id": self._operation_id(name),
|
||
"name": name,
|
||
"title": title,
|
||
"summary": description,
|
||
"input_schema": tool.get("inputSchema") or {"type": "object"},
|
||
"output_schema": tool.get("outputSchema"),
|
||
"annotations": tool.get("annotations"),
|
||
"recommended": is_recommended,
|
||
}
|
||
scored.append(
|
||
(
|
||
0 if is_recommended else 1,
|
||
recommended.get(name, -score),
|
||
name,
|
||
item,
|
||
)
|
||
)
|
||
scored.sort(key=lambda row: (row[0], row[1], row[2]))
|
||
return [row[3] for row in scored[:safe_limit]]
|
||
|
||
def call(
|
||
self,
|
||
operation_id: str,
|
||
arguments: dict[str, Any] | None = None,
|
||
body: Any = None,
|
||
) -> dict[str, Any]:
|
||
if body is not None:
|
||
raise McpConnectorError("MCP 工具参数请全部放入 arguments，不使用 body")
|
||
name = self._remote_name(operation_id)
|
||
catalog = self._catalog()
|
||
if name not in {str(tool.get("name") or "") for tool in catalog["tools"]}:
|
||
raise McpConnectorError("MCP 工具不存在，请先搜索工具目录")
|
||
raw = self._run("call", (name, dict(arguments or {})))
|
||
if raw.get("isError"):
|
||
details = json.dumps(
|
||
_redacted(raw.get("content") or []),
|
||
ensure_ascii=False,
|
||
default=str,
|
||
)
|
||
raise McpConnectorError(f"MCP 工具返回错误: {details[:1000]}")
|
||
data = raw.get("structuredContent")
|
||
if data is None:
|
||
data = {"content": raw.get("content") or []}
|
||
normalized = json.loads(json.dumps(data, ensure_ascii=False, default=str))
|
||
response_bytes = len(
|
||
json.dumps(normalized, ensure_ascii=False, separators=(",", ":")).encode(
|
||
"utf-8"
|
||
)
|
||
)
|
||
return {
|
||
"operation_id": self._operation_id(name),
|
||
"status_code": None,
|
||
"response_bytes": response_bytes,
|
||
"truncated": False,
|
||
"data": normalized,
|
||
}
|