mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(mcp): never exchange the LiteLLM virtual key as the upstream subject token (#39446)
This commit is contained in:
parent
bcd3e2d94d
commit
291d02f8aa
4 changed files with 396 additions and 27 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue