添加MCP客户端超时配置和连接管理改进

添加了MCP客户端的超时配置功能，包括启动超时和工具调用超时参数。改进了HTTP客户端和标准IO客户端的超时处理机制，确保请求在指定时间内完成或取消。增加了对MCP服务器数量的限制，防止配置过多服务器导致系统不稳定。增强了错误处理机制，当连接异常时能够正确清理资源并移除桥接主机。添加了对大型MCP消息的大小验证，防止过大的请求导致系统问题。优化了密钥更改后的处理流程，确保在修改密钥时停用服务器并要求重新测试。
2026-09-03 16:16:58 +08:00
parent 7d5f4023a9
commit 2f7066aa92
11 changed files with 504 additions and 72 deletions
@@ -146,7 +146,7 @@ class McpStdioClient:
        timeout_code: str,
        response_error_code: str = "MCP_TOOL_CALL_FAILED",
    ) -> dict[str, Any]:
-        request_id, pending = self.begin_request(method, params)
+        request_id, pending = self.begin_request(method, params, timeout=timeout)
        return self.wait_response(
            request_id,
            pending,
@@ -156,7 +156,7 @@ class McpStdioClient:
        )

    def begin_request(
-        self, method: str, params: dict[str, Any]
+        self, method: str, params: dict[str, Any], *, timeout: float | None = None
    ) -> tuple[int, _PendingRequest]:
        self._ensure_running()
        with self._pending_lock:
@@ -388,6 +388,7 @@ class McpHttpClient:
        url: str,
        *,
        headers: dict[str, str],
+        startup_timeout_seconds: float = 15,
        on_seen: Callable[[], None],
        on_broken: Callable[[str], None],
        on_tools_changed: Callable[[], None],
@@ -407,6 +408,7 @@ class McpHttpClient:
        self._stream_started = False
        self._last_event_id: str | None = None
        self._stop_event = threading.Event()
+        self._startup_timeout_seconds = startup_timeout_seconds

    def start(self) -> None:
        return
@@ -429,7 +431,7 @@ class McpHttpClient:
        timeout_code: str,
        response_error_code: str = "MCP_TOOL_CALL_FAILED",
    ) -> dict[str, Any]:
-        request_id, pending = self.begin_request(method, params)
+        request_id, pending = self.begin_request(method, params, timeout=timeout)
        return self.wait_response(
            request_id,
            pending,
@@ -439,7 +441,7 @@ class McpHttpClient:
        )

    def begin_request(
-        self, method: str, params: dict[str, Any]
+        self, method: str, params: dict[str, Any], *, timeout: float | None = None
    ) -> tuple[int, _PendingRequest]:
        with self._pending_lock:
            request_id = self._next_id
@@ -454,7 +456,7 @@ class McpHttpClient:
        }
        threading.Thread(
            target=self._dispatch_request,
-            args=(request_id, message),
+            args=(request_id, message, timeout),
            daemon=True,
        ).start()
        return request_id, pending
@@ -526,7 +528,10 @@ class McpHttpClient:
        if self._session_id:
            try:
                request = self._client.build_request(
-                    "DELETE", self.url, headers=self._request_headers()
+                    "DELETE",
+                    self.url,
+                    headers=self._request_headers(),
+                    timeout=min(self._startup_timeout_seconds, 5),
                )
                response = self._client.send(request, stream=True)
                response.close()
@@ -539,9 +544,14 @@ class McpHttpClient:
            )
        )

-    def _dispatch_request(self, request_id: int, message: dict[str, Any]) -> None:
+    def _dispatch_request(
+        self,
+        request_id: int,
+        message: dict[str, Any],
+        timeout: float | None,
+    ) -> None:
        try:
-            response = self._post(message, timeout=None)
+            response = self._post(message, timeout=timeout)
            try:
                self._capture_session(response)
                content_type = response.headers.get("content-type", "").lower()
@@ -588,7 +598,12 @@ class McpHttpClient:

    def _post_notification(self, message: dict[str, Any]) -> None:
        try:
