fix(mcp): never exchange the LiteLLM virtual key as the upstream subject token (#39446)

This commit is contained in:
devin-ai-integration[bot] 2026-09-02 20:32:43 -07:00 committed by GitHub
parent bcd3e2d94d
commit 291d02f8aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 396 additions and 27 deletions

View file

@ -820,20 +820,46 @@ def _should_strip_caller_authorization(
if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate):
return False
normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)}
has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if mcp_server.is_oauth_delegate:
return not has_explicit_litellm_admission_header
admission_consumed_authorization_as_litellm_key: Final = (
user_api_key_auth is not None
and bool(getattr(user_api_key_auth, "api_key", None))
and not has_explicit_litellm_admission_header
)
return admission_consumed_authorization_as_litellm_key or (
return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or (
user_api_key_auth is None and not has_explicit_litellm_admission_header
)
LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-"
def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None:
return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None)
def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool:
"""Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``."""
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
def _authorization_is_litellm_admission_credential(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
"""True when ``Authorization`` carries the LiteLLM key admission validated.
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
"""
if user_api_key_auth is None or not user_api_key_auth.api_key:
return False
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
if not admission_header:
return True
authorization: Final = _raw_header_value(raw_headers, "authorization")
return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(
admission_header, "Bearer"
)
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
"""Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection.
@ -3277,8 +3303,8 @@ class MCPServerManager:
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: dict[str, str] | None,
raw_headers: dict[str, str] | None,
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
) -> str | None:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
@ -3298,10 +3324,29 @@ class MCPServerManager:
return auth_value
return None
@staticmethod
def _extract_subject_token(
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key.
Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix,
any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``.
"""
if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth):
return None
bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers)
if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX):
return None
return bearer
def _obo_subject_token(
self,
server: MCPServer,
raw_headers: dict[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's bearer as the token_exchange (OBO) subject token, for that mode only.
@ -3311,7 +3356,7 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return None
return self._extract_bearer_token(None, raw_headers)
return self._extract_subject_token(None, raw_headers, user_api_key_auth)
def _build_stdio_env(
self,
@ -3566,6 +3611,7 @@ class MCPServerManager:
server: MCPServer,
oauth2_headers: dict[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Run the OBO exchange for a caller-supplied subject at the transport edge.
@ -3577,13 +3623,15 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return
subject_token: Final = self._extract_bearer_token(oauth2_headers, None)
if not subject_token:
if not self._extract_bearer_token(oauth2_headers, None):
return
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
spec: Final = to_server_spec(resolved_server)
if spec is None or not isinstance(spec.config, TokenExchangeConfig):
return
subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if subject_token is None:
raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
@ -3851,7 +3899,7 @@ class MCPServerManager:
# token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes
# never read the inbound bearer, so leave subject_token None to avoid forwarding it.
subject_token: Final = (
self._extract_bearer_token(oauth2_headers, raw_headers)
self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if server.auth_type == MCPAuth.oauth2_token_exchange
else None
)
@ -3931,6 +3979,7 @@ class MCPServerManager:
async def get_prompts_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -3959,7 +4008,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -3982,6 +4031,7 @@ class MCPServerManager:
async def get_resources_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4001,7 +4051,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4024,6 +4074,7 @@ class MCPServerManager:
async def get_resource_templates_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4043,7 +4094,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4068,6 +4119,7 @@ class MCPServerManager:
async def read_resource_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
url: AnyUrl,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
@ -4084,7 +4136,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -4099,6 +4151,7 @@ class MCPServerManager:
async def get_prompt_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
mcp_auth_header: str | dict[str, str] | None = None,
@ -4116,7 +4169,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -5290,7 +5343,7 @@ class MCPServerManager:
MCPAuth.oauth2_token_exchange,
MCPAuth.oauth2_id_jag,
):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
@ -5638,7 +5691,7 @@ class MCPServerManager:
subject_token: str | None = None
if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif isinstance(spec.config, PassthroughConfig):
inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers)
per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header)

View file

@ -2189,6 +2189,7 @@ if MCP_AVAILABLE:
try:
prompts = await global_mcp_server_manager.get_prompts_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2242,6 +2243,7 @@ if MCP_AVAILABLE:
try:
resources = await global_mcp_server_manager.get_resources_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2293,6 +2295,7 @@ if MCP_AVAILABLE:
try:
resource_templates = await global_mcp_server_manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -3211,6 +3214,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.get_prompt_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name=original_prompt_name,
arguments=arguments,
mcp_auth_header=server_auth_header,
@ -3261,6 +3265,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.read_resource_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
url=url,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
@ -3723,6 +3728,7 @@ if MCP_AVAILABLE:
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
allowed_server_ids: set[str] | None = None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Fail fast with HTTP 401 for MCP servers that need user auth but
didn't receive it on this request. Covers both gateway-managed OAuth2
@ -3867,6 +3873,7 @@ if MCP_AVAILABLE:
server=server,
oauth2_headers=oauth2_headers,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
)
# Pass-through OAuth: when the admin has opted a server into
@ -4195,6 +4202,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers. Must run after
@ -4518,6 +4526,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_sse_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers: surface upstream

View file

