From 53318796fd27d7e59bb86850f163836dda7a45e8 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:13:10 -0700 Subject: [PATCH] fix(mcp): separate JWT identity lookup from request authorization --- .../mcp_server/bridge_token_flow.py | 22 +-- litellm/proxy/auth/handle_jwt.py | 71 ++++++++-- .../mcp_server/test_discoverable_endpoints.py | 126 +++++++++++++----- .../proxy/auth/test_handle_jwt.py | 71 +++++++++- 4 files changed, 237 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 2f3e50b803f..6d11a4d2607 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -198,9 +198,7 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No return loaded if isinstance(loaded, str) else None -async def load_active_user_by_id( - user_id: str, *, sso_user_id: str | None = None, user_email: str | None = None -) -> "LiteLLM_UserTable | _KeyResolutionFailure": +async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure": """Load a live litellm user by id, returning the record when the user is active or a precise failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a @@ -234,8 +232,6 @@ async def load_active_user_by_id( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, - sso_user_id=sso_user_id, - user_email=user_email, ) except (ProxyException, HTTPException): return "no_active_key" @@ -247,6 +243,10 @@ async def load_active_user_by_id( return "no_active_key" if user_object is None: return "no_active_key" + return _active_user_record(user_object) + + +def _active_user_record(user_object: "LiteLLM_UserTable") -> "LiteLLM_UserTable | Literal['no_active_key']": if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False: return "no_active_key" return user_object @@ -336,7 +336,7 @@ async def _extract_jwt_user_id(request: Request, token: str) -> str | None: user_api_key_cache, ) - if general_settings.get("enable_jwt_auth") is not True or premium_user is not True: + if general_settings.get("enable_jwt_auth") is not True or premium_user is not True or prisma_client is None: return None try: if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured(): @@ -368,13 +368,13 @@ async def _extract_jwt_user_id(request: Request, token: str) -> str | None: proxy_logging_obj=proxy_logging_obj, request_headers=dict(request.headers), request_method=request.method, + identity_only=True, ) - owner_id: Final = identity["user_id"] - if not owner_id: + resolved_user: Final = identity["user_object"] + if resolved_user is None: return None - # Admin JWTs can return before auth_builder loads the canonical database user. - owner: Final = await load_active_user_by_id(owner_id, sso_user_id=owner_id, user_email=identity["user_email"]) - return None if isinstance(owner, str) else owner.user_id + owner: Final = _active_user_record(resolved_user) + return None if isinstance(owner, str) else identity["user_id"] except Exception as exc: # noqa: BLE001 # public OAuth exchange stays available; unvalidated identities never write credentials verbose_logger.debug("OAuth JWT identity could not be validated (%s)", type(exc).__name__) return None diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4304542fc83..678769b375b 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1674,6 +1674,7 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, route: str, org_alias: str | None = None, + user_id_upsert: bool | None = None, ) -> tuple[ LiteLLM_UserTable | None, LiteLLM_OrganizationTable | None, @@ -1737,7 +1738,11 @@ class JWTAuthManager: user_id=user_id, user_email=user_email, sso_user_id=user_id, - upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email), + upsert=( + jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email) + if user_id_upsert is None + else user_id_upsert + ), ), team_id=team_id, ) @@ -2209,8 +2214,14 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, request_headers: dict | None = None, request_method: str | None = None, + identity_only: bool = False, ) -> JWTAuthBuilderResult: - """Main authentication and authorization builder""" + """Build JWT authentication and authorization context. + + Public OAuth endpoints use identity_only to resolve an existing credential owner + without authorizing the OAuth route or provisioning users/teams. The returned + identity does not grant permission to execute an MCP or model request. + """ # Check if OIDC UserInfo endpoint is enabled, but fall back to standard # JWT auth if the token itself is a well-formed JWT (3-part structure). if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key): @@ -2231,18 +2242,23 @@ class JWTAuthManager: # Check RBAC rbac_role: Final = jwt_handler.get_rbac_role(token=jwt_valid_token) - await JWTAuthManager.check_rbac_role( - jwt_handler, - jwt_valid_token, - general_settings, - request_data, - route, - rbac_role, - ) + if not identity_only: + await JWTAuthManager.check_rbac_role( + jwt_handler, + jwt_valid_token, + general_settings, + request_data, + route, + rbac_role, + ) # Check Scope Based Access scopes: Final = jwt_handler.get_scopes(token=jwt_valid_token) - if jwt_handler.litellm_jwtauth.enforce_scope_based_access and jwt_handler.litellm_jwtauth.scope_mappings: + if ( + not identity_only + and jwt_handler.litellm_jwtauth.enforce_scope_based_access + and jwt_handler.litellm_jwtauth.scope_mappings + ): JWTAuthManager.check_scope_based_access( scope_mappings=jwt_handler.litellm_jwtauth.scope_mappings, scopes=scopes, @@ -2268,6 +2284,39 @@ class JWTAuthManager: elif rbac_role == LitellmUserRoles.INTERNAL_USER: user_id = object_id + if identity_only: + identity_user, _, _, _, identity_user_id = await JWTAuthManager.get_objects( + user_id=user_id, + user_email=user_email, + org_id=None, + end_user_id=None, + team_id=None, + valid_user_email=valid_user_email, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + user_id_upsert=False, + ) + return JWTAuthBuilderResult( + is_proxy_admin=False, + # Admin admission uses the claim ID; other callers use the canonical DB ID. + user_id=user_id if jwt_handler.is_admin(scopes=scopes) else identity_user_id, + user_email=identity_user.user_email if identity_user is not None else user_email, + user_object=identity_user, + team_id=None, + team_object=None, + org_id=None, + org_object=None, + end_user_id=None, + end_user_object=None, + team_membership=None, + token=api_key, + jwt_claims=jwt_valid_token, + ) + # Check admin access admin_result: Final = await JWTAuthManager.check_admin_access( jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index d1ab8809861..bf557667892 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7127,12 +7127,12 @@ async def test_build_oauth_protected_resource_response_obo_end_to_end(): global_mcp_server_manager.registry.clear() -def _token_request(headers): +def _token_request(headers, path="/token"): """A real Starlette request with case-insensitive headers (matches production).""" from starlette.requests import Request raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()] - return Request({"type": "http", "method": "POST", "path": "/token", "headers": raw, "query_string": b""}) + return Request({"type": "http", "method": "POST", "path": path, "headers": raw, "query_string": b""}) @pytest.fixture @@ -11443,10 +11443,12 @@ def _oauth_identity_jwt( @pytest.mark.asyncio @pytest.mark.parametrize("header", ["Authorization", "x-litellm-api-key"]) @pytest.mark.parametrize("policy_allowed", [False, True]) +@pytest.mark.parametrize("admin", [False, True]) async def test_oauth_exchange_stores_token_for_validated_jwt_user( jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], header: str, policy_allowed: bool, + admin: bool, monkeypatch: pytest.MonkeyPatch, ) -> None: import httpx @@ -11456,9 +11458,9 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( from litellm.types.mcp_server.mcp_server_manager import MCPServer handler, signing_key = jwt_oauth_identity - handler.litellm_jwtauth.enforce_team_based_model_access = not policy_allowed - bearer: Final = _oauth_identity_jwt(signing_key) - request: Final = _token_request({header: f"Bearer {bearer}"}) + handler.litellm_jwtauth.custom_validate = lambda claims: policy_allowed + bearer: Final = _oauth_identity_jwt(signing_key, scope="litellm_proxy_admin" if admin else "") + request: Final = _token_request({header: f"Bearer {bearer}"}, path="/jwt-oauth-server/token") server: Final = MCPServer( server_id="jwt-oauth-server", name="jwt-oauth-server", @@ -11534,8 +11536,6 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( "scim_inactive", "custom_validate", "missing_database", - "denied_route", - "required_team", ], ) async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( @@ -11548,7 +11548,6 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request - from litellm.proxy._types import LitellmUserRoles, RoleBasedPermissions handler, signing_key = jwt_oauth_identity key: Final = ( @@ -11573,18 +11572,6 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( ) if rejection == "custom_validate": handler.litellm_jwtauth.custom_validate = lambda claims: False - if rejection == "denied_route": - handler.litellm_jwtauth.enforce_rbac = True - monkeypatch.setattr( - proxy_server, - "general_settings", - { - "enable_jwt_auth": True, - "role_permissions": [RoleBasedPermissions(role=LitellmUserRoles.INTERNAL_USER, routes=["/models"])], - }, - ) - if rejection == "required_team": - handler.litellm_jwtauth.enforce_team_based_model_access = True assert await _extract_user_id_from_request(_token_request({"Authorization": f"Bearer {bearer}"})) is None @@ -11666,7 +11653,7 @@ async def test_oauth_jwt_respects_custom_validation_and_email_policy( @pytest.mark.asyncio @pytest.mark.parametrize("route_allowed", [False, True]) -async def test_oauth_jwt_uses_rbac_user_object_id( +async def test_oauth_jwt_identity_preserves_separate_mcp_route_authorization( jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], monkeypatch: pytest.MonkeyPatch, route_allowed: bool, @@ -11674,6 +11661,7 @@ async def test_oauth_jwt_uses_rbac_user_object_id( from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request from litellm.proxy._types import LitellmUserRoles, RoleBasedPermissions, RoleMapping + from litellm.proxy.auth.handle_jwt import JWTAuthManager handler, signing_key = jwt_oauth_identity handler.litellm_jwtauth.user_id_jwt_field = "sub" @@ -11691,13 +11679,32 @@ async def test_oauth_jwt_uses_rbac_user_object_id( "role_permissions": [ RoleBasedPermissions( role=LitellmUserRoles.INTERNAL_USER, - routes=["/token"] if route_allowed else ["/models"], + routes=["mcp_routes"] if route_allowed else ["/models"], ) ], }, ) - request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}) - assert await _extract_user_id_from_request(request) == ("jwt-owner" if route_allowed else None) + bearer: Final = _oauth_identity_jwt(signing_key) + request: Final = _token_request({"Authorization": f"Bearer {bearer}"}, path="/example/token") + assert await _extract_user_id_from_request(request) == "jwt-owner" + admission: Final = JWTAuthManager.auth_builder( + api_key=bearer, + jwt_handler=handler, + request_data={}, + general_settings=proxy_server.general_settings, + route="/mcp/example", + prisma_client=proxy_server.prisma_client, + user_api_key_cache=handler.user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_server.proxy_logging_obj, + request_method="POST", + ) + if route_allowed: + assert (await admission)["user_id"] == "jwt-owner" + else: + with pytest.raises(HTTPException) as denial: + await admission + assert denial.value.status_code == 403 @pytest.mark.asyncio @@ -11714,11 +11721,12 @@ async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity( from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request + from litellm.proxy.auth.handle_jwt import JWTAuthManager handler, signing_key = jwt_oauth_identity external_id: Final = f"external-{identity}-{inactive}-{admin}" handler.litellm_jwtauth.user_email_jwt_field = "email" - handler.litellm_jwtauth.admin_allowed_routes = ["/token"] + handler.litellm_jwtauth.admin_allowed_routes = ["mcp_routes"] owner: Final = LiteLLM_UserTable( user_id="canonical-oauth-owner", user_email="owner@example.test", @@ -11727,7 +11735,7 @@ async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity( ) database: Final = MagicMock() table: Final = database.db.litellm_usertable - table.find_unique = AsyncMock(side_effect=[None, owner if identity == "sso" else None, owner]) + table.find_unique = AsyncMock(side_effect=[None, owner if identity == "sso" else None]) table.find_first = AsyncMock(return_value=owner) table.update = AsyncMock(return_value=owner) monkeypatch.setattr(proxy_server, "prisma_client", database) @@ -11735,9 +11743,67 @@ async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity( signing_key, owner=external_id, scope="litellm_proxy_admin" if admin else "" ) request: Final = _token_request({"Authorization": f"Bearer {bearer}"}) - assert await _extract_user_id_from_request(request) == (None if inactive else "canonical-oauth-owner") - assert table.find_unique.await_count == (2 if admin else 3) - if not admin: - assert table.find_unique.call_args.kwargs["where"] == {"user_id": "canonical-oauth-owner"} + stored_owner: Final = await _extract_user_id_from_request(request) + assert stored_owner == (None if inactive else external_id if admin else "canonical-oauth-owner") + assert table.find_unique.await_count == 2 if identity == "email": table.find_first.assert_awaited_once() + if not inactive: + admission: Final = await JWTAuthManager.auth_builder( + api_key=bearer, + jwt_handler=handler, + request_data={}, + general_settings=proxy_server.general_settings, + route="/mcp/example", + prisma_client=database, + user_api_key_cache=handler.user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + assert stored_owner == admission["user_id"] + + +@pytest.mark.asyncio +async def test_oauth_jwt_identity_does_not_provision_or_synchronize_teams( + jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], +) -> None: + from litellm.models.user import LiteLLM_UserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request + + handler, signing_key = jwt_oauth_identity + handler.litellm_jwtauth.enforce_team_based_model_access = True + handler.litellm_jwtauth.team_id_default = "new-team" + handler.litellm_jwtauth.team_id_upsert = True + handler.litellm_jwtauth.sync_user_role_and_teams = True + owner: Final = LiteLLM_UserTable(user_id="jwt-owner", teams=["existing-team"]) + handler.user_api_key_cache.set_cache("jwt-owner", owner) + request: Final = _token_request( + {"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}, path="/example/token" + ) + assert await _extract_user_id_from_request(request) == "jwt-owner" + assert owner.teams == ["existing-team"] + proxy_server.prisma_client.db.litellm_teamtable.find_unique.assert_not_called() + proxy_server.prisma_client.db.litellm_teamtable.upsert.assert_not_called() + proxy_server.prisma_client.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["active", "inactive", "missing_database"]) +async def test_oauth_refresh_revalidates_the_same_active_user_rule( + jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], + monkeypatch: pytest.MonkeyPatch, + state: str, +) -> None: + from litellm.models.user import LiteLLM_UserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id + + handler, _ = jwt_oauth_identity + handler.user_api_key_cache.set_cache( + "jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": state != "inactive"}) + ) + if state == "missing_database": + monkeypatch.setattr(proxy_server, "prisma_client", None) + expected: Final = None if state == "active" else "no_active_key" if state == "inactive" else "unresolvable" + assert await _reload_active_user_by_id("jwt-owner") == expected diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 94226b5404d..9e15a442a18 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -2,7 +2,7 @@ import asyncio import re import time from collections.abc import Mapping, Sequence -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException @@ -6786,3 +6786,72 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla } assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == [] assert user.teams == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity_only", [False, True]) +@pytest.mark.parametrize("existing_user", [False, True]) +@pytest.mark.parametrize("model_allowed", [False, True]) +async def test_auth_builder_identity_lookup_does_not_provision_users( + monkeypatch: pytest.MonkeyPatch, identity_only: bool, existing_user: bool, model_allowed: bool +) -> None: + from litellm.proxy._types import ScopeMapping + from litellm.proxy.auth.auth_checks import UserNotFoundError + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + private_key, jwk = _get_rsa_key_and_jwk("identity-mode") + cache: Final = UserApiKeyCache() + cache.set_cache("litellm_jwt_auth_keys_https://identity.example/jwks", [jwk]) + user_id: Final = f"identity-mode-{identity_only}-{existing_user}-{model_allowed}" + user: Final = LiteLLM_UserTable(user_id=user_id, organization_memberships=[]) + if existing_user: + cache.set_cache(user_id, user) + database: Final = MagicMock() + users: Final = database.db.litellm_usertable + users.find_unique = AsyncMock(return_value=None) + users.find_first = AsyncMock(return_value=None) + users.create = AsyncMock(return_value=user) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + user_id_upsert=True, + enforce_scope_based_access=True, + scope_mappings=[ScopeMapping(scope="allowed", models=["allowed-model"])], + ), + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://identity.example/jwks") + monkeypatch.setenv("JWT_ISSUER", "https://identity.example") + monkeypatch.setenv("JWT_AUDIENCE", "gateway") + token: Final = _encode_rsa_jwt( + private_key, "https://identity.example", "gateway", "identity-mode", {"sub": user_id, "scope": "allowed"} + ) + pending: Final = JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=handler, + request_data={"model": "allowed-model" if model_allowed else "forbidden-model"}, + general_settings={}, + route="/example/token" if identity_only else "/mcp/example", + prisma_client=database, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + identity_only=identity_only, + ) + if not identity_only and not model_allowed: + with pytest.raises(HTTPException) as denial: + await pending + assert denial.value.status_code == 403 + users.create.assert_not_awaited() + return + if identity_only and not existing_user: + with pytest.raises(UserNotFoundError): + await pending + else: + result: Final = await pending + assert result["user_id"] == user_id + assert result["user_object"] is not None + assert result["user_object"].user_id == user_id + assert users.create.await_count == (0 if identity_only or existing_user else 1)