-            response = self._post(message, timeout=10)
+            timeout = (
+                self._startup_timeout_seconds
+                if message.get("method") == "notifications/initialized"
+                else 10
+            )
+            response = self._post(message, timeout=timeout)
        except httpx.HTTPError as exc:
            raise McpBridgeError(
                "MCP_HTTP_REQUEST_FAILED",
@@ -616,6 +631,7 @@ class McpHttpClient:
            self.url,
            content=encoded.encode("utf-8"),
            headers=self._request_headers(),
+            timeout=timeout,
        )
        return self._client.send(request, stream=True)

@@ -714,7 +730,7 @@ class McpLegacySseClient(McpHttpClient):
    def start(self) -> None:
        threading.Thread(target=self._event_loop, daemon=True).start()
        try:
-            endpoint = self._endpoint_ready.get(timeout=15)
+            endpoint = self._endpoint_ready.get(timeout=self._startup_timeout_seconds)
        except queue.Empty as exc:
            raise McpBridgeError(
                "MCP_INITIALIZE_FAILED",
@@ -730,9 +746,14 @@ class McpLegacySseClient(McpHttpClient):

        return

-    def _dispatch_request(self, request_id: int, message: dict[str, Any]) -> None:
+    def _dispatch_request(
+        self,
+        request_id: int,
+        message: dict[str, Any],
+        timeout: float | None,
+    ) -> None:
        try:
-            response = self._post(message, timeout=10)
+            response = self._post(message, timeout=timeout)
            try:
                if response.status_code not in {200, 202, 204}:
                    raise McpBridgeError(
@@ -761,6 +782,8 @@ class McpLegacySseClient(McpHttpClient):
                "MCP_INITIALIZE_FAILED", "Legacy MCP endpoint is not ready."
            )
        encoded = json.dumps(message, ensure_ascii=False, separators=(",", ":"))
+        if len(encoded.encode("utf-8")) > MAX_MCP_MESSAGE_BYTES:
+            raise McpBridgeError("MCP_TOOL_CALL_FAILED", "MCP request is too large.")
        request = self._client.build_request(
            "POST",
            self._endpoint,
@@ -770,6 +793,7 @@ class McpLegacySseClient(McpHttpClient):
                "Accept": "application/json, text/event-stream",
                "Content-Type": "application/json",
            },
+            timeout=timeout,
        )
        return self._client.send(request, stream=True)

@@ -801,6 +825,8 @@ class McpLegacySseClient(McpHttpClient):
                        self._endpoint = endpoint
                        continue
                    self._handle_message(_json_rpc_message(data))
+            if not self._stopping:
+                self.on_broken("Legacy MCP SSE stream ended unexpectedly.")
        except (McpBridgeError, httpx.HTTPError) as exc:
            if self._endpoint is None:
                self._endpoint_ready.put(exc)
