scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Merge PR #28008 (gym-cmd/litellm:feat/v1.84.0-mcp-gateway-jwt-auth) into litellm_feat/v1.84.0-mcp-gateway-jwt-auth

Browse source 
# Conflicts:
#	litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
#	litellm/proxy/_experimental/mcp_server/server.py
#	litellm/proxy/management_endpoints/mcp_management_endpoints.py
#	litellm/proxy/proxy_server.py
#	litellm/types/mcp_server/mcp_server_manager.py
#	tests/test_litellm/interactions/test_openapi_compliance.py
#	tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
#	tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-05-21 03:41:08 +00:00
commit 6697fdb03d
No known key found for this signature in database
9 changed files with 622 additions and 184 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -126,6 +126,19 @@ def decode_state_hash(encrypted_state: str) -> dict:
return state_data
def _get_validated_client_redirect_uri(
request: Request, state_data: Dict[str, Any]
) -> str:
"""Return a trusted (same-origin, loopback, or ops-allowlisted)
client redirect URI from OAuth state.
"""
redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url")
if not redirect_uri or not isinstance(redirect_uri, str):
raise HTTPException(status_code=400, detail="Invalid redirect URI")
validate_trusted_redirect_uri(request, redirect_uri)
return redirect_uri
def _append_query_params(url: str, params: Dict[str, str]) -> str:
parsed = urlparse(url)
query_params = parse_qsl(parsed.query, keep_blank_values=True)
@ -677,11 +690,6 @@ async def token_endpoint(
async def callback(request: Request, code: str, state: str):
try:
state_data = decode_state_hash(state)
redirect_uri = state_data.get("client_redirect_uri") or state_data.get(
"base_url"
)
if not redirect_uri or not isinstance(redirect_uri, str):
raise HTTPException(status_code=400, detail="Invalid redirect URI")
original_state = state_data["original_state"]
# Re-validate the client redirect URI at the sink. /authorize
@ -690,7 +698,7 @@ async def callback(request: Request, code: str, state: str):
# expiry and remain valid indefinitely. Validating here blocks
# the open-redirect + code-theft primitive even for pre-fix
# states while permitting same-origin / allowlisted clients.
validate_trusted_redirect_uri(request, redirect_uri)
redirect_uri = _get_validated_client_redirect_uri(request, state_data)
params = {"code": code, "state": original_state}
complete_returned_url = _append_query_params(redirect_uri, params)

70
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2835,6 +2835,30 @@ if MCP_AVAILABLE:
)
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
def _get_passthrough_resource_metadata_url(scope: Scope, server_name: str) -> str:
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
_path = scope.get("_original_path") or scope.get("path", "") or ""
if _path.startswith(f"/{server_name}/mcp"):
return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp"
return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}"
def _get_passthrough_www_authenticate(
scope: Scope,
server_name: str,
invalid_token: bool = False,
) -> str:
resource_metadata_url = _get_passthrough_resource_metadata_url(
scope=scope,
server_name=server_name,
)
params = []
if invalid_token:
params.append('error="invalid_token"')
params.append(f'resource_metadata="{resource_metadata_url}"')
return "Bearer " + ", ".join(params)
async def _raise_preemptive_401_for_unauthenticated_servers(
scope: Scope,
mcp_servers: Optional[List[str]],
@ -2895,18 +2919,10 @@ if MCP_AVAILABLE:
server, oauth2_headers, mcp_server_auth_headers
)
):
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
_path = scope.get("_original_path") or scope.get("path", "") or ""
# Pick the well-known resource-metadata form that matches the inbound
# route pattern so the metadata's `resource` field round-trips to what
# the client actually hit (RFC 9728 §3.2).
if _path.startswith(f"/{server_name}/mcp"):
resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp"
else:
resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}"
www_authenticate = f'Bearer resource_metadata="{resource_metadata_url}"'
www_authenticate = _get_passthrough_www_authenticate(
scope=scope,
server_name=server_name,
)
raise HTTPException(
status_code=401,
detail="Unauthorized",
@ -3045,30 +3061,16 @@ if MCP_AVAILABLE:
for srv in passthrough_servers
]
)
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
_path = scope.get("_original_path") or scope.get("path", "") or ""
for srv, (probe_status, _) in zip(passthrough_servers, probe_results):
if probe_status == 401:
# Token is missing or expired — direct the client at the
# gateway's oauth-protected-resource well-known URL (which
# proxies the upstream IdP's metadata), matching the
# pre-emptive 401 path in
# _raise_preemptive_401_for_unauthenticated_servers. The
# gateway is not the authorization server for pass-through
# servers, so emitting ``authorization_uri=`` would point
# clients at the wrong AS metadata.
if _path.startswith(f"/{srv.name}/mcp"):
resource_metadata_url = (
f"{base_url}/.well-known/oauth-protected-resource/"
f"{srv.name}/mcp"
)
else:
resource_metadata_url = (
f"{base_url}/.well-known/oauth-protected-resource/mcp/"
f"{srv.name}"
)
www_authenticate = f'Bearer resource_metadata="{resource_metadata_url}"'
# Token is missing or expired: keep pass-through clients on the
# protected-resource discovery flow so they re-authorize against
# the upstream IdP metadata proxied by LiteLLM.
www_authenticate = _get_passthrough_www_authenticate(
scope=scope,
server_name=srv.name,
invalid_token=True,
)
raise HTTPException(
status_code=401,
detail="Unauthorized",

1
litellm/proxy/proxy_server.py
View file

@ -15858,6 +15858,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
)
if toolset is not None:
scope = dict(request.scope)
scope["_original_path"] = scope.get("path", "")
scope["path"] = "/mcp"
token = _mcp_active_toolset_id.set(toolset.toolset_id)
try:

1
tests/test_litellm/interactions/test_openapi_compliance.py
View file

@ -166,6 +166,7 @@ class TestResponseCompliance:
"created",
"updated",
"role",
"steps",
"usage",
]

2
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -1,7 +1,7 @@
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, call as mock_call, patch
import pytest
from fastapi.testclient import TestClient

59
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -5,23 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
TRUSTED_PROXY_IP = "10.0.0.5"
TRUSTED_PROXY_RANGES = ["10.0.0.0/8"]
def set_request_from_trusted_proxy(mock_request):
mock_request.client = MagicMock()
mock_request.client.host = TRUSTED_PROXY_IP
@pytest.fixture
def trusted_proxy_origin_headers():
with patch(
"litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy",
return_value=True,
):
yield
# Fixture to mock IP address check for all MCP tests
# This prevents tests from failing due to IP-based access control
@ -63,7 +46,7 @@ def trust_xff():
``test_get_request_base_url_xff_trust_gate``.
"""
with patch(
"litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy",
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy",
return_value=True,
):
yield
@ -592,7 +575,6 @@ async def test_authorize_endpoint_respects_x_forwarded_proto():
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://litellm.example.com/" # HTTP
mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy
set_request_from_trusted_proxy(mock_request)
# Mock the encryption functions
with patch(
@ -661,7 +643,6 @@ async def test_token_endpoint_respects_x_forwarded_proto():
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP
mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy
set_request_from_trusted_proxy(mock_request)
# Mock httpx client response
mock_response = MagicMock()
@ -740,7 +721,6 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto():
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://litellm.example.com/" # HTTP
mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy
set_request_from_trusted_proxy(mock_request)
# Call the endpoint
response = await oauth_protected_resource_mcp(
@ -796,7 +776,6 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto():
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://litellm.example.com/" # HTTP
mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy
set_request_from_trusted_proxy(mock_request)
# Call the endpoint
response = await oauth_authorization_server_mcp(
@ -835,7 +814,6 @@ async def test_register_client_respects_x_forwarded_proto():
mock_request = MagicMock(spec=Request)
mock_request.base_url = "http://proxy.litellm.example/" # HTTP
mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy
set_request_from_trusted_proxy(mock_request)
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
@ -898,7 +876,6 @@ async def test_authorize_endpoint_respects_x_forwarded_host():
"X-Forwarded-Proto": "https",
"X-Forwarded-Host": "proxy.example.com",
}
set_request_from_trusted_proxy(mock_request)
# Mock the encryption functions
with patch(
@ -971,7 +948,6 @@ async def test_token_endpoint_respects_x_forwarded_host():
"X-Forwarded-Proto": "https",
"X-Forwarded-Host": "proxy.example.com",
}
set_request_from_trusted_proxy(mock_request)
# Mock httpx client response
mock_response = MagicMock()
@ -1129,11 +1105,7 @@ async def test_token_endpoint_respects_x_forwarded_host():
],
)
def test_get_request_base_url_comprehensive(
base_url,
x_forwarded_proto,
x_forwarded_host,
x_forwarded_port,
expected_url,
base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url
):
"""Comprehensive test for get_request_base_url with various header combinations.
@ -1152,7 +1124,6 @@ def test_get_request_base_url_comprehensive(
mock_request = MagicMock(spec=Request)
mock_request.base_url = base_url
set_request_from_trusted_proxy(mock_request)
headers = {}
if x_forwarded_proto:
@ -1168,7 +1139,7 @@ def test_get_request_base_url_comprehensive(
mock_request.headers.get = mock_get
with patch(
"litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy",
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy",
return_value=True,
):
result = get_request_base_url(mock_request)
@ -2053,10 +2024,6 @@ async def test_oauth_callback_redirects_with_state():
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
# Mock the state decoding
mock_state_data = {
"base_url": "http://localhost:3000/ui/mcp/oauth/callback",
@ -2103,10 +2070,6 @@ async def test_oauth_callback_preserves_client_redirect_uri_query():
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
) as mock_decode:
@ -2144,10 +2107,6 @@ async def test_oauth_callback_handles_invalid_state():
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
# Mock state decoding to raise an exception
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
@ -2570,10 +2529,6 @@ async def test_callback_revalidates_loopback_on_decoded_base_url():
callback,
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
) as mock_decode:
@ -2600,10 +2555,6 @@ async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri():
callback,
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
) as mock_decode:
@ -2630,10 +2581,6 @@ async def test_callback_rejects_state_missing_redirect_uri():
callback,
)
mock_request = MagicMock()
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
) as mock_decode:

