import asyncio import hashlib import json import sys import threading import time from concurrent.futures import ThreadPoolExecutor import httpx 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 CredentialStoreError, EncryptedCredentialStore SERVER = BACKEND_DIR / "extensions" / "fixtures" / "mcp-echo" / "server.py" def request(**overrides) -> McpServerCreateRequest: values = { "name": "Echo MCP", "command": sys.executable, "args": [str(SERVER)], "permissions": ["notes.read", "secrets.use"], "secret_environment_keys": ["TEST_MCP_SECRET"], } values.update(overrides) return McpServerCreateRequest(**values) def registry(*, launch: bool = True) -> McpServerRegistry: return McpServerRegistry( ToolRegistry(), EncryptedCredentialStore(), get_settings().data_dir, allow_process_launch=launch, ) def test_registry_requires_current_trust_and_never_returns_secret() -> None: service = registry() created = service.create(request()) assert created.trusted is False assert created.secret_environment == {"TEST_MCP_SECRET": False} service.put_secret(created.server_id, "TEST_MCP_SECRET", "do-not-return") configured = service.get(created.server_id) assert configured.secret_environment == {"TEST_MCP_SECRET": True} assert "do-not-return" not in configured.model_dump_json() with pytest.raises(McpRegistryError, match="approve"): service.test(created.server_id) service.trust(created.server_id, created.command_digest) tested = service.test(created.server_id) assert tested.status == "stopped" assert tested.last_test_succeeded is True assert tested.tools_count > 0 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=[])) service.trust(created.server_id, created.command_digest) service.test(created.server_id) enabled = service.enable(created.server_id) assert enabled.enabled is True assert any( item.name.startswith(f"mcp.{created.server_id}.") for item in service.tools.definitions() ) updated = service.update( created.server_id, McpServerUpdateRequest( **request(name="Changed", secret_environment_keys=[]).model_dump(), version=enabled.version, ), ) assert updated.enabled is False assert updated.trusted is False assert not any( item.name.startswith(f"mcp.{created.server_id}.") for item in service.tools.definitions() ) 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) generation = object() service._generations[created.server_id] = generation service._unavailable(created.server_id, generation, "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_old_failure_callback_cannot_stop_replacement_host(monkeypatch) -> None: service = registry() callbacks = [] original_start = service.bridge.start def capture_callback(*args, **kwargs): callbacks.append(args[4]) return original_start(*args, **kwargs) monkeypatch.setattr(service.bridge, "start", capture_callback) created = service.create(request(secret_environment_keys=[])) service.trust(created.server_id, created.command_digest) callback_thread = None try: service.test(created.server_id) service.enable(created.server_id) old_callback = callbacks[-1] callback_started = threading.Event() callback_finished = threading.Event() def delayed_failure(): callback_started.set() old_callback(f"mcp.{created.server_id}", "delayed old failure") callback_finished.set() # 将旧回调排队,而替换者拥有生命周期锁。 with service._lifecycle_lock: callback_thread = threading.Thread(target=delayed_failure, daemon=True) callback_thread.start() assert callback_started.wait(timeout=2) service.disable(created.server_id) service.enable(created.server_id) assert callback_finished.wait(timeout=2) assert service.get(created.server_id).enabled is True assert service.get(created.server_id).status == "ready" assert service.tools.definitions() callbacks[-1](f"mcp.{created.server_id}", "current failure") assert service.get(created.server_id).enabled is False assert service.get(created.server_id).status == "unhealthy" finally: service.shutdown() if callback_thread is not None: callback_thread.join(timeout=2) def test_header_case_only_rename_preserves_secret() -> None: service = registry() config = { "name": "HTTP", "transport": "streamable_http", "url": "https://example.test/mcp", "secret_header_keys": ["Authorization"], } created = service.create(McpServerCreateRequest(**config)) service.put_secret(created.server_id, "Authorization", "synthetic", kind="header") config["secret_header_keys"] = ["authorization"] updated = service.update( created.server_id, McpServerUpdateRequest(**config, version=created.version) ) assert updated.secret_headers == {"authorization": True} assert ( service.credentials.resolve( service._secret_id(created.server_id, "authorization", "header") ) == "synthetic" ) def test_environment_secrets_are_case_sensitive_and_delete_independently() -> None: service = registry() created = service.create(request(secret_environment_keys=["TOKEN", "token"])) service.put_secret(created.server_id, "TOKEN", "upper") service.put_secret(created.server_id, "token", "lower") assert ( service.credentials.resolve(service._secret_id(created.server_id, "TOKEN")) == "upper" ) assert ( service.credentials.resolve(service._secret_id(created.server_id, "token")) == "lower" ) service.delete_secret(created.server_id, "TOKEN") assert service.get(created.server_id).secret_environment == { "TOKEN": False, "token": True, } def test_legacy_environment_credential_migration_is_idempotent() -> None: service = registry() created = service.create(request(secret_environment_keys=["TOKEN"])) suffix = hashlib.sha256(b"environment\0token").hexdigest()[:20] legacy_id = f"mcp.{created.server_id}.{suffix}" service.credentials.put(legacy_id, "legacy-value") service._records[created.server_id]["secret_environment_version"] = 1 service._write() migrated = registry() assert migrated.get(created.server_id).secret_environment == {"TOKEN": True} assert ( migrated.credentials.resolve(migrated._secret_id(created.server_id, "TOKEN")) == "legacy-value" ) assert not migrated.credentials.has(legacy_id) migrated.put_secret(created.server_id, "TOKEN", "new-value") assert ( registry().credentials.resolve(migrated._secret_id(created.server_id, "TOKEN")) == "new-value" ) def test_ambiguous_legacy_credentials_are_not_assigned_to_two_variables() -> None: service = registry() created = service.create(request(secret_environment_keys=["TOKEN", "token"])) suffix = hashlib.sha256(b"environment\0token").hexdigest()[:20] legacy_id = f"mcp.{created.server_id}.{suffix}" service.credentials.put(legacy_id, "cannot-reconstruct-originals") service._records[created.server_id]["secret_environment_version"] = 1 service._write() migrated = registry() current = migrated.get(created.server_id) assert current.secret_environment == {"TOKEN": False, "token": False} assert current.enabled is False assert current.last_test_succeeded is None assert migrated.credentials.has( legacy_id ) # 保持原始密文可恢复。 migrated.put_secret(created.server_id, "TOKEN", "upper") migrated.put_secret(created.server_id, "token", "lower") assert registry().get(created.server_id).secret_environment == { "TOKEN": True, "token": True, } migrated.delete(created.server_id) assert not migrated.credentials.has(legacy_id) def test_credential_id_migration_keeps_new_values_and_is_atomic(monkeypatch) -> None: credentials = EncryptedCredentialStore() credentials.put("mcp.old", "old-value") credentials.put("mcp.new", "new-value") original_write = credentials._write_tokens def fail_write(_tokens): raise CredentialStoreError("synthetic failure") monkeypatch.setattr(credentials, "_write_tokens", fail_write) with pytest.raises(CredentialStoreError): credentials.move_many({"mcp.old": "mcp.new"}) assert credentials.resolve("mcp.old") == "old-value" assert credentials.resolve("mcp.new") == "new-value" monkeypatch.setattr(credentials, "_write_tokens", original_write) credentials.move_many({"mcp.old": "mcp.new"}) assert credentials.resolve("mcp.old") is None assert credentials.resolve("mcp.new") == "new-value" def test_ambiguous_legacy_secret_is_not_resurrected_after_removing_a_key() -> None: service = registry() created = service.create(request(secret_environment_keys=["TOKEN", "token"])) legacy_id = service._legacy_environment_secret_id(created.server_id, "TOKEN") service.credentials.put(legacy_id, "ambiguous-old-value") service._records[created.server_id]["secret_environment_version"] = 1 service._write() migrated = registry() migrated.update( created.server_id, McpServerUpdateRequest( **request(secret_environment_keys=["token"]).model_dump(), version=created.version, ), ) assert registry().get(created.server_id).secret_environment == {"token": False} def test_production_rejects_process_launch_even_after_approval() -> None: service = registry(launch=False) created = service.create(request(secret_environment_keys=[])) service.trust(created.server_id, created.command_digest) with pytest.raises(McpRegistryError) as error: service.enable(created.server_id) assert error.value.code == "MCP_SANDBOX_REQUIRED" @pytest.mark.parametrize("startup,tool", [(120, 300), (1.5, 2.5)]) def test_server_timeouts_survive_bridge_adaptation_and_reload(startup, tool) -> None: service = registry() created = service.create( request( secret_environment_keys=[], startup_timeout_seconds=startup, tool_timeout_seconds=tool, ) ) service.trust(created.server_id, created.command_digest) try: tested = service.test(created.server_id) assert tested.last_test_succeeded is True assert tested.startup_timeout_seconds == startup assert tested.tool_timeout_seconds == tool restored = registry().get(created.server_id) assert restored.startup_timeout_seconds == startup assert restored.tool_timeout_seconds == tool finally: service.shutdown() def test_enable_requires_successful_test_and_update_checks_version() -> None: service = registry() created = service.create(request(secret_environment_keys=[])) service.trust(created.server_id, created.command_digest) with pytest.raises(McpRegistryError) as error: service.enable(created.server_id) assert error.value.code == "MCP_CONNECTION_TEST_REQUIRED" with pytest.raises(McpRegistryError) as error: service.update( created.server_id, McpServerUpdateRequest( **request(secret_environment_keys=[]).model_dump(), version=99 ), ) assert error.value.code == "MCP_SERVER_VERSION_CONFLICT" def test_http_transport_rejects_invalid_cross_transport_fields() -> None: service = registry() with pytest.raises(McpRegistryError) as error: service.create( request( transport="streamable_http", url="https://example.invalid/mcp", secret_environment_keys=[], ) ) assert error.value.code == "MCP_CONFIG_INVALID" def test_registry_rejects_corrupt_persisted_json(tmp_path) -> None: path = tmp_path / "mcp" path.mkdir() (path / "servers.json").write_text("{broken", 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_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( request( command=f'"{sys.executable}" "{SERVER}"', args=[], secret_environment_keys=[], ) ) service.trust(created.server_id, created.command_digest) with pytest.raises(McpRegistryError) as error: service.test(created.server_id) assert error.value.code == "PLUGIN_HOST_START_FAILED" assert service.get(created.server_id).last_test_succeeded is False service.shutdown() def test_enabled_server_is_restored_from_persisted_registry() -> None: first = registry() created = first.create(request(secret_environment_keys=[])) first.trust(created.server_id, created.command_digest) first.test(created.server_id) first.enable(created.server_id) first.shutdown() restored = registry() restored.restore_enabled() current = restored.get(created.server_id) assert current.enabled is True assert current.status == "ready" assert any( item.name.startswith(f"mcp.{created.server_id}.") for item in restored.tools.definitions() ) restored.shutdown() def test_lifecycle_operations_are_serialized_and_tool_names_are_isolated() -> None: service = registry() servers = [ service.create(request(name=f"Echo {index}", secret_environment_keys=[])) for index in range(2) ] for server in servers: service.trust(server.server_id, server.command_digest) service.test(server.server_id) with ThreadPoolExecutor(max_workers=4) as pool: enabled = list( pool.map(lambda item: service.enable(item.server_id), servers * 2) ) assert all(item.enabled for item in enabled) names = [ item.name for item in service.tools.definitions() if item.source == "mcp_server" ] assert len(names) == len(set(names)) assert all( any(name.startswith(f"mcp.{item.server_id}.") for name in names) for item in servers ) with ThreadPoolExecutor(max_workers=4) as pool: list(pool.map(lambda item: service.disable(item.server_id), servers * 2)) assert not any(item.source == "mcp_server" for item in service.tools.definitions()) service.shutdown() def _http_result(request_id: int, result: dict) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "application/json"}, json={"jsonrpc": "2.0", "id": request_id, "result": result}, ) 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) if request_value.method == "GET": return httpx.Response(405) 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"], { "protocolVersion": "2025-11-25", "capabilities": {"tools": {}}, "serverInfo": {"name": "HTTP Fixture", "version": "1"}, }, ) response.headers["MCP-Session-Id"] = "session-test" return response if payload.get("method") == "tools/list": return _http_result( payload["id"], { "tools": [ { "name": "echo", "description": "Echo over HTTP", "inputSchema": {"type": "object", "properties": {}}, } ] }, ) if payload.get("method") == "tools/call": return _http_result( payload["id"], {"structuredContent": {"transport": "http"}} ) return httpx.Response(202) real_client = httpx.Client monkeypatch.setattr( "app.extensions.mcp.httpx.Client", lambda **kwargs: real_client(transport=httpx.MockTransport(handler), **kwargs), ) service = registry() created = service.create( McpServerCreateRequest( name="Remote MCP", transport="streamable_http", url="https://mcp.example.test/mcp", headers={"X-Client": "NotesAgent"}, secret_header_keys=["Authorization"], ) ) service.put_secret( created.server_id, "Authorization", "Bearer hidden", kind="header" ) service.trust(created.server_id, created.command_digest) tested = service.test(created.server_id) assert tested.last_test_succeeded is True assert tested.secret_headers == {"Authorization": True} assert "Bearer hidden" not in tested.model_dump_json() assert service.list_tools(created.server_id)[0].remote_name == "echo" assert any( request.headers.get("mcp-session-id") == "session-test" for request in requests ) assert any( request.headers.get("mcp-protocol-version") == "2025-11-25" for request in requests ) 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( service.tools.execute( ToolCall(tool_call_id="call-1", name=tool_name, arguments={}), ToolExecutionContext(run_id="run-1"), ) ) 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) initialize = { "jsonrpc": "2.0", "id": 1, "result": { "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "Legacy Fixture"}, }, } yield f"data: {json.dumps(initialize)}\n\n".encode() time.sleep(0.1) tools = { "jsonrpc": "2.0", "id": 2, "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=event_stream, ) posted_urls.append(str(request_value.url)) return httpx.Response(202) real_client = httpx.Client monkeypatch.setattr( "app.extensions.mcp.httpx.Client", lambda **kwargs: real_client(transport=httpx.MockTransport(handler), **kwargs), ) service = registry() created = service.create( McpServerCreateRequest( name="Legacy MCP", transport="sse", url="https://legacy.example.test/sse", ) ) service.trust(created.server_id, created.command_digest) tested = service.test(created.server_id) assert tested.last_test_succeeded is True assert posted_urls and all( 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): def __iter__(self): yield b"event: endpoint\ndata: https://attacker.example/messages\n\n" def test_legacy_sse_rejects_cross_origin_message_endpoint(monkeypatch) -> None: def handler(request_value: httpx.Request) -> httpx.Response: assert request_value.method == "GET" return httpx.Response( 200, headers={"content-type": "text/event-stream"}, stream=_CrossOriginLegacyEventStream(), ) real_client = httpx.Client monkeypatch.setattr( "app.extensions.mcp.httpx.Client", lambda **kwargs: real_client(transport=httpx.MockTransport(handler), **kwargs), ) service = registry() created = service.create( McpServerCreateRequest( name="Unsafe legacy MCP", transport="sse", url="https://legacy.example.test/sse", ) ) service.trust(created.server_id, created.command_digest) with pytest.raises(McpRegistryError) as error: service.test(created.server_id) assert error.value.code == "MCP_HTTP_RESPONSE_INVALID" service.shutdown()