from typing import Final import pytest from mcp import Client from mcp.server import Server from mcp.shared.exceptions import MCPError from mcp.types import ( ListPromptsResult, ListResourcesResult, PaginatedRequestParams, PromptsCapability, ResourcesCapability, ServerCapabilities, ToolsCapability, ) from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy.mcp_server.capabilities import ( GATEWAY_OPERATIONS, REVISION_SUPPORT, TRANSLATION_PAIRS, GatewayVersionPolicy, build_discovery, disabled_capabilities, enabled_gateway_operations, ) from litellm.types.mcp import MCPTransport @pytest.mark.parametrize("revision", HANDSHAKE_PROTOCOL_VERSIONS) @pytest.mark.parametrize("transport", tuple(MCPTransport)) def test_discovery_only_exposes_authorized_completed_support(revision, transport): result = build_discovery( configured=(revision, "2026-07-28", "unknown"), revision=revision, transport=transport, authorized_operations=frozenset({"tools/list", "tools/call"}), upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=ServerCapabilities( tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability(), extensions={"io.modelcontextprotocol/ui": {}}, ), client_extensions=frozenset({"io.modelcontextprotocol/ui"}), upstream_extensions=frozenset({"io.modelcontextprotocol/ui"}), ) assert result.supported_versions == ([revision] if transport is MCPTransport.sse else [revision, "2026-07-28"]) assert result.capabilities.tools is not None assert result.capabilities.prompts is None assert result.capabilities.resources is None assert result.capabilities.extensions is None assert result.capabilities.tasks is None assert result.cache_scope == "private" assert result.ttl_ms == 0 @pytest.mark.parametrize("upstream", [frozenset(), frozenset({"unknown"})]) def test_unproven_translation_never_advertises_operations(upstream): result = build_discovery( configured=HANDSHAKE_PROTOCOL_VERSIONS, revision="2025-11-25", transport=MCPTransport.http, authorized_operations=GATEWAY_OPERATIONS, upstream_versions=upstream, capabilities=ServerCapabilities(tools=ToolsCapability()), ) assert result.capabilities.tools is None @pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "2024-11-05"]) def test_unadvertised_revision_never_gains_capabilities(revision): result = build_discovery( configured=("2025-11-25",), revision=revision, transport=MCPTransport.http, authorized_operations=GATEWAY_OPERATIONS, upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=ServerCapabilities(tools=ToolsCapability()), ) assert result.capabilities.tools is None def test_discovery_results_do_not_share_mutable_capabilities(): capabilities = ServerCapabilities( tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability() ) args = dict( configured=HANDSHAKE_PROTOCOL_VERSIONS, revision="2025-11-25", transport=MCPTransport.http, upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=capabilities, ) allowed = build_discovery(**args, authorized_operations=GATEWAY_OPERATIONS) denied = build_discovery(**args, authorized_operations=frozenset()) assert allowed.capabilities.prompts is not None assert allowed.capabilities.resources is not None assert denied.capabilities.model_dump(exclude_none=True) == {} assert allowed.capabilities.tools is not None allowed.capabilities.tools.list_changed = True assert capabilities.tools.list_changed is not True def test_modern_support_excludes_legacy_sse() -> None: modern: Final = REVISION_SUPPORT["2026-07-28"] assert modern.completed is True assert "input_required" in modern.results assert MCPTransport.sse not in modern.transports assert ("2026-07-28", "2026-07-28") in TRANSLATION_PAIRS @pytest.mark.asyncio @pytest.mark.parametrize("versions,accepted", [(("2025-11-25",), True), (("2025-06-18",), False)]) async def test_version_policy_gates_the_actual_sdk_handshake(versions, accepted): server: Final = Server("test-gateway", version="1") server.middleware.append(GatewayVersionPolicy(lambda: versions)) if accepted: async with Client(server, mode="legacy") as client: assert client.protocol_version == "2025-11-25" result = await client.session.send_ping() assert result is not None else: with pytest.RaisesGroup( pytest.RaisesExc(MCPError, match="Unsupported MCP protocol version"), flatten_subgroups=True ): async with Client(server, mode="legacy"): pytest.fail("The excluded revision must not initialize") def test_modern_discovery_requires_opt_in_and_keeps_unsupported_features_disabled() -> None: from pydantic import TypeAdapter from litellm.types.mcp import MCPAdvertisedVersions configured: Final = TypeAdapter(MCPAdvertisedVersions).validate_python(["2026-07-28"]) result: Final = build_discovery( configured=configured, revision="2026-07-28", transport=MCPTransport.http, authorized_operations=GATEWAY_OPERATIONS, upstream_versions=frozenset({"2026-07-28"}), capabilities=ServerCapabilities( tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability() ), ) assert result.supported_versions == ["2026-07-28"] assert result.capabilities.tools is not None assert result.capabilities.prompts is not None assert result.capabilities.resources is not None assert result.capabilities.tasks is None assert result.capabilities.extensions is None @pytest.mark.parametrize("path", ["/mcp/sse", "/mcp/example/sse/"]) def test_modern_protocol_is_rejected_on_legacy_sse_paths(path: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server from litellm.proxy.mcp_server.server import unsupported_protocol_version monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2025-11-25", "2026-07-28"]) assert unsupported_protocol_version({"path": "/mcp", "headers": [(b"mcp-protocol-version", b"2026-07-28")]}) is None assert ( unsupported_protocol_version({"path": path, "headers": [(b"mcp-protocol-version", b"2026-07-28")]}) == "2026-07-28" ) @pytest.mark.parametrize("revision,upstreams,transport,allowed", [ ("2026-07-28", frozenset({"2026-07-28"}), MCPTransport.http, True), ("2025-11-25", frozenset({"2026-07-28"}), MCPTransport.http, True), ("2025-11-25", frozenset({"2025-11-25"}), MCPTransport.sse, True), ("2026-07-28", frozenset({"2025-11-25", "2026-07-28"}), MCPTransport.http, True), ("2026-07-28", frozenset({"2025-11-25"}), MCPTransport.http, True), ("2026-07-28", frozenset({"2026-07-28"}), MCPTransport.stdio, False), ]) def test_subscription_flags_require_a_supported_protocol_pair(revision, upstreams, transport, allowed): result = build_discovery( configured=("2025-11-25", "2026-07-28"), revision=revision, transport=transport, authorized_operations=GATEWAY_OPERATIONS, upstream_versions=upstreams, capabilities=ServerCapabilities( tools=ToolsCapability(list_changed=True), prompts=PromptsCapability(list_changed=True), resources=ResourcesCapability(list_changed=True, subscribe=True), ), ) assert bool(result.capabilities.tools.list_changed) is allowed assert bool(result.capabilities.prompts.list_changed) is allowed assert bool(result.capabilities.resources.list_changed) is allowed assert bool(result.capabilities.resources.subscribe) is allowed assert result.capabilities.tasks is None assert result.capabilities.extensions is None @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["empty", "modern-stdio"]) async def test_subscription_discovery_never_probes_unsupported_targets(mode, monkeypatch): from types import SimpleNamespace from unittest.mock import AsyncMock from litellm.proxy.mcp_server import subscriptions from litellm.proxy.mcp_server.contracts import OperationContext probe = AsyncMock() target = SimpleNamespace( server=SimpleNamespace(protocol_version="2026-07-28", transport="stdio"), client=SimpleNamespace(run_with_session=probe), ) monkeypatch.setattr(subscriptions, "_targets", AsyncMock(return_value=() if mode == "empty" else (target,))) versions, capabilities = await subscriptions.subscription_capabilities(OperationContext(_caller=None)) assert versions assert capabilities == ServerCapabilities() probe.assert_not_awaited() @pytest.mark.asyncio @pytest.mark.parametrize( "disabled,prompts_advertised,resources_advertised", [ (frozenset(), True, True), (frozenset({"prompts"}), False, True), (frozenset({"resources"}), True, False), (frozenset({"prompts", "resources"}), False, False), ], ) async def test_version_policy_hides_disabled_capabilities_from_the_handshake( disabled, prompts_advertised, resources_advertised ) -> None: server: Final = Server("test-gateway", version="1") async def list_prompts(ctx: object, params: PaginatedRequestParams) -> ListPromptsResult: return ListPromptsResult(prompts=[]) async def list_resources(ctx: object, params: PaginatedRequestParams) -> ListResourcesResult: return ListResourcesResult(resources=[]) server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) server.add_request_handler("resources/list", PaginatedRequestParams, list_resources) server.middleware.append(GatewayVersionPolicy(lambda: ("2025-11-25",), lambda: disabled)) async with Client(server, mode="legacy") as client: assert (client.server_capabilities.prompts is not None) is prompts_advertised assert (client.server_capabilities.resources is not None) is resources_advertised @pytest.mark.parametrize( "configured,expected", [ (None, frozenset()), (["prompts"], frozenset({"prompts"})), (["prompts", "resources"], frozenset({"prompts", "resources"})), ], ) def test_disabled_capabilities_follow_general_settings(configured, expected, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", configured) assert disabled_capabilities() == expected def test_enabled_gateway_operations_only_drop_the_disabled_capability() -> None: assert enabled_gateway_operations(frozenset()) == GATEWAY_OPERATIONS remaining: Final = enabled_gateway_operations(frozenset({"resources"})) assert {"tools/list", "tools/call", "prompts/list", "prompts/get"} <= remaining assert not any(operation.startswith("resources/") for operation in remaining)