mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(mcp): report resolved upstream authentication in debug headers
This commit is contained in:
parent
b23995ee29
commit
feeda1a36c
8 changed files with 470 additions and 129 deletions
|
|
@ -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,17 @@ 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
|
||||
) -> Send:
|
||||
"""
|
||||
Return a new ASGI ``send`` callable that injects *debug_headers*
|
||||
into the ``http.response.start`` message.
|
||||
"""
|
||||
|
||||
if resolution is not None:
|
||||
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 +334,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 +354,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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,13 @@ 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)
|
||||
|
||||
# Ensure session managers are initialized
|
||||
if not _SESSION_MANAGERS_INITIALIZED:
|
||||
|
|
|
|||
|
|
@ -1203,3 +1203,69 @@ 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, subject, expected_source, expected_header):
|
||||
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():
|
||||
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():
|
||||
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():
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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,85 @@ class TestWrapSendWithDebugHeaders:
|
|||
asyncio.run(wrapped(body_msg))
|
||||
|
||||
assert captured[0] == body_msg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source", tuple(AuthResolution))
|
||||
async def test_debug_uses_resolution_recorded_after_response_start(source: AuthResolution) -> 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)
|
||||
await wrapped({"type": "http.response.start", "status": 200, "headers": []})
|
||||
assert captured == []
|
||||
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()
|
||||
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)
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -12567,3 +12567,109 @@ 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, extra_headers, expected_source, expected_authorization
|
||||
):
|
||||
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):
|
||||
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):
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue