Merge pull request #40454 from BerriAI/litellm_fix_mcp_debug_auth_resolution_6892

fix(mcp): report resolved upstream authentication in debug headers
This commit is contained in:
joshua-berri 2026-09-09 21:51:22 -07:00 committed by GitHub
commit 17e9b2329e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 564 additions and 216 deletions

View file

@ -22,7 +22,16 @@ Response headers returned (all values are masked for safety):
x-mcp-debug-auth-resolution
Which auth priority was used for the outbound MCP call:
``per-request-header``, ``m2m-client-credentials``, ``static-token``,
``oauth2-passthrough``, or ``no-auth``.
``oauth2-passthrough``, ``stored-user-token``, ``token-exchange``,
``id-jag``, ``aws-sigv4``, ``extra-headers``, or ``no-auth``.
``unresolved`` means no outcome was available before the first response
frame; ``multiple`` means several servers resolved credentials;
``not-applicable`` covers stdio; ``resolution-failed`` is a resolver error.
x-mcp-debug-auth-resolutions
For multiple servers, a JSON map of server IDs to resolution labels.
At most 32 entries are included; x-mcp-debug-auth-resolutions-truncated
is true when additional servers were omitted. No credentials are included.
x-mcp-debug-outbound-url
The upstream MCP server URL that will receive the request.
@ -58,10 +67,16 @@ header is free for OAuth2 discovery::
Symptom: ``x-mcp-debug-oauth2-token`` shows ``(none)`` and
``x-mcp-debug-auth-resolution`` shows ``no-auth``.
This means the client didn't go through the OAuth2 flow. Check that:
1. The ``Authorization`` header is NOT set as a static header in the client config.
2. The ``.well-known/oauth-protected-resource`` endpoint returns valid metadata.
3. The MCP server in LiteLLM config has ``auth_type: oauth2``.
``no-auth`` means the resolved upstream client carries no authentication.
An absent inbound OAuth2 token does not imply the user skipped OAuth: the gateway
can retrieve a stored per-user token, reported as ``stored-user-token``.
``unresolved`` is used when a stream starts before credential resolution, or a
request (such as initialization or a cached tool listing) resolves no credential.
Debug reporting does not fetch credentials or delay a streaming frame to resolve them.
``extra-headers`` identifies supplied headers that won over the resolver or were
the only headers supplied; their values are never inspected to guess a scheme.
``per-request-header`` denotes a legacy credential override, including a BYOK
credential supplied by the gateway; it does not imply a caller-supplied token.
**Common issue: M2M token used instead of user token**
@ -69,8 +84,8 @@ Symptom: ``x-mcp-debug-auth-resolution`` shows ``m2m-client-credentials``.
This means the server has ``client_id``/``client_secret``/``token_url``
configured and LiteLLM is fetching a machine-to-machine token instead of
using the per-user OAuth2 token. If you want per-user tokens, remove the
client credentials from the server config.
using the per-user OAuth2 token. For gateway-stored per-user tokens,
configure ``oauth2_flow: authorization_code``.
Usage from Claude Code::
@ -85,14 +100,16 @@ Usage with curl::
http://localhost:4000/mcp/atlassian_mcp
"""
from typing import TYPE_CHECKING, Final
import json
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final
from starlette.requests import HTTPConnection
from starlette.types import Message, Send
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
# Header the client sends to opt into debug mode
MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
@ -101,6 +118,83 @@ MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
_RESPONSE_HEADER_PREFIX: Final = "x-mcp-debug"
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics"
def record_auth_resolution(server_id: str, source: AuthResolution) -> None:
from mcp.server.lowlevel.server import request_ctx
context: Final[object] = request_ctx.get(None)
request: Final[object] = getattr(context, "request", None)
if isinstance(request, HTTPConnection):
diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY)
if isinstance(diagnostics, MCPAuthDiagnostics):
diagnostics.record(server_id, source)
class MCPAuthDiagnostics:
def __init__(self) -> None:
self._outcomes: tuple[tuple[str, AuthResolution], ...] = ()
def record(self, server_id: str, resolution: AuthResolution) -> None:
self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),)
def resolution(self) -> str:
match self._outcomes:
case ():
return AuthResolution.unresolved.value
case ((_, source),):
return source.value
case _:
return AuthResolution.multiple.value
def headers(self) -> Mapping[str, str]:
if len(self._outcomes) <= 1:
return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()})
return MappingProxyType(
{
"x-mcp-debug-auth-resolution": AuthResolution.multiple.value,
"x-mcp-debug-auth-resolutions": json.dumps(
{
server_id: source.value for server_id, source in self._outcomes[:32]
}, # mutable-ok: JSON encoder requires a concrete dict
separators=(",", ":"),
ensure_ascii=True,
),
**(
MappingProxyType({"x-mcp-debug-auth-resolutions-truncated": "true"})
if len(self._outcomes) > 32
else MappingProxyType({})
),
}
)
class _DiagnosticSend:
def __init__(self, send: Send, headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]]) -> None:
self._send = send
self._headers = headers
self._resolution = resolution
self._start: Message | None = None
async def __call__(self, message: Message) -> None:
if message["type"] == "http.response.start":
self._start = message
return
if self._start is not None:
start: Final = self._start
self._start = None
headers: Final = MappingProxyType({**self._headers, **self._resolution()})
await self._send(
{ # mutable-ok: ASGI send consumes a mutable message mapping
**start,
"headers": tuple(start.get("headers", ()))
+ tuple((key.encode(), value.encode()) for key, value in headers.items()),
}
)
await self._send(message)
class MCPDebug:
"""
Static helper class for MCP OAuth2 debug headers.
@ -144,37 +238,6 @@ class MCPDebug:
return val.strip().lower() in ("true", "1", "yes")
return False
@staticmethod
def resolve_auth_resolution(
server: "MCPServer",
mcp_auth_header: str | None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
oauth2_headers: dict[str, str] | None,
) -> str:
"""
Determine which auth priority will be used for the outbound MCP call.
Returns one of: ``per-request-header``, ``m2m-client-credentials``,
``static-token``, ``oauth2-passthrough``, or ``no-auth``.
"""
from litellm.types.mcp import MCPAuth
has_server_specific: Final = bool(
mcp_server_auth_headers
and (
mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "")
)
)
if has_server_specific or mcp_auth_header:
return "per-request-header"
if server.has_client_credentials:
return "m2m-client-credentials"
if server.authentication_token:
return "static-token"
if oauth2_headers and server.auth_type == MCPAuth.oauth2:
return "oauth2-passthrough"
return "no-auth"
@staticmethod
def build_debug_headers(
*,
@ -244,12 +307,21 @@ class MCPDebug:
return debug
@staticmethod
def wrap_send_with_debug_headers(send: Send, debug_headers: dict[str, str]) -> Send:
def wrap_send_with_debug_headers(
send: Send,
debug_headers: Mapping[str, str],
resolution: Callable[[], Mapping[str, str]] | None = None,
*,
request_method: str | None = None,
) -> Send:
"""
Return a new ASGI ``send`` callable that injects *debug_headers*
into the ``http.response.start`` message.
"""
if resolution is not None and request_method == "POST":
return _DiagnosticSend(send, debug_headers, resolution)
async def _send_with_debug(message: Message) -> None:
if message["type"] == "http.response.start":
headers: Final = list(message.get("headers", []))
@ -266,8 +338,6 @@ class MCPDebug:
raw_headers: dict[str, str] | None,
scope: dict,
mcp_servers: list[str] | None,
mcp_auth_header: str | None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
oauth2_headers: dict[str, str] | None,
client_ip: str | None,
) -> dict[str, str]:
@ -288,16 +358,13 @@ class MCPDebug:
server_url: str | None = None
server_auth_type: str | None = None
auth_resolution = "no-auth"
auth_resolution: Final = AuthResolution.unresolved.value
for server_name in mcp_servers or []:
server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip)
if server:
server_url = server.url
server_auth_type = server.auth_type
auth_resolution = MCPDebug.resolve_auth_resolution(
server, mcp_auth_header, mcp_server_auth_headers, oauth2_headers
)
break
scope_headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope)

