"""管理员托管的 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) tool_count = len(catalog["tools"]) return { "server": catalog["server"], # 对外测试连接响应与 OpenAPI 统一使用 operation_count；保留 # tool_count，兼容已依赖 MCP 连接器原始响应的调用方。 "operation_count": tool_count, "tool_count": tool_count, } 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, }