90
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -109,11 +109,7 @@ async def test_mcp_server_tool_call_body_contains_request_data():
assert body["arguments"] == tool_arguments
def test_prepare_mcp_server_headers_passthrough_omits_authorization():
"""Pass-through servers must not forward the inbound Authorization header
(LiteLLM API key) to the upstream server. The upstream must provide its own
bearer token via a server-specific header.
"""
def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
@ -122,82 +118,9 @@ def test_prepare_mcp_server_headers_passthrough_omits_authorization():
pytest.skip("MCP server not available")
server = MCPServer(
server_id="server-passthrough",
server_id="server-case",
name="server",
transport=MCPTransport.http,
auth_type=MCPAuth.none, # Explicitly none for pass-through
extra_headers=["Authorization"],
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers={"authorization": "Bearer sk-litellm-key"},
)
assert server_auth_header is None
# Authorization should NOT be forwarded for pass-through servers
assert extra_headers is None
def test_prepare_mcp_server_headers_passthrough_forwards_other_headers():
"""Pass-through servers should forward other headers (not Authorization)
from the raw request."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="server-passthrough-headers",
name="server",
transport=MCPTransport.http,
auth_type=MCPAuth.none, # Pass-through mode
extra_headers=["Authorization", "x-request-id", "x-trace-id"],
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers={
"authorization": "Bearer sk-litellm-key",
"x-request-id": "req-123",
"x-trace-id": "trace-456",
},
)
assert server_auth_header is None
# Authorization should be omitted, but other headers forwarded
assert extra_headers == {"x-request-id": "req-123", "x-trace-id": "trace-456"}
def test_prepare_mcp_server_headers_passthrough_forwards_authorization_with_explicit_admission():
"""Transparent OAuth pass-through: when LiteLLM admission used the explicit
`x-litellm-api-key` header, the inbound `Authorization` header is
unambiguously the upstream OAuth bearer and MUST be forwarded.
Regression for EAI-506 V5/V6 — a standards-compliant MCP client (e.g.
OpenCode) completes PKCE against the upstream IdP and sends the resulting
token as plain `Authorization: Bearer <token>` per the MCP spec.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="server-passthrough-explicit-admission",
name="server",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
extra_headers=["Authorization"],
)
@ -208,20 +131,15 @@ def test_prepare_mcp_server_headers_passthrough_forwards_authorization_with_expl
oauth2_headers=None,
raw_headers={
"x-litellm-api-key": "Bearer sk-litellm-key",
"authorization": "Bearer upstream-okta-token",
"authorization": "Bearer token",
},
)
assert server_auth_header is None
assert extra_headers == {"Authorization": "Bearer upstream-okta-token"}
assert extra_headers == {"Authorization": "Bearer token"}
def test_prepare_mcp_server_headers_passthrough_strips_authorization_without_admission_header():
"""Counterpart to the explicit-admission test: without `x-litellm-api-key`,
the inbound `Authorization` may itself be the LiteLLM admission key, so we
strip it to avoid leaking the gateway credential upstream. This preserves
the security guarantee introduced in commit 3753970cc9.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,

561
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -2754,3 +2754,564 @@ def test_build_decode_kwargs_no_warning_when_scoped(
if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert matching == []
def _base64url_encode_int(value: int) -> str:
import base64
value_bytes = value.to_bytes((value.bit_length() + 7) // 8, "big")
return base64.urlsafe_b64encode(value_bytes).decode("utf-8").rstrip("=")
def _get_rsa_key_and_jwk(kid: str):
from cryptography.hazmat.primitives.asymmetric import rsa
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
public_numbers = private_key.public_key().public_numbers()
jwk = {
"kty": "RSA",
"n": _base64url_encode_int(value=public_numbers.n),
"e": _base64url_encode_int(value=public_numbers.e),
"kid": kid,
"alg": "RS256",
"use": "sig",
}
return private_key, jwk
def _encode_rsa_jwt(
private_key,
issuer: str,
audience: str,
kid: str,
extra_claims: Optional[dict] = None,
) -> str:
import time
import jwt
from cryptography.hazmat.primitives import serialization
private_key_pem = private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
current_time = int(time.time())
claims = {
"sub": "test-subject",
"iss": issuer,
"aud": audience,
"iat": current_time,
"exp": current_time + 300,
}
if extra_claims:
claims.update(extra_claims)
return jwt.encode(
claims,
private_key_pem,
algorithm="RS256",
headers={"kid": kid},
)
def _get_jwt_handler_with_issuer_keys(issuers: list, keys_by_url: dict) -> JWTHandler:
from litellm.caching.dual_cache import DualCache
cache = DualCache()
for jwks_url, keys in keys_by_url.items():
cache.set_cache(
key=f"litellm_jwt_auth_keys_{jwks_url}",
value=keys,
)
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(issuers=issuers),
)
return jwt_handler
@pytest.mark.asyncio
async def test_get_public_key_fetches_and_caches_jwks_response():
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.dual_cache import DualCache
jwt_handler = JWTHandler()
cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(public_key_ttl=123),
)
expected_key_id = "cached-key"
_, jwk = _get_rsa_key_and_jwk(kid=expected_key_id)
mock_response = MagicMock()
mock_response.json.return_value = {"keys": [jwk]}
jwt_handler.http_handler.get = AsyncMock(return_value=mock_response)
public_key = await jwt_handler._get_public_key_from_jwks_url(
jwks_url="https://issuer.example.com/keys",
kid=expected_key_id,
)
assert public_key == jwk
cached_keys = await cache.async_get_cache(
key="litellm_jwt_auth_keys_https://issuer.example.com/keys"
)
assert cached_keys == [jwk]
@pytest.mark.asyncio
async def test_get_public_key_tries_next_jwks_url_when_kid_missing(monkeypatch):
from litellm.caching.dual_cache import DualCache
first_jwks_url = "https://first.example.com/keys"
second_jwks_url = "https://second.example.com/keys"
monkeypatch.setenv(
"JWT_PUBLIC_KEY_URL", f"{first_jwks_url}, {second_jwks_url},,"
)
_, first_jwk = _get_rsa_key_and_jwk(kid="first-key")
_, second_jwk = _get_rsa_key_and_jwk(kid="second-key")
cache = DualCache()
cache.set_cache(key=f"litellm_jwt_auth_keys_{first_jwks_url}", value=[first_jwk])
cache.set_cache(
key=f"litellm_jwt_auth_keys_{second_jwks_url}", value=[second_jwk]
)
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(),
)
public_key = await jwt_handler.get_public_key(kid="second-key")
assert public_key == second_jwk
def test_get_jwks_url_for_issuer_falls_back_to_discovery_document():
jwt_handler = JWTHandler()
issuer_config = LiteLLM_JWTAuth(
issuers=[{"issuer": "https://issuer.example.com/tenant/"}]
).issuers[0]
jwks_url = jwt_handler._get_jwks_url_for_issuer(issuer_config=issuer_config)
assert (
jwks_url
== "https://issuer.example.com/tenant/.well-known/openid-configuration"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_validates_selected_issuer_and_maps_claims(
monkeypatch,
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer_one = "https://issuer-one.example.com"
issuer_two = "https://issuer-two.example.com"
issuer_one_jwks_url = f"{issuer_one}/keys"
issuer_two_jwks_url = f"{issuer_two}/keys"
shared_kid = "shared-kid"
_, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
issuer_two_private_key, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer_one,
"jwks_url": issuer_one_jwks_url,
"audience": "audience-one",
"user_id_jwt_field": "email",
"user_email_jwt_field": "email",
},
{
"issuer": issuer_two,
"jwks_url": issuer_two_jwks_url,
"audience": "audience-two",
"user_id_jwt_field": "repository_owner",
"team_id_jwt_field": "repository",
},
],
keys_by_url={
issuer_one_jwks_url: [issuer_one_jwk],
issuer_two_jwks_url: [issuer_two_jwk],
},
)
token = _encode_rsa_jwt(
private_key=issuer_two_private_key,
issuer=issuer_two,
audience="audience-two",
kid=shared_kid,
extra_claims={
"repository_owner": "example-org",
"repository": "example-org/litellm-fork",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert claims[JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == issuer_two
assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-org"
assert jwt_handler.get_team_id(token=claims, default_value=None) == (
"example-org/litellm-fork"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://oidc.eks.eu-west-1.amazonaws.com/id/test-cluster"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="k8s-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": None,
"disable_audience_validation": True,
"user_id_jwt_field": "kubernetes\\.io.namespace",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="kubernetes.default.svc",
kid="k8s-key",
extra_claims={"kubernetes.io": {"namespace": "example-namespace"}},
)
claims = await jwt_handler.auth_jwt(token=token)
assert (
jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_rejects_unknown_issuer(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
configured_issuer = "https://issuer.example.com"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": configured_issuer,
"jwks_url": f"{configured_issuer}/keys",
"audience": "expected-audience",
}
],
keys_by_url={f"{configured_issuer}/keys": [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer="https://unknown-issuer.example.com",
audience="expected-audience",
kid="issuer-key",
)
with pytest.raises(Exception) as exc:
await jwt_handler.auth_jwt(token=token)
assert "Unsupported JWT issuer" in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="wrong-audience",
kid="issuer-key",
)
with pytest.raises(Exception) as exc:
await jwt_handler.auth_jwt(token=token)
assert "Validation fails" in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer_one = "https://issuer-one.example.com"
issuer_two = "https://issuer-two.example.com"
issuer_one_jwks_url = f"{issuer_one}/keys"
issuer_two_jwks_url = f"{issuer_two}/keys"
shared_kid = "shared-kid"
issuer_one_private_key, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
_, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid)
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer_one,
"jwks_url": issuer_one_jwks_url,
"audience": "audience-one",
},
{
"issuer": issuer_two,
"jwks_url": issuer_two_jwks_url,
"audience": "audience-two",
},
],
keys_by_url={
issuer_one_jwks_url: [issuer_one_jwk],
issuer_two_jwks_url: [issuer_two_jwk],
},
)
token = _encode_rsa_jwt(
private_key=issuer_one_private_key,
issuer=issuer_two,
audience="audience-two",
kid=shared_kid,
)
with pytest.raises(Exception) as exc:
await jwt_handler.auth_jwt(token=token)
assert "Validation fails" in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_missing_mapped_claim_fails_closed(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
"user_id_jwt_field": "email",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
)
with pytest.raises(Exception) as exc:
await jwt_handler.auth_jwt(token=token)
assert "missing required mapped claim: email" in str(exc.value)
assert "Validation fails" not in str(exc.value)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled(
monkeypatch,
):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="some-other-client",
kid="issuer-key",
)
with pytest.raises(Exception) as exc:
await jwt_handler.auth_jwt(token=token)
assert "must configure audience" in str(exc.value)
@pytest.mark.asyncio
async def test_global_jwt_ignores_user_supplied_internal_claims(monkeypatch):
from litellm.caching.dual_cache import DualCache
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
jwks_url = "https://global-issuer.example.com/keys"
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
private_key, jwk = _get_rsa_key_and_jwk(kid="global-key")
cache = DualCache()
cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk])
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(
user_id_jwt_field="email",
user_email_jwt_field="email",
team_id_jwt_field="team.id",
team_ids_jwt_field="teams",
org_id_jwt_field="org.id",
end_user_id_jwt_field="end_user.id",
),
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer="https://global-issuer.example.com",
audience="some-other-client",
kid="global-key",
extra_claims={
"email": "real-user@example.com",
"team": {"id": "real-team"},
"teams": ["real-team", "secondary-team"],
"org": {"id": "real-org"},
"end_user": {"id": "real-end-user"},
JWTHandler.LITELLM_JWT_ISSUER_CLAIM: "https://issuer.example.com",
JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user",
JWTHandler.LITELLM_USER_EMAIL_CLAIM: "victim@example.com",
JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team",
JWTHandler.LITELLM_TEAM_IDS_CLAIM: ["victim-team"],
JWTHandler.LITELLM_ORG_ID_CLAIM: "victim-org",
JWTHandler.LITELLM_END_USER_ID_CLAIM: "victim-end-user",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert jwt_handler.get_user_id(token=claims, default_value=None) == (
"real-user@example.com"
)
assert jwt_handler.get_user_email(token=claims, default_value=None) == (
"real-user@example.com"
)
assert jwt_handler.get_team_id(token=claims, default_value=None) == "real-team"
assert jwt_handler.get_team_ids_from_jwt(token=claims) == [
"real-team",
"secondary-team",
]
assert jwt_handler.get_org_id(token=claims, default_value=None) == "real-org"
assert jwt_handler.get_end_user_id(token=claims, default_value=None) == (
"real-end-user"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_strips_unmapped_internal_claims(monkeypatch):
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
"user_email_jwt_field": "email",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
extra_claims={
"email": "real-user@example.com",
JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user",
JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team",
},
)
claims = await jwt_handler.auth_jwt(token=token)
assert JWTHandler.LITELLM_USER_ID_CLAIM not in claims
assert JWTHandler.LITELLM_TEAM_ID_CLAIM not in claims
assert jwt_handler.get_user_id(token=claims, default_value=None) is None
assert jwt_handler.get_team_id(token=claims, default_value=None) is None
assert jwt_handler.get_user_email(token=claims, default_value=None) == (
"real-user@example.com"
)
@pytest.mark.asyncio
async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning(
monkeypatch, caplog
):
import logging
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
monkeypatch.delenv("JWT_ISSUER", raising=False)
monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False)
JWTHandler._unscoped_jwt_warning_emitted = False
issuer = "https://issuer.example.com"
jwks_url = f"{issuer}/keys"
private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key")
jwt_handler = _get_jwt_handler_with_issuer_keys(
issuers=[
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "expected-audience",
}
],
keys_by_url={jwks_url: [jwk]},
)
token = _encode_rsa_jwt(
private_key=private_key,
issuer=issuer,
audience="expected-audience",
kid="issuer-key",
)
with caplog.at_level(logging.WARNING):
await jwt_handler.auth_jwt(token=token)
assert "Tokens minted by any application" not in caplog.text
assert JWTHandler._unscoped_jwt_warning_emitted is False

2
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -2,7 +2,7 @@ import os
import sys
import types
import json
from datetime import datetime, timedelta, timezone
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import List, Optional
from unittest.mock import AsyncMock, MagicMock, patch