diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 9cb929dec92..9463023854e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -941,15 +941,13 @@ async def inject_mcp_jwt_headers_for_upstream( if signer is None or user_api_key_dict is None: return merged - normalized_raw: Final = {k.lower(): v for k, v in (raw_headers or {}).items()} - incoming_bearer_token: str | None = None - auth_hdr: Final = normalized_raw.get("authorization", "") - if auth_hdr.lower().startswith("bearer "): - incoming_bearer_token = auth_hdr[len("bearer ") :] + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + incoming_bearer_token: Final = MCPRequestHandler.guardrail_bearer_from_raw_headers(raw_headers, user_api_key_dict) hook_data: Final = { "mcp_tool_name": "" if for_list_tools else mcp_tool_name, - "incoming_bearer_token": incoming_bearer_token, + "incoming_bearer_token": incoming_bearer_token.get_secret_value() if incoming_bearer_token else None, "extra_headers": merged, } call_type: Final[CallTypesLiteral] = "list_mcp_tools" if for_list_tools else "call_mcp_tool" diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index be24b1630b0..bcb979d720b 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -919,7 +919,6 @@ class MCPRequestHandler: mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, admitted_credential=admitted_credential, - preserve_jwt_subject=bool(validated_user_api_key_auth.jwt_claims), ) return ( @@ -959,49 +958,29 @@ class MCPRequestHandler: mcp_server_auth_headers: dict[str, dict[str, str]] | None, *, admitted_credential: str | None, - preserve_jwt_subject: bool = False, ) -> tuple[dict[str, str] | None, dict[str, str], str | None, dict[str, dict[str, str]] | None]: - """Remove gateway-shaped credentials and the caller's admitted credential from EVERY egress - header context: top-level ``Authorization`` (oauth2 + raw), the deprecated ``x-mcp-auth``, and - per-server ``x-mcp-{alias}-authorization``. Preserve the raw ``x-litellm-api-key`` admission - header; an admitted subject's top-level Authorization is otherwise dropped unconditionally.""" + """Remove gateway-shaped credentials from every egress header context, and the caller's admitted + credential from the implicitly forwarded top-level ``Authorization``. Raw headers (only forwarded + when a server lists them in ``extra_headers``), ``x-mcp-auth`` and ``x-mcp-{alias}-authorization`` + keep the caller's LiteLLM key, since the admin or caller chose to send it there.""" + gateway_shaped: Final = MCPRequestHandler._is_gateway_admission_credential - def cred(value: str | None) -> bool: - return MCPRequestHandler._is_gateway_admission_credential( - value - ) or MCPRequestHandler._is_caller_admission_key(value, admitted_credential) - - # 1. Top-level Authorization → oauth2_headers. authz: Final = oauth2_headers.get("Authorization") if oauth2_headers else None - preserve_subject: Final = bool( - preserve_jwt_subject - and not admitted - and authz - and not _normalize_caller_admission_credential(authz).startswith("sk-") - and not MCPRequestHandler._is_gateway_admission_credential(authz) - and MCPRequestHandler._is_caller_admission_key(authz, admitted_credential) - ) - if admitted or cred(authz): + if admitted or gateway_shaped(authz) or MCPRequestHandler._is_caller_admission_key(authz, admitted_credential): oauth2_headers = None - # 2. raw_headers: drop the admitted subject's Authorization, and ANY header whose value is a - # gateway credential (covers x-mcp-auth and x-mcp-{alias}-authorization in their raw form). raw_headers = { k: v for k, v in raw_headers.items() - if k.lower() == "x-litellm-api-key" - or (preserve_subject and k.lower() == "authorization") - or not ((admitted and k.lower() == "authorization") or cred(v)) + if not ((admitted and k.lower() == "authorization") or gateway_shaped(v)) } - # 3. Deprecated x-mcp-auth value. - if cred(mcp_auth_header): + if gateway_shaped(mcp_auth_header): mcp_auth_header = None - # 4. Per-server x-mcp-{alias}-authorization values (drop the value, then any now-empty server dict). if mcp_server_auth_headers: stripped: Final = { - alias: {h: val for h, val in hdrs.items() if not cred(val)} + alias: {h: val for h, val in hdrs.items() if not gateway_shaped(val)} for alias, hdrs in mcp_server_auth_headers.items() } mcp_server_auth_headers = {alias: hdrs for alias, hdrs in stripped.items() if hdrs} @@ -1892,6 +1871,14 @@ class MCPRequestHandler: return None return bearer + @staticmethod + def guardrail_bearer_from_raw_headers( + raw_headers: Mapping[str, str] | None, auth: UserAPIKeyAuth | None + ) -> SecretStr | None: + if auth is None: + return MCPRequestHandler.get_incoming_bearer_token(raw_headers or {}) + return MCPRequestHandler.get_guardrail_bearer_token(Headers(headers=dict(raw_headers or {})), auth) + @staticmethod def get_oauth2_headers_from_headers(headers: Headers) -> dict[str, str]: """ diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 013336e1d5d..cc02fc43ef3 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -1206,8 +1206,10 @@ def should_strip_caller_authorization( admission credential — a virtual key, an IdP JWT, or an SSO / OIDC / session token whose ``api_key`` is ``None`` — and must never reach the upstream, so it is stripped regardless of the ``api_key`` value. - - **Every other auth-enabled mode**: strip when ``Authorization`` is the - credential admission validated. + - **true_passthrough**: strip when ``Authorization`` is the credential + admission validated. + - **Every other mode**: forward. The caller's ``Authorization`` only reaches + the upstream because the admin listed it in ``extra_headers``. """ if ( user_api_key_auth is not None @@ -1233,6 +1235,8 @@ def should_strip_caller_authorization( mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream ) if not (mcp_server.is_oauth_passthrough or is_delegated_oauth): + if not mcp_server.is_client_forwarded_token: + return False from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode from litellm.proxy.proxy_server import ( general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict @@ -1403,7 +1407,7 @@ def _openapi_forwarded_extra_headers( ) forwarded: Final[dict[str, str]] = {} for header_name in mcp_server.extra_headers: - if not isinstance(header_name, str) or header_name.lower() == "x-litellm-api-key": + if not isinstance(header_name, str): continue if skip_caller_authorization and header_name.lower() == "authorization": continue @@ -6241,7 +6245,7 @@ class MCPServerManager: ) caller_bearer: Final = ( - MCPRequestHandler.get_incoming_bearer_token(raw_headers or {}) + MCPRequestHandler.guardrail_bearer_from_raw_headers(raw_headers, user_api_key_auth) if incoming_bearer_token is ... else incoming_bearer_token ) @@ -6640,7 +6644,7 @@ class MCPServerManager: # For M2M OAuth servers, Authorization must come from token fetch. extra_headers = None else: - extra_headers = oauth2_headers + extra_headers = dict(oauth2_headers) if oauth2_headers else None # Migrated authorization_code: the v2 resolver injects the stored per-user # token, so drop the caller-forwarded Authorization (apply-if-absent would # otherwise let it shadow the resolved token). Delegate keeps it. Centralized @@ -6673,7 +6677,7 @@ class MCPServerManager: ) for header in mcp_server.extra_headers: - if not isinstance(header, str) or header.lower() == "x-litellm-api-key": + if not isinstance(header, str): continue if header.lower() == "authorization" and strip_caller_authorization: continue diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index dfc45417814..8422d516d8c 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -1080,7 +1080,7 @@ def _prepare_mcp_server_headers( ) for header in server.extra_headers: - if not isinstance(header, str) or header.lower() == "x-litellm-api-key": + if not isinstance(header, str): continue if header.lower() == "authorization" and (strip_caller_authorization or withhold_forwarded_authorization): continue diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 2737b567985..815c1d6270c 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -627,7 +627,6 @@ if MCP_AVAILABLE: ) = MCPRequestHandler.scrub_gateway_admission_credentials( admitted=False, admitted_credential=admitted_credential, - preserve_jwt_subject=bool(user_api_key_dict.jwt_claims), oauth2_headers=oauth2_headers_from_request, raw_headers=dict(headers), mcp_auth_header=MCPRequestHandler.get_mcp_auth_header_from_headers(headers), diff --git a/tests/integration/mcp/test_mcp_credentials.py b/tests/integration/mcp/test_mcp_credentials.py index 08dd5234b68..d21e6fea28e 100644 --- a/tests/integration/mcp/test_mcp_credentials.py +++ b/tests/integration/mcp/test_mcp_credentials.py @@ -155,11 +155,11 @@ def test_caller_headers_for_other_servers_and_unknown_headers_never_reach_the_pe @pytest.mark.parametrize( - ("auth_type", "per_server_header"), - (("none", True), ("true_passthrough", False)), + ("auth_type", "per_server_header", "forwarded"), + (("none", True, True), ("true_passthrough", False, False)), ) -def test_callers_own_litellm_key_never_reaches_the_peer_over_rest( - gateway: Gateway, auth_type: str, per_server_header: bool +def test_callers_litellm_key_reaches_the_peer_only_through_a_per_server_header( + gateway: Gateway, auth_type: str, per_server_header: bool, forwarded: bool ) -> None: with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "cred" + uuid.uuid4().hex[:8] @@ -178,19 +178,14 @@ def test_callers_own_litellm_key_never_reaches_the_peer_over_rest( observed: Final = peer.drain() calls: Final = tool_calls(observed) assert len(calls) == 1, calls - header_sets: Final = tuple( - TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed - ) - assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), ( - "caller key appeared in the recorded MCP header set" - ) + assert (_header(calls[0], b"authorization") == f"Bearer {key}".encode()) is forwarded @pytest.mark.parametrize("entry", ("server_mcp", "rest")) def test_empty_primary_header_does_not_expose_authorization_key(gateway: Gateway, entry: EntryPoint) -> None: with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "cred" + uuid.uuid4().hex[:8] - identity: Final = register_mcp(scenario, peer, alias, extra_headers=["x-upstream-token", "x-tenant"]) + identity: Final = register_mcp(scenario, peer, alias, extra_headers=["x-tenant"]) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) caller: Final = McpCaller( gateway, @@ -200,8 +195,6 @@ def test_empty_primary_header_does_not_expose_authorization_key(gateway: Gateway headers={ "x-litellm-api-key": "", "Authorization": f"Bearer {key}", - f"x-mcp-{alias}-authorization": f"Bearer {key}", - "x-upstream-token": key, "x-tenant": "tenant-control", }, ) @@ -216,40 +209,40 @@ def test_empty_primary_header_does_not_expose_authorization_key(gateway: Gateway TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed ) assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), ( - "empty primary header prevented scrubbing the admitted Authorization key" + "admission key reached the peer through a header the server does not list" ) @pytest.mark.parametrize("entry", ENTRY_POINTS) -def test_extra_headers_cannot_forward_gateway_admission_key(gateway: Gateway, entry: EntryPoint) -> None: +def test_extra_headers_forward_listed_admission_key(gateway: Gateway, entry: EntryPoint) -> None: with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "cred" + uuid.uuid4().hex[:8] identity: Final = register_mcp( - scenario, peer, alias, extra_headers=["X-LiteLLM-API-Key", "x-tenant"], + scenario, + peer, + alias, + extra_headers=["X-LiteLLM-API-Key", "x-tenant"], ) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) caller: Final = McpCaller(gateway, key, entry, alias, headers={"x-tenant": "tenant-control"}) peer.drain() outcome: Final = caller.call(f"{alias}-add", ADD, identity if entry != "server_mcp" else None) assert outcome.ok, outcome.raw - observed: Final = peer.drain() - calls: Final = tool_calls(observed) + calls: Final = tool_calls(peer.drain()) assert len(calls) == 1, calls assert _header(calls[0], b"x-tenant") == b"tenant-control" - header_sets: Final = tuple( - TypeAdapter(dict[bytes, bytes]).validate_python(request["headers"]) for request in observed - ) - assert all(all(key.encode() not in value for value in headers.values()) for headers in header_sets), ( - "gateway admission key reached the upstream through Extra Headers" - ) + assert key.encode() in (_header(calls[0], b"x-litellm-api-key") or b"") @pytest.mark.parametrize("entry", ("server_mcp", "rest")) -def test_openapi_extra_headers_cannot_forward_gateway_admission_key(gateway: Gateway, entry: EntryPoint) -> None: +def test_openapi_extra_headers_forward_listed_admission_key(gateway: Gateway, entry: EntryPoint) -> None: with openapi_peer() as peer, gateway.scenario() as scenario: alias: Final = "cred" + uuid.uuid4().hex[:8] identity: Final = register_mcp( - scenario, peer, alias, extra_headers=["X-LiteLLM-API-Key", "x-tenant"], + scenario, + peer, + alias, + extra_headers=["X-LiteLLM-API-Key", "x-tenant"], ) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) caller: Final = McpCaller(gateway, key, entry, alias, headers={"x-tenant": "tenant-control"}) @@ -260,7 +253,7 @@ def test_openapi_extra_headers_cannot_forward_gateway_admission_key(gateway: Gat assert [(request["method"], request["path"]) for request in observed] == [("GET", "/pets/7")] headers: Final = TypeAdapter(dict[bytes, bytes]).validate_python(observed[0]["headers"]) assert headers.get(b"x-tenant") == b"tenant-control" - assert all(key.encode() not in value for value in headers.values()), "gateway admission key reached OpenAPI" + assert key.encode() in headers.get(b"x-litellm-api-key", b"") def test_server_scoped_caller_header_reaches_only_its_server(gateway: Gateway) -> None: diff --git a/tests/unit/proxy/guardrails/test_mcp_jwt_signer.py b/tests/unit/proxy/guardrails/test_mcp_jwt_signer.py index a7b24169398..8b7a420c739 100644 --- a/tests/unit/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/unit/proxy/guardrails/test_mcp_jwt_signer.py @@ -11,7 +11,7 @@ Tests cover: import base64 import time -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import jwt @@ -1292,3 +1292,49 @@ async def test_inject_mcp_jwt_signs_for_tool_call_path(): scopes = set(decoded["scope"].split()) assert "mcp:tools/call" in scopes assert "mcp:tools/search_web:call" in scopes + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("raw_headers", "introspected"), + [ + ({"Authorization": "Bearer sk-admission-key"}, False), + ({"x-litellm-api-key": "sk-admission-key", "Authorization": "Bearer upstream-opaque"}, True), + ], +) +async def test_inject_mcp_jwt_never_introspects_the_admission_key( + raw_headers: dict[str, str], introspected: bool +) -> None: + import httpx + from fastapi import HTTPException + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer( + access_token_discovery_uri="https://idp.example.com/.well-known/openid-configuration", + token_introspection_endpoint="https://idp.example.com/introspect", + ) + inactive_introspection: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"active": False})) + ) + + async def sign() -> dict[str, str]: + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=inactive_introspection, + ): + return await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-admission-key"), user_id="alice"), + raw_headers=raw_headers, + for_list_tools=True, + ) + + if introspected: + with pytest.raises(HTTPException): + await sign() + else: + assert (await sign())["Authorization"].startswith("Bearer ") diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index 7e74e62894f..ea54d456307 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -10141,444 +10141,100 @@ class TestSessionBearerEgressScrub: assert "authorization" not in {k.lower() for k in raw} assert per_server == {"github": {"Authorization": "Bearer gh_injected_upstream"}} - @pytest.mark.parametrize( - "context", - ["per_server", "mcp_auth", "oauth2_authorization", "raw_authorization"], - ) @pytest.mark.parametrize( "caller_value", - ["sk-caller-admission-key-123", "Bearer sk-caller-admission-key-123"], + [ + "sk-caller-admission-key-123", + "Bearer sk-caller-admission-key-123", + "BEARER sk-caller-admission-key-123", + "Bearer sk-caller-admission-key-123", + "Basic sk-caller-admission-key-123", + ], ) - async def test_scrub_removes_caller_admission_key_from_each_egress_context( - self, - context: Literal["per_server", "mcp_auth", "oauth2_authorization", "raw_authorization"], - caller_value: str, - ) -> None: + async def test_scrub_drops_caller_key_only_from_implicit_authorization(self, caller_value: str) -> None: caller_key: Final = "sk-caller-admission-key-123" - upstream_token: Final = "Bearer real-upstream-token" - oauth2_headers: Final = {"Authorization": caller_value if context == "oauth2_authorization" else upstream_token} raw_headers: Final = { "x-litellm-api-key": f"Bearer {caller_key}", - "authorization": caller_value if context == "raw_authorization" else upstream_token, - "x-upstream-token": upstream_token, - } - mcp_auth_header: Final = caller_value if context == "mcp_auth" else upstream_token - mcp_server_auth_headers: Final = { - "echo_srv": { - "Authorization": caller_value if context == "per_server" else upstream_token, - } + "authorization": caller_value, + "x-tenant": caller_value, } oauth2, raw, mcp_auth, per_server = MCPRequestHandler.scrub_gateway_admission_credentials( admitted=False, - oauth2_headers=oauth2_headers, + oauth2_headers={"Authorization": caller_value}, raw_headers=raw_headers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=caller_value, + mcp_server_auth_headers={"echo_srv": {"Authorization": caller_value}}, admitted_credential=caller_key, ) - assert raw["x-litellm-api-key"] == f"Bearer {caller_key}" - assert raw["x-upstream-token"] == upstream_token - if context == "raw_authorization": - assert "authorization" not in raw - else: - assert raw["authorization"] == upstream_token - if context == "oauth2_authorization": - assert oauth2 is None - else: - assert oauth2 == {"Authorization": upstream_token} - if context == "mcp_auth": - assert mcp_auth is None - else: - assert mcp_auth == upstream_token - if context == "per_server": - assert not per_server - else: - assert per_server == {"echo_srv": {"Authorization": upstream_token}} + assert oauth2 is None + assert raw == raw_headers + assert mcp_auth == caller_value + assert per_server == {"echo_srv": {"Authorization": caller_value}} - async def test_scrub_keeps_non_ascii_per_server_token(self) -> None: - upstream_token: Final = "Bearer t\u00f6ken" - _oauth2, _raw, _mcp_auth, per_server = MCPRequestHandler.scrub_gateway_admission_credentials( + async def test_scrub_keeps_distinct_upstream_authorization(self) -> None: + upstream_token: Final = "Bearer real-upstream-token" + oauth2, _raw, _mcp_auth, _per_server = MCPRequestHandler.scrub_gateway_admission_credentials( admitted=False, - admitted_credential="sk-caller-admission-key-123", - oauth2_headers=None, + oauth2_headers={"Authorization": upstream_token}, raw_headers={}, mcp_auth_header=None, - mcp_server_auth_headers={"echo_srv": {"Authorization": upstream_token}}, + mcp_server_auth_headers=None, + admitted_credential="sk-caller-admission-key-123", ) - assert per_server == {"echo_srv": {"Authorization": upstream_token}} - - async def test_scrub_keeps_caller_key_when_admission_credential_is_missing(self) -> None: - caller_key: Final = "Bearer sk-caller-admission-key-123" - oauth2, raw, mcp_auth, per_server = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": caller_key}, - raw_headers={"x-litellm-api-key": caller_key, "authorization": caller_key}, - mcp_auth_header=caller_key, - mcp_server_auth_headers={"echo_srv": {"Authorization": caller_key}}, - admitted_credential=None, - ) - - assert oauth2 == {"Authorization": caller_key} - assert raw == {"x-litellm-api-key": caller_key, "authorization": caller_key} - assert mcp_auth == caller_key - assert per_server == {"echo_srv": {"Authorization": caller_key}} + assert oauth2 == {"Authorization": upstream_token} @pytest.mark.parametrize( - "admission_value,per_server_value", - ( - ("Bearer sk-caller-admission-key-123", "BEARER sk-caller-admission-key-123"), - ("Bearer sk-caller-admission-key-123", "Bearer sk-caller-admission-key-123"), - ("Bearer sk-caller-admission-key-123", "Basic sk-caller-admission-key-123"), - ("Basic sk-caller-admission-key-123", "sk-caller-admission-key-123"), - ("Bearer sk-caller-admission-key-123", "bearer sk-caller-admission-key-123"), - ), + ("headers", "custom_key_header_name", "expected"), + [ + ( + { + "x-litellm-api-key": "Bearer sk-caller", + "authorization": "Bearer upstream", + SpecialHeaders.azure_authorization.value: "provider", + }, + None, + "sk-caller", + ), + ( + {"authorization": "Bearer sk-caller", SpecialHeaders.azure_authorization.value: "provider"}, + None, + "sk-caller", + ), + ( + {"x-custom-api-key": "Bearer sk-caller", "x-litellm-api-key": "Bearer standard"}, + "x-custom-api-key", + "sk-caller", + ), + ({"x-litellm-api-key": "", "authorization": "sk-caller"}, None, "sk-caller"), + ({"authorization": "", "api-key": "sk-caller"}, None, "sk-caller"), + ( + { + SpecialHeaders.azure_authorization.value: "azure-key", + SpecialHeaders.anthropic_authorization.value: "anthropic-key", + SpecialHeaders.google_ai_studio_authorization.value: "google-key", + SpecialHeaders.azure_apim_authorization.value: "apim-key", + }, + None, + "azure-key", + ), + ({SpecialHeaders.anthropic_authorization.value: "anthropic-key"}, None, "anthropic-key"), + ({SpecialHeaders.google_ai_studio_authorization.value: "google-key"}, None, "google-key"), + ({SpecialHeaders.azure_apim_authorization.value: "apim-key"}, None, "apim-key"), + ], ) - async def test_caller_admission_credential_scrubs_bearer_variants_and_basic( - self, - admission_value: str, - per_server_value: str, + async def test_caller_admission_credential_matches_admission_header_precedence( + self, headers: dict[str, str], custom_key_header_name: str | None, expected: str ) -> None: - upstream_authorization: Final = "Bearer unrelated-upstream-token" - headers: Final = Headers( - { - "x-litellm-api-key": admission_value, - "authorization": upstream_authorization, - "x-mcp-auth": per_server_value, - "x-mcp-echo_srv-authorization": per_server_value, - "x-upstream-token": upstream_authorization, - } - ) - user_api_key_auth: Final = UserAPIKeyAuth(api_key="stored-key-hash") admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - user_api_key_auth, - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": upstream_authorization}, - raw_headers=dict(headers), - mcp_auth_header=per_server_value, - mcp_server_auth_headers={ - "echo_srv": {"Authorization": per_server_value}, - "other_srv": {"Authorization": upstream_authorization}, - }, - admitted_credential=admitted_credential, - ) - - assert result == ( - {"Authorization": upstream_authorization}, - { - "x-litellm-api-key": admission_value, - "authorization": upstream_authorization, - "x-upstream-token": upstream_authorization, - }, - None, - {"other_srv": {"Authorization": upstream_authorization}}, - ) - - @pytest.mark.parametrize( - ("empty_header", "admission_header"), - (("x-litellm-api-key", "authorization"), ("authorization", "api-key")), - ) - async def test_empty_admission_header_falls_back_to_authenticated_key( - self, empty_header: str, admission_header: str - ) -> None: - caller_key: Final = "sk-caller-admission-key-123" - headers: Final = Headers({empty_header: "", admission_header: caller_key}) - credential: Final = MCPRequestHandler.caller_admission_credential( - headers, UserAPIKeyAuth(api_key="stored-key-hash"), custom_key_header_name=None - ) - assert credential == caller_key - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": f"Bearer {caller_key}"}, - raw_headers={**dict(headers), "x-upstream-token": caller_key, "x-tenant": "tenant-control"}, - mcp_auth_header=caller_key, - mcp_server_auth_headers={"echo": {"Authorization": caller_key}}, - admitted_credential=credential, - ) - assert result == (None, {empty_header: "", "x-tenant": "tenant-control"}, None, {}) - - @pytest.mark.parametrize("claims", [None, {"sub": "caller"}]) - @pytest.mark.parametrize("bearer", ["validated-jwt", "sk-admission-key"]) - async def test_scrubbing_preserves_only_validated_jwt_admission( - self, claims: dict[str, str] | None, bearer: str - ) -> None: - headers: Final = Headers({"Authorization": f"Bearer {bearer}"}) - auth: Final = UserAPIKeyAuth(api_key="stored-key-hash", jwt_claims=claims) - credential: Final = MCPRequestHandler.caller_admission_credential(headers, auth, custom_key_header_name=None) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": f"Bearer {bearer}"}, - raw_headers={**dict(headers), "x-mcp-echo-authorization": f"Bearer {bearer}"}, - mcp_auth_header=bearer, - mcp_server_auth_headers={"echo": {"Authorization": f"Bearer {bearer}"}}, - admitted_credential=credential, - preserve_jwt_subject=bool(claims), - ) - expected: Final = {"Authorization": f"Bearer {bearer}"} if claims and not bearer.startswith("sk-") else None - assert result == (None, dict(headers) if expected else {}, None, {}) - - async def test_caller_admission_credential_prefers_x_litellm_header(self) -> None: - caller_key: Final = "sk-caller-admission-key-123" - upstream_authorization: Final = "Bearer unrelated-upstream-token" - provider_key: Final = "provider-key-not-admitted" - headers: Final = Headers( - { - "x-litellm-api-key": f"Bearer {caller_key}", - "authorization": upstream_authorization, - SpecialHeaders.azure_authorization.value: provider_key, - "x-mcp-auth": f"Bearer {caller_key}", - "x-mcp-echo_srv-authorization": f"Bearer {caller_key}", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, + Headers(headers), UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": upstream_authorization}, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {caller_key}", - mcp_server_auth_headers={ - "echo_srv": {"Authorization": f"Bearer {caller_key}"}, - "other_srv": {"Authorization": upstream_authorization}, - }, - admitted_credential=admitted_credential, + custom_key_header_name=custom_key_header_name, ) - assert result == ( - {"Authorization": upstream_authorization}, - { - "x-litellm-api-key": f"Bearer {caller_key}", - "authorization": upstream_authorization, - "api-key": provider_key, - }, - None, - {"other_srv": {"Authorization": upstream_authorization}}, - ) - - async def test_caller_admission_credential_prefers_authorization_before_provider_headers(self) -> None: - caller_key: Final = "sk-caller-admission-key-123" - provider_key: Final = "provider-key-not-admitted" - headers: Final = Headers( - { - "authorization": f"Bearer {caller_key}", - SpecialHeaders.azure_authorization.value: provider_key, - "x-mcp-auth": f"Bearer {caller_key}", - "x-mcp-echo_srv-authorization": f"Bearer {caller_key}", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": f"Bearer {caller_key}"}, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {caller_key}", - mcp_server_auth_headers={"echo_srv": {"Authorization": f"Bearer {caller_key}"}}, - admitted_credential=admitted_credential, - ) - - assert result == ( - None, - {"api-key": provider_key}, - None, - {}, - ) - - async def test_caller_admission_credential_uses_configured_custom_header(self) -> None: - caller_key: Final = "sk-caller-admission-key-123" - standard_key: Final = "Bearer standard-key-not-admitted" - upstream_authorization: Final = "Bearer unrelated-upstream-token" - headers: Final = Headers( - { - "x-custom-api-key": f"Bearer {caller_key}", - "x-litellm-api-key": standard_key, - "authorization": upstream_authorization, - "x-mcp-auth": f"Bearer {caller_key}", - "x-mcp-echo_srv-authorization": f"Bearer {caller_key}", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name="x-custom-api-key", - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": upstream_authorization}, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {caller_key}", - mcp_server_auth_headers={"echo_srv": {"Authorization": f"Bearer {caller_key}"}}, - admitted_credential=admitted_credential, - ) - - assert result == ( - {"Authorization": upstream_authorization}, - { - "x-litellm-api-key": standard_key, - "authorization": upstream_authorization, - }, - None, - {}, - ) - - async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent( - self, - ) -> None: - caller_value: Final = "Bearer sk-caller-admission-key-123" - headers: Final = Headers( - { - "x-litellm-api-key": caller_value, - "authorization": caller_value, - "x-mcp-auth": caller_value, - "x-mcp-echo_srv-authorization": caller_value, - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name="x-custom-api-key", - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": caller_value}, - raw_headers=dict(headers), - mcp_auth_header=caller_value, - mcp_server_auth_headers={"echo_srv": {"Authorization": caller_value}}, - admitted_credential=admitted_credential, - ) - - assert admitted_credential == "sk-caller-admission-key-123" - assert result == ( - None, - {"x-litellm-api-key": caller_value}, - None, - {}, - ) - - async def test_caller_admission_credential_uses_master_key_alias_as_admission_gate(self) -> None: - master_key: Final = "sk-" + "1234" - headers: Final = Headers( - { - "x-litellm-api-key": f"Bearer {master_key}", - "authorization": f"Bearer {master_key}", - "x-mcp-auth": f"Bearer {master_key}", - "x-mcp-echo_srv-authorization": f"Bearer {master_key}", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="litellm_proxy_master_key"), - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers={"Authorization": f"Bearer {master_key}"}, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {master_key}", - mcp_server_auth_headers={"echo_srv": {"Authorization": f"Bearer {master_key}"}}, - admitted_credential=admitted_credential, - ) - - assert result == ( - None, - {"x-litellm-api-key": f"Bearer {master_key}"}, - None, - {}, - ) - - @pytest.mark.parametrize( - "provider_header_name", - ( - SpecialHeaders.azure_authorization.value, - SpecialHeaders.anthropic_authorization.value, - SpecialHeaders.google_ai_studio_authorization.value, - SpecialHeaders.azure_apim_authorization.value, - ), - ) - async def test_caller_admission_credential_scrubs_provider_header( - self, - provider_header_name: str, - ) -> None: - provider_key: Final = "provider-admission-key" - headers: Final = Headers( - { - provider_header_name: provider_key, - "x-mcp-auth": f"Bearer {provider_key}", - "x-mcp-echo_srv-authorization": f"Bearer {provider_key}", - "x-upstream-token": "Bearer unrelated-upstream-token", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers=None, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {provider_key}", - mcp_server_auth_headers={ - "echo_srv": {"Authorization": f"Bearer {provider_key}"}, - "other_srv": {"Authorization": "Bearer unrelated-upstream-token"}, - }, - admitted_credential=admitted_credential, - ) - - assert result == ( - None, - {"x-upstream-token": "Bearer unrelated-upstream-token"}, - None, - {"other_srv": {"Authorization": "Bearer unrelated-upstream-token"}}, - ) - - async def test_caller_admission_credential_follows_provider_header_precedence(self) -> None: - azure_key: Final = "azure-admission-key" - headers: Final = Headers( - { - SpecialHeaders.azure_authorization.value: azure_key, - SpecialHeaders.anthropic_authorization.value: "anthropic-not-admitted", - SpecialHeaders.google_ai_studio_authorization.value: "google-not-admitted", - SpecialHeaders.azure_apim_authorization.value: "apim-not-admitted", - "x-mcp-echo_srv-authorization": f"Bearer {azure_key}", - } - ) - admitted_credential: Final = MCPRequestHandler.caller_admission_credential( - headers, - UserAPIKeyAuth(api_key="stored-key-hash"), - custom_key_header_name=None, - ) - result: Final = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, - oauth2_headers=None, - raw_headers=dict(headers), - mcp_auth_header=f"Bearer {azure_key}", - mcp_server_auth_headers={"echo_srv": {"Authorization": f"Bearer {azure_key}"}}, - admitted_credential=admitted_credential, - ) - - assert result == ( - None, - { - "x-api-key": "anthropic-not-admitted", - "x-goog-api-key": "google-not-admitted", - "ocp-apim-subscription-key": "apim-not-admitted", - }, - None, - {}, - ) + assert admitted_credential == expected async def test_caller_admission_credential_returns_none_without_admitted_api_key(self) -> None: caller_value: Final = "Bearer sk-caller-admission-key-123" diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index 13e4704bd90..f7cbf7be8f7 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -321,7 +321,7 @@ def test_validated_admission_jwt_is_not_forwarded_as_an_extra_header(auth_type: "auth_type", [MCPAuth.none, MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.true_passthrough], ) -def test_admission_authorization_is_not_forwarded_for_any_auth_type( +def test_listed_admission_authorization_forwards_except_for_true_passthrough( auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy import proxy_server @@ -341,7 +341,11 @@ def test_admission_authorization_is_not_forwarded_for_any_auth_type( "Authorization": "Bearer sk-" + "1234", "x-tenant": "tenant-control", } - expected_without_admission_key: Final = {"x-tenant": "tenant-control"} + expected_without_admission_key: Final = ( + {"x-tenant": "tenant-control"} + if auth_type == MCPAuth.true_passthrough + else {"x-tenant": "tenant-control", "Authorization": "Bearer sk-" + "1234"} + ) admission_upstream_headers: Final = _prepare_mcp_server_headers( server, mcp_auth_header=None, @@ -414,12 +418,13 @@ def test_no_master_key_preserves_caller_authorization_unless_auth_is_enabled( raw_headers=raw_headers, user_api_key_auth=user, )[1] - assert auth_enabled_headers == {"x-tenant": "t"} - assert _openapi_forwarded_extra_headers(server, raw_headers, user) == {"x-tenant": "t"} + expected_auth_enabled: Final = {"x-tenant": "t"} if auth_type == MCPAuth.true_passthrough else expected_forwarded + assert auth_enabled_headers == expected_auth_enabled + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_auth_enabled @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) -def test_custom_auth_is_stripped_in_no_auth_dev_mode( +def test_custom_auth_is_stripped_in_no_auth_dev_mode_for_true_passthrough( auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy import proxy_server @@ -438,7 +443,11 @@ def test_custom_auth_is_stripped_in_no_auth_dev_mode( raw_headers: Final = {"Authorization": "Bearer custom-tok", "x-tenant": "t"} user: Final = UserAPIKeyAuth(api_key="custom-tok") user.authenticated_by_custom_auth = True - expected_headers: Final = {"x-tenant": "t"} + expected_headers: Final = ( + {"x-tenant": "t"} + if auth_type == MCPAuth.true_passthrough + else {"x-tenant": "t", "Authorization": "Bearer custom-tok"} + ) forwarded_headers: Final = _prepare_mcp_server_headers( server, mcp_auth_header=None, @@ -502,6 +511,11 @@ def test_custom_admission_header_controls_authorization_forwarding( extra_headers=["Authorization", "x-tenant"], ) user: Final = UserAPIKeyAuth(api_key="hashed-key") + expected_headers: Final = ( + expected + if auth_type == MCPAuth.true_passthrough + else {"x-tenant": "t", "Authorization": raw_headers["Authorization"]} + ) upstream_headers: Final = _prepare_mcp_server_headers( server, @@ -511,8 +525,8 @@ def test_custom_admission_header_controls_authorization_forwarding( raw_headers=raw_headers, user_api_key_auth=user, )[1] - assert upstream_headers == expected - assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected + assert upstream_headers == expected_headers + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_headers class TestMCPServerManager: @@ -3856,7 +3870,7 @@ class TestMCPServerManager: assert captured_extra_headers == {"x-request-id": "req-123"} @pytest.mark.asyncio - async def test_call_regular_mcp_tool_never_forwards_admission_key_for_none_auth( + async def test_call_regular_mcp_tool_forwards_listed_admission_key_for_none_auth( self, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy import proxy_server @@ -3907,7 +3921,76 @@ class TestMCPServerManager: assert capture_extra_headers.await_args is not None captured_extra_headers: Final = capture_extra_headers.await_args.args[0] - assert captured_extra_headers == {"x-tenant": "tenant-control"} + assert captured_extra_headers == {"Authorization": "Bearer sk-" + "1234", "x-tenant": "tenant-control"} + + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_listed_key_does_not_reach_a_later_server( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from mcp.types import JSONRPCRequest + + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + tool_call_headers: Final[dict[str, httpx2.Headers]] = {} + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "tools/call": + tool_call_headers[request.url.host] = request.headers + result: Final[dict[str, object]] = {"content": [], "isError": False} + else: + result = { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "upstream", "version": "1"}, + "capabilities": {"tools": {}}, + } + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(respond)) + listing_server: Final = MCPServer( + server_id="lists-key", + name="lists-key", + url="https://lists-key.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + extra_headers=["X-LiteLLM-API-Key"], + ) + later_server: Final = listing_server.model_copy( + update={ + "server_id": "no-key", + "name": "no-key", + "url": "https://no-key.example/mcp", + "extra_headers": None, + } + ) + oauth2_headers: Final = {"Authorization": "Bearer upstream-token"} + raw_headers: Final = { + "x-litellm-api-key": "sk-" + "1234", + "Authorization": "Bearer upstream-token", + } + + for server in (listing_server, later_server): + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=None, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key"), + ) + + assert tool_call_headers["lists-key.example"].get("x-litellm-api-key") == "sk-" + "1234" + assert "x-litellm-api-key" not in tool_call_headers["no-key.example"] @pytest.mark.asyncio async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization( @@ -21576,7 +21659,7 @@ async def test_repeated_stale_discovery_uses_current_callers_endpoint(endpoint: @pytest.mark.parametrize("header_name", ("x-litellm-api-key", "X-LiteLLM-API-Key")) -def test_openapi_extra_headers_exclude_gateway_admission_key(header_name: str) -> None: +def test_openapi_extra_headers_forward_listed_admission_key(header_name: str) -> None: caller_key: Final = "Bearer sk-admission-only" raw_headers: Final = {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} server: Final = MCPServer( @@ -21593,8 +21676,7 @@ def test_openapi_extra_headers_exclude_gateway_admission_key(header_name: str) - None, raw_headers, UserAPIKeyAuth(api_key="sk-admission-only"), - ) == (None, {"X-Tenant": "tenant-control"}, None) - assert raw_headers == {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} + ) == (None, {header_name: caller_key, "X-Tenant": "tenant-control"}, None) @pytest.mark.asyncio @@ -21900,3 +21982,32 @@ def test_id_jag_listed_tool_metadata_is_isolated_between_identity_assertions() - manager.record_listed_tools(server, (first,), callers[0]) manager.record_listed_tools(server, (second,), callers[1]) assert tuple(manager.get_listed_tool(server, "list_notes", caller) for caller in callers) == (first, second, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("raw_headers", "expected_bearer"), + [ + ({"authorization": "Bearer sk-admission-key"}, None), + ({"x-litellm-api-key": "sk-admission-key", "authorization": "Bearer upstream-token"}, "upstream-token"), + ], +) +async def test_pre_call_guardrail_bearer_never_carries_the_admission_key( + raw_headers: dict[str, str], expected_bearer: str | None +) -> None: + from litellm.proxy._types import hash_token + + proxy_logging_obj: Final = _permissive_proxy_logging() + + await MCPServerManager().pre_call_tool_check( + name="tool", + arguments={}, + server_name="guarded", + user_api_key_auth=UserAPIKeyAuth(api_key=hash_token("sk-admission-key")), + proxy_logging_obj=proxy_logging_obj, + server=MCPServer(server_id="guarded", name="guarded", transport=MCPTransport.stdio), + raw_headers=raw_headers, + ) + + hook_kwargs: Final = proxy_logging_obj.create_mcp_request_object_from_kwargs.call_args.args[0] + assert hook_kwargs["incoming_bearer_token"] == expected_bearer diff --git a/tests/unit/proxy/mcp_server/test_operations.py b/tests/unit/proxy/mcp_server/test_operations.py index c482f588521..df2502c6be7 100644 --- a/tests/unit/proxy/mcp_server/test_operations.py +++ b/tests/unit/proxy/mcp_server/test_operations.py @@ -140,7 +140,7 @@ async def test_mapped_caller_bearer_reaches_agent_evaluation( mcp_auth_header=None, mcp_server_auth_headers=None, ) - assert "authorization" not in raw + assert oauth is None provider_requests: Final[asyncio.Queue[str]] = asyncio.Queue() def provider(exchange: httpx.Request) -> httpx.Response: @@ -1917,7 +1917,7 @@ async def test_initial_tool_listing_preserves_legacy_error_fallback( @pytest.mark.parametrize("header_name", ("x-litellm-api-key", "X-LiteLLM-API-Key")) -def test_discovery_extra_headers_exclude_gateway_admission_key(header_name: str) -> None: +def test_discovery_extra_headers_forward_listed_admission_key(header_name: str) -> None: caller_key: Final = "Bearer sk-admission-only" raw_headers: Final = {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} server: Final = MCPServer( @@ -1937,8 +1937,7 @@ def test_discovery_extra_headers_exclude_gateway_admission_key(header_name: str) UserAPIKeyAuth(api_key="sk-admission-only"), ) assert auth_header is None - assert extra_headers == {"X-Tenant": "tenant-control"} - assert raw_headers == {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} + assert extra_headers == {header_name: caller_key, "X-Tenant": "tenant-control"} @pytest.mark.asyncio diff --git a/tests/unit/proxy/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/mcp_server/test_rest_endpoints.py index fdb1ff3cd09..b7eb68cae03 100644 --- a/tests/unit/proxy/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_rest_endpoints.py @@ -2764,7 +2764,7 @@ class TestCallToolRestAPI: assert resolved_server is target_server assert billed.team_id == "team-b" - async def test_extract_mcp_headers_scrubs_admitted_caller_credential(self) -> None: + async def test_extract_mcp_headers_scrubs_caller_key_only_from_implicit_authorization(self) -> None: caller_key: Final = "sk-rest-caller-admission-key-123" upstream_token: Final = "Bearer unrelated-upstream-token" request: Final = _build_request( @@ -2783,9 +2783,14 @@ class TestCallToolRestAPI: assert extracted_headers == ( None, - {"upstream": {"Authorization": upstream_token}}, + { + "echo_srv": {"Authorization": f"Bearer {caller_key}"}, + "upstream": {"Authorization": upstream_token}, + }, { "x-litellm-api-key": f"Bearer {caller_key}", + "authorization": f"Bearer {caller_key}", + "x-mcp-echo_srv-authorization": f"Bearer {caller_key}", "x-mcp-upstream-authorization": upstream_token, }, None, @@ -5097,12 +5102,5 @@ async def test_rest_guardrail_bearer_excludes_admission_key( assert result.is_error is False assert observed.qsize() == 1 assert observed.get_nowait() == expected - if claims is None: - list_headers: Final = rest_endpoints._extract_mcp_headers_from_request( - _build_request(headers), - auth, - )[2] - authorization: Final = list_headers.get("authorization", "") - assert (authorization.removeprefix("Bearer ") or None) == expected finally: global_mcp_tool_registry.unregister_tools_with_prefix("guarded-") diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 24f5173d6a3..d9d61b86590 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1066,7 +1066,7 @@ async def test_mcp_tool_execution_events_creation(): @pytest.mark.asyncio -async def test_aresponses_mcp_call_does_not_forward_admission_authorization( +async def test_aresponses_mcp_call_forwards_listed_admission_authorization( monkeypatch: pytest.MonkeyPatch, ) -> None: import importlib @@ -1165,7 +1165,10 @@ async def test_aresponses_mcp_call_does_not_forward_admission_authorization( assert mcp_client.call_tool.await_args is not None assert manager.create_mcp_client.await_args is not None - assert manager.create_mcp_client.await_args.kwargs["extra_headers"] == {"x-tenant": "tenant-a"} + assert manager.create_mcp_client.await_args.kwargs["extra_headers"] == { + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "tenant-a", + } @pytest.mark.asyncio