@ -914,6 +914,7 @@ async def test_mcp_get_prompt_success():
)
mock_manager.get_prompt_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name="hello",
arguments={"foo": "bar"},
mcp_auth_header={"Authorization": "token"},
@ -976,6 +977,7 @@ async def test_mcp_read_resource_success():
)
mock_manager.read_resource_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
url="https://example.com/resource",
mcp_auth_header={"Authorization": "token"},
extra_headers={"X-Test": "1"},
@ -8268,7 +8270,9 @@ class TestOboPreflightScopedToAllowedServers:
_, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key)
preflight.assert_awaited_once_with(
server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None
)
@pytest.mark.asyncio

View file

@ -30,6 +30,7 @@ from mcp.types import (
TextResourceContents,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm.constants import MCP_METADATA_TIMEOUT
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -2270,7 +2271,9 @@ class TestMCPServerManager:
"""prompts/list on an OBO server must exchange the caller's bearer, not connect with none."""
server = self._token_exchange_server("te-prompts")
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2279,7 +2282,9 @@ class TestMCPServerManager:
"""resources/list on an OBO server must exchange the caller's bearer."""
server = self._token_exchange_server("te-resources")
st = await self._capture_subject_token(
lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_resources_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2290,6 +2295,7 @@ class TestMCPServerManager:
st = await self._capture_subject_token(
lambda m: m.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://up.example.com/r",
raw_headers={"authorization": "Bearer subj-jwt"},
)
@ -2307,7 +2313,9 @@ class TestMCPServerManager:
auth_type=MCPAuth.none,
)
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st is None
@ -3254,7 +3262,7 @@ class TestMCPServerManager:
new_callable=AsyncMock,
return_value=mock_client,
):
prompts = await manager.get_prompts_from_server(server, add_prefix=True)
prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True)
mock_client.list_prompts.assert_awaited_once()
assert len(prompts) == 1
@ -3289,6 +3297,7 @@ class TestMCPServerManager:
):
result = await manager.get_prompt_from_server(
server=server,
user_api_key_auth=None,
prompt_name="hello",
arguments={"tone": "casual"},
)
@ -3334,6 +3343,7 @@ class TestMCPServerManager:
):
result = await manager.get_resources_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
add_prefix=True,
@ -3391,6 +3401,7 @@ class TestMCPServerManager:
):
result = await manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers=None,
add_prefix=False,
@ -3441,6 +3452,7 @@ class TestMCPServerManager:
) as mock_create_client:
result = await manager.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://example.com/resource",
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth:
assert result.isError is True
assert "upstream returned HTTP 503" in result.content[0].text
class TestLitellmAdmissionKeyIsNeverTheSubjectToken:
"""The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the
RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it
present, ``Authorization`` is the caller's own identity token and is exchanged as before."""
_ADMISSION_KEY: Final = "sk-litellm-virtual-key"
_USER_TOKEN: Final = "user-idp-jwt"
@staticmethod
def _token_exchange_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
)
@staticmethod
def _id_jag_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_id_jag,
client_id="cid",
client_secret="csec",
token_exchange_endpoint="https://idp.example.com/token",
id_jag_resource_token_endpoint="https://resource-as.example.com/token",
)
@staticmethod
def _recording_provider() -> MagicMock:
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
provider: Final = MagicMock()
provider.resolve_credentials = AsyncMock(
return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
)
return provider
@staticmethod
def _subjects_seen_by(provider: MagicMock) -> list[str | None]:
return [
call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None
for call in provider.resolve_credentials.call_args_list
]
@staticmethod
def _manager_with_recording_client() -> MCPServerManager:
manager: Final = MCPServerManager()
client: Final = AsyncMock()
client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
client.list_prompts = AsyncMock(return_value=[])
client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[]))
manager._create_mcp_client = AsyncMock(return_value=client)
return manager
@staticmethod
def _subject_token_given_to_client(manager: MCPServerManager) -> str | None:
return manager._create_mcp_client.call_args.kwargs["subject_token"]
async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth):
manager: Final = self._manager_with_recording_client()
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=user_api_key_auth,
)
return self._subject_token_given_to_client(manager)
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag])
async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type):
server = (
self._token_exchange_server("te-call")
if auth_type == MCPAuth.oauth2_token_exchange
else self._id_jag_server("jag-call")
)
subject_token = await self._call_tool_subject(
server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self):
"""The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-rest"),
oauth2_headers=None,
raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self):
"""Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-empty-header"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-same-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={
"x-litellm-api-key": self._ADMISSION_KEY,
"authorization": f"Bearer {self._ADMISSION_KEY}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self):
"""A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-second-key"),
oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": "Bearer sk-another-virtual-key",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-jwt"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
manager._fetch_tools_with_timeout = AsyncMock(return_value=[])
await manager._get_tools_from_server(
server=self._token_exchange_server("te-list-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.get_prompts_from_server(
server=self._token_exchange_server("te-prompts-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-split"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subject_token_given_to_client(manager) == self._USER_TOKEN
@pytest.mark.asyncio
async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
server = MCPServer(
server_id="te-openapi",
name="te_openapi",
server_name="te_openapi",
url=None,
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
spec_path="https://api.example.com/openapi.json",
)
user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice")
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN]
@pytest.mark.asyncio
async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
with pytest.raises(HTTPException) as exc_info:
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert exc_info.value.status_code == 401
headers = exc_info.value.headers or {}
assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "")
assert self._subjects_seen_by(provider) == []
@pytest.mark.asyncio
async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subjects_seen_by(provider) == [self._USER_TOKEN]