View file

@ -80,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
raise_classified_list_failure,
upstream_auth_challenge,
)
from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
MCPPerUserTokenCache,
mcp_per_user_token_cache,
@ -108,12 +109,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import (
LazyPerUserOAuthTokenStore,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import (
build_token_exchanger,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
DEFAULT_CREDENTIAL_HEADER,
AuthorizationCodeConfig,
AuthResolution,
ClientCredentialsConfig,
CredError,
IdJagConfig,
@ -3832,13 +3835,21 @@ class MCPServerManager:
(authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any
other ``CredError`` onto its public HTTP status; it never returns an error as a value.
"""
match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(auth):
match await resolve_credentials_with_source(provider, to_subject(user_api_key_auth, subject_token), spec):
case Ok(credential):
auth: Final = credential.auth
# NoOpAuth has no header_name and so never conflicts.
header_name: Final[str | None] = getattr(auth, "header_name", None)
if header_name is None or not extra_headers:
source: Final = (
AuthResolution.extra_headers
if credential.source == AuthResolution.no_auth and extra_headers
else credential.source
)
record_auth_resolution(server.server_id, source)
return auth, extra_headers
if not has_header(extra_headers, header_name):
record_auth_resolution(server.server_id, credential.source)
return auth, extra_headers
if isinstance(
spec.config,
@ -3853,11 +3864,14 @@ class MCPServerManager:
# one-shot 401 refetch is lost with it). Drop only the header the resolved
# credential is about to occupy, so a static credential the operator aimed at a
# DIFFERENT header still reaches upstream.
record_auth_resolution(server.server_id, credential.source)
return auth, without_header(extra_headers, header_name)
# Other modes: an Authorization already supplied via extra_headers (a forwarded caller
# header or static_headers) is intentional and wins; v1 applies those last.
record_auth_resolution(server.server_id, AuthResolution.extra_headers)
return None, extra_headers
case Error(err):
record_auth_resolution(server.server_id, AuthResolution.failed)
if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig):
# authorization_code's missing per-user token -> the per-server browser-OAuth
# challenge, built here where the full MCPServer is in hand.
@ -3960,6 +3974,7 @@ class MCPServerManager:
Returns:
Configured MCP client instance.
"""
record_auth_resolution(server.server_id, AuthResolution.unresolved)
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
transport: Final = resolved_server.transport or MCPTransport.sse
spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server)
@ -4032,6 +4047,7 @@ class MCPServerManager:
env=resolved_env,
)
record_auth_resolution(server.server_id, AuthResolution.not_applicable)
return MCPClient(
server_url="", # Not used for stdio
transport_type=transport,
@ -4086,6 +4102,20 @@ class MCPServerManager:
aws_session_name=resolved_server.aws_session_name,
)
legacy_source: Final = (
AuthResolution.aws_sigv4
if aws_auth is not None
else AuthResolution.extra_headers
if extra_headers and has_header(extra_headers, auth_header_name or "Authorization")
else AuthResolution.per_request_header
if mcp_auth_header
else AuthResolution.static_token
if auth_value
else AuthResolution.extra_headers
if extra_headers
else AuthResolution.no_auth
)
record_auth_resolution(server.server_id, legacy_source)
return MCPClient(
server_url=server_url,
transport_type=transport,

View file

@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ApiKeyConfig,
AuthorizationCodeConfig,
AuthResolution,
AuthSpecKind,
AwsSigV4Config,
Byok,
@ -76,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
NoneConfig,
PassthroughConfig,
PrivateKeyJwtAuth,
ResolvedCredential,
ServerSpec,
SharedKey,
Subject,
@ -448,3 +450,32 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str:
def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet"))
async def resolve_credentials_with_source(
provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec
) -> Result[ResolvedCredential, CredError]:
match await provider.resolve_credentials(subject, server):
case Error(err):
return Error(err)
case Ok(auth):
if isinstance(auth, NoOpAuth):
return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
match server.config:
case NoneConfig():
return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
case ApiKeyConfig():
return Ok(ResolvedCredential(auth, AuthResolution.static_token))
case PassthroughConfig():
return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough))
case ClientCredentialsConfig():
return Ok(ResolvedCredential(auth, AuthResolution.client_credentials))
case TokenExchangeConfig():
return Ok(ResolvedCredential(auth, AuthResolution.token_exchange))
case IdJagConfig():
return Ok(ResolvedCredential(auth, AuthResolution.id_jag))
case AuthorizationCodeConfig():
return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token))
case AwsSigV4Config():
return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4))
assert_never(server.config)