@@ -827,7 +853,7 @@ class _McpClient(Protocol):
        response_error_code: str = "MCP_TOOL_CALL_FAILED",
    ) -> dict[str, Any]: ...
    def begin_request(
-        self, method: str, params: dict[str, Any]
+        self, method: str, params: dict[str, Any], *, timeout: float | None = None
    ) -> tuple[int, _PendingRequest]: ...
    def wait_response(
        self,
@@ -931,6 +957,7 @@ class McpBridge:
            client = client_type(
                url,
                headers=headers or {},
+                startup_timeout_seconds=backend.startup_timeout_seconds,
                on_seen=seen,
                on_broken=broken,
                on_tools_changed=tools_changed,
@@ -989,6 +1016,12 @@ class McpBridge:
            discovered = self._discover_tools(
                plugin_id, client, backend, declared_permissions, tool_source
            )
+            if status.status == PluginHostState.unhealthy:
+                raise McpBridgeError(
+                    "PLUGIN_HOST_UNAVAILABLE",
+                    status.error or "MCP event stream became unavailable during startup.",
+                    status_code=503,
+                )
            status.status = PluginHostState.ready
            status.tools_count = len(discovered)
            status.last_seen_at = datetime.now(UTC)
@@ -1019,7 +1052,9 @@ class McpBridge:
    ) -> Any:
        host = self._host(plugin_id)
        rpc_id, pending = host.client.begin_request(
-            "tools/call", {"name": remote_name, "arguments": arguments}
+            "tools/call",
+            {"name": remote_name, "arguments": arguments},
+            timeout=host.backend.tool_timeout_seconds,
        )
        call_key = (plugin_id, request_id)
        with self._lock:
@@ -13,12 +13,13 @@ from typing import Any
 from urllib.parse import urlsplit
 from uuid import uuid4

-from pydantic import BaseModel, ConfigDict, create_model
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, create_model

 from app.agent.permissions import KNOWN_PERMISSIONS
 from app.agent.tools import ToolExecutionContext, ToolRegistry
 from app.contracts import (
    McpServer,
+    McpServerConfig,
    McpServerCreateRequest,
    McpServerSecretStatus,
    McpServerTransport,
@@ -40,6 +41,23 @@ _RESERVED_HEADERS = {
    "mcp-protocol-version",
    "mcp-session-id",
 }
+_SERVER_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$")
+_MAX_MCP_SERVERS = 256
+
+
+class _McpServerRecord(McpServerConfig):
+    """Validated on-disk representation with defaults for older C.1 records."""
+
+    version: int = Field(default=1, ge=1)
+    enabled: bool = False
+    approved_digest: str | None = Field(
+        default=None, min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$"
+    )
+    tested_digest: str | None = Field(
+        default=None, min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$"
+    )
+    last_tested_at: datetime | None = None
+    last_test_succeeded: bool | None = None


 class McpRegistryError(RuntimeError):
@@ -105,6 +123,13 @@ class McpServerRegistry:
    @_serialized_lifecycle
    def create(self, request: McpServerCreateRequest) -> McpServer:
        self._validate(request)
+        with self._lock:
+            if len(self._records) >= _MAX_MCP_SERVERS:
+                raise McpRegistryError(
+                    "MCP_SERVER_LIMIT_REACHED",
+                    f"At most {_MAX_MCP_SERVERS} MCP servers can be configured.",
+                    status_code=409,
+                )
        server_id = uuid4().hex[:12]
        record = request.model_dump(mode="json")
        record["name"] = request.name.strip()
@@ -137,8 +162,8 @@ class McpServerRegistry:
        self.disable(server_id)
        with self._lock:
            previous = self._record(server_id)
-            removed = [
-                (kind, key)
+            removed_secret_ids = [
+                self._secret_id(server_id, key, kind)
                for kind, old_keys, new_keys in (
                    (
                        "environment",
@@ -153,6 +178,13 @@ class McpServerRegistry:
                )
                for key in set(old_keys) - set(new_keys)
            ]
+        try:
+            self.credentials.delete_many(removed_secret_ids)
+        except CredentialStoreError as exc:
+            raise McpRegistryError(
+                "MCP_SECRET_STORE_ERROR", str(exc), status_code=500
+            ) from exc
+        with self._lock:
            record = request.model_dump(mode="json", exclude={"version"})
            record["name"] = request.name.strip()
            record["command"] = request.command.strip() if request.command else None
@@ -170,13 +202,6 @@ class McpServerRegistry:
            self._records = updated
            self._last_status.pop(server_id, None)
            self._summaries.pop(server_id, None)
-        for kind, key in removed:
-            try:
-                self.credentials.delete(self._secret_id(server_id, key, kind))
-            except CredentialStoreError as exc:
-                raise McpRegistryError(
-                    "MCP_SECRET_STORE_ERROR", str(exc), status_code=500
-                ) from exc
        return self.get(server_id)

    @_serialized_lifecycle
@@ -192,18 +217,19 @@ class McpServerRegistry:
                )
                for key in keys
            ]
-            updated = dict(self._records)
-            del updated[server_id]
-            self._write(updated)
-            self._records = updated
-            self._last_status.pop(server_id, None)
-            self._summaries.pop(server_id, None)
        try:
            self.credentials.delete_many(secret_ids)
        except CredentialStoreError as exc:
            raise McpRegistryError(
                "MCP_SECRET_STORE_ERROR", str(exc), status_code=500
            ) from exc
+        with self._lock:
+            updated = dict(self._records)
+            del updated[server_id]
+            self._write(updated)
+            self._records = updated
+            self._last_status.pop(server_id, None)
+            self._summaries.pop(server_id, None)
        self.bridge.remove(self._host_id(server_id))

    @_serialized_lifecycle
@@ -236,6 +262,9 @@ class McpServerRegistry:
                    "MCP_SECRET_NOT_DECLARED",
                    "Secret environment key is not declared in this server configuration.",
                )
+        if record.get("enabled"):
+            self.disable(server_id)
+        self._invalidate_test(server_id)
        try:
            self.credentials.put(self._secret_id(server_id, key, kind), secret)
        except CredentialStoreError as exc:
@@ -254,6 +283,9 @@ class McpServerRegistry:
                "MCP_SECRET_NOT_DECLARED",
                "Secret environment key is not declared in this server configuration.",
            )
