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): enable verified stateless protocol capabilities (#46264)

Browse source 
* feat(mcp): advertise verified stateless capabilities

* fix(mcp): honor rate limits during capability discovery

* test(mcp): cover protocol interactions and scoped enablement rollback

* fix(mcp): preserve legacy capabilities and bound discovery probes

* fix(mcp): keep discovery results single-assignment

* ci(mcp): distribute acceptance cases across existing workers

* ci: balance MCP integration cases across idle workers

* ci: isolate MCP interaction acceptance workload

* test(mcp): acknowledge fixture notifications after dispatch

* test(mcp): drain session dispatch before fixture deletion

* test(mcp): preserve deletion of streaming fixture requests

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-10-15 19:18:18 -04:00 • committed by GitHub
parent d489c04554
commit 440744f967
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1082 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/_experimental/mcp_server/capabilities.py
View file

@ -30,6 +30,7 @@ GATEWAY_OPERATIONS: Final = frozenset(
"resources/list",
"resources/read",
"resources/templates/list",
"subscriptions/listen",
}
)
_CAPABILITY_OPERATIONS: Final[Mapping[MCPDisableableCapability, frozenset[str]]] = MappingProxyType(
@ -124,13 +125,30 @@ def build_discovery(
if operations and revision_support is not None
else frozenset()
)
notifications: Final = revision in HANDSHAKE_PROTOCOL_VERSIONS or (
"subscriptions/listen" in operations and revision == "2026-07-28" and transport is MCPTransport.http
)
caller_capabilities: Final = capabilities.model_copy(deep=True)
tools: Final = caller_capabilities.tools
prompts: Final = caller_capabilities.prompts
resources: Final = caller_capabilities.resources
return DiscoverResult(
supported_versions=list(supported),
capabilities=ServerCapabilities(
tools=caller_capabilities.tools if {"tools/list", "tools/call"} <= operations else None,
prompts=caller_capabilities.prompts if {"prompts/list", "prompts/get"} <= operations else None,
resources=caller_capabilities.resources if {"resources/list", "resources/read"} <= operations else None,
tools=tools.model_copy(update={"list_changed": tools.list_changed if notifications else None})
if tools is not None and {"tools/list", "tools/call"} <= operations
else None,
prompts=prompts.model_copy(update={"list_changed": prompts.list_changed if notifications else None})
if prompts is not None and {"prompts/list", "prompts/get"} <= operations
else None,
resources=resources.model_copy(
update={
"list_changed": resources.list_changed if notifications else None,
"subscribe": resources.subscribe if notifications else None,
}
)
if resources is not None and {"resources/list", "resources/read"} <= operations
else None,
extensions={
key: value for key, value in (caller_capabilities.extensions or {}).items() if key in extensions
}

51
litellm/proxy/_experimental/mcp_server/operations.py
View file

@ -13,6 +13,7 @@ from functools import partial
from types import EllipsisType
from typing import Any, Final, NoReturn, TypeAlias, overload
import anyio
from fastapi import HTTPException
from mcp import MCPError, ReadResourceResult, Resource
from mcp.types import (
@ -52,6 +53,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_logger
from litellm.constants import (
MAXIMUM_TRACEBACK_LINES_TO_LOG,
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
@ -3501,6 +3503,7 @@ class GatewayOperations:
async def _execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult:
match operation:
case DiscoverRequest():
deadline: Final = time.monotonic() + MCP_TOOL_LISTING_TIMEOUT
disabled: Final = disabled_capabilities()
listings: Final = (
()
@ -3524,6 +3527,17 @@ class GatewayOperations:
)
try:
results: Final = await asyncio.gather(*tasks)
from litellm.proxy._experimental.mcp_server.subscriptions import subscription_capabilities
async def discover_notifications() -> tuple[frozenset[str], ServerCapabilities]:
remaining: Final = deadline - time.monotonic()
if context.protocol_version == "2026-07-28" and not context.mcp_proxy_mode and remaining > 0:
# Optional probes share the catalog budget, including target preparation.
with anyio.move_on_after(remaining):
return await subscription_capabilities(context)
return frozenset(MCP_LEGACY_VERSIONS), ServerCapabilities()
upstream_versions, notifications = await discover_notifications()
finally:
try:
for task in tasks:
@ -3536,21 +3550,30 @@ class GatewayOperations:
revision=context.protocol_version or "2025-11-25",
transport=MCPTransport.http,
authorized_operations=enabled_gateway_operations(disabled),
upstream_versions=frozenset(MCP_LEGACY_VERSIONS),
upstream_versions=upstream_versions,
capabilities=ServerCapabilities(
tools=ToolsCapability()
if any(isinstance(result, ListToolsResult) and result.tools for result in results)
else None,
prompts=PromptsCapability()
if any(isinstance(result, ListPromptsResult) and result.prompts for result in results)
else None,
resources=ResourcesCapability()
if any(
(isinstance(result, ListResourcesResult) and result.resources)
or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates)
for result in results
)
else None,
tools=notifications.tools
or (
ToolsCapability()
if any(isinstance(result, ListToolsResult) and result.tools for result in results)
else None
),
prompts=notifications.prompts
or (
PromptsCapability()
if any(isinstance(result, ListPromptsResult) and result.prompts for result in results)
else None
),
resources=notifications.resources
or (
ResourcesCapability()
if any(
(isinstance(result, ListResourcesResult) and result.resources)
or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates)
for result in results
)
else None
),
),
)
case AuthorizedToolCall():

52
litellm/proxy/_experimental/mcp_server/subscriptions.py
View file

@ -32,7 +32,9 @@ from mcp.types import (
ListToolsRequest,
ListToolsResult,
PaginatedRequestParams,
PromptsCapability,
RequestId,
ResourcesCapability,
ServerCapabilities,
ServerNotification,
SubscribeRequest,
@ -40,6 +42,7 @@ from mcp.types import (
SubscriptionFilter,
SubscriptionsListenRequestParams,
SubscriptionsListenResult,
ToolsCapability,
UnsubscribeRequest,
UnsubscribeRequestParams,
)
@ -52,6 +55,7 @@ from litellm.proxy._experimental.mcp_server.catalog import (
)
from litellm.proxy._experimental.mcp_server.contracts import OperationContext
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCP_LEGACY_VERSIONS
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -219,6 +223,50 @@ async def _targets(context: OperationContext) -> tuple[SubscriptionTarget, ...]:
return tuple([await prepare(server) for server in allowed])
async def subscription_capabilities(context: OperationContext) -> tuple[frozenset[str], ServerCapabilities]:
try:
targets: Final = await _targets(context)
except Exception: # noqa: BLE001 # missing optional upstream credentials must not erase usable catalog results
return frozenset(MCP_LEGACY_VERSIONS), ServerCapabilities()
configured: Final = frozenset(target.server.protocol_version for target in targets)
versions: Final = (
configured | frozenset(MCP_LEGACY_VERSIONS) if not configured or "auto" in configured else configured
)
if not targets or any(
target.server.protocol_version == "2026-07-28" and target.server.transport != "http" for target in targets
):
return versions, ServerCapabilities()
async def probe(target: SubscriptionTarget) -> ServerCapabilities:
async def read(session: ClientSession) -> ServerCapabilities:
return session.server_capabilities or ServerCapabilities()
try:
if await operations._mcp_server_rate_limit_rejection(target.server, context.user_api_key_auth) is not None: # pyright: ignore[reportPrivateUsage] # reuse request-scoped catalog admission
return ServerCapabilities()
with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT):
return await target.client.run_with_session(read, quiet_on_error=True)
except Exception: # noqa: BLE001 # an optional notification probe must not disable an otherwise usable catalog
return ServerCapabilities()
declared: Final = tuple(await asyncio.gather(*(probe(target) for target in targets)))
return versions, ServerCapabilities(
tools=ToolsCapability(list_changed=True)
if all(item.tools and item.tools.list_changed for item in declared)
else None,
prompts=PromptsCapability(list_changed=True)
if all(item.prompts and item.prompts.list_changed for item in declared)
else None,
resources=ResourcesCapability(
list_changed=all(item.resources and item.resources.list_changed for item in declared) or None,
subscribe=(len(targets) == 1 and any(item.resources and item.resources.subscribe for item in declared))
or None,
)
if any(item.resources for item in declared)
else None,
)
async def _listing(event: ServerEvent, context: OperationContext) -> str | None:
if isinstance(event, ResourceUpdated):
return None
@ -432,6 +480,10 @@ class GatewaySubscriptions:
) -> SubscriptionsListenResult:
context: Final = await authorize()
targets: Final = await self._resolve_targets(context)
if any(
target.server.protocol_version == "2026-07-28" and target.server.transport != "http" for target in targets
):
raise MCPError(code=-32602, message="Modern upstream subscriptions require HTTP")
if not targets:
raise HTTPException(status_code=403, detail="No authorized MCP subscription targets")
if params.notifications.resource_subscriptions and len(targets) != 1:

84
tests/integration/_support/mcp.py
View file

@ -26,7 +26,7 @@ from pydantic import BaseModel, JsonValue
from sse_starlette.sse import AppStatus
from starlette.requests import Request as StarletteRequest
from starlette.responses import Response
from starlette.types import Message, Receive, Scope, Send
from starlette.types import ASGIApp, Message, Receive, Scope, Send
Transport = Literal["http", "sse", "stdio"]
STDIO_PEER: Final = Path(__file__).with_name("mcp_stdio_peer.py")
@ -140,6 +140,70 @@ def _capturing(app: Callable[[Scope, Receive, Send], object], observed: queue.Qu
return capture
def _acknowledge_notifications_after_dispatch(app: ASGIApp) -> ASGIApp:
in_flight: Final[dict[bytes, set[asyncio.Event]]] = {}
async def dispatch(scope: Scope, receive: Receive, send: Send, dispatched: asyncio.Event | None = None) -> None:
if scope["type"] != "http" or scope["method"] != "POST":
await app(scope, receive, send)
return
pending: Final[list[Message]] = []
async def capture(message: Message) -> None:
if not pending and message["type"] == "http.response.start" and message["status"] == 202:
pending.append(message)
elif pending:
if (
len(pending) != 1
or message["type"] != "http.response.body"
or message.get("body", b"")
or message.get("more_body", False)
):
raise RuntimeError("MCP notification acknowledgement must be a single empty response")
pending.append(message)
else:
if dispatched is not None and message["type"] == "http.response.start":
dispatched.set()
await send(message)
try:
await app(scope, receive, capture)
finally:
if dispatched is not None:
dispatched.set()
if pending:
if len(pending) != 2:
raise RuntimeError("MCP notification acknowledgement did not complete")
for message in pending:
await send(message)
async def ordered(scope: Scope, receive: Receive, send: Send) -> None:
session: Final = next(
(value for name, value in scope.get("headers", ()) if name.lower() == b"mcp-session-id"), None
)
if scope["type"] != "http" or session is None:
await dispatch(scope, receive, send)
return
if scope["method"] == "DELETE":
while pending := tuple(done for done in in_flight.get(session, ()) if not done.is_set()):
await asyncio.gather(*(done.wait() for done in pending))
if scope["method"] != "POST":
await dispatch(scope, receive, send)
return
done: Final = asyncio.Event()
active: Final = in_flight.setdefault(session, set())
active.add(done)
try:
await dispatch(scope, receive, send, done)
finally:
active.remove(done)
if not active:
del in_flight[session]
done.set()
return ordered
def _drain_sse_streams() -> None:
AppStatus.should_exit = True
@ -565,6 +629,7 @@ class McpCaller:
entry: EntryPoint
alias: str | None = None
headers: Mapping[str, str] = field(default_factory=dict)
protocol_version: str | None = None
def _path(self) -> str:
if self.entry == "server_mcp":
@ -582,10 +647,23 @@ class McpCaller:
def rpc(self, method: str, params: JsonRpc | None = None) -> httpx.Response:
if self.entry == "sse":
return _legacy_sse_rpc(self.gateway, self._headers(), method, params)
modern: Final = self.protocol_version == "2026-07-28"
arguments: Final = dict(params or {})
return self.gateway.client.post(
self._path(),
json={"jsonrpc": "2.0", "id": 1, "method": method, "params": dict(params or {})},
headers=self._headers(),
json={"jsonrpc": "2.0", "id": 1, "method": method, "params": {
**arguments,
**({"_meta": {
"io.modelcontextprotocol/protocolVersion": self.protocol_version,
"io.modelcontextprotocol/clientInfo": {"name": "integration-caller", "version": "1"},
"io.modelcontextprotocol/clientCapabilities": {},
}} if modern else {}),
}},
headers={**self._headers(), **({
"MCP-Protocol-Version": self.protocol_version,
"Mcp-Method": method,
"Mcp-Name": str(arguments.get("name", arguments.get("uri", ""))),
} if modern else {})},
)
def initialize(self) -> Outcome:

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

@ -1,5 +1,6 @@
import json
import os
from contextlib import contextmanager, nullcontext
from pathlib import Path
from typing import Final
@ -7,6 +8,7 @@ import httpx
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.mcp import McpCaller
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
@ -85,6 +87,122 @@ def interaction_peer(request: Request) -> Reply:
)
@contextmanager
def sdk_interaction_peer(*, sampling: bool = False):
import asyncio
import queue
from contextlib import asynccontextmanager
from mcp.server import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.server.transport_security import TransportSecuritySettings
from mcp.types import (
CallToolResult,
CreateMessageRequest,
CreateMessageRequestParams,
SamplingMessage,
ElicitRequest,
ElicitRequestFormParams,
GetPromptResult,
InputRequiredResult,
ListToolsResult,
ListPromptsResult,
ListResourcesResult,
Prompt,
PromptMessage,
ReadResourceResult,
Resource,
TextContent,
TextResourceContents,
Tool,
)
from starlette.applications import Starlette
from starlette.routing import Mount
from integration._support.asgi import asgi_server
from integration._support.mcp import McpPeer, _capturing, _drain_sse_streams
from sse_starlette.sse import AppStatus
async def tools(ctx, params):
return ListToolsResult(tools=[Tool(name="confirm", input_schema={"type": "object"})])
async def prompts(ctx, params):
return ListPromptsResult(prompts=[Prompt(name="confirm")])
async def resources(ctx, params):
return ListResourcesResult(resources=[Resource(name="confirm", uri="test://confirm")])
async def interact(ctx, params):
if sampling:
if not params.request_state:
return InputRequiredResult(
request_state="sample-state",
input_requests={
"sample": CreateMessageRequest(
params=CreateMessageRequestParams(
messages=[
SamplingMessage(role="user", content=TextContent(type="text", text="sample prompt"))
],
max_tokens=8,
)
)
},
)
assert params.request_state == "sample-state"
sampled = params.input_responses["sample"]
assert "sampled response" in sampled.model_dump_json()
return CallToolResult(content=[TextContent(type="text", text="sample confirmed")])
if not params.request_state:
return InputRequiredResult(
request_state="opaque:" + ctx.method,
input_requests={
"consent": ElicitRequest(
params=ElicitRequestFormParams(
message="Confirm", requested_schema={"type": "object", "properties": {}}
)
)
},
)
assert params.request_state == "opaque:" + ctx.method
assert params.input_responses["consent"].action == "accept"
if ctx.method == "tools/call":
return CallToolResult(content=[TextContent(type="text", text="confirmed")])
if ctx.method == "prompts/get":
return GetPromptResult(
messages=[PromptMessage(role="user", content=TextContent(type="text", text="confirmed"))]
)
return ReadResourceResult(contents=[TextResourceContents(uri="test://confirm", text="confirmed")])
service = Server(
"modern-interactions",
on_list_tools=tools,
on_list_prompts=prompts,
on_list_resources=resources,
on_call_tool=interact,
on_get_prompt=interact,
on_read_resource=interact,
)
manager = StreamableHTTPSessionManager(
service,
stateless=True,
json_response=True,
security_settings=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
@asynccontextmanager
async def lifespan(app):
AppStatus.should_exit = False
async with manager.run():
yield
_drain_sse_streams()
watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro()))
await asyncio.wait_for(asyncio.gather(*watchers), 2)
app = Starlette(routes=[Mount("/mcp", app=manager.handle_request)], lifespan=lifespan)
observed = queue.Queue()
with asgi_server(_capturing(app, observed)) as url:
yield McpPeer(url + "/mcp/", observed)
def rpc(gateway: Gateway, method: str, params: dict[str, JsonValue], *, status: int = 200) -> dict:
response: Final = gateway.client.post(
"/mcp/",
@ -101,11 +219,12 @@ def rpc(gateway: Gateway, method: str, params: dict[str, JsonValue], *, status:
return response.json()
@pytest.mark.parametrize("sdk_upstream", [False, True])
@pytest.mark.parametrize("changed_target", [False, True])
def test_continuations_resume_on_another_replica_and_reject_changed_operations(
tmp_path: Path, changed_target: bool
tmp_path: Path, changed_target: bool, sdk_upstream: bool
) -> None:
with wire_server(interaction_peer) as peer, httpx.Client() as client:
with (sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as peer, httpx.Client() as client:
config: Final = tmp_path / "proxy.yaml"
config.write_text(
yaml.safe_dump(
@ -182,7 +301,8 @@ def test_continuations_resume_on_another_replica_and_reject_changed_operations(
assert peer.drain() == (), "Rejected continuation must not contact the upstream"
def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tmp_path: Path) -> None:
@pytest.mark.parametrize("sdk_upstream", [False, True])
def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tmp_path: Path, sdk_upstream: bool) -> None:
auth_module: Final = tmp_path / "interaction_auth.py"
auth_module.write_text(
"from pathlib import Path\n"
@ -201,7 +321,11 @@ def test_continuation_reauthenticates_caller_and_rechecks_revoked_permissions(tm
}
auth_state: Final = auth_module.with_suffix(".json")
auth_state.write_text(json.dumps(identity))
with wire_server(interaction_peer) as peer, wire_server(interaction_peer) as other_peer, httpx.Client() as client:
with (
(sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as peer,
(sdk_interaction_peer() if sdk_upstream else wire_server(interaction_peer)) as other_peer,
httpx.Client() as client,
):
config: Final = tmp_path / "proxy.yaml"
config.write_text(
yaml.safe_dump(
@ -397,6 +521,9 @@ def test_missing_continuation_key_reports_configuration_for_each_carrier(tmp_pat
"event_kind,access",
[
("tools", "normal"),
("tools", "rollback"),
("tools", "legacy-upstream-sse"),
("updated", "legacy-upstream-sse"),
("tools", "modern-boundary"),
("prompts", "normal"),
("resources", "normal"),
@ -435,6 +562,12 @@ def test_missing_continuation_key_reports_configuration_for_each_carrier(tmp_pat
("tools", "renewal"),
("tools", "renewal-scopes"),
("tools", "renewal-failed"),
("tools", "mixed-capabilities"),
("tools", "probe-failure"),
("tools", "missing-credentials"),
("tools", "concurrent-probes"),
("tools", "probe-timeout"),
("tools", "rate-limited"),
],
)
def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
@ -447,7 +580,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
import httpx2
from integration._support.asgi import asgi_server
from integration._support.mcp import _drain_sse_streams
from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams
from mcp import ClientSession, MCPError
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
@ -494,9 +627,16 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
}
)
probe_fails: Final = Event()
probe_failed: Final = Event()
listener_credentials = []
catalog_credentials = []
setup_changed = Event()
probe_arrivals = []
catalog_reads = []
limited_requests = []
probe_catalog_goal = []
probe_barrier = asyncio.Event()
active_listeners = set()
@ -540,7 +680,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
"user_role": "internal_user",
"object_permission": {
"object_permission_id": "subscription-permission",
"mcp_servers": ["subscription-server", "second-server"] if access == "fanout" else ["subscription-server"],
"mcp_servers": ["subscription-server", "second-server"] if access in ("fanout", "mixed-capabilities", "missing-credentials", "concurrent-probes", "rate-limited") else ["subscription-server"],
},
}
auth_state: Final = auth_module.with_suffix(".json")
@ -565,13 +705,26 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
return EmptyResult()
if getattr(ctx.request, "headers", {}).get("x-probe-target") == "limited":
limited_requests.append(ctx.method)
if ctx.method == "tools/list":
catalog_credentials.append(ctx.request.headers.get("authorization"))
catalog_credentials.append(getattr(ctx.request, "headers", {}).get("authorization"))
if access == "setup-race" and connected.is_set() and not setup_changed.is_set():
setup_changed.set()
await change_tools()
if ctx.method == "server/discover" and probe_fails.is_set() and len(catalog_reads) >= probe_catalog_goal[0]:
probe_failed.set()
if access == "concurrent-probes":
probe_arrivals.append(ctx.method)
if len(probe_arrivals) == 2:
probe_barrier.set()
await probe_barrier.wait()
elif access == "probe-timeout":
await probe_barrier.wait()
else:
raise MCPError(code=-32603, message="Discovery unavailable")
if ctx.method == "subscriptions/listen":
listener_credentials.append(ctx.request.headers.get("authorization"))
listener_credentials.append(getattr(ctx.request, "headers", {}).get("authorization"))
if access == "legacy-caller-no-notifications":
raise MCPError(code=-32601, message="Upstream has no subscription support")
return await upstream_listener(ctx, SubscriptionsListenRequestParams.model_validate(ctx.params))
@ -600,6 +753,8 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
)
}
)
if ctx.method in ("tools/list", "prompts/list", "resources/list", "resources/templates/list"):
catalog_reads.append(ctx.method)
if access == "pages" and ctx.method == "tools/list":
result = ListToolsResult.model_validate(result)
offset = int((ctx.params or {}).get("cursor") or "0")
@ -693,7 +848,11 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
assert {tool.name for tool in (await other.list_tools()).tools} == {"peer-change_tools"}
if access.startswith("legacy-caller") or access in ("legacy-both", "legacy-sse"):
await session.initialize()
initialized = await session.initialize()
assert initialized.capabilities.tools.list_changed
assert initialized.capabilities.prompts.list_changed
assert initialized.capabilities.resources.list_changed
assert initialized.capabilities.resources.subscribe
before = await session.list_tools()
assert "peer-added" not in {tool.name for tool in before.tools}
if event_kind == "updated":
@ -772,10 +931,52 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
with pytest.raises(TimeoutError):
await asyncio.wait_for(received.get(), 0.5)
return
session.adopt(DiscoverResult.model_validate(await session.send_discover("2026-07-28")))
discovery = DiscoverResult.model_validate(await session.send_discover("2026-07-28"))
session.adopt(discovery)
if access in ("mixed-capabilities", "missing-credentials"):
assert not discovery.capabilities.tools.list_changed
else:
assert discovery.capabilities.tools.list_changed
assert discovery.capabilities.resources.list_changed
assert bool(discovery.capabilities.resources.subscribe) is (access not in ("fanout", "concurrent-probes", "rate-limited"))
if access == "rate-limited":
assert limited_requests, "Admitted discovery never reached the limited upstream"
limited_requests.clear()
limited = DiscoverResult.model_validate(await session.send_discover("2026-07-28"))
assert limited_requests == [], "Rejected discovery reached the rate-limited upstream"
assert limited.capabilities.tools is not None
assert not limited.capabilities.tools.list_changed
assert not (limited.capabilities.resources and limited.capabilities.resources.subscribe)
return
if access in ("probe-failure", "concurrent-probes", "probe-timeout"):
probe_catalog_goal.append(len(catalog_reads) + 4 * len(identity["object_permission"]["mcp_servers"]))
probe_fails.set()
try:
degraded = DiscoverResult.model_validate(
await asyncio.wait_for(session.send_discover("2026-07-28"), 5)
)
finally:
probe_barrier.set()
assert probe_failed.is_set(), ("Controlled discovery probe was not exercised", catalog_reads, probe_catalog_goal)
assert degraded.capabilities.tools is not None
if access == "concurrent-probes":
assert len(probe_arrivals) == 2
assert degraded.capabilities.tools.list_changed
else:
assert not (degraded.capabilities.tools and degraded.capabilities.tools.list_changed)
assert not (degraded.capabilities.resources and degraded.capabilities.resources.subscribe)
probe_fails.clear()
visible = await session.list_tools()
assert "peer-change_tools" in {tool.name for tool in visible.tools}
return
before = await session.list_tools()
assert "peer-change_tools" in {tool.name for tool in before.tools}
assert "peer-added" not in {tool.name for tool in before.tools}
if access in ("mixed-capabilities", "missing-credentials"):
assert not discovery.capabilities.resources.list_changed
assert not discovery.capabilities.resources.subscribe
assert not connected.is_set()
return
if access == "modern-boundary":
for method, params in (
("notifications/initialized", {}),
@ -862,11 +1063,19 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
await asyncio.sleep(0.2)
assert token_grants[0] == token_grants[-1]
if access != "setup-race":
async with streamable_http_client(upstream_url + "/mcp") as direct_streams:
direct_connection = (
sse_client(upstream_url + "/sse")
if access == "legacy-upstream-sse"
else streamable_http_client(upstream_url + "/mcp")
)
async with direct_connection as direct_streams:
async with ClientSession(direct_streams[0], direct_streams[1]) as direct:
direct.adopt(
DiscoverResult.model_validate(await direct.send_discover("2026-07-28"))
)
if access == "legacy-upstream-sse":
await direct.initialize()
else:
direct.adopt(
DiscoverResult.model_validate(await direct.send_discover("2026-07-28"))
)
changed = await direct.call_tool("change_tools", {})
assert not changed.is_error
if access in (
@ -929,7 +1138,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
await asyncio.wait_for(anext(subscription), 0.5)
assert await asyncio.to_thread(disconnected.wait, 3), "Upstream listener survived disconnect"
app: Final = upstream.streamable_http_app(
app: Final = (upstream.sse_app if access == "legacy-upstream-sse" else upstream.streamable_http_app)(
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False)
)
@ -942,10 +1151,13 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro()))
await asyncio.wait_for(asyncio.gather(*watchers), 2)
hosted: Final = Starlette(
routes=[Route("/token", token_endpoint, methods=["POST"]), Mount("/", app=app)], lifespan=lifespan
)
with asgi_server(hosted) as upstream_url, httpx.Client() as client:
transport_app: Final = app if access == "legacy-upstream-sse" else _acknowledge_notifications_after_dispatch(app)
hosted: Final = Starlette(routes=[Route("/token", token_endpoint, methods=["POST"]), Mount("/", app=transport_app)], lifespan=lifespan)
with (
asgi_server(hosted) as upstream_url,
(sdk_interaction_peer() if access == "mixed-capabilities" else nullcontext(None)) as plain,
httpx.Client() as client,
):
config: Final = tmp_path / "proxy.yaml"
config.write_text(
yaml.safe_dump(
@ -966,10 +1178,12 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
if access.startswith("renewal")
else {}
),
"url": upstream_url + "/mcp",
"transport": "http",
"protocol_version": "2025-11-25" if access in ("legacy", "legacy-both") else "2026-07-28",
"url": upstream_url + ("/sse" if access == "legacy-upstream-sse" else "/mcp"),
"transport": "sse" if access == "legacy-upstream-sse" else "http",
"protocol_version": "2025-11-25" if access in ("legacy", "legacy-both", "legacy-upstream-sse") else "2026-07-28",
**({"allowed_tools": ["change_tools"]} if access == "hidden" else {}),
**({"rpm": 1, "static_headers": {"X-Probe-Target": "limited"}} if access == "rate-limited" else {}),
**({"timeout": 2} if access == "concurrent-probes" else {}),
**(
{"auth_type": "oauth_delegate", "extra_headers": ["Authorization"]}
if access == "fanout"
@ -980,14 +1194,22 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
{
"second": {
"server_id": "second-server",
"url": upstream_url + "/mcp",
"url": plain.url if plain is not None else upstream_url + "/mcp",
"transport": "http",
**({"timeout": 2} if access == "concurrent-probes" else {}),
"protocol_version": "2026-07-28",
"auth_type": "oauth_delegate",
"extra_headers": ["Authorization"],
**({
"auth_type": "oauth2",
"oauth2_flow": "authorization_code",
"client_id": "subscription-client",
"authorization_url": upstream_url + "/authorize",
"token_url": upstream_url + "/token",
} if access == "missing-credentials" else {}),
}
}
if access == "fanout"
if access in ("fanout", "mixed-capabilities", "missing-credentials", "concurrent-probes", "rate-limited")
else {}
),
},
@ -1004,6 +1226,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
environment: Final = {
"PYTHONPATH": os.pathsep.join((str(source), str(source / "tests"))),
"LITELLM_SALT_KEY": "subscription-pagination-key",
**({"LITELLM_MCP_TOOL_LISTING_TIMEOUT": "2"} if access == "probe-timeout" else {}),
"STORE_MODEL_IN_DB": "False",
"DISABLE_SCHEMA_UPDATE": "true",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
@ -1030,6 +1253,28 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools(
eventually(lambda: len(active_listeners), lambda count: count == 0, seconds=10)
assert not active_listeners, "Upstream listener survived client disconnect"
if access == "rollback":
rolled_back = yaml.safe_load(config.read_text())
rolled_back["general_settings"]["mcp_advertised_versions"] = ["2025-11-25"]
config.write_text(yaml.safe_dump(rolled_back))
async def legacy_after_rollback(gateway: Gateway) -> None:
async with httpx2.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client:
async with streamable_http_client(str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client) as streams:
async with ClientSession(streams[0], streams[1]) as session:
await session.initialize()
assert "peer-added" in {tool.name for tool in (await session.list_tools()).tools}
result = await session.call_tool("peer-added", {})
assert not result.is_error
assert result.content[0].text == "new tool"
with owned_proxy(seed, tmp_path / "rollback", environment, config=config, database_setup=(), remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "INTEGRATION_PROXY_DATABASE_URL", "INTEGRATION_PROXY_READ_REPLICA_URL")) as rollback:
rejected = rpc(rollback, "tools/list", {}, status=400)
assert "error" in rejected
asyncio.run(legacy_after_rollback(rollback))
eventually(lambda: len(active_listeners), lambda count: count == 0, seconds=10)
assert tool_executions == ["change_tools", "added"], "Rollback duplicated or restarted upstream tool execution"
def test_subscription_partial_setup_releases_upstream_capacity() -> None:
import asyncio
@ -1130,3 +1375,232 @@ def test_subscription_partial_setup_releases_upstream_capacity() -> None:
with asgi_server(hosted) as url:
asyncio.run(exercise(url))
@pytest.mark.parametrize("enabled", [False, True])
def test_modern_sampling_obeys_opt_in_and_completes_upstream_interaction(tmp_path: Path, enabled: bool) -> None:
def provider(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return Reply(body=b'{"object":"list","data":[]}')
assert request.method == "POST" and request.target == "/v1/chat/completions"
body = json.loads(request.body)
assert "sample prompt" in json.dumps(body["messages"])
return Reply(
body=json.dumps(
{
"id": "sample-response",
"object": "chat.completion",
"created": 1,
"model": "sampling-model",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "sampled response"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4},
}
).encode()
)
with sdk_interaction_peer(sampling=True) as peer, wire_server(provider) as model, httpx.Client() as client:
config = tmp_path / "sampling.yaml"
config.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": "sampling-model",
"litellm_params": {
"model": "openai/sampling-model",
"api_base": model.url + "/v1",
"api_key": "synthetic-provider-key",
},
}
],
"mcp_servers": {
"peer": {
"url": peer.url,
"transport": "http",
"protocol_version": "2026-07-28",
"allow_sampling": enabled,
}
},
"general_settings": {
"master_key": _KEY,
"store_model_in_db": False,
"mcp_advertised_versions": ["2026-07-28"],
},
"litellm_settings": {"default_mcp_sampling_model": "sampling-model"},
}
)
)
environment = {
"PYTHONPATH": _PROXY_PYTHONPATH,
"STORE_MODEL_IN_DB": "False",
"DISABLE_SCHEMA_UPDATE": "true",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"LITELLM_SALT_KEY": "sampling-test-salt",
}
with owned_proxy(
Gateway(client, _KEY, peer.url),
tmp_path / "gateway",
environment,
config=config,
database_setup=(),
remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "REDIS_URL", "REDIS_HOST"),
) as gateway:
reply = McpCaller(gateway, gateway.key, "root", protocol_version="2026-07-28").rpc(
"tools/call", {"name": "peer-confirm", "arguments": {}}
)
assert reply.status_code == 200, reply.text
response = reply.json()
calls = tuple(request for request in model.drain() if request.method == "POST")
if enabled:
assert response["result"]["content"][0]["text"] == "sample confirmed", response
assert len(calls) == 1
else:
assert response["result"]["isError"], response
assert "sampling" in json.dumps(response).lower(), response
assert calls == (), "Disabled sampling reached the model provider"
@pytest.mark.parametrize("downstream", ["2025-11-25", "2026-07-28"])
@pytest.mark.parametrize("upstream", ["2025-11-25", "2026-07-28"])
def test_input_interaction_pairs_preserve_consent_and_never_restart_legacy_work(
tmp_path: Path, downstream: str, upstream: str
) -> None:
import asyncio
import httpx2
from integration._support.mcp import mcp_peer, tool_calls
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, DiscoverResult, ElicitResult, InputRequiredResult
modern_upstream = upstream == "2026-07-28"
peer_context = wire_server(interaction_peer) if modern_upstream else mcp_peer("sse", rich=True)
with peer_context as peer, httpx.Client() as http_client:
config = tmp_path / "proxy.yaml"
config.write_text(yaml.safe_dump({
"model_list": [],
"mcp_servers": {"peer": {"url": peer.url, "transport": "http" if modern_upstream else "sse", "protocol_version": upstream, "allow_elicitation": True}},
"general_settings": {"master_key": _KEY, "store_model_in_db": False, "mcp_advertised_versions": ["2025-11-25", "2026-07-28"]},
}))
consent_requests = []
async def consent(context, params):
consent_requests.append(params.message)
return ElicitResult(action="accept", content=None if modern_upstream else {"confirmed": True})
async def exercise(gateway):
async with httpx2.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client:
async with streamable_http_client(str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client) as streams:
async with ClientSession(streams[0], streams[1], elicitation_callback=consent) as session:
if downstream == "2026-07-28":
session.adopt(DiscoverResult.model_validate(await session.send_discover(downstream)))
else:
await session.initialize()
name = "peer-confirm" if modern_upstream else "peer-elicit"
args = {} if modern_upstream else {"question": "Confirm"}
await session.list_tools()
result = await session.call_tool(name, args, allow_input_required=True)
if downstream == "2026-07-28" and modern_upstream:
assert isinstance(result, InputRequiredResult)
assert set(result.input_requests) == {"consent"}
result = await session.call_tool(name, args, request_state=result.request_state, input_responses={"consent": ElicitResult(action="accept")})
assert isinstance(result, CallToolResult)
if downstream == "2026-07-28" and not modern_upstream:
assert result.is_error
assert "legacy upstream cannot resume input" in result.content[0].text.lower()
assert "partially completed" in result.content[0].text.lower()
assert consent_requests == []
else:
assert not result.is_error, result
assert result.content[0].text == ("confirmed" if modern_upstream else "elicited:accept")
assert len(consent_requests) == (0 if downstream == "2026-07-28" else 1)
seed = Gateway(http_client, _KEY, peer.url)
with owned_proxy(seed, tmp_path / "gateway", {"PYTHONPATH": _PROXY_PYTHONPATH, "STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true", "LITELLM_SALT_KEY": "interaction-matrix-test"}, config=config, database_setup=(), remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "REDIS_URL", "REDIS_HOST")) as gateway:
asyncio.run(exercise(gateway))
observed = peer.drain()
calls = [json.loads(item.body) for item in observed if item.method == "POST" and json.loads(item.body).get("method") == "tools/call"] if modern_upstream else tool_calls(observed)
assert len(calls) == (2 if modern_upstream else 1), "Gateway restarted work instead of returning or resuming the original request"
def test_legacy_notification_dispatch_finishes_before_session_deletion() -> None:
import asyncio
from contextlib import asynccontextmanager
from threading import Event
from integration._support.asgi import asgi_server
from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams
from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from sse_starlette.sse import AppStatus
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Message, Receive, Scope, Send
peer = MCPServer("notification-order")
sdk_app = peer.streamable_http_app(
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False)
)
@asynccontextmanager
async def lifespan(_: Starlette):
AppStatus.should_exit = False
async with sdk_app.router.lifespan_context(sdk_app):
yield
_drain_sse_streams()
watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro()))
await asyncio.wait_for(asyncio.gather(*watchers), 2)
hosted = Starlette(routes=[Mount("/", app=sdk_app)], lifespan=lifespan)
ordered = _acknowledge_notifications_after_dispatch(hosted)
deleted = Event()
async def pause_after_acceptance(scope: Scope, receive: Receive, send: Send) -> None:
accepted = False
async def capture(message: Message) -> None:
nonlocal accepted
if message["type"] == "http.response.start":
accepted = message["status"] == 202
await send(message)
if accepted and message["type"] == "http.response.body" and not message.get("more_body", False):
assert await asyncio.to_thread(deleted.wait, 5), "Client did not complete session deletion"
await ordered(scope, receive, capture)
with asgi_server(pause_after_acceptance) as url, httpx.Client(timeout=8) as client:
headers = {
"Accept": "application/json, text/event-stream",
"MCP-Protocol-Version": "2025-11-25",
}
initialized = client.post(
url + "/mcp",
headers=headers,
json={
"jsonrpc": "2.0", "id": 1, "method": "initialize",
"params": {
"protocolVersion": "2025-11-25", "capabilities": {},
"clientInfo": {"name": "notification-order", "version": "1"},
},
},
)
assert initialized.status_code == 200
headers["Mcp-Session-Id"] = initialized.headers["mcp-session-id"]
try:
accepted = client.post(
url + "/mcp", headers=headers,
json={"jsonrpc": "2.0", "method": "notifications/initialized"},
)
assert accepted.status_code == 202
assert accepted.content == b""
terminated = client.delete(url + "/mcp", headers=headers)
assert terminated.status_code == 200
finally:
deleted.set()

