From 4a4a0e8d3c6832c53560a19d370f84f8bfbb2047 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Thu, 15 Oct 2026 16:16:49 -0400 Subject: [PATCH] feat(mcp): relay authorized modern upstream subscriptions (#46153) * feat(mcp): relay authorized modern upstream subscriptions * fix(mcp): close subscriptions when caller identity changes * fix(mcp): preserve subscription scope and setup changes * fix(mcp): preserve subscriptions across managed token renewal * fix(mcp): forward legacy notifications with subscription-owned authorization * test(mcp): verify modern method boundaries and unsupported legacy resources * test(mcp): model SDK context fields in cleanup regressions --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 47 +- .../client_credentials.py | 10 + .../outbound_credentials/resolver.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 115 ++- .../_experimental/mcp_server/subscriptions.py | 542 +++++++++++++ tests/integration/mcp/test_interactions.py | 743 +++++++++++++++++- .../test_mcp_client.py | 53 ++ .../test_client_credentials.py | 41 + .../outbound_credentials/test_resolver.py | 18 + .../mcp_server/test_operations.py | 4 +- .../mcp_server/test_subscriptions.py | 357 +++++++++ 11 files changed, 1904 insertions(+), 28 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/subscriptions.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 50553b59365..bc3ba796ee3 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -186,6 +186,23 @@ def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]: } +async def _receive_session_message( + message: ServerNotification | Exception, + *, + stream_error: asyncio.Future[Exception], + on_notification: Callable[[ServerNotification], Awaitable[None]] | None, +) -> None: + if not isinstance(message, Exception): + if on_notification is not None: + await on_notification(message) + return + if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)): + return + if not stream_error.done(): + stream_error.set_result(message) + raise RuntimeError("MCP response stream failed") + + def _first_non_cancelled_cause(exc: BaseException, cleanup_errors: tuple[Exception, ...] = ()) -> BaseException | None: queue: Final[list[BaseException]] = [exc] while queue: @@ -534,6 +551,17 @@ class MCPClient: async def discovery_auth_fingerprint(self) -> str: return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity()) + async def subscription_auth_fingerprint(self) -> str: + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + + request: Final = await self.prepare_request_auth() + auth: Final = self._resolved_auth + if isinstance(auth, ClientCredentialsBearerAuth) and auth.subscription_identity is not None: + request.headers[auth.header_name] = "managed-grant:" + auth.subscription_identity + return self._hash_discovery_auth(request, self._aws_identity()) + def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None: return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None @@ -697,6 +725,7 @@ class MCPClient: *, on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None, on_cleanup: Callable[[], None] | None = None, + on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None, ) -> TSessionResult: """ Execute an operation within a transport and session context. @@ -717,16 +746,6 @@ class MCPClient: stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() _notify_stream_error(on_stream_error, stream_error) - async def receive_message( - message: ServerNotification | Exception, - ) -> None: - if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)): - return - if not stream_error.done(): - stream_error.set_result(message) - # The SDK closes pending requests when its message handler raises. - raise RuntimeError("MCP response stream failed") - session_kwargs: Final = ( self._session_callbacks_override if self._session_callbacks_override is not None @@ -738,7 +757,9 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=self.timeout, - message_handler=receive_message, + message_handler=partial( + _receive_session_message, stream_error=stream_error, on_notification=on_notification + ), sampling_callback=session_kwargs["sampling_callback"], elicitation_callback=session_kwargs["elicitation_callback"], logging_callback=session_kwargs["logging_callback"], @@ -805,6 +826,7 @@ class MCPClient: quiet_on_error: bool = False, on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None, on_cleanup: Callable[[], None] | None = None, + on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None, ) -> TSessionResult: """Open a session, run the provided coroutine, and clean up. @@ -823,8 +845,9 @@ class MCPClient: http_client=http_client, on_stream_error=on_stream_error, on_cleanup=on_cleanup, + on_notification=on_notification, ) - if on_stream_error is not None or on_cleanup is not None + if on_stream_error is not None or on_cleanup is not None or on_notification is not None else self._execute_session_operation(transport_ctx, operation, http_client=http_client) ) result: Final = await execution diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index e628b3880b3..1652f5d8bd4 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -28,6 +28,7 @@ from __future__ import annotations import asyncio import hashlib +import json import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator from dataclasses import dataclass @@ -352,11 +353,20 @@ class ClientCredentialsBearerAuth(httpx2.Auth): access_token: str, refetch: Callable[[str], Awaitable[str | None]], carrier: HeaderCarrier, + *, + granted_scopes: tuple[str, ...] | None = None, ) -> None: self._carrier = carrier self.header_name = carrier.header_name self._access_token = SecretStr(access_token) self._refetch = refetch + self.subscription_identity: str | None = ( + hashlib.sha256( + json.dumps((_identity_key(carrier), sorted(set(granted_scopes)), carrier.value_prefix)).encode() + ).hexdigest() + if isinstance(carrier, ClientCredentialsConfig) and granted_scopes is not None + else None + ) async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: token: Final = self._access_token.get_secret_value() diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 981843f2234..66b44e91fd7 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -327,7 +327,7 @@ class UpstreamCredentialProvider: match await self._client_credentials_source.get(server_id, config): case Ok(token): refetch: Final = partial(self._client_credentials_source.refetch, server_id, config) - return Ok(ClientCredentialsBearerAuth(token.access_token, refetch, config)) + return Ok(ClientCredentialsBearerAuth(token.access_token, refetch, config, granted_scopes=token.scopes)) case Error(err): return Error(err) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c6ea55e9642..b0e6d08e649 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -511,6 +511,7 @@ if MCP_AVAILABLE: from mcp.server.auth.provider import AccessToken from mcp.server.context import ServerRequestContext from mcp.server.lowlevel.server import NotificationOptions + from mcp.server.lowlevel.server import Server as MCPProtocolServer from mcp.server.models import InitializationOptions from mcp.shared.exceptions import MCPError from mcp.shared.inbound import InboundLadderRejection, classify_inbound_request, find_duplicated_routing_header @@ -581,7 +582,12 @@ if MCP_AVAILABLE: ) -> InitializationOptions: base_options: Final = Server.create_initialization_options( self, - notification_options=notification_options, + notification_options=notification_options + or NotificationOptions( + tools_changed=not _mcp_proxy_mode.get(), + prompts_changed=not _mcp_proxy_mode.get(), + resources_changed=not _mcp_proxy_mode.get(), + ), experimental_capabilities=experimental_capabilities or {}, extensions=extensions, ) @@ -606,7 +612,7 @@ if MCP_AVAILABLE: ######################################################## ############ Initialize the MCP Server ################# ######################################################## - server: Final[Server] = Server( + server: Final = MCPProtocolServer[object]( name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) @@ -828,6 +834,8 @@ if MCP_AVAILABLE: @contextlib.asynccontextmanager async def _legacy_operation_context(ctx: ServerRequestContext, *, trace: bool) -> AsyncGenerator[OperationContext]: + if ctx.protocol_version != "2026-07-28" and ctx.method in ("tools/list", "prompts/list", "resources/list"): + await subscriptions.wait_initialized(ctx) with contextlib.ExitStack() as cleanup: cleanup.callback(active_mcp_request_ctx_var.reset, active_mcp_request_ctx_var.set(ctx)) cleanup.callback(active_mcp_session_var.reset, active_mcp_session_var.set(ctx.session)) @@ -993,6 +1001,88 @@ if MCP_AVAILABLE: async with _legacy_operation_context(ctx, trace=False) as context: return await operations.GatewayOperations().execute(DiscoverRequest(params=params), context) + from mcp.types import ( + EmptyResult, + NotificationParams, + SubscribeRequestParams, + SubscriptionsListenRequestParams, + SubscriptionsListenResult, + UnsubscribeRequestParams, + ) + + from litellm.proxy._experimental.mcp_server.subscriptions import GatewaySubscriptions, SubscriptionAuthorization + + subscriptions: Final = GatewaySubscriptions() + legacy_subscription_request: Final[contextvars.ContextVar[StarletteRequest | None]] = contextvars.ContextVar( + "legacy_subscription_request", default=None + ) + + @contextlib.asynccontextmanager + async def _subscription_authorization( + ctx: ServerRequestContext[object, StarletteRequest], + ) -> AsyncGenerator[SubscriptionAuthorization]: + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() + request: Final[StarletteRequest | None] = ( + ctx.request if isinstance(ctx.request, StarletteRequest) else legacy_subscription_request.get() + ) + if not isinstance(request, StarletteRequest): + raise HTTPException(status_code=400, detail="Subscriptions require an HTTP request") + async with _legacy_operation_context(ctx, trace=False) as initial: + + async def authorize() -> OperationContext: + auth, token, servers, server_headers, oauth_headers, headers = await extract_mcp_auth_context( + request.scope, request.url.path + ) + return operations.prepare_context( + auth, + token, + servers, + server_headers, + oauth_headers, + headers, + initial.client_ip, + False, + initial.wire_compat, + initial.protocol_version, + incoming_bearer_token=MCPRequestHandler.get_guardrail_bearer_token(request.headers, auth), + ) + + authorization: Final = SubscriptionAuthorization(initial, authorize, _owner_fingerprint_for) + yield authorization + + async def listen_subscriptions( + ctx: ServerRequestContext, params: SubscriptionsListenRequestParams + ) -> SubscriptionsListenResult: + async with _subscription_authorization(ctx) as authorization: + return await subscriptions.listen(ctx, params, authorization, authorization.owner) + + async def legacy_notifications(ctx: ServerRequestContext, params: NotificationParams) -> None: + if _mcp_proxy_mode.get(): + return + async with subscriptions.initializing(ctx) as start: + if start: + async with _subscription_authorization(ctx) as authorization: + await subscriptions.legacy_session(ctx, authorization) + + async def subscribe_resource(ctx: ServerRequestContext, params: SubscribeRequestParams) -> EmptyResult: + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() + await subscriptions.legacy_resource(ctx, str(params.uri), subscribe=True) + return EmptyResult() + + async def unsubscribe_resource(ctx: ServerRequestContext, params: UnsubscribeRequestParams) -> EmptyResult: + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() + await subscriptions.legacy_resource(ctx, str(params.uri), subscribe=False) + return EmptyResult() + + server.add_notification_handler("notifications/initialized", NotificationParams, legacy_notifications) + server.add_request_handler("resources/subscribe", SubscribeRequestParams, subscribe_resource) + server.add_request_handler("resources/unsubscribe", UnsubscribeRequestParams, unsubscribe_resource) + + server.add_request_handler("subscriptions/listen", SubscriptionsListenRequestParams, listen_subscriptions) + server.add_request_handler("server/discover", RequestParams, discover) server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools) server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call) @@ -1214,7 +1304,7 @@ if MCP_AVAILABLE: def _owner_fingerprint_for( user_api_key_auth: UserAPIKeyAuth | None, - oauth2_headers: dict[str, str] | None = None, + oauth2_headers: Mapping[str, str] | None = None, client_ip: str | None = None, ) -> str: """ @@ -2678,13 +2768,18 @@ if MCP_AVAILABLE: sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream), server.lifespan(server) as lifespan_state, ): - await serve_loop( - server, - read_stream, - write_stream, - lifespan_state=lifespan_state, - init_options=server.create_initialization_options(), - ) + request_token: Final = legacy_subscription_request.set(StarletteRequest(scope)) + try: + await serve_loop( + server, + read_stream, + write_stream, + lifespan_state=lifespan_state, + init_options=server.create_initialization_options(), + ) + finally: + legacy_subscription_request.reset(request_token) + except MCPUpstreamAuthError as e: # Upstream delegated auth returned 401; surface it to the client so # standards-compliant MCP clients trigger the upstream OAuth flow. diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py new file mode 100644 index 00000000000..bc22e29ef96 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -0,0 +1,542 @@ +import asyncio +import hashlib +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from contextlib import AsyncExitStack, asynccontextmanager +from dataclasses import dataclass, replace +from itertools import chain +from typing import Final, NoReturn + +import anyio +from anyio.abc import TaskGroup, TaskStatus +from fastapi import HTTPException +from mcp import ClientSession, MCPError +from mcp.client.subscriptions import listen +from mcp.server.context import ServerRequestContext +from mcp.server.session import ServerSession +from mcp.server.subscriptions import InMemorySubscriptionBus, ListenHandler +from mcp.shared.subscriptions import ( + PromptsListChanged, + ResourcesListChanged, + ResourceUpdated, + ServerEvent, + ToolsListChanged, + event_from_wire, + event_matches, + event_to_notification, +) +from mcp.types import ( + EmptyResult, + ListPromptsRequest, + ListPromptsResult, + ListResourcesRequest, + ListToolsRequest, + ListToolsResult, + PaginatedRequestParams, + RequestId, + ServerCapabilities, + ServerNotification, + SubscribeRequest, + SubscribeRequestParams, + SubscriptionFilter, + SubscriptionsListenRequestParams, + SubscriptionsListenResult, + UnsubscribeRequest, + UnsubscribeRequestParams, +) + +from litellm.constants import MCP_TOOL_LISTING_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT +from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server import operations +from litellm.proxy._experimental.mcp_server.catalog import ( + _configuration_identity as configuration_identity, # pyright: ignore[reportPrivateUsage] # reuse catalog-owned configuration identity without treating discovered metadata as configuration +) +from litellm.proxy._experimental.mcp_server.contracts import OperationContext +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class SubscriptionAuthorization: + def __init__( + self, + initial: OperationContext, + refresh: Callable[[], Awaitable[OperationContext]], + fingerprint: Callable[[UserAPIKeyAuth | None, Mapping[str, str] | None, str | None], str], + ) -> None: + self._refresh = refresh + self._fingerprint = fingerprint + self.owner = self._owner(initial) + self._identity = self._caller(initial) + + def _owner(self, context: OperationContext) -> str: + return self._fingerprint(context.user_api_key_auth, dict(context.oauth2_headers or {}), context.client_ip) + + @staticmethod + def _caller(context: OperationContext) -> tuple[str | None, str | None, str | None, str | None] | None: + caller: Final = context.user_api_key_auth + return (caller.user_id, caller.team_id, caller.org_id, caller.end_user_id) if caller is not None else None + + async def __call__(self) -> OperationContext: + fresh: Final = await self._refresh() + if self._caller(fresh) != self._identity or self._owner(fresh) != self.owner: + raise HTTPException(status_code=403, detail="Subscription caller changed") + return fresh + + +def _supported_notifications(requested: SubscriptionFilter, capabilities: ServerCapabilities) -> SubscriptionFilter: + return SubscriptionFilter( + tools_list_changed=bool(requested.tools_list_changed and capabilities.tools and capabilities.tools.list_changed) + or None, + prompts_list_changed=bool( + requested.prompts_list_changed and capabilities.prompts and capabilities.prompts.list_changed + ) + or None, + resources_list_changed=bool( + requested.resources_list_changed and capabilities.resources and capabilities.resources.list_changed + ) + or None, + resource_subscriptions=list(requested.resource_subscriptions or ()) + if capabilities.resources and capabilities.resources.subscribe + else None, + ) + + +class LegacyNotifications: + def __init__(self, *, max_buffered_events: int = 128) -> None: + self.honored = SubscriptionFilter() + self._session: ClientSession | None = None + self._events: Final[asyncio.Queue[ServerEvent | MCPError]] = asyncio.Queue(max_buffered_events) + + async def receive(self, notification: ServerNotification) -> None: + event: Final = event_from_wire( + notification.method, + notification.params.model_dump(by_alias=True) if notification.params is not None else None, + ) + if event is None or not event_matches( + self.honored, frozenset(self.honored.resource_subscriptions or ()), event + ): + return + self.publish(event) + + def publish(self, event: ServerEvent) -> None: + try: + self._events.put_nowait(event) + except asyncio.QueueFull: + self._events.get_nowait() + self._events.put_nowait(MCPError(code=-32001, message="Legacy notification buffer overflow; refresh")) + + def __aiter__(self) -> "LegacyNotifications": + return self + + async def __anext__(self) -> ServerEvent: + while True: + with anyio.move_on_after(15 if self._session is not None else float("inf")) as idle: + event = await self._events.get() + if isinstance(event, MCPError): + raise event + return event + if idle.cancel_called and self._session is not None: + await self._session.send_ping() + + @asynccontextmanager + async def listen( + self, session: ClientSession, requested: SubscriptionFilter + ) -> AsyncGenerator["LegacyNotifications"]: + self._session = session + capabilities: Final = session.server_capabilities or ServerCapabilities() + self.honored = _supported_notifications(requested, capabilities) + cleanup: Final = AsyncExitStack() + try: + for uri in self.honored.resource_subscriptions or (): + await session.send_request(SubscribeRequest(params=SubscribeRequestParams(uri=uri)), EmptyResult) + cleanup.push_async_callback( + session.send_request, UnsubscribeRequest(params=UnsubscribeRequestParams(uri=uri)), EmptyResult + ) + yield self + finally: + with anyio.move_on_after(5, shield=True): + await cleanup.aclose() + self._session = None + + +@dataclass(frozen=True) +class SubscriptionTarget: + server: MCPServer + client: MCPClient + credential: str + + +class _AuthorizedSession(ServerSession): + def __init__(self, session: ServerSession, authorize: Callable[[], Awaitable[OperationContext]]) -> None: + self._session = session + self._authorize = authorize + + async def send_notification( + self, notification: ServerNotification, related_request_id: RequestId | None = None + ) -> None: + await self._authorize() + await self._session.send_notification(notification, related_request_id) + + +class _ListenerBus(InMemorySubscriptionBus): + def __init__(self) -> None: + super().__init__() + self.ready = asyncio.Event() + + def subscribe(self, listener: Callable[[ServerEvent], None]) -> Callable[[], None]: + unsubscribe: Final = super().subscribe(listener) + self.ready.set() + return unsubscribe + + +async def _targets(context: OperationContext) -> tuple[SubscriptionTarget, ...]: + manager: Final = operations.global_mcp_server_manager + auth, token, servers, server_headers, oauth_headers, headers, ip = context.legacy_auth() + async with manager.catalog.operation(): + allowed: Final = await operations._get_allowed_mcp_servers(auth, servers, ip) + + async def prepare(server: MCPServer) -> SubscriptionTarget: + server_auth, extra = operations._prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=server_headers, + mcp_auth_header=token, + oauth2_headers=oauth_headers, + raw_headers=headers, + user_api_key_auth=auth, + scope_servers=list(allowed), + ) + client: Final = await manager.create_mcp_client( + server, + mcp_auth_header=server_auth, + extra_headers={**(extra or {}), **(server.static_headers or {})}, + stdio_env=manager.build_stdio_env(server, headers), + subject_token=manager._obo_subject_token(server, headers, auth), # pyright: ignore[reportPrivateUsage] # reuse the manager-owned OBO policy + user_api_key_auth=auth, + raw_headers=headers, + client_ip=ip, + ) + return SubscriptionTarget(server, client, await client.subscription_auth_fingerprint()) + + return tuple([await prepare(server) for server in allowed]) + + +async def _listing(event: ServerEvent, context: OperationContext) -> str | None: + if isinstance(event, ResourceUpdated): + return None + operation: Final = ( + ListToolsRequest() + if isinstance(event, ToolsListChanged) + else ListPromptsRequest() + if isinstance(event, PromptsListChanged) + else ListResourcesRequest() + ) + gateway: Final = operations.GatewayOperations() + digest: Final = hashlib.sha256() + cursor: str | None = None # rebind-ok: advance the signed gateway cursor for each catalog page + with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): + for _ in range(MCP_TOOL_LISTING_MAX_PAGES): + result = await gateway.execute( + operation.model_copy(update={"params": PaginatedRequestParams(cursor=cursor)}), context + ) + items = ( + result.tools + if isinstance(result, ListToolsResult) + else result.prompts + if isinstance(result, ListPromptsResult) + else result.resources + ) + for item in items: + digest.update(item.model_dump_json().encode()) + cursor = result.next_cursor + if cursor is None: + return digest.hexdigest() + raise MCPError(code=-32603, message="Subscription catalog exceeds the pagination limit") + + +async def _recheck_authorization(check: Callable[[], Awaitable[OperationContext]]) -> None: + while True: + await asyncio.sleep(15) + await check() + + +async def _deliver_legacy( + ctx: ServerRequestContext, + bus: _ListenerBus, + check: Callable[[], Awaitable[OperationContext]], + notifications: SubscriptionFilter, + ready: Callable[[SubscriptionFilter], None] | None, +) -> NoReturn: + events: Final = LegacyNotifications() + unsubscribe: Final = bus.subscribe(events.publish) + try: + if ready is not None: + ready(notifications) + while True: + event = await anext(events) + await check() + await ctx.session.send_notification(event_to_notification(event, {})) + finally: + unsubscribe() + + +@dataclass +class LegacySubscriptionSession: + group: TaskGroup + authorization: SubscriptionAuthorization + resources: Mapping[str, anyio.CancelScope] + lock: asyncio.Lock + + +class GatewaySubscriptions: + def __init__( + self, + *, + max_listeners: int = 128, + max_per_caller: int = 8, + resolve_targets: Callable[[OperationContext], Awaitable[tuple[SubscriptionTarget, ...]]] = _targets, + ) -> None: + self._resolve_targets = resolve_targets + self._slots = anyio.CapacityLimiter(max_listeners) + self._upstream_slots = anyio.CapacityLimiter(max_listeners) + self._max_per_caller = max_per_caller + self._owners: tuple[str, ...] = () + self._legacy: Mapping[int, LegacySubscriptionSession] = {} + self._legacy_started = asyncio.Event() + self._initializing: Mapping[int, asyncio.Event] = {} + + async def listen( + self, + ctx: ServerRequestContext, + params: SubscriptionsListenRequestParams, + authorize: Callable[[], Awaitable[OperationContext]], + owner: str, + ) -> SubscriptionsListenResult: + if ctx.protocol_version != "2026-07-28": + raise MCPError(code=-32601, message="Subscriptions require the modern MCP protocol") + async with self._capacity(owner): + return await self._listen(ctx, params, authorize) + + @asynccontextmanager + async def _capacity(self, owner: str) -> AsyncGenerator[None]: + if self._owners.count(owner) >= self._max_per_caller: + raise MCPError(code=-32603, message="Caller subscription limit reached") + try: + self._slots.acquire_nowait() + except anyio.WouldBlock as error: + raise MCPError(code=-32603, message="Subscription limit reached") from error + self._owners += (owner,) + try: + yield + finally: + position: Final = self._owners.index(owner) + self._owners = self._owners[:position] + self._owners[position + 1 :] + self._slots.release() + + @asynccontextmanager + async def initializing(self, ctx: ServerRequestContext) -> AsyncGenerator[bool]: + key: Final = id(ctx.session.client_params) + if key in self._initializing: + yield False + return + ready: Final = asyncio.Event() + self._initializing = {**self._initializing, key: ready} + try: + yield True + finally: + ready.set() + self._initializing = { + registered: event for registered, event in self._initializing.items() if registered != key + } + + async def wait_initialized(self, ctx: ServerRequestContext) -> None: + ready: Final = self._initializing.get(id(ctx.session.client_params)) + if ready is not None: + await ready.wait() + + async def legacy_session(self, ctx: ServerRequestContext, authorization: SubscriptionAuthorization) -> None: + key: Final = id(ctx.session.client_params) + if ctx.session.client_params is None or key in self._legacy: + return + async with anyio.create_task_group() as group: + state: Final = LegacySubscriptionSession(group, authorization, {}, asyncio.Lock()) + self._legacy = {**self._legacy, key: state} + self._legacy_started.set() + try: + await group.start( + self._legacy_listener, + ctx, + SubscriptionFilter(tools_list_changed=True, prompts_list_changed=True, resources_list_changed=True), + authorization, + ) + if (ready := self._initializing.get(key)) is not None: + ready.set() + await anyio.sleep_forever() + finally: + self._legacy = {registered: state for registered, state in self._legacy.items() if registered != key} + group.cancel_scope.cancel() + + async def _legacy_listener( + self, + ctx: ServerRequestContext, + notifications: SubscriptionFilter, + authorization: SubscriptionAuthorization, + scope: anyio.CancelScope | None = None, + *, + task_status: TaskStatus[None], + ) -> None: + with scope if scope is not None else anyio.CancelScope(): + + def ready(honored: SubscriptionFilter) -> None: + if ( + notifications.resource_subscriptions + and honored.resource_subscriptions != notifications.resource_subscriptions + ): + raise MCPError(code=-32602, message="Upstream does not support this resource subscription") + task_status.started() + + async with self._capacity(authorization.owner): + await self._listen( + ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, ready + ) + + async def legacy_resource(self, ctx: ServerRequestContext, uri: str, *, subscribe: bool) -> None: + if ctx.session.client_params is None: + raise MCPError(code=-32602, message="Resource subscriptions require a live legacy session") + key: Final = id(ctx.session.client_params) + with anyio.fail_after(10): + while key not in self._legacy: + self._legacy_started.clear() + await self._legacy_started.wait() + state: Final = self._legacy[key] + async with state.lock: + await state.authorization() + if subscribe and uri not in state.resources: + scope: Final = anyio.CancelScope() + await state.group.start( + self._legacy_listener, + ctx, + SubscriptionFilter(resource_subscriptions=[uri]), + state.authorization, + scope, + ) + state.resources = {**state.resources, uri: scope} + elif not subscribe and uri in state.resources: + state.resources[uri].cancel() + state.resources = {resource: scope for resource, scope in state.resources.items() if resource != uri} + + async def _listen( + self, + ctx: ServerRequestContext, + params: SubscriptionsListenRequestParams, + authorize: Callable[[], Awaitable[OperationContext]], + ready: Callable[[SubscriptionFilter], None] | None = None, + ) -> SubscriptionsListenResult: + context: Final = await authorize() + targets: Final = await self._resolve_targets(context) + if not targets: + raise HTTPException(status_code=403, detail="No authorized MCP subscription targets") + if params.notifications.resource_subscriptions and len(targets) != 1: + raise MCPError(code=-32602, message="Resource subscriptions require exactly one allowed server") + bus: Final = _ListenerBus() + handler: Final = ListenHandler(bus, max_subscriptions=1, max_buffered_events=128) + + async def check() -> OperationContext: + fresh: Final = await authorize() + current: Final = await self._resolve_targets(fresh) + if tuple((configuration_identity(item.server), item.credential) for item in current) != tuple( + (configuration_identity(item.server), item.credential) for item in targets + ): + raise MCPError(code=-32001, message="Subscription authorization or upstream configuration changed") + return fresh + + async def relay(target: SubscriptionTarget, *, task_status: TaskStatus[SubscriptionFilter]) -> None: + await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, target.server) # pyright: ignore[reportPrivateUsage] # same admission policy as other gateway operations + + legacy: Final = LegacyNotifications() + + async def connected(session: ClientSession) -> None: + requested: Final = params.notifications + if ctx.protocol_version != "2026-07-28" and not any( + _supported_notifications(requested, session.server_capabilities or ServerCapabilities()) + .model_dump() + .values() + ): + task_status.started(SubscriptionFilter()) + return + requested_events: Final = tuple( + event + for event in (ToolsListChanged(), PromptsListChanged(), ResourcesListChanged()) + if event_matches(requested, frozenset(), event) + ) + baseline: Final = {event: await _listing(event, context) for event in requested_events} + stream: Final = ( + listen( + session, + tools_list_changed=bool(requested.tools_list_changed), + prompts_list_changed=bool(requested.prompts_list_changed), + resources_list_changed=bool(requested.resources_list_changed), + resource_subscriptions=requested.resource_subscriptions or (), + ) + if session.protocol_version == "2026-07-28" + else legacy.listen(session, requested) + ) + async with stream as subscription: + events: Final = tuple( + event for event in requested_events if event_matches(subscription.honored, frozenset(), event) + ) + setup_context: Final = await check() + catalog: Final = operations.global_mcp_server_manager + catalog.invalidate_discovery_lists(target.server.server_id) + catalog._drop_listed_tools(target.server.server_id) # pyright: ignore[reportPrivateUsage] # refresh the setup snapshot through the same catalog owner as event delivery + previous: Final = {event: await _listing(event, setup_context) for event in events} + setup_changes: Final = tuple(event for event in events if baseline[event] != previous[event]) + task_status.started(subscription.honored) + await bus.ready.wait() + for event in setup_changes: + await bus.publish(event) + async for event in subscription: + fresh = await check() + manager = operations.global_mcp_server_manager + manager.invalidate_discovery_lists(target.server.server_id) + manager._drop_listed_tools(target.server.server_id) # pyright: ignore[reportPrivateUsage] # evict served catalogs without invalidating OAuth metadata + listing = await _listing(event, fresh) + if not isinstance(event, ResourceUpdated): + if previous.get(event) == listing: + continue + previous[event] = listing + await bus.publish(event) + raise MCPError(code=-32001, message="Upstream subscription ended; listen again and refresh") + + try: + self._upstream_slots.acquire_nowait() + except anyio.WouldBlock as error: + raise MCPError(code=-32603, message="Upstream subscription limit reached") from error + try: + await target.client.run_with_session(connected, on_notification=legacy.receive) + finally: + self._upstream_slots.release() + + async with anyio.create_task_group() as group: + honored: Final = tuple( + [SubscriptionFilter.model_validate(await group.start(relay, target)) for target in targets] + ) + notifications: Final = SubscriptionFilter( + tools_list_changed=any(item.tools_list_changed for item in honored) or None, + prompts_list_changed=any(item.prompts_list_changed for item in honored) or None, + resources_list_changed=any(item.resources_list_changed for item in honored) or None, + resource_subscriptions=list( + dict.fromkeys(chain.from_iterable(item.resource_subscriptions or () for item in honored)) + ) + or None, + ) + group.start_soon(_recheck_authorization, check) + try: + if ctx.protocol_version != "2026-07-28": + await _deliver_legacy(ctx, bus, check, notifications, ready) + return await handler( + replace(ctx, session=_AuthorizedSession(ctx.session, check)), + params.model_copy(update={"notifications": notifications}), + ) + finally: + group.cancel_scope.cancel() + + raise MCPError(code=-32001, message="Subscription ended before completion") diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index 42af8cf5dba..c254139c109 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -6,12 +6,10 @@ from typing import Final import httpx import pytest import yaml -from pydantic import JsonValue - from integration._support.client import Gateway from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server - +from pydantic import JsonValue _KEY: Final = "sk-interaction-test" _PROXY_PYTHONPATH: Final = os.pathsep.join( @@ -393,3 +391,742 @@ def test_missing_continuation_key_reports_configuration_for_each_carrier(tmp_pat response: Final = rpc(gateway, method, params, status=400) assert response["error"]["code"] == -32602, response assert "LITELLM_SALT_KEY" in response["error"]["message"], response + + +@pytest.mark.parametrize( + "event_kind,access", + [ + ("tools", "normal"), + ("tools", "modern-boundary"), + ("prompts", "normal"), + ("resources", "normal"), + ("updated", "normal"), + ("tools", "revoked"), + ("tools", "identity"), + ("tools", "team"), + ("tools", "org"), + ("tools", "end-user"), + ("tools", "idle"), + ("tools", "hidden"), + ("tools", "upstream-ended"), + ("tools", "legacy-caller-isolation"), + ("tools", "legacy-caller-no-notifications"), + ("updated", "legacy-caller-no-notifications"), + ("tools", "legacy-caller-revoked"), + ("tools", "legacy-sse"), + ("updated", "legacy-sse"), + ("prompts", "legacy-caller"), + ("resources", "legacy-caller"), + ("prompts", "legacy-both"), + ("resources", "legacy-both"), + ("tools", "legacy-caller"), + ("updated", "legacy-caller"), + ("tools", "legacy-both"), + ("updated", "legacy-both"), + ("tools", "legacy"), + ("prompts", "legacy"), + ("resources", "legacy"), + ("updated", "legacy"), + ("tools", "pages"), + ("tools", "fanout"), + ("tools", "setup-race"), + ("tools", "metadata"), + ("tools", "config-change"), + ("tools", "renewal"), + ("tools", "renewal-scopes"), + ("tools", "renewal-failed"), + ], +) +def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( + tmp_path: Path, event_kind: str, access: str +) -> None: + import asyncio + from collections.abc import Callable + from contextlib import AsyncExitStack, asynccontextmanager + from threading import Event + + import httpx2 + from integration._support.asgi import asgi_server + from integration._support.mcp import _drain_sse_streams + from mcp import ClientSession, MCPError + from mcp.client.sse import sse_client + from mcp.client.streamable_http import streamable_http_client + from mcp.client.subscriptions import SubscriptionLost, listen + from mcp.server.mcpserver import MCPServer + from mcp.server.subscriptions import ( + InMemorySubscriptionBus, + ListenHandler, + PromptsListChanged, + ResourcesListChanged, + ResourceUpdated, + ToolsListChanged, + ) + from mcp.server.transport_security import TransportSecuritySettings + from mcp.shared.subscriptions import ServerEvent + from mcp.types import DiscoverResult, ListToolsResult, PaginatedRequestParams, SubscriptionsListenRequestParams + from sse_starlette.sse import AppStatus + from starlette.applications import Starlette + from starlette.responses import JSONResponse + from starlette.routing import Mount, Route + + disconnected: Final = Event() + connected: Final = Event() + tool_executions = [] + token_grants = [] + renewal_started = Event() + + async def token_endpoint(request): + from urllib.parse import parse_qs + + form = parse_qs((await request.body()).decode()) + assert form["grant_type"] == ["client_credentials"] + assert form["client_id"] == ["subscription-service"] + assert form["scope"] == ["read write"] + token_grants.append(form) + if renewal_started.is_set() and access == "renewal-failed": + return JSONResponse({"error": "invalid_client"}, status_code=401) + return JSONResponse( + { + "access_token": "synthetic-renewal-" + str(len(token_grants)), + "token_type": "Bearer", + "expires_in": 70, + "scope": "read" if renewal_started.is_set() and access == "renewal-scopes" else "read write", + } + ) + + listener_credentials = [] + catalog_credentials = [] + setup_changed = Event() + + active_listeners = set() + + class ObservedBus(InMemorySubscriptionBus): + def subscribe(self, listener: Callable[[ServerEvent], None]) -> Callable[[], None]: + unsubscribe = super().subscribe(listener) + marker = object() + active_listeners.add(marker) + connected.set() + + def close() -> None: + unsubscribe() + active_listeners.discard(marker) + disconnected.set() + + return close + + auth_module: Final = tmp_path / "subscription_auth.py" + auth_module.write_text( + "from pathlib import Path\n" + "from fastapi import HTTPException, Request\n" + "from litellm.proxy._types import UserAPIKeyAuth\n" + "async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth:\n" + " if api_key not in ('sk-subscriptions-test', 'sk-subscriptions-other'):\n" + " raise HTTPException(status_code=401, detail='Unknown caller')\n" + " import json\n" + " data = json.loads(Path(__file__).with_suffix('.json').read_text())\n" + " change = data.pop('subscription_test_update', None)\n" + " if change:\n" + " from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager as manager\n" + " server = manager.config_mcp_servers['subscription-server']\n" + " fields = {'authorization_url': 'https://authorization.example/authorize'} if change == 'metadata' else {'url': server.url.rstrip('/') + '/changed'}\n" + " manager.config_mcp_servers[server.server_id] = server.model_copy(update=fields)\n" + " if api_key == 'sk-subscriptions-other':\n" + " data = {**data, 'api_key': api_key, 'user_id': 'bob', 'object_permission': {**data['object_permission'], 'mcp_tool_permissions': {'subscription-server': ['change_tools']}}}\n" + " return UserAPIKeyAuth.model_validate(data)\n" + ) + identity: Final = { + "api_key": "sk-subscriptions-test", + "user_id": "alice", + "user_role": "internal_user", + "object_permission": { + "object_permission_id": "subscription-permission", + "mcp_servers": ["subscription-server", "second-server"] if access == "fanout" else ["subscription-server"], + }, + } + auth_state: Final = auth_module.with_suffix(".json") + auth_state.write_text(json.dumps(identity)) + bus: Final = ObservedBus() + upstream_listener: Final = ListenHandler(bus) + + async def control_listen(ctx, call_next): + if ctx.protocol_version != "2026-07-28" and ctx.method == "notifications/initialized": + from mcp.shared.subscriptions import event_to_notification + + queue = asyncio.Queue() + unsubscribe = bus.subscribe(queue.put_nowait) + connected.set() + try: + while True: + await ctx.session.send_notification(event_to_notification(await queue.get(), {})) + finally: + unsubscribe() + if ctx.method in ("resources/subscribe", "resources/unsubscribe"): + from mcp.types import EmptyResult + + return EmptyResult() + + if ctx.method == "tools/list": + catalog_credentials.append(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 == "subscriptions/listen": + listener_credentials.append(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)) + result = await call_next(ctx) + if access == "legacy-caller-no-notifications" and ctx.method in ("initialize", "server/discover"): + from mcp.types import InitializeResult, ServerCapabilities + + parsed = (DiscoverResult if ctx.method == "server/discover" else InitializeResult).model_validate(result) + return parsed.model_copy(update={"capabilities": ServerCapabilities()}) + if ctx.method == "initialize": + from mcp.types import ( + InitializeResult, + PromptsCapability, + ResourcesCapability, + ServerCapabilities, + ToolsCapability, + ) + + result = InitializeResult.model_validate(result) + return result.model_copy( + update={ + "capabilities": ServerCapabilities( + tools=ToolsCapability(list_changed=True), + prompts=PromptsCapability(list_changed=True), + resources=ResourcesCapability(list_changed=True, subscribe=True), + ) + } + ) + if access == "pages" and ctx.method == "tools/list": + result = ListToolsResult.model_validate(result) + offset = int((ctx.params or {}).get("cursor") or "0") + return result.model_copy( + update={ + "tools": result.tools[offset : offset + 1], + "next_cursor": str(offset + 1) if offset + 1 < len(result.tools) else None, + } + ) + return result + + upstream: Final = MCPServer( + "subscription-peer", + subscriptions=None if access == "legacy-caller-no-notifications" else bus, + middleware=[control_listen], + ) + + @upstream.resource("test://value") + def value() -> str: + return "value" + + @upstream.tool() + async def change_tools() -> str: + tool_executions.append("change_tools") + + async def added() -> str: + tool_executions.append("added") + return "new tool" + + if access == "upstream-ended": + upstream_listener.close() + return "ended" + if event_kind == "tools": + upstream.add_tool(added) + await bus.publish(ToolsListChanged()) + elif event_kind == "prompts": + upstream.prompt()(added) + await bus.publish(PromptsListChanged()) + elif event_kind == "resources": + upstream.resource("test://added")(added) + await bus.publish(ResourcesListChanged()) + else: + await bus.publish(ResourceUpdated(uri="test://value")) + return "changed" + + async def exercise(gateway: Gateway) -> None: + headers = {"Authorization": "Bearer " + gateway.key} + if access == "fanout": + headers = {"x-litellm-api-key": gateway.key, "Authorization": "Bearer upstream-only-test-token"} + async with httpx2.AsyncClient(headers=headers, timeout=30) as client: + gateway_url = str(gateway.client.base_url).rstrip("/") + connection = ( + sse_client(gateway_url + "/mcp/sse", headers=headers) + if access == "legacy-sse" + else streamable_http_client(gateway_url + "/mcp/", http_client=client) + ) + async with connection as streams: + received = asyncio.Queue() + + async def receive(notification, queue=received): + from mcp.shared.subscriptions import event_from_wire + + if not isinstance(notification, Exception): + event = event_from_wire( + notification.method, + notification.params.model_dump(by_alias=True) if notification.params else None, + ) + if event is not None: + await queue.put(event) + + async with ( + ClientSession(streams[0], streams[1], message_handler=receive) as session, + AsyncExitStack() as additional, + ): + other_received = asyncio.Queue() + if access == "legacy-caller-isolation": + other_client = await additional.enter_async_context( + httpx2.AsyncClient(headers={"Authorization": "Bearer sk-subscriptions-other"}) + ) + other_streams = await additional.enter_async_context( + streamable_http_client(gateway_url + "/mcp/", http_client=other_client) + ) + + async def receive_other(notification): + await receive(notification, other_received) + + other = await additional.enter_async_context( + ClientSession(other_streams[0], other_streams[1], message_handler=receive_other) + ) + await other.initialize() + 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() + before = await session.list_tools() + assert "peer-added" not in {tool.name for tool in before.tools} + if event_kind == "updated": + if access == "legacy-caller-no-notifications": + with pytest.raises(MCPError): + await asyncio.wait_for(session.subscribe_resource("test://value"), 3) + assert listener_credentials == [] + assert not active_listeners + assert not (await session.call_tool("peer-change_tools", {})).is_error + return + await session.subscribe_resource("test://value") + if access != "legacy-caller-no-notifications": + assert await asyncio.to_thread(connected.wait, 5) + if access == "legacy-caller-revoked": + auth_state.write_text( + json.dumps( + { + **identity, + "object_permission": { + "object_permission_id": "subscription-permission", + "mcp_servers": ["denied"], + }, + } + ) + ) + async with streamable_http_client(upstream_url + "/mcp") 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")) + ) + assert not (await direct.call_tool("change_tools", {})).is_error + with pytest.raises(TimeoutError): + await asyncio.wait_for(received.get(), 0.5) + assert await asyncio.to_thread(disconnected.wait, 3) + return + changed = await session.call_tool("peer-change_tools", {}) + assert not changed.is_error + if access == "legacy-caller-no-notifications": + with pytest.raises(TimeoutError): + await asyncio.wait_for(received.get(), 0.5) + assert "peer-added" in {tool.name for tool in (await session.list_tools()).tools} + assert listener_credentials == [] + return + event = await asyncio.wait_for(received.get(), 10) + assert ( + event + == ( + { + "tools": ToolsListChanged(), + "prompts": PromptsListChanged(), + "resources": ResourcesListChanged(), + "updated": ResourceUpdated(uri="test://value"), + }[event_kind] + ) + ) + if access == "legacy-caller-isolation": + with pytest.raises(TimeoutError): + await asyncio.wait_for(other_received.get(), 0.5) + assert {tool.name for tool in (await other.list_tools()).tools} == {"peer-change_tools"} + before_denied = tuple(tool_executions) + denied = await other.call_tool("peer-added", {}) + assert denied.is_error + assert tuple(tool_executions) == before_denied + if event_kind == "tools": + assert "peer-added" in {tool.name for tool in (await session.list_tools()).tools} + elif event_kind == "prompts": + assert "peer-added" in {prompt.name for prompt in (await session.list_prompts()).prompts} + elif event_kind == "resources": + assert "test://added" in { + str(resource.uri) for resource in (await session.list_resources()).resources + } + else: + await session.unsubscribe_resource("test://value") + changed = await session.call_tool("peer-change_tools", {}) + assert not changed.is_error + 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"))) + 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 == "modern-boundary": + for method, params in ( + ("notifications/initialized", {}), + ("initialize", {"protocolVersion": "2026-07-28", "clientInfo": {"name": "boundary", "version": "1"}, "capabilities": {}}), + ("resources/subscribe", {"uri": "test://value"}), + ("resources/unsubscribe", {"uri": "test://value"}), + ): + notification = method == "notifications/initialized" + response = await client.post( + gateway_url + "/mcp/", + headers={ + "MCP-Protocol-Version": "2026-07-28", + "Mcp-Method": method, + "Mcp-Name": str(params.get("uri", "")), + "Accept": "application/json, text/event-stream", + }, + json={ + "jsonrpc": "2.0", + **({} if notification else {"id": "unsupported"}), + "method": method, + "params": {**params, "_meta": _META}, + }, + timeout=5, + ) + if notification: + assert response.status_code == 202 + else: + assert response.status_code == 404, response.text + assert response.json()["error"]["code"] == -32601 + assert listener_credentials == [], "Removed methods opened an implicit upstream listener" + assert not active_listeners + async with listen( + session, + tools_list_changed=event_kind == "tools", + prompts_list_changed=event_kind == "prompts", + resources_list_changed=event_kind == "resources", + resource_subscriptions=["test://value"] if event_kind == "updated" else (), + ) as subscription: + if access == "fanout": + assert listener_credentials == [None, None], ( + "Unscoped caller credential reached multiple upstreams" + ) + assert catalog_credentials and all(value is None for value in catalog_credentials), ( + "Unscoped caller credential reached catalog listings" + ) + if access in ("metadata", "config-change"): + auth_state.write_text(json.dumps({**identity, "subscription_test_update": access})) + if access in ("revoked", "identity", "team", "org", "end-user", "idle"): + changed_identity = ( + { + **identity, + { + "identity": "user_id", + "team": "team_id", + "org": "org_id", + "end-user": "end_user_id", + }[access]: "bob", + } + if access in ("identity", "team", "org", "end-user") + else { + **identity, + "object_permission": { + "object_permission_id": "subscription-permission", + "mcp_servers": ["denied"], + }, + } + ) + auth_state.write_text(json.dumps(changed_identity)) + if access == "idle": + with pytest.raises((SubscriptionLost, StopAsyncIteration)): + await asyncio.wait_for(anext(subscription), 20) + assert await asyncio.to_thread(disconnected.wait, 3) + return + if access.startswith("renewal"): + assert token_grants + admitted_grants = len(token_grants) + renewal_started.set() + async with asyncio.timeout(25): + while len(token_grants) == admitted_grants: + try: + await session.list_tools() + except MCPError: + assert access == "renewal-failed" + 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: + async with ClientSession(direct_streams[0], direct_streams[1]) as direct: + 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 ( + "revoked", + "identity", + "team", + "org", + "end-user", + "idle", + "upstream-ended", + "config-change", + "renewal-scopes", + "renewal-failed", + ): + with pytest.raises((SubscriptionLost, StopAsyncIteration)): + await asyncio.wait_for(anext(subscription), 5) + assert await asyncio.to_thread(disconnected.wait, 3) + return + if access == "hidden": + after = await session.list_tools() + assert {tool.name for tool in after.tools} == {"peer-change_tools"} + with pytest.raises(TimeoutError): + await asyncio.wait_for(anext(subscription), 0.5) + return + event = await asyncio.wait_for(anext(subscription), 5) + expected = { + "tools": ToolsListChanged(), + "prompts": PromptsListChanged(), + "resources": ResourcesListChanged(), + "updated": ResourceUpdated(uri="test://value"), + } + assert event == expected[event_kind] + if access == "fanout": + assert all(value is None for value in catalog_credentials), ( + "Unscoped caller credential reached refreshed catalogs" + ) + if event_kind == "tools": + assert subscription.honored.tools_list_changed + after = await session.list_tools() + if access == "pages": + assert after.next_cursor + assert "peer-added" not in {tool.name for tool in after.tools} + after = await session.list_tools( + params=PaginatedRequestParams(cursor=after.next_cursor) + ) + assert "peer-added" in {tool.name for tool in after.tools} + elif event_kind == "prompts": + assert subscription.honored.prompts_list_changed + prompts = await session.list_prompts() + assert "peer-added" in {prompt.name for prompt in prompts.prompts} + elif event_kind == "resources": + assert subscription.honored.resources_list_changed + resources = await session.list_resources() + assert "test://added" in {str(resource.uri) for resource in resources.resources} + else: + assert subscription.honored.resource_subscriptions == ["test://value"] + if access == "modern-boundary": + assert len(listener_credentials) == 1 + with pytest.raises(TimeoutError): + 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( + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False) + ) + + @asynccontextmanager + async def lifespan(_: Starlette): + AppStatus.should_exit = False + async with app.router.lifespan_context(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: 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: + config: Final = tmp_path / "proxy.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [], + "mcp_servers": { + "peer": { + "server_id": "subscription-server", + **( + { + "auth_type": "oauth2", + "oauth2_flow": "client_credentials", + "client_id": "subscription-service", + "client_secret": "synthetic-client-secret", + "token_url": upstream_url + "/token", + "scopes": ["read", "write"], + } + 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", + **({"allowed_tools": ["change_tools"]} if access == "hidden" else {}), + **( + {"auth_type": "oauth_delegate", "extra_headers": ["Authorization"]} + if access == "fanout" + else {} + ), + }, + **( + { + "second": { + "server_id": "second-server", + "url": upstream_url + "/mcp", + "transport": "http", + "protocol_version": "2026-07-28", + "auth_type": "oauth_delegate", + "extra_headers": ["Authorization"], + } + } + if access == "fanout" + else {} + ), + }, + "general_settings": { + "master_key": "sk-subscriptions-test", + "custom_auth": "subscription_auth.authenticate", + "store_model_in_db": False, + "mcp_advertised_versions": ["2025-11-25", "2026-07-28"], + }, + } + ) + ) + source: Final = Path(__file__).resolve().parents[3] + environment: Final = { + "PYTHONPATH": os.pathsep.join((str(source), str(source / "tests"))), + "LITELLM_SALT_KEY": "subscription-pagination-key", + "STORE_MODEL_IN_DB": "False", + "DISABLE_SCHEMA_UPDATE": "true", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true", + } + seed: Final = Gateway(client, "sk-subscriptions-test", upstream_url) + with owned_proxy( + seed, + tmp_path / "gateway", + 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 gateway: + asyncio.run(exercise(gateway)) + from integration._support.client import eventually + + eventually(lambda: len(active_listeners), lambda count: count == 0, seconds=10) + assert not active_listeners, "Upstream listener survived client disconnect" + + +def test_subscription_partial_setup_releases_upstream_capacity() -> None: + import asyncio + from contextlib import asynccontextmanager + from threading import Event + + import httpx2 + from integration._support.asgi import asgi_server + from integration._support.mcp import _drain_sse_streams + from mcp import ClientSession, MCPError + from mcp.client.streamable_http import streamable_http_client + from mcp.client.subscriptions import listen + from mcp.server.mcpserver import MCPServer + from mcp.server.subscriptions import InMemorySubscriptionBus + from mcp.server.transport_security import TransportSecuritySettings + from mcp.types import DiscoverResult, SubscriptionsListenRequestParams + from sse_starlette.sse import AppStatus + from starlette.applications import Starlette + from starlette.routing import Mount + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.contracts import OperationContext + from litellm.proxy._experimental.mcp_server.subscriptions import GatewaySubscriptions, SubscriptionTarget + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer as RegisteredServer + + disconnected = Event() + + class ObservedBus(InMemorySubscriptionBus): + def subscribe(self, listener): + unsubscribe = super().subscribe(listener) + + def close(): + unsubscribe() + disconnected.set() + + return close + + selected = [] + + async def resolve(context): + return tuple(selected) + + async def authorize(): + return OperationContext(_caller=None) + + service = GatewaySubscriptions(max_listeners=1, resolve_targets=resolve) + + async def dispatch(ctx, call_next): + if ctx.method == "subscriptions/listen": + return await service.listen( + ctx, SubscriptionsListenRequestParams.model_validate(ctx.params), authorize, "alice" + ) + return await call_next(ctx) + + security = TransportSecuritySettings(enable_dns_rebinding_protection=False) + upstream = MCPServer("capacity-upstream", subscriptions=ObservedBus()).streamable_http_app( + transport_security=security + ) + gateway = MCPServer("capacity-gateway", middleware=[dispatch]).streamable_http_app(transport_security=security) + + @asynccontextmanager + async def lifespan(app): + AppStatus.should_exit = False + async with upstream.router.lifespan_context(upstream), gateway.router.lifespan_context(gateway): + 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("/upstream", app=upstream), Mount("/gateway", app=gateway)], lifespan=lifespan) + + async def exercise(url): + client = MCPClient( + server_url=url + "/upstream/mcp", transport_type=MCPTransport.http, protocol_version="2026-07-28" + ) + target = SubscriptionTarget( + RegisteredServer( + server_id="capacity-peer", name="capacity-peer", transport="http", protocol_version="2026-07-28" + ), + client, + "uncredentialed", + ) + selected.extend((target, target)) + async with httpx2.AsyncClient(timeout=30) as http: + async with streamable_http_client(url + "/gateway/mcp", http_client=http) as streams: + async with ClientSession(streams[0], streams[1]) as session: + session.adopt(DiscoverResult.model_validate(await session.send_discover("2026-07-28"))) + with pytest.raises(MCPError): + async with listen(session): + pytest.fail("Partial upstream setup was acknowledged") + assert await asyncio.to_thread(disconnected.wait, 3), "Partial listener leaked its upstream" + selected.pop() + disconnected.clear() + async with listen(session) as subscription: + assert not subscription.honored.tools_list_changed + assert await asyncio.to_thread(disconnected.wait, 3), "Retry leaked its upstream" + + with asgi_server(hosted) as url: + asyncio.run(exercise(url)) diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 5aba5d5dfd6..418738a9dca 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -2821,6 +2821,8 @@ async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: repeated: Final = await client("private-original-credential").discovery_auth_fingerprint() replaced: Final = await client("private-replaced-credential").discovery_auth_fingerprint() assert original == repeated + assert await client("private-original-credential").subscription_auth_fingerprint() == original + assert await client("private-replaced-credential").subscription_auth_fingerprint() == replaced assert original != replaced assert len(original) == 64 assert "private-original-credential" not in original @@ -2843,6 +2845,8 @@ async def test_discovery_auth_fingerprint_tracks_aws_sigv4_credential_identity() repeated: Final = await client("AKIAORIGINAL").discovery_auth_fingerprint() rotated: Final = await client("AKIAROTATED").discovery_auth_fingerprint() assert original == repeated + assert await client("AKIAORIGINAL").subscription_auth_fingerprint() == original + assert await client("AKIAROTATED").subscription_auth_fingerprint() == rotated assert original != rotated assert "AKIAORIGINAL" not in original @@ -5088,3 +5092,52 @@ async def test_cancelling_persistent_operation_stops_work_and_keeps_session_usab session.close() await asyncio.wait_for(session.wait_closed(), 5) await asyncio.gather(waiter, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("header_name", ["Authorization", "x-upstream-token"]) +async def test_subscription_fingerprint_preserves_context_across_managed_renewal(header_name): + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ClientCredentialsConfig + + async def refetch(token: str) -> str | None: + return None + + config = ClientCredentialsConfig(client_id="service", header_name=header_name) + + def client(token, *, scopes=("read",), url="https://upstream.example/mcp", tenant="alice"): + return MCPClient( + server_url=url, + extra_headers={"x-tenant": tenant}, + resolved_auth=ClientCredentialsBearerAuth(token, refetch, config, granted_scopes=scopes), + ) + + original = client("original") + fingerprint = await original.subscription_auth_fingerprint() + assert fingerprint == await client("renewed").subscription_auth_fingerprint() + assert fingerprint != await client("renewed", scopes=("write",)).subscription_auth_fingerprint() + assert fingerprint != await client("renewed", url="https://other.example/mcp").subscription_auth_fingerprint() + assert fingerprint != await client("renewed", tenant="bob").subscription_auth_fingerprint() + assert await original.discovery_auth_fingerprint() != await client("renewed").discovery_auth_fingerprint() + assert fingerprint != await client("renewed", scopes=None).subscription_auth_fingerprint() + assert await client("original", scopes=None).subscription_auth_fingerprint() != await client("renewed", scopes=None).subscription_auth_fingerprint() + assert (await original.prepare_request_auth()).headers[header_name] == "Bearer original" + + +@pytest.mark.asyncio +async def test_notification_callback_keeps_stream_errors_separate_from_protocol_errors() -> None: + from litellm.experimental_mcp_client.client import _receive_session_message + + stream_error = asyncio.get_running_loop().create_future() + receiver = AsyncMock() + await _receive_session_message( + MCPError(code=-32602, message="Unsupported request"), stream_error=stream_error, on_notification=receiver + ) + assert not stream_error.done() + receiver.assert_not_awaited() + first = OSError("Original transport failure") + for failure in (first, ValueError("Later parser failure")): + with pytest.raises(RuntimeError, match="MCP response stream failed"): + await _receive_session_message(failure, stream_error=stream_error, on_notification=receiver) + assert stream_error.result() is first + receiver.assert_not_awaited() diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index 1cad9a1fccb..77cbd8fa464 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -523,3 +523,44 @@ async def test_token_exchange_failure_diagnostics(mode, monkeypatch, caplog): else: assert "POST https://idp/ -> HTTP" in caplog.text assert {"denied":"denied", "invalid":"invalid response", "missing":"no access token"}[mode] in caplog.text + + +@pytest.mark.parametrize("changed", [ + {"client_id": "other-client"}, + {"client_secret": SecretStr("other-secret")}, + {"token_url": "https://other.example/token"}, + {"scopes": ("read",)}, + {"audience": "other-audience"}, + {"upstream_resource": "https://other.example/mcp"}, + {"token_endpoint_auth_method": "client_secret_basic"}, + {"value_prefix": "Token"}, +]) +def test_subscription_identity_tracks_grant_configuration(changed): + config = ClientCredentialsConfig( + client_id="client", client_secret=SecretStr("secret"), token_url="https://idp.example/token", + scopes=("read", "write"), audience="audience", upstream_resource="https://upstream.example/mcp", + token_endpoint_auth_method="client_secret_post", + ) + + async def refetch(token: str) -> str | None: + return None + + original = ClientCredentialsBearerAuth("first", refetch, config, granted_scopes=("read", "write")) + renewed = ClientCredentialsBearerAuth("second", refetch, config, granted_scopes=("write", "read")) + changed_auth = ClientCredentialsBearerAuth("second", refetch, config.model_copy(update=changed), granted_scopes=("read", "write")) + reduced = ClientCredentialsBearerAuth("second", refetch, config, granted_scopes=("read",)) + assert original.subscription_identity is not None + assert original.subscription_identity == renewed.subscription_identity + assert original.subscription_identity != changed_auth.subscription_identity + assert original.subscription_identity != reduced.subscription_identity + assert ClientCredentialsBearerAuth("first", refetch, config).subscription_identity is None + + +def test_subscription_identity_requires_managed_grant_metadata(): + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import HeaderCarrier + + async def refetch(token: str) -> str | None: + return None + + auth = ClientCredentialsBearerAuth("token", refetch, HeaderCarrier(), granted_scopes=("read",)) + assert auth.subscription_identity is None diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 0eee7334a26..283d0754940 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -1315,3 +1315,21 @@ async def test_id_jag_source_describes_final_token_after_both_exchanges() -> Non assert result.ok.source.value == "id-jag" assert _emitted(result.ok.auth)["Authorization"] == "Bearer resource-token" assert len(endpoint.calls) == 2 + + +@pytest.mark.asyncio +async def test_client_credentials_subscription_identity_preserves_resolved_grant_scopes(): + async def resolve(token): + source = _FakeM2MSource(Ok(token)) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials(_SUBJECT, _spec(_M2M)) + assert isinstance(result, Ok) + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + assert isinstance(result.ok, ClientCredentialsBearerAuth) + return result.ok.subscription_identity + + original = await resolve(OAuthToken(access_token="original", scopes=("read", "write"))) + renewed = await resolve(OAuthToken(access_token="renewed", scopes=("write", "read"))) + reduced = await resolve(OAuthToken(access_token="renewed", scopes=("read",))) + assert original is not None + assert original == renewed + assert original != reduced diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index b037d3b5d25..e96664aae84 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -276,7 +276,7 @@ async def test_legacy_adapter_cleans_context_after_cancelled_operation(): previous_session = server.active_mcp_session_var.get() previous_request = active_mcp_request_ctx_var.get() - request = SimpleNamespace(session=object(), protocol_version="2025-06-18") + request = SimpleNamespace(session=object(), protocol_version="2025-06-18", method="tools/call") auth = (None, None, None, None, None, None, None) async def cancelled_operation(): @@ -303,7 +303,7 @@ async def test_legacy_adapter_cleans_context_when_trace_setup_fails(): previous_session = server.active_mcp_session_var.get() previous_request = active_mcp_request_ctx_var.get() - request = SimpleNamespace(session=object()) + request = SimpleNamespace(session=object(), protocol_version="2025-06-18", method="tools/call") async def enter_operation(): async with server._legacy_operation_context(request, trace=True): diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py new file mode 100644 index 00000000000..e44aa88c56c --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -0,0 +1,357 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi import HTTPException +from mcp import MCPError +from mcp.server.context import ServerRequestContext +from mcp.server.session import ServerSession +from mcp.server.subscriptions import ListenHandler +from mcp.shared.subscriptions import PromptsListChanged, ToolsListChanged +from mcp.types import SubscriptionFilter, SubscriptionsListenRequestParams + +from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server import subscriptions +from litellm.proxy._experimental.mcp_server.contracts import OperationContext +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def context(session: ServerSession, revision: str = "2026-07-28") -> ServerRequestContext: + return ServerRequestContext( + session=session, + lifespan_context=None, + protocol_version=revision, + method="subscriptions/listen", + request_id="listener", + ) + + +@pytest.mark.asyncio +async def test_queued_notification_rechecks_authority_before_delivery() -> None: + bus: Final = subscriptions._ListenerBus() + sender: Final = Mock(spec=ServerSession) + first_delivery: Final = asyncio.Event() + release: Final = asyncio.Event() + revoked: Final = asyncio.Event() + observed: Final[list[str]] = [] + + async def send(notification, related_request_id=None) -> None: + assert related_request_id == "listener" + observed.append(notification.method) + if notification.method == "notifications/tools/list_changed": + first_delivery.set() + await release.wait() + + async def authorize() -> OperationContext: + if revoked.is_set(): + raise HTTPException(status_code=403, detail="Revoked caller") + return OperationContext(_caller=UserAPIKeyAuth(user_id="alice")) + + sender.send_notification = AsyncMock(side_effect=send) + handler: Final = ListenHandler(bus, max_buffered_events=2) + request: Final = SubscriptionsListenRequestParams( + notifications=SubscriptionFilter(tools_list_changed=True, prompts_list_changed=True) + ) + task: Final = asyncio.create_task(handler(context(subscriptions._AuthorizedSession(sender, authorize)), request)) + try: + await asyncio.wait_for(bus.ready.wait(), 1) + await bus.publish(ToolsListChanged()) + await asyncio.wait_for(first_delivery.wait(), 1) + await bus.publish(PromptsListChanged()) + revoked.set() + release.set() + with pytest.raises(HTTPException, match="Revoked caller"): + await asyncio.wait_for(task, 1) + assert observed == ["notifications/subscriptions/acknowledged", "notifications/tools/list_changed"] + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ["2024-11-05", "2025-11-25"]) +async def test_legacy_listener_rejected_before_upstream_work(revision: str) -> None: + authorize: Final = AsyncMock() + service: Final = subscriptions.GatewaySubscriptions() + with pytest.raises(MCPError, match="modern MCP protocol"): + await service.listen( + context(Mock(spec=ServerSession), revision), + SubscriptionsListenRequestParams(notifications=SubscriptionFilter(tools_list_changed=True)), + authorize, + "alice", + ) + authorize.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["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", + protocol_version="2025-11-25" if mode == "legacy" else "2026-07-28", + ) + target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") + targets: Final = () if mode == "denied" else (target, target) if mode == "multiple-resources" else (target,) + + async def resolve(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + return targets + + service: Final = subscriptions.GatewaySubscriptions(resolve_targets=resolve) + request: Final = SubscriptionsListenRequestParams( + notifications=SubscriptionFilter(resource_subscriptions=["test://value"]) + ) + error: Final = HTTPException if mode == "denied" else MCPError + expected: Final = { + "multiple-resources": "exactly one allowed server", + "denied": "No authorized MCP subscription targets", + }[mode] + with pytest.raises(error, match=expected): + await service.listen( + context(Mock(spec=ServerSession)), request, AsyncMock(return_value=OperationContext(_caller=None)), "alice" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("second_owner", ["alice", "bob"]) +async def test_listener_limits_release_capacity_after_cancellation(second_owner: str) -> None: + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def prepare(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + entered.set() + await release.wait() + return () + + prepare_mock: Final = AsyncMock(side_effect=prepare) + service: Final = subscriptions.GatewaySubscriptions(max_listeners=1, max_per_caller=1, resolve_targets=prepare_mock) + request: Final = SubscriptionsListenRequestParams(notifications=SubscriptionFilter(tools_list_changed=True)) + authorize: Final = AsyncMock(return_value=OperationContext(_caller=None)) + ctx: Final = context(Mock(spec=ServerSession)) + first: Final = asyncio.create_task(service.listen(ctx, request, authorize, "alice")) + try: + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(MCPError, match=r"[Ss]ubscription limit reached"): + await service.listen(ctx, request, authorize, second_owner) + assert prepare_mock.await_count == 1 + finally: + first.cancel() + await asyncio.gather(first, return_exceptions=True) + release.set() + with pytest.raises(HTTPException, match="No authorized MCP subscription targets"): + await service.listen(ctx, request, authorize, second_owner) + assert prepare_mock.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_mode", [False, True]) +async def test_subscription_requires_a_supported_http_gateway_route(proxy_mode: bool) -> None: + from litellm.proxy._experimental.mcp_server import server + + token = server._mcp_proxy_mode.set(proxy_mode) + try: + error = MCPError if proxy_mode else HTTPException + with pytest.raises(error) as rejected: + await server.listen_subscriptions( + context(Mock(spec=ServerSession)), + SubscriptionsListenRequestParams(notifications=SubscriptionFilter(tools_list_changed=True)), + ) + if proxy_mode: + assert rejected.value.code == -32601 + assert "/mcp/proxy" in str(rejected.value) + else: + assert rejected.value.status_code == 400 + assert "HTTP request" in rejected.value.detail + finally: + server._mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("changed", [None, "user_id", "team_id", "org_id", "end_user_id", "api_key"]) +async def test_subscription_authorization_preserves_original_caller(changed: str | None) -> None: + from dataclasses import replace + + from litellm.proxy._experimental.mcp_server.server import _owner_fingerprint_for + + initial = OperationContext(_caller=UserAPIKeyAuth(api_key="first", user_id="alice", team_id="team")) + current = replace( + initial, + _caller=initial.user_api_key_auth.model_copy(update={changed: "different"}) + if changed + else initial.user_api_key_auth, + ) + + async def refresh() -> OperationContext: + return current + + authorization = subscriptions.SubscriptionAuthorization(initial, refresh, _owner_fingerprint_for) + if changed: + with pytest.raises(HTTPException, match="Subscription caller changed"): + await authorization() + else: + assert await authorization() is current + assert authorization.owner == _owner_fingerprint_for(initial.user_api_key_auth) + + +@pytest.mark.asyncio +async def test_legacy_events_filter_resource_updates_and_unsubscribe_on_exit() -> None: + from mcp.shared.subscriptions import ResourceUpdated + from mcp.types import ( + ResourcesCapability, + ResourceUpdatedNotification, + ResourceUpdatedNotificationParams, + ServerCapabilities, + ToolListChangedNotification, + ToolsCapability, + ) + + session = Mock() + session.server_capabilities = ServerCapabilities( + tools=ToolsCapability(list_changed=True), resources=ResourcesCapability(subscribe=True) + ) + session.send_request = AsyncMock() + events = subscriptions.LegacyNotifications() + requested = SubscriptionFilter( + tools_list_changed=True, prompts_list_changed=True, resource_subscriptions=["test://value"] + ) + async with events.listen(session, requested) as listener: + assert listener.honored == SubscriptionFilter(tools_list_changed=True, resource_subscriptions=["test://value"]) + await events.receive(ResourceUpdatedNotification(params=ResourceUpdatedNotificationParams(uri="test://other"))) + await events.receive(ToolListChangedNotification()) + await events.receive(ResourceUpdatedNotification(params=ResourceUpdatedNotificationParams(uri="test://value"))) + assert await anext(listener) == ToolsListChanged() + assert await anext(listener) == ResourceUpdated(uri="test://value") + assert [(call.args[0].method, str(call.args[0].params.uri)) for call in session.send_request.await_args_list] == [ + ("resources/subscribe", "test://value"), + ("resources/unsubscribe", "test://value"), + ] + + +@pytest.mark.asyncio +async def test_legacy_notification_overflow_closes_instead_of_silently_losing_changes() -> None: + from mcp.types import ServerCapabilities, ToolListChangedNotification, ToolsCapability + + session = Mock() + session.server_capabilities = ServerCapabilities(tools=ToolsCapability(list_changed=True)) + events = subscriptions.LegacyNotifications(max_buffered_events=1) + async with events.listen(session, SubscriptionFilter(tools_list_changed=True)) as listener: + await events.receive(ToolListChangedNotification()) + await events.receive(ToolListChangedNotification()) + with pytest.raises(MCPError, match="overflow"): + await anext(listener) + + +@pytest.mark.asyncio +async def test_partial_legacy_resource_setup_only_unsubscribes_accepted_resources() -> None: + from mcp.types import ResourcesCapability, ServerCapabilities + + session = Mock() + session.server_capabilities = ServerCapabilities(resources=ResourcesCapability(subscribe=True)) + session.send_request = AsyncMock(side_effect=[None, MCPError(code=-32602, message="Denied resource"), None, None]) + events = subscriptions.LegacyNotifications() + with pytest.raises(MCPError, match="Denied resource"): + async with events.listen( + session, SubscriptionFilter(resource_subscriptions=["test://allowed", "test://denied"]) + ): + pytest.fail("Rejected setup must not acknowledge a listener") + assert [(call.args[0].method, str(call.args[0].params.uri)) for call in session.send_request.await_args_list] == [ + ("resources/subscribe", "test://allowed"), + ("resources/subscribe", "test://denied"), + ("resources/unsubscribe", "test://allowed"), + ] + + +@pytest.mark.asyncio +async def test_legacy_initialization_waiters_are_session_local_and_released_on_failure() -> None: + import anyio + + service = subscriptions.GatewaySubscriptions() + first = context(Mock(spec=ServerSession), "2025-11-25") + second = context(Mock(spec=ServerSession), "2025-11-25") + async with service.initializing(first) as started: + assert started + async with service.initializing(first) as duplicate: + assert not duplicate + waiter = asyncio.create_task(service.wait_initialized(first)) + await anyio.lowlevel.checkpoint() + assert not waiter.done() + await service.wait_initialized(second) + assert not waiter.done() + await waiter + async with service.initializing(first) as restarted: + assert restarted + + +@pytest.mark.asyncio +async def test_sessionless_legacy_resource_subscription_is_rejected() -> None: + session = Mock(spec=ServerSession) + session.client_params = None + service = subscriptions.GatewaySubscriptions() + with pytest.raises(MCPError, match="live legacy session"): + await service.legacy_resource(context(session, "2025-11-25"), "test://value", subscribe=True) + + +@pytest.mark.asyncio +async def test_legacy_upstream_without_capabilities_does_not_receive_subscribe() -> None: + session = Mock() + session.server_capabilities = None + session.send_request = AsyncMock() + events = subscriptions.LegacyNotifications() + async with events.listen( + session, SubscriptionFilter(tools_list_changed=True, resource_subscriptions=["test://value"]) + ) as listener: + assert listener.honored == SubscriptionFilter() + session.send_request.assert_not_awaited() + + +@pytest.mark.parametrize("proxy_mode", [False, True]) +def test_legacy_notification_capabilities_require_a_notification_route(proxy_mode: bool) -> None: + from litellm.proxy._experimental.mcp_server import server + + token = server._mcp_proxy_mode.set(proxy_mode) + try: + capabilities = server.server.create_initialization_options().capabilities + assert capabilities.tools.list_changed is (not proxy_mode) + if proxy_mode: + assert capabilities.resources is None + finally: + server._mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +async def test_idle_legacy_listener_detects_closed_upstream_without_restarting_work(monkeypatch) -> None: + import anyio + + session = Mock() + session.server_capabilities = None + session.send_ping = AsyncMock(side_effect=OSError("Upstream closed")) + deadline = anyio.move_on_after + monkeypatch.setattr(subscriptions.anyio, "move_on_after", lambda seconds, **kw: deadline(0.001, **kw)) + events = subscriptions.LegacyNotifications() + with pytest.raises(OSError, match="Upstream closed"): + async with events.listen(session, SubscriptionFilter()) as listener: + await anext(listener) + session.send_ping.assert_awaited_once() + session.send_request.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["subscribe_resource", "unsubscribe_resource"]) +async def test_proxy_route_rejects_legacy_resource_subscription_before_setup(method: str) -> None: + from mcp.types import SubscribeRequestParams, UnsubscribeRequestParams + + from litellm.proxy._experimental.mcp_server import server + + params = (SubscribeRequestParams if method == "subscribe_resource" else UnsubscribeRequestParams)(uri="test://value") + token = server._mcp_proxy_mode.set(True) + try: + with pytest.raises(MCPError, match="/mcp/proxy"): + await getattr(server, method)(context(Mock(spec=ServerSession)), params) + finally: + server._mcp_proxy_mode.reset(token)