View file

@ -26,10 +26,11 @@ union (see `result.py`), not `expression.Result`.
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from dataclasses import dataclass, field
from enum import Enum
from typing import Annotated, Final, Literal
import httpx
from expression import case, tag, tagged_union
from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
from typing_extensions import assert_never
@ -46,6 +47,29 @@ from litellm.types.mcp import (
)
class AuthResolution(str, Enum):
no_auth = "no-auth"
stored_user_token = "stored-user-token"
static_token = "static-token"
per_request_header = "per-request-header"
oauth2_passthrough = "oauth2-passthrough"
client_credentials = "m2m-client-credentials"
token_exchange = "token-exchange"
id_jag = "id-jag"
aws_sigv4 = "aws-sigv4"
extra_headers = "extra-headers"
not_applicable = "not-applicable"
unresolved = "unresolved"
failed = "resolution-failed"
multiple = "multiple"
@dataclass(frozen=True, slots=True)
class ResolvedCredential:
auth: httpx.Auth = field(repr=False)
source: AuthResolution
class AuthSpecKind(str, Enum):
"""The server's statically-declared upstream-auth mode — derived from its `config`.

View file

@ -49,7 +49,11 @@ from litellm.proxy._experimental.mcp_server.mcp_context import (
_mcp_gateway_server_name,
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode
)
from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug
from litellm.proxy._experimental.mcp_server.mcp_debug import (
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY,
MCPAuthDiagnostics,
MCPDebug,
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
_redact_mcp_resource_url,
get_passthrough_www_authenticate,
@ -4472,13 +4476,15 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
scope=dict(scope),
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
client_ip=_client_ip,
)
if _debug_headers:
send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers)
diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None
if diagnostics is not None:
scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics
send = MCPDebug.wrap_send_with_debug_headers(
send, _debug_headers, diagnostics.headers, request_method=scope.get("method")
)
# Ensure session managers are initialized
if not _SESSION_MANAGERS_INITIALIZED:

View file

@ -19,6 +19,7 @@ from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.outbound_credentials import (
ApiKeyConfig,
AuthConfig,
AuthorizationCodeConfig,
AwsSigV4Config,
Byok,
@ -1203,3 +1204,71 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot():
assert isinstance(result, Ok)
headers, _ = await _emitted_async(result.ok)
assert headers["Authorization"] == "caller-token"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("config", "subject", "expected_source", "expected_header"),
[
(NoneConfig(), _SUBJECT, "no-auth", None),
(PassthroughConfig(), _SUBJECT, "no-auth", None),
(PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"),
(ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"),
(AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"),
],
)
async def test_resolved_source_matches_the_credential_sent_upstream(
config: AuthConfig, subject: Subject, expected_source: str, expected_header: str | None
) -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")})
provider = UpstreamCredentialProvider(oauth_token_store=store)
result = await resolve_credentials_with_source(provider, subject, _spec(config))
assert isinstance(result, Ok)
assert result.ok.source.value == expected_source
assert _emitted(result.ok.auth).get("Authorization") == expected_header
assert "stored-alice" not in repr(result.ok)
assert "static-key" not in repr(result.ok)
@pytest.mark.asyncio
async def test_resolved_source_preserves_missing_user_token_error() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig()))
assert isinstance(result, Error)
assert result.error.tag == "unauthorized"
@pytest.mark.asyncio
async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at")))
m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M))
assert isinstance(m2m, Ok)
headers, _ = await _emitted_async(m2m.ok.auth)
assert headers["Authorization"] == "Bearer m2m-at"
assert m2m.ok.source.value == "m2m-client-credentials"
assert source.gets == ["s"]
exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at")))
exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO))
assert isinstance(exchanged, Ok)
assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at"
assert exchanged.ok.source.value == "token-exchange"
assert len(exchanger.calls) == 1
@pytest.mark.asyncio
async def test_id_jag_source_describes_final_token_after_both_exchanges() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token"))
provider = UpstreamCredentialProvider(token_endpoint=endpoint)
result = await resolve_credentials_with_source(provider, _with_inbound("identity-token"), _spec(_id_jag_config()))
assert isinstance(result, Ok)
assert result.ok.source.value == "id-jag"
assert _emitted(result.ok.auth)["Authorization"] == "Bearer resource-token"
assert len(endpoint.calls) == 2

View file

@ -3,11 +3,17 @@ Tests for MCPDebug — MCP OAuth2 debug response headers.
"""
import asyncio
from unittest.mock import MagicMock
from typing import Final
import pytest
from starlette.types import Message
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
from litellm.proxy._experimental.mcp_server.mcp_debug import (
MCP_DEBUG_REQUEST_HEADER,
MCPDebug,
MCPAuthDiagnostics,
)
@ -166,77 +172,6 @@ class TestBuildDebugHeaders:
assert set(headers.keys()) == expected_keys
class TestResolveAuthResolution:
def _make_server(self, **kwargs):
server = MagicMock()
server.alias = kwargs.get("alias", "test")
server.server_name = kwargs.get("server_name", "test")
server.has_client_credentials = kwargs.get("has_client_credentials", False)
server.authentication_token = kwargs.get("authentication_token", None)
server.auth_type = kwargs.get("auth_type", None)
return server
def test_per_request_header(self):
server = self._make_server()
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header="Bearer xxx",
mcp_server_auth_headers=None,
oauth2_headers=None,
)
assert result == "per-request-header"
def test_server_specific_header(self):
server = self._make_server(alias="atlas")
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header=None,
mcp_server_auth_headers={"atlas": {"Authorization": "Bearer xxx"}},
oauth2_headers=None,
)
assert result == "per-request-header"
def test_m2m(self):
server = self._make_server(has_client_credentials=True)
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
)
assert result == "m2m-client-credentials"
def test_static_token(self):
server = self._make_server(authentication_token="static-tok")
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
)
assert result == "static-token"
def test_oauth2_passthrough(self):
server = self._make_server(auth_type="oauth2")
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer eyJ..."},
)
assert result == "oauth2-passthrough"
def test_no_auth(self):
server = self._make_server()
result = MCPDebug.resolve_auth_resolution(
server,
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
)
assert result == "no-auth"
class TestWrapSendWithDebugHeaders:
def test_injects_headers(self):
captured = []
@ -269,3 +204,90 @@ class TestWrapSendWithDebugHeaders:
asyncio.run(wrapped(body_msg))
assert captured[0] == body_msg
@pytest.mark.asyncio
@pytest.mark.parametrize("source", tuple(AuthResolution))
@pytest.mark.parametrize("method", ("GET", "DELETE", "POST"))
async def test_debug_defers_resolution_until_first_frame_only_for_post(source: AuthResolution, method: str) -> None:
captured: Final[list[Message]] = []
diagnostics: Final = MCPAuthDiagnostics()
async def send(message: Message) -> None:
captured.append(message)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(
send, diagnostics.headers(), diagnostics.headers, request_method=method
)
await wrapped({"type": "http.response.start", "status": 200, "headers": []})
assert len(captured) == (0 if method == "POST" else 1)
diagnostics.record("s1", source)
body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True}
await wrapped(body)
assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == (
source.value.encode() if method == "POST" else b"unresolved"
)
assert captured[1] == body
@pytest.mark.asyncio
async def test_early_stream_frame_reports_unresolved_without_waiting() -> None:
captured: Final[list[Message]] = []
diagnostics: Final = MCPAuthDiagnostics()
async def send(message: Message) -> None:
captured.append(message)
wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers, request_method="POST")
await wrapped({"type": "http.response.start", "status": 200, "headers": []})
await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True})
diagnostics.record("s1", AuthResolution.stored_user_token)
await wrapped({"type": "http.response.body", "body": b"data: pong\n\n", "more_body": False})
assert len(captured) == 3
assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == b"unresolved"
def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers() -> None:
alice: Final = MCPAuthDiagnostics()
bob: Final = MCPAuthDiagnostics()
alice.record("s1", AuthResolution.stored_user_token)
assert bob.resolution() == "unresolved"
alice.record("s1", AuthResolution.token_exchange)
assert alice.resolution() == "token-exchange"
alice.record("s2", AuthResolution.static_token)
assert alice.resolution() == "multiple"
assert alice.headers()["x-mcp-debug-auth-resolutions"] == '{"s1":"token-exchange","s2":"static-token"}'
@pytest.mark.asyncio
async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None:
from unittest.mock import MagicMock
from mcp.server.lowlevel.server import request_ctx
from mcp.shared.context import RequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.mcp_debug import (
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY,
record_auth_resolution,
)
session: Final = MagicMock()
first: Final = MCPAuthDiagnostics()
second: Final = MCPAuthDiagnostics()
async def record(diagnostics: MCPAuthDiagnostics, source: AuthResolution) -> None:
context: Final = RequestContext(
request_id=1, meta=None, session=session, lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
)
token: Final = request_ctx.set(context)
try:
await asyncio.sleep(0)
record_auth_resolution("same-server", source)
finally:
request_ctx.reset(token)
await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header))
assert first.resolution() == "stored-user-token"
assert second.resolution() == "per-request-header"

