diff --git a/litellm/proxy/_experimental/mcp_server/capabilities.py b/litellm/proxy/_experimental/mcp_server/capabilities.py index e12c2346d1a..92c8004ff6f 100644 --- a/litellm/proxy/_experimental/mcp_server/capabilities.py +++ b/litellm/proxy/_experimental/mcp_server/capabilities.py @@ -30,6 +30,7 @@ GATEWAY_OPERATIONS: Final = frozenset( "resources/list", "resources/read", "resources/templates/list", + "subscriptions/listen", } ) _CAPABILITY_OPERATIONS: Final[Mapping[MCPDisableableCapability, frozenset[str]]] = MappingProxyType( @@ -124,13 +125,30 @@ def build_discovery( if operations and revision_support is not None else frozenset() ) + notifications: Final = revision in HANDSHAKE_PROTOCOL_VERSIONS or ( + "subscriptions/listen" in operations and revision == "2026-07-28" and transport is MCPTransport.http + ) caller_capabilities: Final = capabilities.model_copy(deep=True) + tools: Final = caller_capabilities.tools + prompts: Final = caller_capabilities.prompts + resources: Final = caller_capabilities.resources return DiscoverResult( supported_versions=list(supported), capabilities=ServerCapabilities( - tools=caller_capabilities.tools if {"tools/list", "tools/call"} <= operations else None, - prompts=caller_capabilities.prompts if {"prompts/list", "prompts/get"} <= operations else None, - resources=caller_capabilities.resources if {"resources/list", "resources/read"} <= operations else None, + tools=tools.model_copy(update={"list_changed": tools.list_changed if notifications else None}) + if tools is not None and {"tools/list", "tools/call"} <= operations + else None, + prompts=prompts.model_copy(update={"list_changed": prompts.list_changed if notifications else None}) + if prompts is not None and {"prompts/list", "prompts/get"} <= operations + else None, + resources=resources.model_copy( + update={ + "list_changed": resources.list_changed if notifications else None, + "subscribe": resources.subscribe if notifications else None, + } + ) + if resources is not None and {"resources/list", "resources/read"} <= operations + else None, extensions={ key: value for key, value in (caller_capabilities.extensions or {}).items() if key in extensions } diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 43435a6fcb2..2de54e8ce1b 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -13,6 +13,7 @@ from functools import partial from types import EllipsisType from typing import Any, Final, NoReturn, TypeAlias, overload +import anyio from fastapi import HTTPException from mcp import MCPError, ReadResourceResult, Resource from mcp.types import ( @@ -52,6 +53,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.constants import ( MAXIMUM_TRACEBACK_LINES_TO_LOG, + MCP_TOOL_LISTING_TIMEOUT, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -3501,6 +3503,7 @@ class GatewayOperations: async def _execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: case DiscoverRequest(): + deadline: Final = time.monotonic() + MCP_TOOL_LISTING_TIMEOUT disabled: Final = disabled_capabilities() listings: Final = ( () @@ -3524,6 +3527,17 @@ class GatewayOperations: ) try: results: Final = await asyncio.gather(*tasks) + from litellm.proxy._experimental.mcp_server.subscriptions import subscription_capabilities + + async def discover_notifications() -> tuple[frozenset[str], ServerCapabilities]: + remaining: Final = deadline - time.monotonic() + if context.protocol_version == "2026-07-28" and not context.mcp_proxy_mode and remaining > 0: + # Optional probes share the catalog budget, including target preparation. + with anyio.move_on_after(remaining): + return await subscription_capabilities(context) + return frozenset(MCP_LEGACY_VERSIONS), ServerCapabilities() + + upstream_versions, notifications = await discover_notifications() finally: try: for task in tasks: @@ -3536,21 +3550,30 @@ class GatewayOperations: revision=context.protocol_version or "2025-11-25", transport=MCPTransport.http, authorized_operations=enabled_gateway_operations(disabled), - upstream_versions=frozenset(MCP_LEGACY_VERSIONS), + upstream_versions=upstream_versions, capabilities=ServerCapabilities( - tools=ToolsCapability() - if any(isinstance(result, ListToolsResult) and result.tools for result in results) - else None, - prompts=PromptsCapability() - if any(isinstance(result, ListPromptsResult) and result.prompts for result in results) - else None, - resources=ResourcesCapability() - if any( - (isinstance(result, ListResourcesResult) and result.resources) - or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates) - for result in results - ) - else None, + tools=notifications.tools + or ( + ToolsCapability() + if any(isinstance(result, ListToolsResult) and result.tools for result in results) + else None + ), + prompts=notifications.prompts + or ( + PromptsCapability() + if any(isinstance(result, ListPromptsResult) and result.prompts for result in results) + else None + ), + resources=notifications.resources + or ( + ResourcesCapability() + if any( + (isinstance(result, ListResourcesResult) and result.resources) + or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates) + for result in results + ) + else None + ), ), ) case AuthorizedToolCall(): diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py index bc22e29ef96..dcf24f974fe 100644 --- a/litellm/proxy/_experimental/mcp_server/subscriptions.py +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -32,7 +32,9 @@ from mcp.types import ( ListToolsRequest, ListToolsResult, PaginatedRequestParams, + PromptsCapability, RequestId, + ResourcesCapability, ServerCapabilities, ServerNotification, SubscribeRequest, @@ -40,6 +42,7 @@ from mcp.types import ( SubscriptionFilter, SubscriptionsListenRequestParams, SubscriptionsListenResult, + ToolsCapability, UnsubscribeRequest, UnsubscribeRequestParams, ) @@ -52,6 +55,7 @@ from litellm.proxy._experimental.mcp_server.catalog import ( ) from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp import MCP_LEGACY_VERSIONS from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -219,6 +223,50 @@ async def _targets(context: OperationContext) -> tuple[SubscriptionTarget, ...]: return tuple([await prepare(server) for server in allowed]) +async def subscription_capabilities(context: OperationContext) -> tuple[frozenset[str], ServerCapabilities]: + try: + targets: Final = await _targets(context) + except Exception: # noqa: BLE001 # missing optional upstream credentials must not erase usable catalog results + return frozenset(MCP_LEGACY_VERSIONS), ServerCapabilities() + configured: Final = frozenset(target.server.protocol_version for target in targets) + versions: Final = ( + configured | frozenset(MCP_LEGACY_VERSIONS) if not configured or "auto" in configured else configured + ) + if not targets or any( + target.server.protocol_version == "2026-07-28" and target.server.transport != "http" for target in targets + ): + return versions, ServerCapabilities() + + async def probe(target: SubscriptionTarget) -> ServerCapabilities: + async def read(session: ClientSession) -> ServerCapabilities: + return session.server_capabilities or ServerCapabilities() + + try: + if await operations._mcp_server_rate_limit_rejection(target.server, context.user_api_key_auth) is not None: # pyright: ignore[reportPrivateUsage] # reuse request-scoped catalog admission + return ServerCapabilities() + with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): + return await target.client.run_with_session(read, quiet_on_error=True) + except Exception: # noqa: BLE001 # an optional notification probe must not disable an otherwise usable catalog + return ServerCapabilities() + + declared: Final = tuple(await asyncio.gather(*(probe(target) for target in targets))) + return versions, ServerCapabilities( + tools=ToolsCapability(list_changed=True) + if all(item.tools and item.tools.list_changed for item in declared) + else None, + prompts=PromptsCapability(list_changed=True) + if all(item.prompts and item.prompts.list_changed for item in declared) + else None, + resources=ResourcesCapability( + list_changed=all(item.resources and item.resources.list_changed for item in declared) or None, + subscribe=(len(targets) == 1 and any(item.resources and item.resources.subscribe for item in declared)) + or None, + ) + if any(item.resources for item in declared) + else None, + ) + + async def _listing(event: ServerEvent, context: OperationContext) -> str | None: if isinstance(event, ResourceUpdated): return None @@ -432,6 +480,10 @@ class GatewaySubscriptions: ) -> SubscriptionsListenResult: context: Final = await authorize() targets: Final = await self._resolve_targets(context) + if any( + target.server.protocol_version == "2026-07-28" and target.server.transport != "http" for target in targets + ): + raise MCPError(code=-32602, message="Modern upstream subscriptions require HTTP") if not targets: raise HTTPException(status_code=403, detail="No authorized MCP subscription targets") if params.notifications.resource_subscriptions and len(targets) != 1: diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py index a4032bcb5f3..f69079ce90b 100644 --- a/tests/integration/_support/mcp.py +++ b/tests/integration/_support/mcp.py @@ -26,7 +26,7 @@ from pydantic import BaseModel, JsonValue from sse_starlette.sse import AppStatus from starlette.requests import Request as StarletteRequest from starlette.responses import Response -from starlette.types import Message, Receive, Scope, Send +from starlette.types import ASGIApp, Message, Receive, Scope, Send Transport = Literal["http", "sse", "stdio"] STDIO_PEER: Final = Path(__file__).with_name("mcp_stdio_peer.py") @@ -140,6 +140,70 @@ def _capturing(app: Callable[[Scope, Receive, Send], object], observed: queue.Qu return capture +def _acknowledge_notifications_after_dispatch(app: ASGIApp) -> ASGIApp: + in_flight: Final[dict[bytes, set[asyncio.Event]]] = {} + + async def dispatch(scope: Scope, receive: Receive, send: Send, dispatched: asyncio.Event | None = None) -> None: + if scope["type"] != "http" or scope["method"] != "POST": + await app(scope, receive, send) + return + pending: Final[list[Message]] = [] + + async def capture(message: Message) -> None: + if not pending and message["type"] == "http.response.start" and message["status"] == 202: + pending.append(message) + elif pending: + if ( + len(pending) != 1 + or message["type"] != "http.response.body" + or message.get("body", b"") + or message.get("more_body", False) + ): + raise RuntimeError("MCP notification acknowledgement must be a single empty response") + pending.append(message) + else: + if dispatched is not None and message["type"] == "http.response.start": + dispatched.set() + await send(message) + + try: + await app(scope, receive, capture) + finally: + if dispatched is not None: + dispatched.set() + if pending: + if len(pending) != 2: + raise RuntimeError("MCP notification acknowledgement did not complete") + for message in pending: + await send(message) + + async def ordered(scope: Scope, receive: Receive, send: Send) -> None: + session: Final = next( + (value for name, value in scope.get("headers", ()) if name.lower() == b"mcp-session-id"), None + ) + if scope["type"] != "http" or session is None: + await dispatch(scope, receive, send) + return + if scope["method"] == "DELETE": + while pending := tuple(done for done in in_flight.get(session, ()) if not done.is_set()): + await asyncio.gather(*(done.wait() for done in pending)) + if scope["method"] != "POST": + await dispatch(scope, receive, send) + return + done: Final = asyncio.Event() + active: Final = in_flight.setdefault(session, set()) + active.add(done) + try: + await dispatch(scope, receive, send, done) + finally: + active.remove(done) + if not active: + del in_flight[session] + done.set() + + return ordered + + def _drain_sse_streams() -> None: AppStatus.should_exit = True @@ -565,6 +629,7 @@ class McpCaller: entry: EntryPoint alias: str | None = None headers: Mapping[str, str] = field(default_factory=dict) + protocol_version: str | None = None def _path(self) -> str: if self.entry == "server_mcp": @@ -582,10 +647,23 @@ class McpCaller: def rpc(self, method: str, params: JsonRpc | None = None) -> httpx.Response: if self.entry == "sse": return _legacy_sse_rpc(self.gateway, self._headers(), method, params) + modern: Final = self.protocol_version == "2026-07-28" + arguments: Final = dict(params or {}) return self.gateway.client.post( self._path(), - json={"jsonrpc": "2.0", "id": 1, "method": method, "params": dict(params or {})}, - headers=self._headers(), + json={"jsonrpc": "2.0", "id": 1, "method": method, "params": { + **arguments, + **({"_meta": { + "io.modelcontextprotocol/protocolVersion": self.protocol_version, + "io.modelcontextprotocol/clientInfo": {"name": "integration-caller", "version": "1"}, + "io.modelcontextprotocol/clientCapabilities": {}, + }} if modern else {}), + }}, + headers={**self._headers(), **({ + "MCP-Protocol-Version": self.protocol_version, + "Mcp-Method": method, + "Mcp-Name": str(arguments.get("name", arguments.get("uri", ""))), + } if modern else {})}, ) def initialize(self) -> Outcome: diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index c254139c109..8a4bdb8e8ec 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -1,5 +1,6 @@ import json import os +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Final @@ -7,6 +8,7 @@ import httpx import pytest import yaml from integration._support.client import Gateway +from integration._support.mcp import McpCaller from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -85,6 +87,122 @@ def interaction_peer(request: Request) -> Reply: ) +@contextmanager +def sdk_interaction_peer(*, sampling: bool = False): + import asyncio + import queue + from contextlib import asynccontextmanager + + from mcp.server import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.server.transport_security import TransportSecuritySettings + from mcp.types import ( + CallToolResult, + CreateMessageRequest, + CreateMessageRequestParams, + SamplingMessage, + ElicitRequest, + ElicitRequestFormParams, + GetPromptResult, + InputRequiredResult, + ListToolsResult, + ListPromptsResult, + ListResourcesResult, + Prompt, + PromptMessage, + ReadResourceResult, + Resource, + TextContent, + TextResourceContents, + Tool, + ) + from starlette.applications import Starlette + from starlette.routing import Mount + from integration._support.asgi import asgi_server + from integration._support.mcp import McpPeer, _capturing, _drain_sse_streams + from sse_starlette.sse import AppStatus + + async def tools(ctx, params): + return ListToolsResult(tools=[Tool(name="confirm", input_schema={"type": "object"})]) + + async def prompts(ctx, params): + return ListPromptsResult(prompts=[Prompt(name="confirm")]) + + async def resources(ctx, params): + return ListResourcesResult(resources=[Resource(name="confirm", uri="test://confirm")]) + + async def interact(ctx, params): + if sampling: + if not params.request_state: + return InputRequiredResult( + request_state="sample-state", + input_requests={ + "sample": CreateMessageRequest( + params=CreateMessageRequestParams( + messages=[ + SamplingMessage(role="user", content=TextContent(type="text", text="sample prompt")) + ], + max_tokens=8, + ) + ) + }, + ) + assert params.request_state == "sample-state" + sampled = params.input_responses["sample"] + assert "sampled response" in sampled.model_dump_json() + return CallToolResult(content=[TextContent(type="text", text="sample confirmed")]) + if not params.request_state: + return InputRequiredResult( + request_state="opaque:" + ctx.method, + input_requests={ + "consent": ElicitRequest( + params=ElicitRequestFormParams( + message="Confirm", requested_schema={"type": "object", "properties": {}} + ) + ) + }, + ) + assert params.request_state == "opaque:" + ctx.method + assert params.input_responses["consent"].action == "accept" + if ctx.method == "tools/call": + return CallToolResult(content=[TextContent(type="text", text="confirmed")]) + if ctx.method == "prompts/get": + return GetPromptResult( + messages=[PromptMessage(role="user", content=TextContent(type="text", text="confirmed"))] + ) + return ReadResourceResult(contents=[TextResourceContents(uri="test://confirm", text="confirmed")]) + + service = Server( + "modern-interactions", + on_list_tools=tools, + on_list_prompts=prompts, + on_list_resources=resources, + on_call_tool=interact, + on_get_prompt=interact, + on_read_resource=interact, + ) + manager = StreamableHTTPSessionManager( + service, + stateless=True, + json_response=True, + security_settings=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + + @asynccontextmanager + async def lifespan(app): + AppStatus.should_exit = False + async with manager.run(): + yield + _drain_sse_streams() + watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())) + await asyncio.wait_for(asyncio.gather(*watchers), 2) + + app = Starlette(routes=[Mount("/mcp", app=manager.handle_request)], lifespan=lifespan) + observed = queue.Queue() + with asgi_server(_capturing(app, observed)) as url: + yield McpPeer(url + "/mcp/", observed) + + def rpc(gateway: Gateway, method: str, params: dict[str, JsonValue], *, status: int = 200) -> dict: response: Final = gateway.client.post( "/mcp/", @@ -101,11 +219,12 @@ def rpc(gateway: Gateway, method: str, params: dict[str, JsonValue], *, status: return response.json() +@pytest.mark.parametrize("sdk_upstream", [False, True]) @pytest.mark.parametrize("changed_target", [False, True]) def test_continuations_resume_on_another_replica_and_reject_changed_operations( - tmp_path: Path, changed_target: bool + tmp_path: Path, changed_target: bool, sdk_upstream: bool ) -> None: - with wire_server(interaction_peer) as peer, httpx.Client() as client: + with (sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as peer, httpx.Client() as client: config: Final = tmp_path / "proxy.yaml" config.write_text( yaml.safe_dump( @@ -182,7 +301,8 @@ def test_continuations_resume_on_another_replica_and_reject_changed_operations( assert peer.drain() == (), "Rejected continuation must not contact the upstream" -def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tmp_path: Path) -> None: +@pytest.mark.parametrize("sdk_upstream", [False, True]) +def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tmp_path: Path, sdk_upstream: bool) -> None: auth_module: Final = tmp_path / "interaction_auth.py" auth_module.write_text( "from pathlib import Path\n" @@ -201,7 +321,11 @@ def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tm } auth_state: Final = auth_module.with_suffix(".json") auth_state.write_text(json.dumps(identity)) - with wire_server(interaction_peer) as peer, wire_server(interaction_peer) as other_peer, httpx.Client() as client: + with ( + (sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as peer, + (sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as other_peer, + httpx.Client() as client, + ): config: Final = tmp_path / "proxy.yaml" config.write_text( yaml.safe_dump( @@ -397,6 +521,9 @@ def test_missing_continuation_key_reports_configuration_for_each_carrier(tmp_pat "event_kind,access", [ ("tools", "normal"), + ("tools", "rollback"), + ("tools", "legacy-upstream-sse"), + ("updated", "legacy-upstream-sse"), ("tools", "modern-boundary"), ("prompts", "normal"), ("resources", "normal"), @@ -435,6 +562,12 @@ def test_missing_continuation_key_reports_configuration_for_each_carrier(tmp_pat ("tools", "renewal"), ("tools", "renewal-scopes"), ("tools", "renewal-failed"), + ("tools", "mixed-capabilities"), + ("tools", "probe-failure"), + ("tools", "missing-credentials"), + ("tools", "concurrent-probes"), + ("tools", "probe-timeout"), + ("tools", "rate-limited"), ], ) def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( @@ -447,7 +580,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( import httpx2 from integration._support.asgi import asgi_server - from integration._support.mcp import _drain_sse_streams + from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams from mcp import ClientSession, MCPError from mcp.client.sse import sse_client from mcp.client.streamable_http import streamable_http_client @@ -494,9 +627,16 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( } ) + probe_fails: Final = Event() + probe_failed: Final = Event() listener_credentials = [] catalog_credentials = [] setup_changed = Event() + probe_arrivals = [] + catalog_reads = [] + limited_requests = [] + probe_catalog_goal = [] + probe_barrier = asyncio.Event() active_listeners = set() @@ -540,7 +680,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( "user_role": "internal_user", "object_permission": { "object_permission_id": "subscription-permission", - "mcp_servers": ["subscription-server", "second-server"] if access == "fanout" else ["subscription-server"], + "mcp_servers": ["subscription-server", "second-server"] if access in ("fanout", "mixed-capabilities", "missing-credentials", "concurrent-probes", "rate-limited") else ["subscription-server"], }, } auth_state: Final = auth_module.with_suffix(".json") @@ -565,13 +705,26 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( return EmptyResult() + if getattr(ctx.request, "headers", {}).get("x-probe-target") == "limited": + limited_requests.append(ctx.method) if ctx.method == "tools/list": - catalog_credentials.append(ctx.request.headers.get("authorization")) + catalog_credentials.append(getattr(ctx.request, "headers", {}).get("authorization")) if access == "setup-race" and connected.is_set() and not setup_changed.is_set(): setup_changed.set() await change_tools() + if ctx.method == "server/discover" and probe_fails.is_set() and len(catalog_reads) >= probe_catalog_goal[0]: + probe_failed.set() + if access == "concurrent-probes": + probe_arrivals.append(ctx.method) + if len(probe_arrivals) == 2: + probe_barrier.set() + await probe_barrier.wait() + elif access == "probe-timeout": + await probe_barrier.wait() + else: + raise MCPError(code=-32603, message="Discovery unavailable") if ctx.method == "subscriptions/listen": - listener_credentials.append(ctx.request.headers.get("authorization")) + listener_credentials.append(getattr(ctx.request, "headers", {}).get("authorization")) if access == "legacy-caller-no-notifications": raise MCPError(code=-32601, message="Upstream has no subscription support") return await upstream_listener(ctx, SubscriptionsListenRequestParams.model_validate(ctx.params)) @@ -600,6 +753,8 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( ) } ) + if ctx.method in ("tools/list", "prompts/list", "resources/list", "resources/templates/list"): + catalog_reads.append(ctx.method) if access == "pages" and ctx.method == "tools/list": result = ListToolsResult.model_validate(result) offset = int((ctx.params or {}).get("cursor") or "0") @@ -693,7 +848,11 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( assert {tool.name for tool in (await other.list_tools()).tools} == {"peer-change_tools"} if access.startswith("legacy-caller") or access in ("legacy-both", "legacy-sse"): - await session.initialize() + initialized = await session.initialize() + assert initialized.capabilities.tools.list_changed + assert initialized.capabilities.prompts.list_changed + assert initialized.capabilities.resources.list_changed + assert initialized.capabilities.resources.subscribe before = await session.list_tools() assert "peer-added" not in {tool.name for tool in before.tools} if event_kind == "updated": @@ -772,10 +931,52 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( with pytest.raises(TimeoutError): await asyncio.wait_for(received.get(), 0.5) return - session.adopt(DiscoverResult.model_validate(await session.send_discover("2026-07-28"))) + discovery = DiscoverResult.model_validate(await session.send_discover("2026-07-28")) + session.adopt(discovery) + if access in ("mixed-capabilities", "missing-credentials"): + assert not discovery.capabilities.tools.list_changed + else: + assert discovery.capabilities.tools.list_changed + assert discovery.capabilities.resources.list_changed + assert bool(discovery.capabilities.resources.subscribe) is (access not in ("fanout", "concurrent-probes", "rate-limited")) + if access == "rate-limited": + assert limited_requests, "Admitted discovery never reached the limited upstream" + limited_requests.clear() + limited = DiscoverResult.model_validate(await session.send_discover("2026-07-28")) + assert limited_requests == [], "Rejected discovery reached the rate-limited upstream" + assert limited.capabilities.tools is not None + assert not limited.capabilities.tools.list_changed + assert not (limited.capabilities.resources and limited.capabilities.resources.subscribe) + return + if access in ("probe-failure", "concurrent-probes", "probe-timeout"): + probe_catalog_goal.append(len(catalog_reads) + 4 * len(identity["object_permission"]["mcp_servers"])) + probe_fails.set() + try: + degraded = DiscoverResult.model_validate( + await asyncio.wait_for(session.send_discover("2026-07-28"), 5) + ) + finally: + probe_barrier.set() + assert probe_failed.is_set(), ("Controlled discovery probe was not exercised", catalog_reads, probe_catalog_goal) + assert degraded.capabilities.tools is not None + if access == "concurrent-probes": + assert len(probe_arrivals) == 2 + assert degraded.capabilities.tools.list_changed + else: + assert not (degraded.capabilities.tools and degraded.capabilities.tools.list_changed) + assert not (degraded.capabilities.resources and degraded.capabilities.resources.subscribe) + probe_fails.clear() + visible = await session.list_tools() + assert "peer-change_tools" in {tool.name for tool in visible.tools} + return before = await session.list_tools() assert "peer-change_tools" in {tool.name for tool in before.tools} assert "peer-added" not in {tool.name for tool in before.tools} + if access in ("mixed-capabilities", "missing-credentials"): + assert not discovery.capabilities.resources.list_changed + assert not discovery.capabilities.resources.subscribe + assert not connected.is_set() + return if access == "modern-boundary": for method, params in ( ("notifications/initialized", {}), @@ -862,11 +1063,19 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( await asyncio.sleep(0.2) assert token_grants[0] == token_grants[-1] if access != "setup-race": - async with streamable_http_client(upstream_url + "/mcp") as direct_streams: + direct_connection = ( + sse_client(upstream_url + "/sse") + if access == "legacy-upstream-sse" + else streamable_http_client(upstream_url + "/mcp") + ) + async with direct_connection as direct_streams: async with ClientSession(direct_streams[0], direct_streams[1]) as direct: - direct.adopt( - DiscoverResult.model_validate(await direct.send_discover("2026-07-28")) - ) + if access == "legacy-upstream-sse": + await direct.initialize() + else: + direct.adopt( + DiscoverResult.model_validate(await direct.send_discover("2026-07-28")) + ) changed = await direct.call_tool("change_tools", {}) assert not changed.is_error if access in ( @@ -929,7 +1138,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( await asyncio.wait_for(anext(subscription), 0.5) assert await asyncio.to_thread(disconnected.wait, 3), "Upstream listener survived disconnect" - app: Final = upstream.streamable_http_app( + app: Final = (upstream.sse_app if access == "legacy-upstream-sse" else upstream.streamable_http_app)( transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False) ) @@ -942,10 +1151,13 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())) await asyncio.wait_for(asyncio.gather(*watchers), 2) - hosted: Final = Starlette( - routes=[Route("/token", token_endpoint, methods=["POST"]), Mount("/", app=app)], lifespan=lifespan - ) - with asgi_server(hosted) as upstream_url, httpx.Client() as client: + transport_app: Final = app if access == "legacy-upstream-sse" else _acknowledge_notifications_after_dispatch(app) + hosted: Final = Starlette(routes=[Route("/token", token_endpoint, methods=["POST"]), Mount("/", app=transport_app)], lifespan=lifespan) + with ( + asgi_server(hosted) as upstream_url, + (sdk_interaction_peer() if access == "mixed-capabilities" else nullcontext(None)) as plain, + httpx.Client() as client, + ): config: Final = tmp_path / "proxy.yaml" config.write_text( yaml.safe_dump( @@ -966,10 +1178,12 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( if access.startswith("renewal") else {} ), - "url": upstream_url + "/mcp", - "transport": "http", - "protocol_version": "2025-11-25" if access in ("legacy", "legacy-both") else "2026-07-28", + "url": upstream_url + ("/sse" if access == "legacy-upstream-sse" else "/mcp"), + "transport": "sse" if access == "legacy-upstream-sse" else "http", + "protocol_version": "2025-11-25" if access in ("legacy", "legacy-both", "legacy-upstream-sse") else "2026-07-28", **({"allowed_tools": ["change_tools"]} if access == "hidden" else {}), + **({"rpm": 1, "static_headers": {"X-Probe-Target": "limited"}} if access == "rate-limited" else {}), + **({"timeout": 2} if access == "concurrent-probes" else {}), **( {"auth_type": "oauth_delegate", "extra_headers": ["Authorization"]} if access == "fanout" @@ -980,14 +1194,22 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( { "second": { "server_id": "second-server", - "url": upstream_url + "/mcp", + "url": plain.url if plain is not None else upstream_url + "/mcp", "transport": "http", + **({"timeout": 2} if access == "concurrent-probes" else {}), "protocol_version": "2026-07-28", "auth_type": "oauth_delegate", "extra_headers": ["Authorization"], + **({ + "auth_type": "oauth2", + "oauth2_flow": "authorization_code", + "client_id": "subscription-client", + "authorization_url": upstream_url + "/authorize", + "token_url": upstream_url + "/token", + } if access == "missing-credentials" else {}), } } - if access == "fanout" + if access in ("fanout", "mixed-capabilities", "missing-credentials", "concurrent-probes", "rate-limited") else {} ), }, @@ -1004,6 +1226,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( environment: Final = { "PYTHONPATH": os.pathsep.join((str(source), str(source / "tests"))), "LITELLM_SALT_KEY": "subscription-pagination-key", + **({"LITELLM_MCP_TOOL_LISTING_TIMEOUT": "2"} if access == "probe-timeout" else {}), "STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true", "LITELLM_LOCAL_MODEL_COST_MAP": "True", @@ -1030,6 +1253,28 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( eventually(lambda: len(active_listeners), lambda count: count == 0, seconds=10) assert not active_listeners, "Upstream listener survived client disconnect" + if access == "rollback": + rolled_back = yaml.safe_load(config.read_text()) + rolled_back["general_settings"]["mcp_advertised_versions"] = ["2025-11-25"] + config.write_text(yaml.safe_dump(rolled_back)) + + async def legacy_after_rollback(gateway: Gateway) -> None: + async with httpx2.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client: + async with streamable_http_client(str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client) as streams: + async with ClientSession(streams[0], streams[1]) as session: + await session.initialize() + assert "peer-added" in {tool.name for tool in (await session.list_tools()).tools} + result = await session.call_tool("peer-added", {}) + assert not result.is_error + assert result.content[0].text == "new tool" + + with owned_proxy(seed, tmp_path / "rollback", environment, config=config, database_setup=(), remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "INTEGRATION_PROXY_DATABASE_URL", "INTEGRATION_PROXY_READ_REPLICA_URL")) as rollback: + rejected = rpc(rollback, "tools/list", {}, status=400) + assert "error" in rejected + asyncio.run(legacy_after_rollback(rollback)) + eventually(lambda: len(active_listeners), lambda count: count == 0, seconds=10) + assert tool_executions == ["change_tools", "added"], "Rollback duplicated or restarted upstream tool execution" + def test_subscription_partial_setup_releases_upstream_capacity() -> None: import asyncio @@ -1130,3 +1375,232 @@ def test_subscription_partial_setup_releases_upstream_capacity() -> None: with asgi_server(hosted) as url: asyncio.run(exercise(url)) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_modern_sampling_obeys_opt_in_and_completes_upstream_interaction(tmp_path: Path, enabled: bool) -> None: + def provider(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/v1/chat/completions" + body = json.loads(request.body) + assert "sample prompt" in json.dumps(body["messages"]) + return Reply( + body=json.dumps( + { + "id": "sample-response", + "object": "chat.completion", + "created": 1, + "model": "sampling-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "sampled response"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4}, + } + ).encode() + ) + + with sdk_interaction_peer(sampling=True) as peer, wire_server(provider) as model, httpx.Client() as client: + config = tmp_path / "sampling.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "sampling-model", + "litellm_params": { + "model": "openai/sampling-model", + "api_base": model.url + "/v1", + "api_key": "synthetic-provider-key", + }, + } + ], + "mcp_servers": { + "peer": { + "url": peer.url, + "transport": "http", + "protocol_version": "2026-07-28", + "allow_sampling": enabled, + } + }, + "general_settings": { + "master_key": _KEY, + "store_model_in_db": False, + "mcp_advertised_versions": ["2026-07-28"], + }, + "litellm_settings": {"default_mcp_sampling_model": "sampling-model"}, + } + ) + ) + environment = { + "PYTHONPATH": _PROXY_PYTHONPATH, + "STORE_MODEL_IN_DB": "False", + "DISABLE_SCHEMA_UPDATE": "true", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_SALT_KEY": "sampling-test-salt", + } + with owned_proxy( + Gateway(client, _KEY, peer.url), + tmp_path / "gateway", + environment, + config=config, + database_setup=(), + remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "REDIS_URL", "REDIS_HOST"), + ) as gateway: + reply = McpCaller(gateway, gateway.key, "root", protocol_version="2026-07-28").rpc( + "tools/call", {"name": "peer-confirm", "arguments": {}} + ) + assert reply.status_code == 200, reply.text + response = reply.json() + calls = tuple(request for request in model.drain() if request.method == "POST") + if enabled: + assert response["result"]["content"][0]["text"] == "sample confirmed", response + assert len(calls) == 1 + else: + assert response["result"]["isError"], response + assert "sampling" in json.dumps(response).lower(), response + assert calls == (), "Disabled sampling reached the model provider" + + + +@pytest.mark.parametrize("downstream", ["2025-11-25", "2026-07-28"]) +@pytest.mark.parametrize("upstream", ["2025-11-25", "2026-07-28"]) +def test_input_interaction_pairs_preserve_consent_and_never_restart_legacy_work( + tmp_path: Path, downstream: str, upstream: str +) -> None: + import asyncio + + import httpx2 + from integration._support.mcp import mcp_peer, tool_calls + from mcp import ClientSession + from mcp.client.streamable_http import streamable_http_client + from mcp.types import CallToolResult, DiscoverResult, ElicitResult, InputRequiredResult + + modern_upstream = upstream == "2026-07-28" + peer_context = wire_server(interaction_peer) if modern_upstream else mcp_peer("sse", rich=True) + with peer_context as peer, httpx.Client() as http_client: + config = tmp_path / "proxy.yaml" + config.write_text(yaml.safe_dump({ + "model_list": [], + "mcp_servers": {"peer": {"url": peer.url, "transport": "http" if modern_upstream else "sse", "protocol_version": upstream, "allow_elicitation": True}}, + "general_settings": {"master_key": _KEY, "store_model_in_db": False, "mcp_advertised_versions": ["2025-11-25", "2026-07-28"]}, + })) + consent_requests = [] + + async def consent(context, params): + consent_requests.append(params.message) + return ElicitResult(action="accept", content=None if modern_upstream else {"confirmed": True}) + + async def exercise(gateway): + async with httpx2.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client: + async with streamable_http_client(str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client) as streams: + async with ClientSession(streams[0], streams[1], elicitation_callback=consent) as session: + if downstream == "2026-07-28": + session.adopt(DiscoverResult.model_validate(await session.send_discover(downstream))) + else: + await session.initialize() + name = "peer-confirm" if modern_upstream else "peer-elicit" + args = {} if modern_upstream else {"question": "Confirm"} + await session.list_tools() + result = await session.call_tool(name, args, allow_input_required=True) + if downstream == "2026-07-28" and modern_upstream: + assert isinstance(result, InputRequiredResult) + assert set(result.input_requests) == {"consent"} + result = await session.call_tool(name, args, request_state=result.request_state, input_responses={"consent": ElicitResult(action="accept")}) + assert isinstance(result, CallToolResult) + if downstream == "2026-07-28" and not modern_upstream: + assert result.is_error + assert "legacy upstream cannot resume input" in result.content[0].text.lower() + assert "partially completed" in result.content[0].text.lower() + assert consent_requests == [] + else: + assert not result.is_error, result + assert result.content[0].text == ("confirmed" if modern_upstream else "elicited:accept") + assert len(consent_requests) == (0 if downstream == "2026-07-28" else 1) + + seed = Gateway(http_client, _KEY, peer.url) + with owned_proxy(seed, tmp_path / "gateway", {"PYTHONPATH": _PROXY_PYTHONPATH, "STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true", "LITELLM_SALT_KEY": "interaction-matrix-test"}, config=config, database_setup=(), remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "REDIS_URL", "REDIS_HOST")) as gateway: + asyncio.run(exercise(gateway)) + observed = peer.drain() + calls = [json.loads(item.body) for item in observed if item.method == "POST" and json.loads(item.body).get("method") == "tools/call"] if modern_upstream else tool_calls(observed) + assert len(calls) == (2 if modern_upstream else 1), "Gateway restarted work instead of returning or resuming the original request" + + +def test_legacy_notification_dispatch_finishes_before_session_deletion() -> None: + import asyncio + from contextlib import asynccontextmanager + from threading import Event + + from integration._support.asgi import asgi_server + from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams + from mcp.server.mcpserver import MCPServer + from mcp.server.transport_security import TransportSecuritySettings + from sse_starlette.sse import AppStatus + from starlette.applications import Starlette + from starlette.routing import Mount + from starlette.types import Message, Receive, Scope, Send + + peer = MCPServer("notification-order") + sdk_app = peer.streamable_http_app( + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False) + ) + + @asynccontextmanager + async def lifespan(_: Starlette): + AppStatus.should_exit = False + async with sdk_app.router.lifespan_context(sdk_app): + yield + _drain_sse_streams() + watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())) + await asyncio.wait_for(asyncio.gather(*watchers), 2) + + hosted = Starlette(routes=[Mount("/", app=sdk_app)], lifespan=lifespan) + ordered = _acknowledge_notifications_after_dispatch(hosted) + deleted = Event() + + async def pause_after_acceptance(scope: Scope, receive: Receive, send: Send) -> None: + accepted = False + + async def capture(message: Message) -> None: + nonlocal accepted + if message["type"] == "http.response.start": + accepted = message["status"] == 202 + await send(message) + if accepted and message["type"] == "http.response.body" and not message.get("more_body", False): + assert await asyncio.to_thread(deleted.wait, 5), "Client did not complete session deletion" + + await ordered(scope, receive, capture) + + with asgi_server(pause_after_acceptance) as url, httpx.Client(timeout=8) as client: + headers = { + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2025-11-25", + } + initialized = client.post( + url + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2025-11-25", "capabilities": {}, + "clientInfo": {"name": "notification-order", "version": "1"}, + }, + }, + ) + assert initialized.status_code == 200 + headers["Mcp-Session-Id"] = initialized.headers["mcp-session-id"] + try: + accepted = client.post( + url + "/mcp", headers=headers, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + ) + assert accepted.status_code == 202 + assert accepted.content == b"" + terminated = client.delete(url + "/mcp", headers=headers) + assert terminated.status_code == 200 + finally: + deleted.set() diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index cdb3e17bea8..9d15ea0afa9 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -8,7 +8,7 @@ import textwrap import time import uuid from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from pathlib import Path from typing import Final @@ -1991,13 +1991,43 @@ def test_standard_pattern_discovery_ignores_forwarded_headers_from_an_untrusted_ @contextmanager -def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, Wire]]: - with scripted_peer( +def _sdk_scope_peer() -> Iterator[McpPeer]: + import queue + from mcp.server.mcpserver import MCPServer + from mcp.server.transport_security import TransportSecuritySettings + from integration._support.asgi import asgi_server + from integration._support.mcp import _capturing + + service = MCPServer("scope-operations") + + @service.tool() + def read_op() -> str: + return "read completed" + + @service.tool() + def write_op() -> str: + return "write completed" + + @service.tool() + def forbid_op() -> str: + return "must not execute" + + app = service.streamable_http_app( + stateless_http=True, json_response=True, + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + observed = queue.Queue() + with asgi_server(_capturing(app, observed)) as url: + yield McpPeer(url + "/mcp", observed) + + +@contextmanager +def _scope_enforcing_peer(auth: AuthorizationServer, *, modern: bool = False) -> Iterator[tuple[McpPeer, Wire]]: + with (_sdk_scope_peer() if modern else scripted_peer( ScriptedTool("read_op", lambda _: text_result("read completed")), ScriptedTool("write_op", lambda _: text_result("write completed")), ScriptedTool("forbid_op", lambda _: text_result("must not execute")), - ) as inner: - + )) as inner: def enforce(request: Request) -> Reply: body: Final = json.loads(request.body) if request.body else {} if body.get("method") == "tools/call": @@ -2026,51 +2056,55 @@ def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, yield McpPeer(wire.url + "/mcp", inner.calls), wire +@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"]) @pytest.mark.parametrize("entry", ["rest", "server_mcp"]) -def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: Gateway, entry: EntryPoint) -> None: +def test_managed_scope_step_up_preserves_read_and_write_without_retry( + gateway: Gateway, entry: EntryPoint, tmp_path: Path, protocol: str +) -> None: + modern = protocol == "2026-07-28" + config = tmp_path / "scope-proxy.yaml" + config.write_text(yaml.safe_dump({ + "model_list": [], + "general_settings": { + "master_key": gateway.key, "store_model_in_db": True, + "mcp_advertised_versions": ["2025-11-25", "2026-07-28"], + }, + })) with ( + (owned_proxy(gateway, tmp_path / "gateway", {"DISABLE_SCHEMA_UPDATE": "true"}, config=config, database_setup=()) + if modern else nullcontext(gateway)) as active, oauth_server(scopes=("tools.read", "tools.write")) as auth, - _scope_enforcing_peer(auth) as (peer, wire), - gateway.scenario() as scenario, + _scope_enforcing_peer(auth, modern=modern) as (peer, wire), + active.scenario() as scenario, ): alias: Final = "scope" + uuid.uuid4().hex[:8] identity: Final = _register_oauth( - scenario, - peer, - auth, - alias, - auth_type="oauth2", - oauth2_flow="authorization_code", + scenario, peer, auth, alias, auth_type="oauth2", oauth2_flow="authorization_code", protocol_version=protocol, credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]}, ) owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) pkce: Final = _Pkce(secrets.token_urlsafe(32)) - code, _ = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read") - first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce) + code, _ = _authorize_through_gateway(active, auth, alias, owner, "scope-client", pkce, "tools.read") + first: Final = _redeem(active, alias, owner, "scope-client", code, pkce) assert first.status_code == 200, first.text - caller: Final = McpCaller(gateway, owner, entry, alias) + caller: Final = McpCaller(active, owner, entry, alias, protocol_version=protocol) read: Final = caller.call(f"{alias}-read_op", {}, identity) assert read.ok and read.text == "read completed", read.raw peer.drain() wire.drain() denied: Final = ( - call_tool(gateway, owner, identity, f"{alias}-write_op", {}) - if entry == "rest" + call_tool(active, owner, identity, f"{alias}-write_op", {}) if entry == "rest" else caller.rpc("tools/call", {"name": f"{alias}-write_op", "arguments": {}}) ) assert denied.status_code == 403, denied.text challenge: Final = denied.headers["www-authenticate"] assert 'error="insufficient_scope"' in challenge and 'scope="tools.write"' in challenge - assert f'resource_metadata="{_base(gateway)}/.well-known/oauth-protected-resource/' in challenge - attempts: Final = tuple( - request - for request in wire.drain() - if request.body and json.loads(request.body).get("method") == "tools/call" - ) + assert f'resource_metadata="{_base(active)}/.well-known/oauth-protected-resource/' in challenge + attempts: Final = tuple(request for request in wire.drain() if request.body and json.loads(request.body).get("method") == "tools/call") assert len(attempts) == 1 and tool_calls(peer.drain()) == (), "denied operation was retried or executed" - second_code, _ = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.write") - second: Final = _redeem(gateway, alias, owner, "scope-client", second_code, pkce) + second_code, _ = _authorize_through_gateway(active, auth, alias, owner, "scope-client", pkce, "tools.write") + second: Final = _redeem(active, alias, owner, "scope-client", second_code, pkce) assert second.status_code == 200, second.text with auth.lock: granted: Final = auth.access_tokens[_issued_token(second.json())]["scope"] @@ -2080,7 +2114,7 @@ def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: G assert result.ok and result.text == f"{operation} completed", result.raw forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity) assert forbidden.status == 200 and not forbidden.ok, forbidden.raw - unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity) + unauthorized: Final = McpCaller(active, stranger, entry, alias, protocol_version=protocol).call(f"{alias}-read_op", {}, identity) assert unauthorized.status == 401, unauthorized.raw diff --git a/tests/integration/mcp/test_pagination.py b/tests/integration/mcp/test_pagination.py index 267da498c92..e52757ec19a 100644 --- a/tests/integration/mcp/test_pagination.py +++ b/tests/integration/mcp/test_pagination.py @@ -8,7 +8,7 @@ import pytest import yaml from mcp import ClientSession, MCPError from mcp.client.streamable_http import streamable_http_client -from mcp.types import CallToolRequest, CallToolRequestParams, CallToolResult, PaginatedRequestParams +from mcp.types import CallToolRequest, CallToolRequestParams, CallToolResult, DiscoverResult, PaginatedRequestParams from integration._support.client import Gateway from integration._support.mcp import McpPeer, paginated_mcp_peer @@ -18,24 +18,27 @@ from litellm.types.mcp import MCPTransport @asynccontextmanager -async def catalog_session(gateway): +async def catalog_session(gateway, protocol: str = "2025-11-25"): async with httpx.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client: async with streamable_http_client( str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client ) as streams: async with ClientSession(streams[0], streams[1]) as session: - await session.initialize() + if protocol == "2026-07-28": + session.adopt(DiscoverResult.model_validate(await session.send_discover(protocol))) + else: + await session.initialize() yield session -def config_file(directory: Path, upstream: str) -> Path: +def config_file(directory: Path, upstream: str, protocol: str | None = None) -> Path: config = directory / "proxy.yaml" config.write_text( yaml.safe_dump( { "model_list": [], - "mcp_servers": {"pages": {"url": upstream, "transport": "http"}}, - "general_settings": {"master_key": "sk-pagination-test", "store_model_in_db": False}, + "mcp_servers": {"pages": {"url": upstream, "transport": "http", **({"protocol_version": protocol} if protocol else {})}}, + "general_settings": {"master_key": "sk-pagination-test", "store_model_in_db": False, **({"mcp_advertised_versions": ["2025-11-25", "2026-07-28"]} if protocol == "2026-07-28" else {})}, } ) ) @@ -45,7 +48,8 @@ def config_file(directory: Path, upstream: str) -> Path: REMOVE_DATABASE = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "LITELLM_LICENSE", "LITELLM_LICENSE_PATH") -def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path): +@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"]) +def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path, protocol: str): async def exercise(first_replica, second_replica): for method, field in ( ("list_tools", "tools"), @@ -53,10 +57,10 @@ def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path): ("list_resources", "resources"), ("list_resource_templates", "resource_templates"), ): - async with catalog_session(first_replica) as first_session: + async with catalog_session(first_replica, protocol) as first_session: first = await getattr(first_session, method)() assert first.next_cursor - async with catalog_session(second_replica) as second_session: + async with catalog_session(second_replica, protocol) as second_session: second = await getattr(second_session, method)(params=PaginatedRequestParams(cursor=first.next_cursor)) replay = await getattr(second_session, method)(params=PaginatedRequestParams(cursor=first.next_cursor)) assert getattr(second, field) == getattr(replay, field) @@ -65,14 +69,14 @@ def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path): names = [item.name for page in (first, second, third) for item in getattr(page, field)] assert len(names) == len(set(names)) == 3 assert third.next_cursor is None - async with catalog_session(second_replica) as session: + async with catalog_session(second_replica, protocol) as session: called = await session.call_tool("pages-add2", {"a": 3, "b": 4}) assert not called.is_error assert called.content[0].text == "7" with paginated_mcp_peer() as peer, httpx.Client() as client: seed = Gateway(client, "sk-pagination-test", peer.url) - config = config_file(tmp_path, peer.url) + config = config_file(tmp_path, peer.url, protocol) options = dict(config=config, database_setup=(), remove_environment=REMOVE_DATABASE) environment = { "STORE_MODEL_IN_DB": "False", @@ -128,7 +132,8 @@ def test_no_salt_preserves_complete_lists_and_calls_but_rejects_continuations(tm @pytest.mark.parametrize("grant", ["direct", "access_group"]) -def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, monkeypatch, grant: str): +@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"]) +def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, monkeypatch, grant: str, protocol: str): import os import subprocess import sys @@ -144,15 +149,15 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m stranger_b = Gateway(b.client, stranger, peer.url) first_pages = {} for method in ("list_tools", "list_prompts", "list_resources", "list_resource_templates"): - async with catalog_session(owner_a) as session: + async with catalog_session(owner_a, protocol) as session: first_pages[method] = await getattr(session, method)() assert first_pages[method].next_cursor - async with catalog_session(owner_b) as session: + async with catalog_session(owner_b, protocol) as session: continued = await getattr(session, method)( params=PaginatedRequestParams(cursor=first_pages[method].next_cursor) ) assert continued.next_cursor - async with catalog_session(stranger_b) as session: + async with catalog_session(stranger_b, protocol) as session: peer.drain() with pytest.raises(MCPError, match="fresh listing"): await getattr(session, method)( @@ -172,7 +177,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m }, ) for method, first in first_pages.items(): - async with catalog_session(owner_b) as session: + async with catalog_session(owner_b, protocol) as session: peer.drain() with pytest.raises(MCPError, match="fresh listing"): await getattr(session, method)(params=PaginatedRequestParams(cursor=first.next_cursor)) @@ -181,7 +186,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m changed = a.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "new catalog generation"}) assert changed.status_code == 202, changed.text for method, first in first_pages.items(): - async with catalog_session(owner_b) as session: + async with catalog_session(owner_b, protocol) as session: peer.drain() with pytest.raises(MCPError, match="fresh listing"): await getattr(session, method)(params=PaginatedRequestParams(cursor=first.next_cursor)) @@ -212,7 +217,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m config = tmp_path / "database-proxy.yaml" config.write_text( yaml.safe_dump( - {"model_list": [], "general_settings": {"master_key": seed.key, "store_model_in_db": True}} + {"model_list": [], "general_settings": {"master_key": seed.key, "store_model_in_db": True, "mcp_advertised_versions": ["2025-11-25", "2026-07-28"]}} ) ) environment = { @@ -230,8 +235,8 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m owned_proxy(seed, tmp_path / "b", environment, **options) as b, a.scenario() as scenario, ): - identity = register_mcp(scenario, peer, "pages") - spare: Final = register_mcp(scenario, spare_peer, "spare") + identity = register_mcp(scenario, peer, "pages", protocol_version=protocol) + spare: Final = register_mcp(scenario, spare_peer, "spare", protocol_version=protocol) group = a.request( "POST", "/v1/access_group", diff --git a/tests/unit/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py index 91498f0879c..0cbcf1b7ee1 100644 --- a/tests/unit/integration_support/test_upstream.py +++ b/tests/unit/integration_support/test_upstream.py @@ -1,3 +1,4 @@ +import asyncio from collections import deque from pathlib import Path from typing import Final @@ -50,3 +51,207 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error ("WEBSOCKET", {"query": [["model", "synthetic"]]}), ("WEBSOCKET_FRAME", {"type": "response.create"}), ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method,status", (("POST", 200), ("POST", 400), ("DELETE", 200))) +async def test_mcp_notification_ordering_preserves_other_responses( + method: str, status: int, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.mcp import _acknowledge_notifications_after_dispatch + from starlette.types import Receive, Send + + outgoing: Final[list[Message]] = [] + messages: Final[tuple[Message, ...]] = ( + {"type": "http.response.start", "status": status, "headers": []}, + {"type": "http.response.body", "body": b"result", "more_body": True}, + {"type": "http.response.body", "body": b"", "more_body": False}, + ) + + async def receive() -> Message: + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + outgoing.append(message) + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + for index, message in enumerate(messages): + await send(message) + assert outgoing == list(messages[: index + 1]), "Ordinary response streaming was buffered" + + await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": method}, receive, send) + assert outgoing == list(messages) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", (RuntimeError("dispatch failed"), asyncio.CancelledError())) +async def test_mcp_notification_dispatch_failure_is_not_acknowledged( + failure: BaseException, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.mcp import _acknowledge_notifications_after_dispatch + from starlette.types import Receive, Send + + outgoing: Final[list[Message]] = [] + + async def receive() -> Message: + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + outgoing.append(message) + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + if scope["method"] == "DELETE": + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + return + await send({"type": "http.response.start", "status": 202, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + raise failure + + ordered: Final = _acknowledge_notifications_after_dispatch(app) + scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"failed-session")]} + with pytest.raises(type(failure)) as caught: + await ordered({**scope, "method": "POST"}, receive, send) + assert caught.value is failure + assert outgoing == [], "Failed dispatch must not send an acknowledgement" + await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2) + assert outgoing[0]["status"] == 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "responses", + ( + (), + ({"type": "http.response.body", "body": b"unexpected"},), + ({"type": "http.response.body", "body": b"", "more_body": True},), + ({"type": "http.response.start", "status": 500, "headers": []},), + ({"type": "http.response.body", "body": b""}, {"type": "http.response.body", "body": b""}), + ), +) +async def test_mcp_notification_ordering_rejects_invalid_acknowledgements( + responses: tuple[Message, ...], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.mcp import _acknowledge_notifications_after_dispatch + from starlette.types import Receive, Send + + outgoing: Final[list[Message]] = [] + + async def receive() -> Message: + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + outgoing.append(message) + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 202, "headers": []}) + for message in responses: + await send(message) + + with pytest.raises(RuntimeError, match="MCP notification acknowledgement"): + await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": "POST"}, receive, send) + assert outgoing == [], "Invalid acknowledgement must not be reported as accepted" + + +@pytest.mark.asyncio +async def test_mcp_session_deletion_waits_for_concurrent_notification_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.mcp import _acknowledge_notifications_after_dispatch + from starlette.types import Receive, Send + + started: Final = [asyncio.Event(), asyncio.Event()] + release: Final = asyncio.Event() + deletion_started: Final = asyncio.Event() + deleted: Final = asyncio.Event() + other_deleted: Final = asyncio.Event() + completed: Final[list[str]] = [] + + async def receive() -> Message: + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + return None + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + if scope["method"] == "POST": + started[int(scope["path"])].set() + await release.wait() + completed.append(scope["path"]) + await send({"type": "http.response.start", "status": 202, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + else: + if scope["headers"][0][1] == b"other-session": + other_deleted.set() + else: + deleted.set() + completed.append("DELETE") + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + ordered: Final = _acknowledge_notifications_after_dispatch(app) + scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"owned-session")]} + posts: Final = [ + asyncio.create_task(ordered({**scope, "method": "POST", "path": str(index)}, receive, send)) + for index in range(2) + ] + + async def delete() -> None: + deletion_started.set() + await ordered({**scope, "method": "DELETE"}, receive, send) + + deletion = None + try: + await asyncio.wait_for(asyncio.gather(*(event.wait() for event in started)), 2) + deletion = asyncio.create_task(delete()) + await deletion_started.wait() + await asyncio.wait_for( + ordered({**scope, "method": "DELETE", "headers": [(b"mcp-session-id", b"other-session")]}, receive, send), 2 + ) + assert other_deleted.is_set(), "Unrelated session deletion was blocked" + assert not deleted.is_set(), "Session deletion overtook in-flight notification dispatch" + finally: + release.set() + await asyncio.wait_for(asyncio.gather(*posts, *([deletion] if deletion is not None else [])), 2) + assert set(completed[:2]) == {"0", "1"} + assert completed[-1] == "DELETE" + + +@pytest.mark.asyncio +async def test_mcp_session_deletion_can_terminate_streaming_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.mcp import _acknowledge_notifications_after_dispatch + from starlette.types import Receive, Send + + streaming: Final = asyncio.Event() + deleted: Final = asyncio.Event() + + async def receive() -> Message: + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + return None + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + if scope["method"] == "POST": + streaming.set() + await deleted.wait() + else: + deleted.set() + await send({"type": "http.response.body", "body": b""}) + + ordered: Final = _acknowledge_notifications_after_dispatch(app) + scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"streaming-session")]} + post: Final = asyncio.create_task(ordered({**scope, "method": "POST"}, receive, send)) + try: + await asyncio.wait_for(streaming.wait(), 2) + await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2) + assert deleted.is_set(), "Session deletion must terminate an ordinary streaming request" + finally: + deleted.set() + await asyncio.wait_for(post, 2) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py b/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py index 39a6b1fcf15..e5104e7941f 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py @@ -164,6 +164,51 @@ def test_modern_protocol_is_rejected_on_legacy_sse_paths(path: str, monkeypatch: ) +@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._experimental.mcp_server import subscriptions + from litellm.proxy._experimental.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", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 47578344f77..3957f486c4c 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1943,3 +1943,56 @@ async def test_guardrail_context_never_recovers_a_previous_request_bearer( assert capture.data["incoming_bearer_token"] is None finally: auth_context_var.reset(previous) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delay", ["target", "probe", "exhausted"]) +async def test_discovery_optional_probes_use_remaining_listing_budget(delay, monkeypatch): + from types import SimpleNamespace + from mcp.types import ( + DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, + ListResourceTemplatesResult, ToolsCapability, Tool, + ) + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import subscriptions + + entered = asyncio.Event() + cleaned = asyncio.Event() + + async def blocked(*args, **kwargs): + entered.set() + try: + await asyncio.Event().wait() + finally: + cleaned.set() + + async def tools(*args, **kwargs): + if delay == "exhausted": + await asyncio.sleep(0.08) + return ListToolsResult(tools=[Tool(name="usable", input_schema={"type": "object"})]) + + target = SimpleNamespace( + server=SimpleNamespace(protocol_version="2026-07-28", transport="http"), + client=SimpleNamespace(run_with_session=blocked), + ) + monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2026-07-28"]) + # The existing listing budget includes optional target preparation and handshakes. + monkeypatch.setattr(operations, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False) + with ( + patch.object(operations, "_execute_handle_list_tools", side_effect=tools), + patch.object(operations, "_execute_list_prompts", return_value=ListPromptsResult(prompts=[])), + patch.object(operations, "_execute_list_resources", return_value=ListResourcesResult(resources=[])), + patch.object(operations, "_execute_list_resource_templates", return_value=ListResourceTemplatesResult(resource_templates=[])), + patch.object(operations, "_mcp_server_rate_limit_rejection", return_value=None), + patch.object(subscriptions, "_targets", side_effect=blocked if delay == "target" else None, return_value=(target,)) as targets, + ): + context = prepare_context(UserAPIKeyAuth(user_id="scoped"), protocol_version="2026-07-28") + result = await asyncio.wait_for(GatewayOperations().execute(DiscoverRequest(), context), 0.5) + assert result.capabilities.tools == ToolsCapability() + assert result.capabilities.resources is None + if delay == "exhausted": + targets.assert_not_awaited() + assert not entered.is_set() + else: + assert entered.is_set() + assert cleaned.is_set() diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py index e44aa88c56c..a5d42e58b99 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -87,14 +87,14 @@ async def test_legacy_listener_rejected_before_upstream_work(revision: str) -> N @pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["multiple-resources", "denied"]) +@pytest.mark.parametrize("mode", ["stdio", "multiple-resources", "denied"]) async def test_unsupported_targets_do_not_open_upstream_listeners(mode: str) -> None: client: Final = MCPClient(server_url="http://127.0.0.1:1/mcp", transport_type=MCPTransport.http) server: Final = MCPServer( server_id="upstream", name="upstream", url="http://localhost/mcp", - transport="http", + transport="stdio" if mode == "stdio" else "http", protocol_version="2025-11-25" if mode == "legacy" else "2026-07-28", ) target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") @@ -109,6 +109,7 @@ async def test_unsupported_targets_do_not_open_upstream_listeners(mode: str) -> ) error: Final = HTTPException if mode == "denied" else MCPError expected: Final = { + "stdio": "Modern upstream subscriptions require HTTP", "multiple-resources": "exactly one allowed server", "denied": "No authorized MCP subscription targets", }[mode]