+        if record.get("enabled"):
+            self.disable(server_id)
+        self._invalidate_test(server_id)
        try:
            self.credentials.delete(self._secret_id(server_id, key, kind))
        except CredentialStoreError as exc:
@@ -465,17 +497,27 @@ class McpServerRegistry:
        self.tools.register(definition, arguments_model, executor)

    def _unavailable(self, server_id: str, message: str) -> None:
-        with self._lock:
-            for name in self._registered.pop(server_id, []):
-                self.tools.unregister(name)
-            record = self._records.get(server_id)
-            if record is not None:
-                self._records[server_id] = {**record, "enabled": False}
-                self._last_status[server_id] = {
-                    "status": PluginHostState.unhealthy,
-                    "error": message,
-                }
-                self._write()
+        # A failure may race with enable(). Waiting for the lifecycle mutation makes
+        # sure tools registered immediately before the callback are also removed.
+        with self._lifecycle_lock:
+            try:
+                with self._lock:
+                    record = self._records.get(server_id)
+                    registered = self._registered.pop(server_id, [])
+                    for name in registered:
+                        self.tools.unregister(name)
+                    if record is not None and (record.get("enabled") or registered):
+                        self._records[server_id] = {**record, "enabled": False}
+                        self._last_status[server_id] = {
+                            "status": PluginHostState.unhealthy,
+                            "error": message,
+                        }
+                        self._write()
+            finally:
+                # broken() can run on the client's reader/event thread. stop() does
+                # not join that thread, and setting _stopping before closing the
+                # transport prevents the close itself from reporting another failure.
+                self.bridge.remove(self._host_id(server_id))

    def _require_launch_allowed(
        self, record: dict[str, Any], *, require_test: bool
@@ -767,6 +809,23 @@ class McpServerRegistry:
                status_code=404,
            ) from exc

+    def _invalidate_test(self, server_id: str) -> None:
+        """Make credential changes safe before touching the encrypted store."""
+
+        with self._lock:
+            record = self._record(server_id)
+            invalidated = {
+                **record,
+                "tested_digest": None,
+                "last_tested_at": None,
+                "last_test_succeeded": None,
+            }
+            updated = {**self._records, server_id: invalidated}
+            self._write(updated)
+            self._records = updated
+            self._last_status.pop(server_id, None)
+            self._summaries.pop(server_id, None)
+
    @property
    def _path(self) -> Path:
        return self.data_dir / "mcp" / "servers.json"
@@ -788,7 +847,29 @@ class McpServerRegistry:
                "MCP server registry has an invalid format.",
                status_code=500,
            )
-        return value
+        if len(value) > _MAX_MCP_SERVERS:
+            raise McpRegistryError(
+                "MCP_REGISTRY_INVALID",
+                "MCP server registry contains too many records.",
+                status_code=500,
+            )
+        normalized: dict[str, dict[str, Any]] = {}
+        config_fields = set(McpServerConfig.model_fields)
+        try:
+            for server_id, raw in value.items():
+                if not isinstance(server_id, str) or not _SERVER_ID.fullmatch(server_id):
+                    raise ValueError("invalid server id")
+                record = _McpServerRecord.model_validate(raw)
+                config = record.model_dump(mode="json", include=config_fields)
+                self._validate(McpServerCreateRequest.model_validate(config))
+                normalized[server_id] = record.model_dump(mode="json")
+        except (McpRegistryError, ValidationError, ValueError, TypeError) as exc:
+            raise McpRegistryError(
+                "MCP_REGISTRY_INVALID",
+                "MCP server registry contains an invalid record.",
+                status_code=500,
+            ) from exc
+        return normalized

    def _write(self, records: dict[str, dict[str, Any]] | None = None) -> None:
        temporary = self._path.with_suffix(".tmp")