diff --git a/litellm/proxy/mcp_server/catalog.py b/litellm/proxy/mcp_server/catalog.py index d2ff322ba5c..e30f26abf41 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -826,6 +826,20 @@ async def prepare_tool_listing_headers( return server_auth_header, extra_headers, catalog_auth_header +def challenge_rejected_credential(server: MCPServer, outcome: ServerOutcome, listed: Sequence[MCPServer]) -> None: + """Answer a listing of this one server with the gateway's 401 when the caller can fix the rejection.""" + from litellm.proxy.mcp_server.legacy_callbacks import get_listing_scope_response + + scope_response: Final = get_listing_scope_response() + if ( + scope_response is not None + and len(listed) == 1 + and outcome.tag == "auth_required" + and server.reauthorizes_in_client + ): + scope_response.deny_reauth() + + async def get_filtered_server_tools( server: MCPServer, *, @@ -943,14 +957,11 @@ async def get_filtered_server_tools( ) return page.model_copy(update={"tools": filtered_tools}), ServerListOk(tool_count=len(filtered_tools)) except MCPUpstreamAuthError as e: - # Absorb so one unauthenticated server does not empty every other server's - # tools. Surfacing the upstream 401 to the client as a re-auth challenge is - # intentionally not done here: raising from this list handler cannot produce a - # 401 + WWW-Authenticate (the MCP session manager serializes it as a JSON-RPC - # error). Single-server routes surface it via the request-scope preemptive - # check in _raise_preemptive_401_for_unauthenticated_servers instead. + # Absorbed so peers keep their tools; a raise here would reach the client as a JSON-RPC error, not a 401 verbose_logger.debug("MCP list_tools: omitting %s; it needs upstream auth", server.name) - return ListToolsResult(tools=[]), classify_list_exception(e) + outcome: Final = classify_list_exception(e) + challenge_rejected_credential(server, outcome, allowed_mcp_servers) + return ListToolsResult(tools=[]), outcome except Exception as e: verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) return ListToolsResult(tools=[]), classify_list_exception(e) @@ -1199,6 +1210,8 @@ async def list_gateway_catalog( update={"meta": {key: value for key, value in page_meta.items() if key != SERVER_OUTCOMES_META_KEY}} ) except Exception as error: + outcome: Final = classify_list_exception(error) + challenge_rejected_credential(server, outcome, allowed) if cursor is not None: raise MCPError( code=INVALID_PARAMS, message="Upstream continuation failed; start a fresh listing" @@ -1207,11 +1220,7 @@ async def list_gateway_catalog( request, (), None, - { - "litellm.ai/server_outcomes": { - _aggregate_server_key(server): classify_list_exception(error).model_dump(mode="json") - } - }, + {"litellm.ai/server_outcomes": {_aggregate_server_key(server): outcome.model_dump(mode="json")}}, ) pages, next_cursor, outcomes = await paginate_catalog( diff --git a/litellm/proxy/mcp_server/exceptions.py b/litellm/proxy/mcp_server/exceptions.py index 8be3ae78d15..db9bc4fc0ee 100644 --- a/litellm/proxy/mcp_server/exceptions.py +++ b/litellm/proxy/mcp_server/exceptions.py @@ -1,10 +1,11 @@ """Exceptions raised by the LiteLLM MCP proxy.""" from typing import Final +from urllib.parse import quote from fastapi import HTTPException -from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope +from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope, well_known_root_suffix TOKEN_EXCHANGE_SUBJECT_ERROR: Final = ( "Missing or invalid subject token. Authenticate with your identity provider and send its token " @@ -101,12 +102,14 @@ class MCPUpstreamAuthError(Exception): challenge: str | None = self.www_authenticate if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url: prefix: Final = base_url.rstrip("/") - if aggregate_step_up: - resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource" - elif request_path and request_path.startswith(f"/{self.server_name}/mcp"): - resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp" - else: - resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}" + name: Final = quote(self.server_name, safe="/:@!$&'()*+,;=") + resource_metadata_url: Final = ( + f"{prefix}/.well-known/oauth-protected-resource" + if aggregate_step_up + else f"{prefix}/.well-known/oauth-protected-resource{well_known_root_suffix()}/{name}/mcp" + if request_path and request_path.startswith(f"/{self.server_name}/mcp") + else f"{prefix}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{name}" + ) challenge = f'Bearer resource_metadata="{resource_metadata_url}"' scoped_challenge: Final = ( (f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"') diff --git a/litellm/proxy/mcp_server/legacy_callbacks.py b/litellm/proxy/mcp_server/legacy_callbacks.py index 3d2f97492bb..82d98f39604 100644 --- a/litellm/proxy/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/mcp_server/legacy_callbacks.py @@ -17,6 +17,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server.contracts import OperationContext from litellm.proxy.mcp_server.mcp_context import get_active_mcp_request_ctx from litellm.proxy.mcp_server.scope_challenge import ( + LISTING_METHODS, OAuthScopeResponse, is_tool_call_request, scope_response_from_request, @@ -113,6 +114,12 @@ def get_scope_response() -> OAuthScopeResponse | None: return scope_response_from_request(context.request) if context is not None else None +def get_listing_scope_response() -> OAuthScopeResponse | None: + """The response gate of the active request, when the SDK parsed that request as a listing.""" + context: Final = get_active_mcp_request_ctx() + return get_scope_response() if context is not None and context.method in LISTING_METHODS else None + + async def record_upstream_tool_authorization(response: httpx2.Response, *, pending: OAuthScopeResponse | None) -> None: if pending is None or not pending.pending or not response.is_success or response.request.method != "POST": return diff --git a/litellm/proxy/mcp_server/scope_challenge.py b/litellm/proxy/mcp_server/scope_challenge.py index 128b27b0f91..927cb38c42e 100644 --- a/litellm/proxy/mcp_server/scope_challenge.py +++ b/litellm/proxy/mcp_server/scope_challenge.py @@ -19,18 +19,31 @@ SCOPE_RESPONSE_KEY: Final = "litellm.mcp.scope_response" _REQUEST: Final = TypeAdapter[dict[str, object]](dict[str, object]) -def is_tool_call_request(body: bytes) -> bool: +LISTING_METHODS: Final = frozenset({"tools/list", "prompts/list", "resources/list", "resources/templates/list"}) + + +def _request_method(body: bytes) -> object: try: - return _REQUEST.validate_json(body).get("method") == "tools/call" + return _REQUEST.validate_json(body).get("method") except ValidationError: - return False + return None + + +def is_tool_call_request(body: bytes) -> bool: + return _request_method(body) == "tools/call" + + +def is_listing_request(body: bytes) -> bool: + method: Final = _request_method(body) + return isinstance(method, str) and method in LISTING_METHODS class OAuthScopeResponse: - def __init__(self, send: Send, base_url: str, request_path: str | None) -> None: + def __init__(self, send: Send, base_url: str, request_path: str | None, reauthorize_as: str | None = None) -> None: self._send = send self._base_url = base_url self._request_path = request_path + self._reauthorize_as = reauthorize_as self._ready = asyncio.Event() self._lock = asyncio.Lock() self._challenge: MCPUpstreamAuthError | None = None @@ -44,7 +57,15 @@ class OAuthScopeResponse: self._ready.set() def deny(self, error: MCPUpstreamAuthError) -> None: - if self.pending and error.status_code == 403 and error.required_scope is not None: + if error.status_code == 403 and error.required_scope is not None: + self._refuse(error) + + def deny_reauth(self) -> None: + if self._reauthorize_as is not None: + self._refuse(MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=self._reauthorize_as)) + + def _refuse(self, error: MCPUpstreamAuthError) -> None: + if self.pending: self._challenge = error self._ready.set() diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index bbbbd09e527..c255eddb214 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -10,6 +10,7 @@ import contextvars import hashlib import json import os +import re import time import types from collections import Counter @@ -81,9 +82,11 @@ from litellm.proxy.mcp_server.oauth_utils import ( well_known_root_suffix, ) from litellm.proxy.mcp_server.scope_challenge import ( + LISTING_METHODS, SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response, + is_listing_request, is_tool_call_request, ) from litellm.proxy.mcp_server.ui_session_utils import ( @@ -225,16 +228,13 @@ _SESSION_MANAGERS_INITIALIZED = False _INITIALIZATION_LOCK: Final = asyncio.Lock() -def _jsonrpc_text_has_top_level_method(text: str) -> bool: - """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at - the root object's top level. +def _jsonrpc_text_top_level_key_end(text: str, key: str) -> int | None: + """Where the root object's top-level ``key`` ends in a (possibly truncated) JSON-RPC envelope. - Used to tell a request/notification (carries ``method``) apart from a - response (carries ``result``/``error`` and no top-level ``method``). A - response payload can itself nest a ``method`` field, so only keys at the + A response payload can itself nest a ``method`` field, so only keys at the root object's depth are inspected rather than searching the whole string. - Returns ``True`` only when a top-level ``method`` key is positively found; - truncation that hides it yields ``False``. + Returns an index only when the key is positively found; truncation that + hides it yields ``None``. """ depth = 0 in_string = False @@ -243,7 +243,7 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: reading_key = False expect_key = False key_chars: list[str] = [] - for ch in text: + for index, ch in enumerate(text): if in_string: if escaped: escaped = False @@ -251,8 +251,8 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: escaped = True elif ch == '"': in_string = False - if reading_key and depth == 1 and "".join(key_chars) == "method": - return True + if reading_key and depth == 1 and "".join(key_chars) == key: + return index + 1 elif reading_key: key_chars.append(ch) continue @@ -276,7 +276,32 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: expect_key = bool(in_object) and in_object[-1] elif ch == ":": expect_key = False - return False + return None + + +def _jsonrpc_text_has_top_level_method(text: str) -> bool: + """Whether a (possibly truncated) JSON-RPC envelope has a top-level ``method`` key. + + Used to tell a request/notification (carries ``method``) apart from a + response (carries ``result``/``error`` and no top-level ``method``). + """ + return _jsonrpc_text_top_level_key_end(text, "method") is not None + + +_JSONRPC_STRING_VALUE: Final = re.compile(r'\s*:\s*("(?:[^"\\]|\\.)*")') + + +def _jsonrpc_text_method(text: str) -> str | None: + """The top-level ``method`` of a (possibly truncated) JSON-RPC envelope, when the cut leaves it whole.""" + key_end: Final = _jsonrpc_text_top_level_key_end(text, "method") + value: Final = _JSONRPC_STRING_VALUE.match(text, key_end) if key_end is not None else None + if value is None: + return None + try: + method: Final = json.loads(value.group(1)) + except ValueError: + return None + return method if isinstance(method, str) else None def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: @@ -2470,6 +2495,27 @@ if MCP_AVAILABLE: ), ) + def _is_listing_body(body: bytes) -> bool: + """A listing request, including one the routing peek cut off after its method.""" + if is_listing_request(body): + return True + peek_cut: Final = len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES + return peek_cut and _jsonrpc_text_method(body.decode("utf-8", errors="replace")) in LISTING_METHODS + + def _server_to_reauthorize( + mcp_servers: Sequence[str] | None, scoped_server_endpoint: bool, body: bytes, client_ip: str | None + ) -> str | None: + """The server a client addressed alone for a listing it can fix by authenticating again.""" + if not scoped_server_endpoint or mcp_servers is None or len(mcp_servers) != 1 or not _is_listing_body(body): + return None + server: Final = operations.global_mcp_server_manager.get_mcp_server_by_identifier( + mcp_servers[0], client_ip=client_ip + ) + reauthorizes: Final = ( + server is not None and MCPServerManager.resolve_oauth2_flow_for_request(server).reauthorizes_in_client + ) + return mcp_servers[0] if reauthorizes else None + async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: @@ -2638,13 +2684,15 @@ if MCP_AVAILABLE: await rejection(scope, receive, send) return + reauthorize_as: Final = _server_to_reauthorize(mcp_servers, scoped_server_endpoint, body, _client_ip) scope_response: Final = ( OAuthScopeResponse( send, get_request_base_url(StarletteRequest(scope)), get_route_relative_request_path(scope), + reauthorize_as=reauthorize_as, ) - if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES + if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES or reauthorize_as is not None else None ) if scope_response is not None: diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 1d3d5911eb1..ed0f050a76a 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -383,6 +383,12 @@ class MCPServer(LiteLLMBaseModel): door and is never a precondition for opening a session.""" return self.is_true_passthrough or self.is_oauth_delegate + @property + def reauthorizes_in_client(self) -> bool: + """True when a caller can fix an upstream 401 by authenticating again: the client-forwarded + token modes and gateway-vaulted per-user OAuth.""" + return self.is_client_forwarded_token or (self.is_gateway_managed_oauth2 and self.needs_user_oauth_token) + @property def is_dcr_bridge(self) -> bool: """True when this client-forwarded-token server serves the gateway-hosted DCR front door diff --git a/tests/unit/proxy/mcp_server/test_catalog.py b/tests/unit/proxy/mcp_server/test_catalog.py index 45c648b81d2..3c78e8f2561 100644 --- a/tests/unit/proxy/mcp_server/test_catalog.py +++ b/tests/unit/proxy/mcp_server/test_catalog.py @@ -2,9 +2,11 @@ import asyncio from collections.abc import Mapping, Sequence from types import SimpleNamespace from typing import Final, Literal -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock +import httpx2 import pytest +from mcp.server.context import ServerRequestContext from mcp.shared.exceptions import MCPError from mcp.types import ( ListPromptsRequest, @@ -18,17 +20,24 @@ from mcp.types import ( PaginatedRequestParams, Tool, ) +from starlette.requests import Request from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.mcp_server import catalog +from litellm.proxy.mcp_server import operations as mcp_operations from litellm.proxy.mcp_server.contracts import OperationContext +from litellm.proxy.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy.mcp_server.faults.list_outcomes import ( SERVER_OUTCOMES_META_KEY, AggregateToolListing, + ServerListFault, ServerListOk, ) -from litellm.types.mcp import MCPTransport +from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var, mcp_gateway_server_name +from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse +from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer CatalogKind = Literal["tools", "prompts", "resources", "templates"] @@ -772,3 +781,293 @@ def test_partial_optional_catalog_never_advertises_freshness() -> None: assert result.ttl_ms == 0 assert result.cache_scope == "private" + +_GATEWAY: Final = "https://gateway.example" +_UNAUTHORIZED: Final = ServerListFault(tag="auth_required", status_code=401) + + +def _server( + auth_type: MCPAuth, + name: str = "docs", + oauth2_flow: str | None = None, + delegate_auth_to_upstream: bool = False, + authentication_token: str | None = None, + alias: str | None = None, +) -> MCPServer: + return MCPServer( + server_id=f"{name}-id", + name=name, + alias=alias or name, + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + oauth2_flow=oauth2_flow, + delegate_auth_to_upstream=delegate_auth_to_upstream, + authentication_token=authentication_token, + ) + + +_VAULTED: Final = _server(MCPAuth.oauth2, oauth2_flow="authorization_code") +_PEER: Final = _server(MCPAuth.none, name="wiki") +_STATIC: Final = _server(MCPAuth.api_key, authentication_token="k") +_REAUTHORIZABLE: Final = [ + pytest.param(_VAULTED, id="vaulted-per-user-oauth"), + pytest.param(_server(MCPAuth.true_passthrough), id="true-passthrough"), + pytest.param(_server(MCPAuth.oauth_delegate), id="oauth-delegate"), +] +_NOT_REAUTHORIZABLE: Final = [ + pytest.param(_server(MCPAuth.none), id="no-auth"), + pytest.param(_STATIC, id="static-api-key"), + pytest.param(_server(MCPAuth.bearer_token, authentication_token="k"), id="static-bearer"), + pytest.param(_server(MCPAuth.oauth2, oauth2_flow="client_credentials"), id="client-credentials"), + pytest.param(_server(MCPAuth.oauth2, delegate_auth_to_upstream=True), id="legacy-upstream-delegate"), +] + + +def _request_context(gate: OAuthScopeResponse, method: str) -> ServerRequestContext: + return ServerRequestContext( + session=MagicMock(), + lifespan_context={}, + protocol_version="2025-11-25", + method=method, + request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}), + ) + + +def _challenged( + server: MCPServer, + outcome: ServerListFault | ServerListOk, + listed: Sequence[MCPServer] | None = None, + reauthorize_as: str | None = "docs", + method: str = "tools/list", +) -> bool: + gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as=reauthorize_as) + token: Final = active_mcp_request_ctx_var.set(_request_context(gate, method)) + try: + catalog.challenge_rejected_credential(server, outcome, [server] if listed is None else listed) + finally: + active_mcp_request_ctx_var.reset(token) + return not gate.pending + + +@pytest.mark.parametrize("server", _REAUTHORIZABLE) +def test_a_rejected_credential_the_caller_can_renew_is_challenged(server: MCPServer) -> None: + assert _challenged(server, _UNAUTHORIZED) + + +@pytest.mark.parametrize("server", _NOT_REAUTHORIZABLE) +def test_a_rejected_credential_the_caller_cannot_renew_stays_absorbed(server: MCPServer) -> None: + assert not _challenged(server, _UNAUTHORIZED) + + +@pytest.mark.parametrize( + "outcome", + [ + ServerListOk(tool_count=3), + ServerListFault(tag="forbidden", status_code=403), + ServerListFault(tag="upstream_error", status_code=500), + ServerListFault(tag="unreachable"), + ], +) +def test_only_an_unauthorized_outcome_is_challenged(outcome: ServerListFault | ServerListOk) -> None: + assert not _challenged(_VAULTED, outcome) + + +def test_a_listing_that_spans_several_servers_is_not_challenged() -> None: + assert not _challenged(_VAULTED, _UNAUTHORIZED, [_VAULTED, _PEER]) + + +def test_a_request_that_is_not_a_single_server_listing_is_not_challenged() -> None: + assert not _challenged(_VAULTED, _UNAUTHORIZED, reauthorize_as=None) + + +@pytest.mark.parametrize("method", ["tools/call", "server/discover", "initialize"]) +def test_a_request_the_sdk_parsed_as_something_other_than_a_listing_is_not_challenged(method: str) -> None: + assert not _challenged(_VAULTED, _UNAUTHORIZED, method=method) + + +@pytest.mark.asyncio +async def test_a_rejected_stored_token_is_answered_with_the_gateways_own_challenge() -> None: + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, _GATEWAY, "/docs/mcp", reauthorize_as="docs") + token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list")) + try: + started: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) + catalog.challenge_rejected_credential(_VAULTED, _UNAUTHORIZED, [_VAULTED]) + await asyncio.wait_for(started, 1) + finally: + active_mcp_request_ctx_var.reset(token) + + start: Final = send.await_args_list[0].args[0] + assert start["status"] == 401 + assert dict(start["headers"])[b"www-authenticate"].decode() == ( + 'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp"' + ) + + +def _upstream_rejecting_the_token(_request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(401, headers={"www-authenticate": 'Bearer error="invalid_token"'}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("server", "peers", "challenged"), + [ + pytest.param(_server(MCPAuth.true_passthrough), (), True, id="true-passthrough"), + pytest.param(_server(MCPAuth.oauth_delegate), (), True, id="oauth-delegate"), + pytest.param(_STATIC, (), False, id="static-api-key"), + pytest.param(_server(MCPAuth.true_passthrough), (_PEER,), False, id="selected-group-of-two"), + ], +) +async def test_tools_listing_challenges_a_single_server_whose_token_the_upstream_rejected( + monkeypatch: pytest.MonkeyPatch, server: MCPServer, peers: Sequence[MCPServer], challenged: bool +) -> None: + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(_upstream_rejecting_the_token)) + monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager) + gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs") + context: Final = mcp_operations.prepare_context( + UserAPIKeyAuth(token="h", user_id="u1"), oauth2_headers={"Authorization": "Bearer rejected"} + ) + request_token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list")) + name_token: Final = mcp_gateway_server_name.set("docs") + try: + listing, outcome = await catalog.get_filtered_server_tools( + server, context=context, allowed_mcp_servers=[server, *peers], prefetched_oauth_creds={} + ) + finally: + mcp_gateway_server_name.reset(name_token) + active_mcp_request_ctx_var.reset(request_token) + + assert listing.tools == [] + assert outcome.tag == "auth_required" + assert (not gate.pending) is challenged + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_name", ["ListPromptsRequest", "ListResourcesRequest", "ListResourceTemplatesRequest"]) +@pytest.mark.parametrize( + ("server", "challenged"), + [ + pytest.param(_VAULTED, True, id="vaulted-per-user-oauth"), + pytest.param(_STATIC, False, id="static-api-key"), + ], +) +async def test_optional_listing_challenges_a_server_whose_stored_token_was_rejected( + monkeypatch: pytest.MonkeyPatch, request_name: str, server: MCPServer, challenged: bool +) -> None: + from mcp import types + + from litellm.proxy import proxy_server + + monkeypatch.setenv("LITELLM_SALT_KEY", "listing-reauth-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(mcp_operations.global_mcp_server_manager, "registry", {server.server_id: server}) + monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])) + monkeypatch.setattr( + catalog, "fetch_optional_catalog_page", AsyncMock(side_effect=MCPUpstreamAuthError(401, None, "docs")) + ) + gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs") + token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "prompts/list")) + try: + result: Final = await catalog.list_gateway_catalog( + mcp_operations.prepare_context(), getattr(types, request_name)() + ) + finally: + active_mcp_request_ctx_var.reset(token) + + assert next(iter(result.meta[SERVER_OUTCOMES_META_KEY].values())) == {"status": "auth_required", "http_status": 401} + assert (not gate.pending) is challenged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_name", "result_name", "items"), + [ + ("ListPromptsRequest", "ListPromptsResult", "prompts"), + ("ListResourcesRequest", "ListResourcesResult", "resources"), + ("ListResourceTemplatesRequest", "ListResourceTemplatesResult", "resource_templates"), + ], +) +async def test_optional_listing_challenges_when_the_token_is_rejected_on_a_later_page( + monkeypatch: pytest.MonkeyPatch, request_name: str, result_name: str, items: str +) -> None: + from mcp import types + + from litellm.proxy import proxy_server + + monkeypatch.setenv("LITELLM_SALT_KEY", "listing-reauth-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(mcp_operations.global_mcp_server_manager, "registry", {_VAULTED.server_id: _VAULTED}) + monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[_VAULTED])) + first_page: Final = getattr(types, result_name)(**{items: [], "next_cursor": "next"}) + monkeypatch.setattr( + catalog, + "fetch_optional_catalog_page", + AsyncMock(side_effect=[first_page, MCPUpstreamAuthError(401, None, "docs")]), + ) + gate: Final = OAuthScopeResponse(AsyncMock(), _GATEWAY, "/docs/mcp", reauthorize_as="docs") + token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "prompts/list")) + request: Final = getattr(types, request_name) + try: + first: Final = await catalog.list_gateway_catalog(mcp_operations.prepare_context(), request()) + assert first.next_cursor + assert gate.pending + with pytest.raises(MCPError, match="Upstream continuation failed"): + await catalog.list_gateway_catalog( + mcp_operations.prepare_context(), + request(params=types.PaginatedRequestParams(cursor=first.next_cursor)), + ) + finally: + active_mcp_request_ctx_var.reset(token) + + assert not gate.pending + + +@pytest.mark.parametrize( + ("request_path", "resource"), + [("/mcp/docs", "litellm/mcp/docs"), ("/docs/mcp", "litellm/docs/mcp")], +) +def test_the_challenge_names_the_metadata_route_under_a_server_root_path( + monkeypatch: pytest.MonkeyPatch, request_path: str, resource: str +) -> None: + monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") + + challenge: Final = MCPUpstreamAuthError(401, None, "docs").to_http_exception(_GATEWAY, request_path) + + assert (challenge.headers or {})["www-authenticate"] == ( + f'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/{resource}"' + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_path", "addressed", "resource"), + [ + pytest.param("/docs/mcp", "docs", "docs/mcp", id="per-server-path"), + pytest.param("/Docs/mcp", "Docs", "Docs/mcp", id="other-case"), + pytest.param("/docs_internal/mcp", "docs_internal", "docs_internal/mcp", id="name-not-alias"), + pytest.param("/mcp/Docs", "Docs", "mcp/Docs", id="standard-path-other-case"), + pytest.param("/mcp/a:b@c+d,e", "a:b@c+d,e", "mcp/a:b@c+d,e", id="path-legal-characters-stay-as-typed"), + pytest.param("/mcp/a b%c", "a b%c", "mcp/a%20b%25c", id="path-illegal-characters-are-encoded"), + pytest.param("/mcp/日本", "日本", "mcp/%E6%97%A5%E6%9C%AC", id="non-ascii-name"), + pytest.param("/日本/mcp", "日本", "%E6%97%A5%E6%9C%AC/mcp", id="non-ascii-name-legacy-path"), + ], +) +async def test_the_challenge_names_the_url_the_client_addressed(request_path: str, addressed: str, resource: str) -> None: + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, _GATEWAY, request_path, reauthorize_as=addressed) + aliased: Final = _server(MCPAuth.oauth2, name="docs_internal", oauth2_flow="authorization_code", alias="docs") + token: Final = active_mcp_request_ctx_var.set(_request_context(gate, "tools/list")) + try: + started: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) + catalog.challenge_rejected_credential(aliased, _UNAUTHORIZED, [aliased]) + await asyncio.wait_for(started, 1) + finally: + active_mcp_request_ctx_var.reset(token) + + start: Final = send.await_args_list[0].args[0] + assert dict(start["headers"])[b"www-authenticate"].decode() == ( + f'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/{resource}"' + ) diff --git a/tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py index 03a7373c908..5054af4a2d4 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -25,6 +25,7 @@ from litellm.proxy.mcp_server.mcp_server_manager import ( MCPServerManager, _extract_upstream_auth_failure, ) +from litellm.proxy.mcp_server.scope_challenge import OAuthScopeResponse, is_listing_request from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -375,12 +376,10 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server(): @pytest.mark.asyncio -async def test_single_server_route_also_absorbs_upstream_auth_error(): +async def test_single_server_route_also_absorbs_upstream_auth_error() -> None: """A single-server route (//mcp) absorbs an upstream-auth error just like the aggregate: the failing server contributes no tools and an "auth_required" outcome rather than re-raising. - Surfacing it to the client as a 401 + WWW-Authenticate challenge cannot be done from this list - handler — the MCP session manager serializes a raise into a JSON-RPC error, not an HTTP 401 — so - re-auth surfacing is handled by a request-scope preemptive check, tracked separately.""" + The 401 a single-server request receives is sent through its scope response, not raised from here.""" from unittest.mock import patch from litellm.proxy._types import UserAPIKeyAuth @@ -836,3 +835,73 @@ def test_managed_step_up_metadata_matches_the_sdk_configured_resource(path): requested_resource=f"https://gateway.test{path}", configured_resource=f"https://gateway.test{resource_path}", ), "the SDK must accept the advertised resource before it can reauthorize" + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + (b'{"jsonrpc":"2.0","id":1,"method":"tools/list"}', True), + (b'{"jsonrpc":"2.0","id":1,"method":"prompts/list"}', True), + (b'{"jsonrpc":"2.0","id":1,"method":"resources/list"}', True), + (b'{"jsonrpc":"2.0","id":1,"method":"resources/templates/list"}', True), + (b'{"jsonrpc":"2.0","id":1,"method":"tools/call"}', False), + (b'{"jsonrpc":"2.0","id":1,"method":"initialize"}', False), + (b'{"jsonrpc":"2.0","id":1,"method":["tools/list"]}', False), + (b'{"jsonrpc":"2.0","id":1}', False), + (b"not json", False), + ], +) +def test_listing_requests_are_recognised_by_method(body: bytes, expected: bool) -> None: + assert is_listing_request(body) is expected + + +async def _pending_listing_response(gate: OAuthScopeResponse, send: AsyncMock) -> "asyncio.Task[None]": + task: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) + send.assert_not_awaited() + return task + + +@pytest.mark.asyncio +async def test_reauth_denial_replaces_the_pending_listing_response_with_the_gateway_challenge() -> None: + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp", reauthorize_as="docs") + started: Final = await _pending_listing_response(gate, send) + + gate.deny_reauth() + await asyncio.wait_for(started, 1) + + start, body = [call.args[0] for call in send.await_args_list] + assert start["status"] == 401 + assert dict(start["headers"])[b"www-authenticate"].decode() == ( + 'Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp"' + ) + assert body["body"] == b'{"detail":"Unauthorized"}' + + +@pytest.mark.asyncio +async def test_a_gate_that_does_not_reauthorize_ignores_a_reauth_denial() -> None: + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") + started: Final = await _pending_listing_response(gate, send) + + gate.deny_reauth() + assert gate.pending + gate.allow() + await asyncio.wait_for(started, 1) + + assert [call.args[0]["status"] for call in send.await_args_list] == [200] + + +@pytest.mark.asyncio +async def test_scope_step_up_denial_still_ignores_an_unauthorized_response() -> None: + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp", reauthorize_as="docs") + started: Final = await _pending_listing_response(gate, send) + + gate.deny(MCPUpstreamAuthError(401, None, "docs")) + assert gate.pending + gate.allow() + await asyncio.wait_for(started, 1) + + assert [call.args[0]["status"] for call in send.await_args_list] == [200] diff --git a/tests/unit/proxy/mcp_server/test_mcp_server.py b/tests/unit/proxy/mcp_server/test_mcp_server.py index 456b1994e2d..b42f913f078 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server.py @@ -2,6 +2,8 @@ import asyncio import json import os +import queue +from collections.abc import Sequence from contextlib import asynccontextmanager, nullcontext from typing import Final from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -17,6 +19,7 @@ from litellm.proxy.mcp_server.mcp_server_manager import ( MCPServerManager, MCPTransport, ) +from litellm.types.mcp import MCPAuth mcp_server_manager = MCPServerManager() @@ -3209,6 +3212,140 @@ async def test_mcp_routing_stashes_only_method_for_auth_and_replays_full_body(me ) +def _docs_server(auth_type: str, legacy_client_credentials: bool = False) -> MCPServer: + return MCPServer( + server_id="docs-id", + name="docs", + alias="docs", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth(auth_type), + token_url="https://idp.example/token" if legacy_client_credentials else None, + client_id="client" if legacy_client_credentials else None, + client_secret="secret" if legacy_client_credentials else None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "method", "shape", "selected_servers", "docs", "expected"), + [ + ("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/docs", "prompts/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/docs", "resources/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/docs", "resources/templates/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/DOCS", "tools/list", "plain", ["DOCS"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/DOCS"), + ("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("api_key"), "none"), + ("/mcp/docs", "tools/list", "plain", ["docs"], _docs_server("oauth2", legacy_client_credentials=True), "none"), + ("/mcp/docs", "tools/call", "plain", ["docs"], _docs_server("oauth_delegate"), "holds"), + ("/mcp/docs", "tools/list", "plain", ["docs", "wiki"], _docs_server("oauth_delegate"), "none"), + ("/mcp/docs", "tools/list", "plain", ["unregistered_group"], _docs_server("oauth_delegate"), "none"), + ("/mcp/docs", "tools/list", "plain", None, _docs_server("oauth_delegate"), "none"), + ("/mcp/docs", "ping", "plain", ["docs"], _docs_server("oauth_delegate"), "none"), + ("/docs/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/docs/mcp"), + ("/Docs/mcp", "tools/list", "plain", ["Docs"], _docs_server("oauth_delegate"), "reauthorizes:/Docs/mcp"), + ("/mcp/docs-id", "tools/list", "plain", ["docs-id"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs-id"), + ("/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "none"), + ("/toolset/team/mcp", "tools/list", "plain", ["docs"], _docs_server("oauth_delegate"), "none"), + ("/mcp/docs", "tools/list", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/docs", "prompts/list", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "reauthorizes:/mcp/docs"), + ("/mcp/docs", "tools/list", "cursor-before-method", ["docs"], _docs_server("oauth_delegate"), "holds"), + ("/mcp/docs", "tools/call", "cursor-after-method", ["docs"], _docs_server("oauth_delegate"), "holds"), + ("/mcp/docs", "ping", "duplicate-method", ["docs"], _docs_server("oauth_delegate"), "none"), + ], +) +async def test_mcp_routing_holds_a_single_servers_listing_response_for_a_reauth_challenge( + path: str, method: str, shape: str, selected_servers: Sequence[str] | None, docs: MCPServer, expected: str +) -> None: + try: + from litellm.proxy.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse + from litellm.proxy.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + async def gate_kind(gate: object) -> str: + if not isinstance(gate, OAuthScopeResponse): + return "none" + gate.deny_reauth() + if gate.pending: + return "holds" + await gate({"type": "http.response.start", "status": 200, "headers": []}) + challenge: Final = dict(sent.await_args_list[0].args[0]["headers"])[b"www-authenticate"].decode() + return "reauthorizes:" + challenge.split("oauth-protected-resource")[1].rstrip('"') + + from litellm.proxy.mcp_server import operations + + envelope: Final = {"jsonrpc": "2.0", "id": 2, "method": method, "params": {"cursor": "c" * 6000}} + bodies: Final = { + "plain": json.dumps({**envelope, "params": {}}).encode(), + "cursor-after-method": json.dumps(envelope).encode(), + "cursor-before-method": json.dumps({"params": envelope["params"], **envelope}).encode(), + "duplicate-method": b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{},"method":"ping"}', + } + jsonrpc_body: Final = bodies[shape] + scope: Final = { + "type": "http", + "method": "POST", + "path": path, + "scheme": "https", + "server": ("gateway.example", 443), + "headers": [ + (b"host", b"gateway.example"), + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": jsonrpc_body, "more_body": False}]) + held: Final = queue.SimpleQueue[object]() + sent: Final = AsyncMock() + + async def stateless_handle(handled_scope: Scope, _receive: Receive, _send: Send) -> None: + held.put(handled_scope.get(SCOPE_RESPONSE_KEY)) + + with ( + patch( # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam + "litellm.proxy.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, selected_servers, None, None, None), + ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), + patch( # test-quality-ok: init flag is a module global; no injection seam + "litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam + session_manager_stateless, "handle_request", side_effect=stateless_handle + ), + patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam + session_manager_stateful, "handle_request", side_effect=stateless_handle + ), + patch.object( # test-quality-ok: the registry is a module-level singleton's attribute; the suite's only seam + operations.global_mcp_server_manager, "registry", {docs.server_id: docs} + ), + patch( # test-quality-ok: the missing-token challenge runs before the gate and is not under test here + "litellm.proxy.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam + session_manager_stateless, "_server_instances", {} + ), + patch.object( # test-quality-ok: session managers are module-level singletons; the suite's only seam + session_manager_stateful, "_server_instances", {} + ), + ): + await handle_streamable_http_mcp(scope, receive, sent) + + assert await gate_kind(held.get_nowait()) == expected + + @pytest.mark.asyncio async def test_mcp_routing_batch_body_is_not_stashed_for_auth(): try: diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index 621a7b65005..4b08fa918cf 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -4190,6 +4190,29 @@ def test_jsonrpc_text_has_top_level_method_ignores_nested_method(): assert _jsonrpc_text_has_top_level_method(truncated) is False +@pytest.mark.parametrize( + ("text", "method"), + [ + pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', "tools/list", id="complete"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="cut-after-method"), + pytest.param('{"jsonrpc":"2.0","params":{"cursor":"' + "c" * 5000, None, id="cut-before-method"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools/li', None, id="cut-inside-method"), + pytest.param('{"jsonrpc":"2.0","params":{"method":"tools/list"},"method":"ping"}', "ping", id="nested-method-first"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":"tools\\/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="escaped-slash"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":"\\u0074ools/list","params":{"cursor":"' + "c" * 5000, "tools/list", id="unicode-escape"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":"bad\\x","params":{}}', None, id="invalid-escape"), + pytest.param('{"jsonrpc":"2.0","id":1,"method":7,"params":{}}', None, id="not-a-string"), + pytest.param('{"jsonrpc":"2.0","id":1,"result":{"method":"tools/list"}}', None, id="response"), + ], +) +def test_jsonrpc_text_method_reads_the_top_level_method_when_the_cut_leaves_it_whole( + text: str, method: str | None +) -> None: + from litellm.proxy.mcp_server.server import _jsonrpc_text_method + + assert _jsonrpc_text_method(text) == method + + @pytest.mark.asyncio @pytest.mark.parametrize("response_field", ("result", "error")) @pytest.mark.parametrize(("character", "bytes_before_cap"), (("", 0), ("x", 0), ("é", 1), ("中", 2), ("😀", 3)))