90
tests/integration/mcp/test_mcp_oauth_flows.py
View file

@ -8,7 +8,7 @@ import textwrap
import time
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from pathlib import Path
from typing import Final
@ -1991,13 +1991,43 @@ def test_standard_pattern_discovery_ignores_forwarded_headers_from_an_untrusted_
@contextmanager
def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, Wire]]:
with scripted_peer(
def _sdk_scope_peer() -> Iterator[McpPeer]:
import queue
from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from integration._support.asgi import asgi_server
from integration._support.mcp import _capturing
service = MCPServer("scope-operations")
@service.tool()
def read_op() -> str:
return "read completed"
@service.tool()
def write_op() -> str:
return "write completed"
@service.tool()
def forbid_op() -> str:
return "must not execute"
app = service.streamable_http_app(
stateless_http=True, json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
observed = queue.Queue()
with asgi_server(_capturing(app, observed)) as url:
yield McpPeer(url + "/mcp", observed)
@contextmanager
def _scope_enforcing_peer(auth: AuthorizationServer, *, modern: bool = False) -> Iterator[tuple[McpPeer, Wire]]:
with (_sdk_scope_peer() if modern else scripted_peer(
ScriptedTool("read_op", lambda _: text_result("read completed")),
ScriptedTool("write_op", lambda _: text_result("write completed")),
ScriptedTool("forbid_op", lambda _: text_result("must not execute")),
) as inner:
)) as inner:
def enforce(request: Request) -> Reply:
body: Final = json.loads(request.body) if request.body else {}
if body.get("method") == "tools/call":
@ -2026,51 +2056,55 @@ def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer,
yield McpPeer(wire.url + "/mcp", inner.calls), wire
@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"])
@pytest.mark.parametrize("entry", ["rest", "server_mcp"])
def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: Gateway, entry: EntryPoint) -> None:
def test_managed_scope_step_up_preserves_read_and_write_without_retry(
gateway: Gateway, entry: EntryPoint, tmp_path: Path, protocol: str
) -> None:
modern = protocol == "2026-07-28"
config = tmp_path / "scope-proxy.yaml"
config.write_text(yaml.safe_dump({
"model_list": [],
"general_settings": {
"master_key": gateway.key, "store_model_in_db": True,
"mcp_advertised_versions": ["2025-11-25", "2026-07-28"],
},
}))
with (
(owned_proxy(gateway, tmp_path / "gateway", {"DISABLE_SCHEMA_UPDATE": "true"}, config=config, database_setup=())
if modern else nullcontext(gateway)) as active,
oauth_server(scopes=("tools.read", "tools.write")) as auth,
_scope_enforcing_peer(auth) as (peer, wire),
gateway.scenario() as scenario,
_scope_enforcing_peer(auth, modern=modern) as (peer, wire),
active.scenario() as scenario,
):
alias: Final = "scope" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="authorization_code",
scenario, peer, auth, alias, auth_type="oauth2", oauth2_flow="authorization_code", protocol_version=protocol,
credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]},
)
owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
pkce: Final = _Pkce(secrets.token_urlsafe(32))
code, _ = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read")
first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce)
code, _ = _authorize_through_gateway(active, auth, alias, owner, "scope-client", pkce, "tools.read")
first: Final = _redeem(active, alias, owner, "scope-client", code, pkce)
assert first.status_code == 200, first.text
caller: Final = McpCaller(gateway, owner, entry, alias)
caller: Final = McpCaller(active, owner, entry, alias, protocol_version=protocol)
read: Final = caller.call(f"{alias}-read_op", {}, identity)
assert read.ok and read.text == "read completed", read.raw
peer.drain()
wire.drain()
denied: Final = (
call_tool(gateway, owner, identity, f"{alias}-write_op", {})
if entry == "rest"
call_tool(active, owner, identity, f"{alias}-write_op", {}) if entry == "rest"
else caller.rpc("tools/call", {"name": f"{alias}-write_op", "arguments": {}})
)
assert denied.status_code == 403, denied.text
challenge: Final = denied.headers["www-authenticate"]
assert 'error="insufficient_scope"' in challenge and 'scope="tools.write"' in challenge
assert f'resource_metadata="{_base(gateway)}/.well-known/oauth-protected-resource/' in challenge
attempts: Final = tuple(
request
for request in wire.drain()
if request.body and json.loads(request.body).get("method") == "tools/call"
)
assert f'resource_metadata="{_base(active)}/.well-known/oauth-protected-resource/' in challenge
attempts: Final = tuple(request for request in wire.drain() if request.body and json.loads(request.body).get("method") == "tools/call")
assert len(attempts) == 1 and tool_calls(peer.drain()) == (), "denied operation was retried or executed"
second_code, _ = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.write")
second: Final = _redeem(gateway, alias, owner, "scope-client", second_code, pkce)
second_code, _ = _authorize_through_gateway(active, auth, alias, owner, "scope-client", pkce, "tools.write")
second: Final = _redeem(active, alias, owner, "scope-client", second_code, pkce)
assert second.status_code == 200, second.text
with auth.lock:
granted: Final = auth.access_tokens[_issued_token(second.json())]["scope"]
@ -2080,7 +2114,7 @@ def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: G
assert result.ok and result.text == f"{operation} completed", result.raw
forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity)
assert forbidden.status == 200 and not forbidden.ok, forbidden.raw
unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity)
unauthorized: Final = McpCaller(active, stranger, entry, alias, protocol_version=protocol).call(f"{alias}-read_op", {}, identity)
assert unauthorized.status == 401, unauthorized.raw

