scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138

feat(mcp): relay authorized modern upstream subscriptions (#46153)

Browse source 
* 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>
This commit is contained in:
joshua-berri 2026-10-15 16:16:49 -04:00 • committed by GitHub
parent fbd078bcbe
commit 4a4a0e8d3c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1904 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/experimental_mcp_client/client.py
View file

@ -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

10
litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py
View file

@ -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()

2
litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py
View file

@ -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)

115
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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.

542
litellm/proxy/_experimental/mcp_server/subscriptions.py Normal file
View file

@ -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")

743
tests/integration/mcp/test_interactions.py
View file

@ -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))

53
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

@ -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()

41
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py
View file

@ -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

18
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py
View file

@ -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

4
tests/unit/proxy/_experimental/mcp_server/test_operations.py
View file

@ -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):

357
tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py Normal file
View file

@ -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)