litellm/tests/unit/proxy/mcp_server/test_capabilities.py

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)