From 61e3b5ddae1fdbbf12f6fa087b8bf0ee3e318d12 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:31:07 -0700 Subject: [PATCH] fix(mcp): persist OAuth credentials for rowless JWT admins --- .../mcp_server/bridge_token_flow.py | 5 ++- litellm/proxy/auth/handle_jwt.py | 36 +++++++++++-------- .../mcp_server/test_discoverable_endpoints.py | 18 +++++++++- 3 files changed, 40 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 6d11a4d2607..b693cc046c8 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -371,10 +371,9 @@ async def _extract_jwt_user_id(request: Request, token: str) -> str | None: identity_only=True, ) resolved_user: Final = identity["user_object"] - if resolved_user is None: + if resolved_user is not None and isinstance(_active_user_record(resolved_user), str): return None - owner: Final = _active_user_record(resolved_user) - return None if isinstance(owner, str) else identity["user_id"] + return 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 d5fcc578167..2d8cfb614ce 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -62,6 +62,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.user_repository import UserRepository from litellm.types.agents import AgentResponse +from litellm.types.proxy.auth.auth_checks import UserNotFoundError from .auth_checks import ( _allowed_routes_check, @@ -2343,21 +2344,26 @@ class JWTAuthManager: ) 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, - ) + try: + 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, + ) + except UserNotFoundError: + if not jwt_handler.is_admin(scopes=scopes): + raise + identity_user, identity_user_id = None, user_id return JWTAuthBuilderResult( is_proxy_admin=False, # Admin admission uses the claim ID; other callers use the canonical DB ID. 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 bf557667892..f6dc7932e2f 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 @@ -11444,11 +11444,13 @@ def _oauth_identity_jwt( @pytest.mark.parametrize("header", ["Authorization", "x-litellm-api-key"]) @pytest.mark.parametrize("policy_allowed", [False, True]) @pytest.mark.parametrize("admin", [False, True]) +@pytest.mark.parametrize("owner_state", ["active", "missing", "inactive", "database_error"]) async def test_oauth_exchange_stores_token_for_validated_jwt_user( jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], header: str, policy_allowed: bool, admin: bool, + owner_state: str, monkeypatch: pytest.MonkeyPatch, ) -> None: import httpx @@ -11475,6 +11477,7 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server + from litellm.models.user import LiteLLM_UserTable from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.types.llms.custom_http import httpxSpecialProvider @@ -11485,6 +11488,18 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( return httpx.Response(200, json={"access_token": "upstream-token", "token_type": "Bearer"}) 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() + if owner_state in ("missing", "database_error"): + handler.user_api_key_cache.delete_cache("jwt-owner") + if owner_state == "database_error": + users.find_unique.side_effect = RuntimeError("database unavailable") + if owner_state == "inactive": + handler.user_api_key_cache.set_cache( + "jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": False}) + ) table: Final = database.db.litellm_mcpusercredentials table.find_unique = AsyncMock(return_value=None) table.upsert = AsyncMock() @@ -11509,7 +11524,8 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( ) assert response.status_code == 200 assert json.loads(response.body)["access_token"] == "upstream-token" - if not policy_allowed: + users.create.assert_not_awaited() + if not policy_allowed or owner_state in ("inactive", "database_error") or (owner_state == "missing" and not admin): table.upsert.assert_not_awaited() return table.upsert.assert_awaited_once()