Kronecker/NotesAgentic
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 5 Wiki Activity

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

Browse Source 
添加了MCP客户端的超时配置功能，包括启动超时和工具调用超时参数。
改进了HTTP客户端和标准IO客户端的超时处理机制，确保请求在指定时间内完成或取消。
增加了对MCP服务器数量的限制，防止配置过多服务器导致系统不稳定。
增强了错误处理机制，当连接异常时能够正确清理资源并移除桥接主机。
添加了对大型MCP消息的大小验证，防止过大的请求导致系统问题。
优化了密钥更改后的处理流程，确保在修改密钥时停用服务器并要求重新测试。
This commit is contained in:
admin
2026-09-03 16:16:58 +08:00
parent 7d5f4023a9
commit 2f7066aa92
11 changed files with 504 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/tests/test_api.py
+69 -20
View File
@@ -1,26 +1,10 @@
import asyncio
import threading
from types import SimpleNamespace
from app.main import health, service_status
from app.routes import (
get_index_status,
list_notes,
list_plugins,
list_provider_presets,
list_providers,
list_skills,
)
from app.routes import (
create_provider,
create_task,
delete_provider,
delete_task,
get_provider,
get_task,
list_tasks,
update_provider,
update_task,
)
from app.contracts import (
McpServerSecretStatus,
McpServerSecretWriteRequest,
ProviderCreateRequest,
ProviderType,
ProviderUpdateRequest,
@@ -28,6 +12,63 @@ from app.contracts import (
TaskStatus,
TaskUpdateRequest,
)
from app.main import health, service_status
from app.routes import (
create_provider,
create_task,
delete_provider,
delete_task,
get_index_status,
get_provider,
get_task,
list_notes,
list_plugins,
list_provider_presets,
list_providers,
list_skills,
list_tasks,
update_provider,
update_task,
)
def test_mcp_secret_routes_offload_blocking_lifecycle_work(monkeypatch) -> None:
from app import routes
caller_thread = threading.get_ident()
worker_threads: list[int] = []
class FakeMcpRegistry:
def put_secret(self, server_id, key, secret, *, kind):
worker_threads.append(threading.get_ident())
return McpServerSecretStatus(key=key, configured=True)
def delete_secret(self, server_id, key, *, kind):
worker_threads.append(threading.get_ident())
return McpServerSecretStatus(key=key, configured=False)
monkeypatch.setattr(
routes,
"container",
SimpleNamespace(mcp_servers=FakeMcpRegistry()),
)
written = asyncio.run(
routes.put_mcp_server_secret(
"server-1",
"TOKEN",
McpServerSecretWriteRequest(secret="hidden"),
kind="environment",
)
)
deleted = asyncio.run(
routes.delete_mcp_server_secret(
"server-1", "TOKEN", kind="environment"
)
)
assert written.configured is True
assert deleted.configured is False
assert worker_threads and all(item != caller_thread for item in worker_threads)
def test_health() -> None:
@@ -101,6 +142,14 @@ def test_openapi_contains_documented_frontend_interfaces() -> None:
"/api/plugins/{plugin_id}/settings/{key}/secret",
"/api/plugins/{plugin_id}/enable",
"/api/plugins/{plugin_id}/disable",
"/api/mcp/servers",
"/api/mcp/servers/{server_id}",
"/api/mcp/servers/{server_id}/tools",
"/api/mcp/servers/{server_id}/trust",
"/api/mcp/servers/{server_id}/test",
"/api/mcp/servers/{server_id}/enable",
"/api/mcp/servers/{server_id}/disable",
"/api/mcp/servers/{server_id}/secrets/{key}",
"/api/providers/test",
"/api/providers/presets",
"/api/credentials/{credential_id}",
backend/tests/test_mcp_registry.py
+169 -2
View File
@@ -1,6 +1,7 @@
import asyncio
import json
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
@@ -10,8 +11,9 @@ import pytest
from app.agent.tools import ToolExecutionContext, ToolRegistry
from app.config import BACKEND_DIR, get_settings
from app.contracts import McpServerCreateRequest, McpServerUpdateRequest, ToolCall
from app.extensions.mcp import McpLegacySseClient
from app.extensions.mcp_registry import McpRegistryError, McpServerRegistry
from app.providers.credentials import EncryptedCredentialStore
from app.providers.credentials import CredentialStoreError, EncryptedCredentialStore
SERVER = BACKEND_DIR / "extensions" / "fixtures" / "mcp-echo" / "server.py"
@@ -59,6 +61,28 @@ def test_registry_requires_current_trust_and_never_returns_secret() -> None:
service.shutdown()
def test_secret_change_disables_server_and_requires_a_new_connection_test() -> None:
service = registry()
created = service.create(request())
service.put_secret(created.server_id, "TEST_MCP_SECRET", "first")
service.trust(created.server_id, created.command_digest)
service.test(created.server_id)
service.enable(created.server_id)
service.put_secret(created.server_id, "TEST_MCP_SECRET", "second")
current = service.get(created.server_id)
assert current.enabled is False
assert current.last_test_succeeded is None
assert not any(
item.name.startswith(f"mcp.{created.server_id}.")
for item in service.tools.definitions()
)
with pytest.raises(McpRegistryError) as error:
service.enable(created.server_id)
assert error.value.code == "MCP_CONNECTION_TEST_REQUIRED"
service.shutdown()
def test_update_disables_server_and_revokes_command_trust() -> None:
service = registry()
created = service.create(request(secret_environment_keys=[]))
@@ -87,6 +111,70 @@ def test_update_disables_server_and_revokes_command_trust() -> None:
service.shutdown()
def test_update_remains_retryable_when_removed_secret_cleanup_fails(monkeypatch) -> None:
service = registry()
created = service.create(request())
service.put_secret(created.server_id, "TEST_MCP_SECRET", "keep-until-retry")
def fail_delete_many(_secret_ids: list[str]) -> set[str]:
raise CredentialStoreError("credential store unavailable")
monkeypatch.setattr(service.credentials, "delete_many", fail_delete_many)
with pytest.raises(McpRegistryError) as error:
service.update(
created.server_id,
McpServerUpdateRequest(
**request(secret_environment_keys=[]).model_dump(),
version=created.version,
),
)
current = service.get(created.server_id)
assert error.value.code == "MCP_SECRET_STORE_ERROR"
assert current.version == created.version
assert current.secret_environment == {"TEST_MCP_SECRET": True}
service.shutdown()
def test_delete_keeps_server_retryable_when_secret_cleanup_fails(monkeypatch) -> None:
service = registry()
created = service.create(request())
service.put_secret(created.server_id, "TEST_MCP_SECRET", "keep-until-retry")
def fail_delete_many(_secret_ids: list[str]) -> set[str]:
raise CredentialStoreError("credential store unavailable")
monkeypatch.setattr(service.credentials, "delete_many", fail_delete_many)
with pytest.raises(McpRegistryError) as error:
service.delete(created.server_id)
current = service.get(created.server_id)
assert error.value.code == "MCP_SECRET_STORE_ERROR"
assert current.server_id == created.server_id
assert current.secret_environment == {"TEST_MCP_SECRET": True}
service.shutdown()
def test_unavailable_server_removes_bridge_host(monkeypatch) -> None:
service = registry()
created = service.create(request(secret_environment_keys=[]))
with service._lock:
service._records[created.server_id] = {
**service._records[created.server_id],
"enabled": True,
}
removed: list[str] = []
monkeypatch.setattr(service.bridge, "remove", removed.append)
service._unavailable(created.server_id, "connection lost")
current = service.get(created.server_id)
assert removed == [f"mcp.{created.server_id}"]
assert current.enabled is False
assert current.status == "unhealthy"
service.shutdown()
def test_production_rejects_process_launch_even_after_approval() -> None:
service = registry(launch=False)
created = service.create(request(secret_environment_keys=[]))
@@ -141,6 +229,36 @@ def test_registry_rejects_corrupt_persisted_json(tmp_path) -> None:
assert error.value.code == "MCP_REGISTRY_INVALID"
def test_registry_rejects_structurally_invalid_record(tmp_path) -> None:
path = tmp_path / "mcp"
path.mkdir()
(path / "servers.json").write_text(
json.dumps({"server-1": {"name": "Broken", "transport": "stdio"}}),
encoding="utf-8",
)
with pytest.raises(McpRegistryError) as error:
McpServerRegistry(
ToolRegistry(),
EncryptedCredentialStore(),
tmp_path,
allow_process_launch=True,
)
assert error.value.code == "MCP_REGISTRY_INVALID"
def test_registry_rejects_create_before_exceeding_persisted_limit(
monkeypatch,
) -> None:
service = registry()
service.create(request(name="Only server"))
monkeypatch.setattr("app.extensions.mcp_registry._MAX_MCP_SERVERS", 1)
with pytest.raises(McpRegistryError) as error:
service.create(request(name="One too many"))
assert error.value.code == "MCP_SERVER_LIMIT_REACHED"
assert len(service.list()) == 1
service.shutdown()
def test_stdio_command_is_not_parsed_as_a_shell_string() -> None:
service = registry()
created = service.create(
@@ -217,6 +335,7 @@ def test_streamable_http_supports_session_headers_secrets_and_tool_summary(
monkeypatch,
) -> None:
requests: list[httpx.Request] = []
request_timeouts: dict[str, float] = {}
def handler(request_value: httpx.Request) -> httpx.Response:
requests.append(request_value)
@@ -225,6 +344,9 @@ def test_streamable_http_supports_session_headers_secrets_and_tool_summary(
if request_value.method == "DELETE":
return httpx.Response(405)
payload = json.loads(request_value.content)
timeout = request_value.extensions.get("timeout", {}).get("read")
if isinstance(timeout, (int, float)):
request_timeouts[payload.get("method", "notification")] = float(timeout)
if payload.get("method") == "initialize":
response = _http_result(
payload["id"],
@@ -290,6 +412,9 @@ def test_streamable_http_supports_session_headers_secrets_and_tool_summary(
assert all(
request.headers.get("authorization") == "Bearer hidden" for request in requests
)
assert request_timeouts["initialize"] == 15
assert request_timeouts["notifications/initialized"] == 15
assert request_timeouts["tools/list"] == 15
enabled = service.enable(created.server_id)
tool_name = service.list_tools(created.server_id)[0].name
result = asyncio.run(
@@ -301,11 +426,15 @@ def test_streamable_http_supports_session_headers_secrets_and_tool_summary(
assert enabled.enabled is True
assert result.success is True
assert result.output == {"transport": "http"}
assert request_timeouts["tools/call"] == 30
service.disable(created.server_id)
service.shutdown()
class _LegacyEventStream(httpx.SyncByteStream):
def __init__(self) -> None:
self.closed = threading.Event()
def __iter__(self):
yield b"event: endpoint\ndata: /messages\n\n"
time.sleep(0.1)
@@ -326,17 +455,22 @@ class _LegacyEventStream(httpx.SyncByteStream):
"result": {"tools": []},
}
yield f"data: {json.dumps(tools)}\n\n".encode()
self.closed.wait()
def close(self) -> None:
self.closed.set()
def test_legacy_sse_uses_same_origin_endpoint(monkeypatch) -> None:
posted_urls: list[str] = []
event_stream = _LegacyEventStream()
def handler(request_value: httpx.Request) -> httpx.Response:
if request_value.method == "GET":
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
stream=_LegacyEventStream(),
stream=event_stream,
)
posted_urls.append(str(request_value.url))
return httpx.Response(202)
@@ -361,6 +495,39 @@ def test_legacy_sse_uses_same_origin_endpoint(monkeypatch) -> None:
url == "https://legacy.example.test/messages" for url in posted_urls
)
service.shutdown()
event_stream.close()
class _EndingLegacyEventStream(httpx.SyncByteStream):
def __iter__(self):
yield b"event: endpoint\ndata: /messages\n\n"
def test_legacy_sse_eof_marks_client_unavailable(monkeypatch) -> None:
def handler(_request_value: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
headers={"content-type": "text/event-stream"},
stream=_EndingLegacyEventStream(),
)
real_client = httpx.Client
monkeypatch.setattr(
"app.extensions.mcp.httpx.Client",
lambda **kwargs: real_client(transport=httpx.MockTransport(handler), **kwargs),
)
broken = threading.Event()
client = McpLegacySseClient(
"https://legacy.example.test/sse",
headers={},
startup_timeout_seconds=1,
on_seen=lambda: None,
on_broken=lambda _message: broken.set(),
on_tools_changed=lambda: None,
)
client.start()
assert broken.wait(timeout=1)
client.stop()
class _CrossOriginLegacyEventStream(httpx.SyncByteStream):