45
tests/integration/mcp/test_pagination.py
View file

@ -8,7 +8,7 @@ import pytest
import yaml
from mcp import ClientSession, MCPError
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolRequest, CallToolRequestParams, CallToolResult, PaginatedRequestParams
from mcp.types import CallToolRequest, CallToolRequestParams, CallToolResult, DiscoverResult, PaginatedRequestParams
from integration._support.client import Gateway
from integration._support.mcp import McpPeer, paginated_mcp_peer
@ -18,24 +18,27 @@ from litellm.types.mcp import MCPTransport
@asynccontextmanager
async def catalog_session(gateway):
async def catalog_session(gateway, protocol: str = "2025-11-25"):
async with httpx.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client:
async with streamable_http_client(
str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client
) as streams:
async with ClientSession(streams[0], streams[1]) as session:
await session.initialize()
if protocol == "2026-07-28":
session.adopt(DiscoverResult.model_validate(await session.send_discover(protocol)))
else:
await session.initialize()
yield session
def config_file(directory: Path, upstream: str) -> Path:
def config_file(directory: Path, upstream: str, protocol: str | None = None) -> Path:
config = directory / "proxy.yaml"
config.write_text(
yaml.safe_dump(
{
"model_list": [],
"mcp_servers": {"pages": {"url": upstream, "transport": "http"}},
"general_settings": {"master_key": "sk-pagination-test", "store_model_in_db": False},
"mcp_servers": {"pages": {"url": upstream, "transport": "http", **({"protocol_version": protocol} if protocol else {})}},
"general_settings": {"master_key": "sk-pagination-test", "store_model_in_db": False, **({"mcp_advertised_versions": ["2025-11-25", "2026-07-28"]} if protocol == "2026-07-28" else {})},
}
)
)
@ -45,7 +48,8 @@ def config_file(directory: Path, upstream: str) -> Path:
REMOVE_DATABASE = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "LITELLM_LICENSE", "LITELLM_LICENSE_PATH")
def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path):
@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"])
def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path, protocol: str):
async def exercise(first_replica, second_replica):
for method, field in (
("list_tools", "tools"),
@ -53,10 +57,10 @@ def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path):
("list_resources", "resources"),
("list_resource_templates", "resource_templates"),
):
async with catalog_session(first_replica) as first_session:
async with catalog_session(first_replica, protocol) as first_session:
first = await getattr(first_session, method)()
assert first.next_cursor
async with catalog_session(second_replica) as second_session:
async with catalog_session(second_replica, protocol) as second_session:
second = await getattr(second_session, method)(params=PaginatedRequestParams(cursor=first.next_cursor))
replay = await getattr(second_session, method)(params=PaginatedRequestParams(cursor=first.next_cursor))
assert getattr(second, field) == getattr(replay, field)
@ -65,14 +69,14 @@ def test_catalog_pages_are_portable_repeatable_and_complete(tmp_path: Path):
names = [item.name for page in (first, second, third) for item in getattr(page, field)]
assert len(names) == len(set(names)) == 3
assert third.next_cursor is None
async with catalog_session(second_replica) as session:
async with catalog_session(second_replica, protocol) as session:
called = await session.call_tool("pages-add2", {"a": 3, "b": 4})
assert not called.is_error
assert called.content[0].text == "7"
with paginated_mcp_peer() as peer, httpx.Client() as client:
seed = Gateway(client, "sk-pagination-test", peer.url)
config = config_file(tmp_path, peer.url)
config = config_file(tmp_path, peer.url, protocol)
options = dict(config=config, database_setup=(), remove_environment=REMOVE_DATABASE)
environment = {
"STORE_MODEL_IN_DB": "False",
@ -128,7 +132,8 @@ def test_no_salt_preserves_complete_lists_and_calls_but_rejects_continuations(tm
@pytest.mark.parametrize("grant", ["direct", "access_group"])
def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, monkeypatch, grant: str):
@pytest.mark.parametrize("protocol", ["2025-11-25", "2026-07-28"])
def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, monkeypatch, grant: str, protocol: str):
import os
import subprocess
import sys
@ -144,15 +149,15 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m
stranger_b = Gateway(b.client, stranger, peer.url)
first_pages = {}
for method in ("list_tools", "list_prompts", "list_resources", "list_resource_templates"):
async with catalog_session(owner_a) as session:
async with catalog_session(owner_a, protocol) as session:
first_pages[method] = await getattr(session, method)()
assert first_pages[method].next_cursor
async with catalog_session(owner_b) as session:
async with catalog_session(owner_b, protocol) as session:
continued = await getattr(session, method)(
params=PaginatedRequestParams(cursor=first_pages[method].next_cursor)
)
assert continued.next_cursor
async with catalog_session(stranger_b) as session:
async with catalog_session(stranger_b, protocol) as session:
peer.drain()
with pytest.raises(MCPError, match="fresh listing"):
await getattr(session, method)(
@ -172,7 +177,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m
},
)
for method, first in first_pages.items():
async with catalog_session(owner_b) as session:
async with catalog_session(owner_b, protocol) as session:
peer.drain()
with pytest.raises(MCPError, match="fresh listing"):
await getattr(session, method)(params=PaginatedRequestParams(cursor=first.next_cursor))
@ -181,7 +186,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m
changed = a.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "new catalog generation"})
assert changed.status_code == 202, changed.text
for method, first in first_pages.items():
async with catalog_session(owner_b) as session:
async with catalog_session(owner_b, protocol) as session:
peer.drain()
with pytest.raises(MCPError, match="fresh listing"):
await getattr(session, method)(params=PaginatedRequestParams(cursor=first.next_cursor))
@ -212,7 +217,7 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m
config = tmp_path / "database-proxy.yaml"
config.write_text(
yaml.safe_dump(
{"model_list": [], "general_settings": {"master_key": seed.key, "store_model_in_db": True}}
{"model_list": [], "general_settings": {"master_key": seed.key, "store_model_in_db": True, "mcp_advertised_versions": ["2025-11-25", "2026-07-28"]}}
)
)
environment = {
@ -230,8 +235,8 @@ def test_continuations_reauthorize_and_reject_registry_changes(tmp_path: Path, m
owned_proxy(seed, tmp_path / "b", environment, **options) as b,
a.scenario() as scenario,
):
identity = register_mcp(scenario, peer, "pages")
spare: Final = register_mcp(scenario, spare_peer, "spare")
identity = register_mcp(scenario, peer, "pages", protocol_version=protocol)
spare: Final = register_mcp(scenario, spare_peer, "spare", protocol_version=protocol)
group = a.request(
"POST",
"/v1/access_group",

205
tests/unit/integration_support/test_upstream.py
View file

@ -1,3 +1,4 @@
import asyncio
from collections import deque
from pathlib import Path
from typing import Final
@ -50,3 +51,207 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error
("WEBSOCKET", {"query": [["model", "synthetic"]]}),
("WEBSOCKET_FRAME", {"type": "response.create"}),
]
@pytest.mark.asyncio
@pytest.mark.parametrize("method,status", (("POST", 200), ("POST", 400), ("DELETE", 200)))
async def test_mcp_notification_ordering_preserves_other_responses(
method: str, status: int, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _acknowledge_notifications_after_dispatch
from starlette.types import Receive, Send
outgoing: Final[list[Message]] = []
messages: Final[tuple[Message, ...]] = (
{"type": "http.response.start", "status": status, "headers": []},
{"type": "http.response.body", "body": b"result", "more_body": True},
{"type": "http.response.body", "body": b"", "more_body": False},
)
async def receive() -> Message:
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
outgoing.append(message)
async def app(scope: Scope, receive: Receive, send: Send) -> None:
for index, message in enumerate(messages):
await send(message)
assert outgoing == list(messages[: index + 1]), "Ordinary response streaming was buffered"
await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": method}, receive, send)
assert outgoing == list(messages)
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", (RuntimeError("dispatch failed"), asyncio.CancelledError()))
async def test_mcp_notification_dispatch_failure_is_not_acknowledged(
failure: BaseException, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _acknowledge_notifications_after_dispatch
from starlette.types import Receive, Send
outgoing: Final[list[Message]] = []
async def receive() -> Message:
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
outgoing.append(message)
async def app(scope: Scope, receive: Receive, send: Send) -> None:
if scope["method"] == "DELETE":
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
return
await send({"type": "http.response.start", "status": 202, "headers": []})
await send({"type": "http.response.body", "body": b""})
raise failure
ordered: Final = _acknowledge_notifications_after_dispatch(app)
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"failed-session")]}
with pytest.raises(type(failure)) as caught:
await ordered({**scope, "method": "POST"}, receive, send)
assert caught.value is failure
assert outgoing == [], "Failed dispatch must not send an acknowledgement"
await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2)
assert outgoing[0]["status"] == 200
@pytest.mark.asyncio
@pytest.mark.parametrize(
"responses",
(
(),
({"type": "http.response.body", "body": b"unexpected"},),
({"type": "http.response.body", "body": b"", "more_body": True},),
({"type": "http.response.start", "status": 500, "headers": []},),
({"type": "http.response.body", "body": b""}, {"type": "http.response.body", "body": b""}),
),
)
async def test_mcp_notification_ordering_rejects_invalid_acknowledgements(
responses: tuple[Message, ...], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _acknowledge_notifications_after_dispatch
from starlette.types import Receive, Send
outgoing: Final[list[Message]] = []
async def receive() -> Message:
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
outgoing.append(message)
async def app(scope: Scope, receive: Receive, send: Send) -> None:
await send({"type": "http.response.start", "status": 202, "headers": []})
for message in responses:
await send(message)
with pytest.raises(RuntimeError, match="MCP notification acknowledgement"):
await _acknowledge_notifications_after_dispatch(app)({"type": "http", "method": "POST"}, receive, send)
assert outgoing == [], "Invalid acknowledgement must not be reported as accepted"
@pytest.mark.asyncio
async def test_mcp_session_deletion_waits_for_concurrent_notification_dispatch(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _acknowledge_notifications_after_dispatch
from starlette.types import Receive, Send
started: Final = [asyncio.Event(), asyncio.Event()]
release: Final = asyncio.Event()
deletion_started: Final = asyncio.Event()
deleted: Final = asyncio.Event()
other_deleted: Final = asyncio.Event()
completed: Final[list[str]] = []
async def receive() -> Message:
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
return None
async def app(scope: Scope, receive: Receive, send: Send) -> None:
if scope["method"] == "POST":
started[int(scope["path"])].set()
await release.wait()
completed.append(scope["path"])
await send({"type": "http.response.start", "status": 202, "headers": []})
await send({"type": "http.response.body", "body": b""})
else:
if scope["headers"][0][1] == b"other-session":
other_deleted.set()
else:
deleted.set()
completed.append("DELETE")
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
ordered: Final = _acknowledge_notifications_after_dispatch(app)
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"owned-session")]}
posts: Final = [
asyncio.create_task(ordered({**scope, "method": "POST", "path": str(index)}, receive, send))
for index in range(2)
]
async def delete() -> None:
deletion_started.set()
await ordered({**scope, "method": "DELETE"}, receive, send)
deletion = None
try:
await asyncio.wait_for(asyncio.gather(*(event.wait() for event in started)), 2)
deletion = asyncio.create_task(delete())
await deletion_started.wait()
await asyncio.wait_for(
ordered({**scope, "method": "DELETE", "headers": [(b"mcp-session-id", b"other-session")]}, receive, send), 2
)
assert other_deleted.is_set(), "Unrelated session deletion was blocked"
assert not deleted.is_set(), "Session deletion overtook in-flight notification dispatch"
finally:
release.set()
await asyncio.wait_for(asyncio.gather(*posts, *([deletion] if deletion is not None else [])), 2)
assert set(completed[:2]) == {"0", "1"}
assert completed[-1] == "DELETE"
@pytest.mark.asyncio
async def test_mcp_session_deletion_can_terminate_streaming_request(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
from integration._support.mcp import _acknowledge_notifications_after_dispatch
from starlette.types import Receive, Send
streaming: Final = asyncio.Event()
deleted: Final = asyncio.Event()
async def receive() -> Message:
return {"type": "http.disconnect"}
async def send(message: Message) -> None:
return None
async def app(scope: Scope, receive: Receive, send: Send) -> None:
await send({"type": "http.response.start", "status": 200, "headers": []})
if scope["method"] == "POST":
streaming.set()
await deleted.wait()
else:
deleted.set()
await send({"type": "http.response.body", "body": b""})
ordered: Final = _acknowledge_notifications_after_dispatch(app)
scope: Final[Scope] = {"type": "http", "headers": [(b"mcp-session-id", b"streaming-session")]}
post: Final = asyncio.create_task(ordered({**scope, "method": "POST"}, receive, send))
try:
await asyncio.wait_for(streaming.wait(), 2)
await asyncio.wait_for(ordered({**scope, "method": "DELETE"}, receive, send), 2)
assert deleted.is_set(), "Session deletion must terminate an ordinary streaming request"
finally:
deleted.set()
await asyncio.wait_for(post, 2)

45
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
View file

@ -164,6 +164,51 @@ def test_modern_protocol_is_rejected_on_legacy_sse_paths(path: str, monkeypatch:
)
@pytest.mark.parametrize("revision,upstreams,transport,allowed", [
("2026-07-28", frozenset({"2026-07-28"}), MCPTransport.http, True),
("2025-11-25", frozenset({"2026-07-28"}), MCPTransport.http, True),
("2025-11-25", frozenset({"2025-11-25"}), MCPTransport.sse, True),
("2026-07-28", frozenset({"2025-11-25", "2026-07-28"}), MCPTransport.http, True),
("2026-07-28", frozenset({"2025-11-25"}), MCPTransport.http, True),
("2026-07-28", frozenset({"2026-07-28"}), MCPTransport.stdio, False),
])
def test_subscription_flags_require_a_supported_protocol_pair(revision, upstreams, transport, allowed):
result = build_discovery(
configured=("2025-11-25", "2026-07-28"), revision=revision, transport=transport,
authorized_operations=GATEWAY_OPERATIONS, upstream_versions=upstreams,
capabilities=ServerCapabilities(
tools=ToolsCapability(list_changed=True), prompts=PromptsCapability(list_changed=True),
resources=ResourcesCapability(list_changed=True, subscribe=True),
),
)
assert bool(result.capabilities.tools.list_changed) is allowed
assert bool(result.capabilities.prompts.list_changed) is allowed
assert bool(result.capabilities.resources.list_changed) is allowed
assert bool(result.capabilities.resources.subscribe) is allowed
assert result.capabilities.tasks is None
assert result.capabilities.extensions is None
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["empty", "modern-stdio"])
async def test_subscription_discovery_never_probes_unsupported_targets(mode, monkeypatch):
from types import SimpleNamespace
from unittest.mock import AsyncMock
from litellm.proxy._experimental.mcp_server import subscriptions
from litellm.proxy._experimental.mcp_server.contracts import OperationContext
probe = AsyncMock()
target = SimpleNamespace(
server=SimpleNamespace(protocol_version="2026-07-28", transport="stdio"),
client=SimpleNamespace(run_with_session=probe),
)
monkeypatch.setattr(subscriptions, "_targets", AsyncMock(return_value=() if mode == "empty" else (target,)))
versions, capabilities = await subscriptions.subscription_capabilities(OperationContext(_caller=None))
assert versions
assert capabilities == ServerCapabilities()
probe.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"disabled,prompts_advertised,resources_advertised",

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