View file

@ -3,6 +3,7 @@ import contextvars
import os
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -1866,96 +1867,85 @@ async def test_streamable_http_session_manager_is_stateless():
@pytest.mark.asyncio
async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless():
"""
Test that routing correctly sends:
- initialize (no mcp-session-id) stateful manager (so client gets mcp-session-id)
- tools/list (no mcp-session-id) stateless manager (curl, Inspector)
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
handle_streamable_http_mcp,
session_manager_stateful,
session_manager_stateless,
@pytest.mark.parametrize("debug", (False, True))
@pytest.mark.parametrize(
("method", "request_body", "stateful"),
(
("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True),
("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False),
("GET", b"", False),
("DELETE", b"", False),
),
)
async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(
debug: bool, method: str, request_body: bytes, stateful: bool
) -> None:
from mcp.server.lowlevel.server import request_ctx
from mcp.shared.context import RequestContext
from starlette.requests import Request
from starlette.types import Message, Receive, Scope, Send
from litellm.proxy._experimental.mcp_server import server as mcp_module
from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
scope: Final[Scope] = {"type": "http", "method": method, "path": "/mcp", "headers": []}
receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False})
send: Final = AsyncMock()
observe_start: Final = AsyncMock()
body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True}
async def handle_request(request_scope: Scope, receive: Receive, outgoing: Send) -> None:
await outgoing({"type": "http.response.start", "status": 200, "headers": []})
await observe_start(send.await_count)
context: Final = RequestContext(
request_id=1, meta=None, session=MagicMock(), lifespan_context=None, request=Request(request_scope)
)
except ImportError:
pytest.skip("MCP server not available")
token: Final = request_ctx.set(context)
try:
record_auth_resolution("s1", AuthResolution.stored_user_token)
finally:
request_ctx.reset(token)
await outgoing(body)
async def make_request(method_body: bytes, path: str = "/mcp/progress_test"):
scope = {
"type": "http",
"method": "POST",
"path": path,
"headers": [
(b"content-type", b"application/json"),
(b"authorization", b"Bearer test-key"),
],
}
receive = AsyncMock(
return_value={
"type": "http.request",
"body": method_body,
"more_body": False,
}
)
send = AsyncMock()
stateless_called = []
stateful_called = []
async def stateless_handle(s, r, se):
stateless_called.append(1)
async def stateful_handle(s, r, se):
stateful_called.append(1)
with (
patch(
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
new_callable=AsyncMock,
return_value=(MagicMock(), None, ["progress_test"], None, None, None),
stateless_handle: Final = AsyncMock(side_effect=handle_request)
stateful_handle: Final = AsyncMock(side_effect=handle_request)
with (
patch(
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
new_callable=AsyncMock,
return_value=(
UserAPIKeyAuth(user_id="debug-user"),
None,
None,
None,
None,
{"x-litellm-mcp-debug": "true"} if debug else {},
),
patch(
"litellm.proxy._experimental.mcp_server.server.set_auth_context",
),
patch(
"litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
True,
),
patch.object(
session_manager_stateless,
"handle_request",
side_effect=stateless_handle,
),
patch.object(
session_manager_stateful,
"handle_request",
side_effect=stateful_handle,
),
patch.object(
session_manager_stateless,
"_server_instances",
{},
),
patch.object(
session_manager_stateful,
"_server_instances",
{},
),
):
await handle_streamable_http_mcp(scope, receive, send)
),
patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True),
patch(
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
SimpleNamespace(handle_request=stateless_handle),
),
patch(
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
SimpleNamespace(handle_request=stateful_handle),
),
):
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
return bool(stateless_called), bool(stateful_called)
# initialize → stateful
init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}'
stateless_called, stateful_called = await make_request(init_body)
assert stateful_called and not stateless_called, "initialize (no session) should route to stateful, not stateless"
# tools/list → stateless
tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
stateless_called, stateful_called = await make_request(tools_body)
assert stateless_called and not stateful_called, "tools/list (no session) should route to stateless, not stateful"
assert stateful_handle.await_count == (1 if stateful else 0)
assert stateless_handle.await_count == (0 if stateful else 1)
observe_start.assert_awaited_once_with(0 if debug and method == "POST" else 1)
assert send.await_count == 2
assert send.call_args_list[0].args[0]["status"] == 200
assert send.call_args_list[1].args[0] == body
headers: Final = dict(send.call_args_list[0].args[0]["headers"])
if debug:
assert headers[b"x-mcp-debug-auth-resolution"] == (b"stored-user-token" if method == "POST" else b"unresolved")
else:
assert not any(name.startswith(b"x-mcp-debug") for name in headers)
@pytest.mark.asyncio

View file

@ -5,7 +5,7 @@ import logging
import os
import sys
from datetime import datetime
from typing import Any, Dict, Final, Optional
from typing import Any, Dict, Final, Literal, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -12567,3 +12567,112 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch,
with pytest.raises(HTTPException) as exc_info:
await call
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
@pytest.mark.parametrize(
("config", "extra_headers", "expected_source", "expected_authorization"),
[
("stored", None, "stored-user-token", "Bearer stored-token"),
("stored", {"aUtHoRiZaTiOn": "Bearer injected"}, "stored-user-token", "Bearer stored-token"),
("static", None, "static-token", "Bearer static-token"),
("static", {"authorization": "Bearer injected"}, "extra-headers", "Bearer injected"),
("none", None, "no-auth", None),
("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"),
],
)
async def test_debug_resolution_matches_final_header_conflict_winner(
config: Literal["stored", "static", "none"],
extra_headers: dict[str, str] | None,
expected_source: str,
expected_authorization: str | None,
) -> None:
from mcp.server.lowlevel.server import request_ctx
from mcp.shared.context import RequestContext
from starlette.requests import Request
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser
from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics
from litellm.proxy._experimental.mcp_server.outbound_credentials import (
ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
from litellm.types.mcp_server.mcp_server_manager import MCPServer
class Store:
def __init__(self) -> None:
self.calls = 0
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
self.calls += 1
return OAuthToken(access_token="stored-token") if user_id == "alice" else None
store = Store()
context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice"))
diagnostics = MCPAuthDiagnostics()
token = request_ctx.set(RequestContext(
request_id=1, meta=None, session=MagicMock(), lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
))
selected = {
"stored": AuthorizationCodeConfig(),
"static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))),
"none": NoneConfig(),
}[config]
try:
auth, remaining = await MCPServerManager()._resolve_v2_auth(
server=MCPServer(
server_id="s", name="s", transport="http", url="https://up.example/mcp",
static_headers={"Authorization": "Bearer configured"},
),
spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected),
provider=UpstreamCredentialProvider(oauth_token_store=store),
subject_token=None,
user_api_key_auth=context.user_api_key_auth,
extra_headers=extra_headers,
)
request = httpx.Request("GET", "https://up.example/mcp", headers=remaining)
if auth is not None:
next(auth.auth_flow(request))
assert diagnostics.resolution() == expected_source
assert request.headers.get("Authorization") == expected_authorization
assert store.calls == (1 if config == "stored" else 0)
finally:
request_ctx.reset(token)
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "stdio"])
async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None:
from mcp.server.lowlevel.server import request_ctx
from mcp.shared.context import RequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics
from litellm.types.mcp_server.mcp_server_manager import MCPServer
diagnostics = MCPAuthDiagnostics()
token = request_ctx.set(RequestContext(
request_id=1, meta=None, session=MagicMock(), lifespan_context=None,
request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}),
))
try:
server = MCPServer(
server_id="signed", name="signed", transport=transport,
url="https://up.example/mcp", auth_type="aws_sigv4",
aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret",
aws_region_name="us-east-1", aws_service_name="execute-api",
command="python", args=["-c", "pass"],
)
client = await MCPServerManager()._create_mcp_client(server)
if transport == "stdio":
assert diagnostics.resolution() == "not-applicable"
else:
assert diagnostics.resolution() == "aws-sigv4"
request = httpx.Request("POST", "https://up.example/mcp", content=b"{}")
next(client._aws_auth.auth_flow(request))
assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ")
assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"]
finally:
request_ctx.reset(token)