@ -1943,3 +1943,56 @@ async def test_guardrail_context_never_recovers_a_previous_request_bearer(
assert capture.data["incoming_bearer_token"] is None
finally:
auth_context_var.reset(previous)
@pytest.mark.asyncio
@pytest.mark.parametrize("delay", ["target", "probe", "exhausted"])
async def test_discovery_optional_probes_use_remaining_listing_budget(delay, monkeypatch):
from types import SimpleNamespace
from mcp.types import (
DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult,
ListResourceTemplatesResult, ToolsCapability, Tool,
)
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import subscriptions
entered = asyncio.Event()
cleaned = asyncio.Event()
async def blocked(*args, **kwargs):
entered.set()
try:
await asyncio.Event().wait()
finally:
cleaned.set()
async def tools(*args, **kwargs):
if delay == "exhausted":
await asyncio.sleep(0.08)
return ListToolsResult(tools=[Tool(name="usable", input_schema={"type": "object"})])
target = SimpleNamespace(
server=SimpleNamespace(protocol_version="2026-07-28", transport="http"),
client=SimpleNamespace(run_with_session=blocked),
)
monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2026-07-28"])
# The existing listing budget includes optional target preparation and handshakes.
monkeypatch.setattr(operations, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False)
with (
patch.object(operations, "_execute_handle_list_tools", side_effect=tools),
patch.object(operations, "_execute_list_prompts", return_value=ListPromptsResult(prompts=[])),
patch.object(operations, "_execute_list_resources", return_value=ListResourcesResult(resources=[])),
patch.object(operations, "_execute_list_resource_templates", return_value=ListResourceTemplatesResult(resource_templates=[])),
patch.object(operations, "_mcp_server_rate_limit_rejection", return_value=None),
patch.object(subscriptions, "_targets", side_effect=blocked if delay == "target" else None, return_value=(target,)) as targets,
):
context = prepare_context(UserAPIKeyAuth(user_id="scoped"), protocol_version="2026-07-28")
result = await asyncio.wait_for(GatewayOperations().execute(DiscoverRequest(), context), 0.5)
assert result.capabilities.tools == ToolsCapability()
assert result.capabilities.resources is None
if delay == "exhausted":
targets.assert_not_awaited()
assert not entered.is_set()
else:
assert entered.is_set()
assert cleaned.is_set()

5
tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py
View file

@ -87,14 +87,14 @@ async def test_legacy_listener_rejected_before_upstream_work(revision: str) -> N
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["multiple-resources", "denied"])
@pytest.mark.parametrize("mode", ["stdio", "multiple-resources", "denied"])
async def test_unsupported_targets_do_not_open_upstream_listeners(mode: str) -> None:
client: Final = MCPClient(server_url="http://127.0.0.1:1/mcp", transport_type=MCPTransport.http)
server: Final = MCPServer(
server_id="upstream",
name="upstream",
url="http://localhost/mcp",
transport="http",
transport="stdio" if mode == "stdio" else "http",
protocol_version="2025-11-25" if mode == "legacy" else "2026-07-28",
)
target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint")
@ -109,6 +109,7 @@ async def test_unsupported_targets_do_not_open_upstream_listeners(mode: str) ->
)
error: Final = HTTPException if mode == "denied" else MCPError
expected: Final = {
"stdio": "Modern upstream subscriptions require HTTP",
"multiple-resources": "exactly one allowed server",
"denied": "No authorized MCP subscription targets",
}[mode]