From 167edf2769b28f76870b654014f99312843cc327 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:10:03 -0700 Subject: [PATCH] fix(proxy): read the database user row only in the credential mint The token exchange mint keeps reading the user row from the database, since JWT auth caches the user it creates before adding it to the JWT's team and a mint off that cached row refused the first exchange for a new user. Introspection and the refresh revalidation go back to the cache read, so a resource server calling /introspect per request pays no database read. --- .../mcp_server/bridge_token_flow.py | 18 ++++++--- .../mcp_server/proxy_api_credentials.py | 4 +- .../mcp_server/test_discoverable_endpoints.py | 39 +++++++++++++++++-- .../mcp_server/test_proxy_api_credentials.py | 30 +++++++++++++- 4 files changed, 78 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 4235471f2d9..f19cb87ae18 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -262,7 +262,12 @@ 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) -> "LiteLLM_UserTable | _KeyResolutionFailure": +UserRowSource = Literal["cache", "database"] + + +async def load_active_user_by_id( + user_id: str, source: UserRowSource = "cache" +) -> "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 @@ -273,11 +278,12 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. The - row is read from the database, never the cache: JWT auth caches the user it creates before it adds + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. + ``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the + cache for the requests the credential makes next: JWT auth caches the user it creates before it adds that user to the JWT's team and adding a member never evicts the cached row, so a credential minted - off the cache would refuse the very first exchange as not a member. The fresh row replaces the cached - one.""" + off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache + read, so introspection, which a resource server may call per request, stays off the database.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -300,7 +306,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, - check_db_only=True, + check_db_only=source == "database", ) except (ProxyException, HTTPException): return "no_active_key" diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index a34119edf10..2f7fcaef645 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -42,7 +42,7 @@ async def mint_proxy_credential( user_id: str, team_id: str | None ) -> MintedProxyCredential | ProxyCredentialMintFailure: """Mint the ``lite login`` credential for a consented grant. Membership is checked - live, so a team the user left between consent and redemption (or between refreshes) + live against the database row, so a team the user left between consent and redemption (or between refreshes) refuses the grant instead of minting a credential attributed to a team they are no longer on. The team is exactly the one the consent page sealed into the grant; nothing is picked on the user's behalf here, so a refresh can never move the credential, and a @@ -54,7 +54,7 @@ async def mint_proxy_credential( the minter's own first-team fallback stays inert. The credential carries the role the proxy already enforces for the user on every request, so a row with no role (JWT auth's upsert writes none) mints as an internal user instead of being refused.""" - user: Final = await load_active_user_by_id(user_id) + user: Final = await load_active_user_by_id(user_id, source="database") if isinstance(user, str): return user if team_id is not None and team_id not in user.teams: 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 30b179f1a26..6965b3b4ebe 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 @@ -7572,8 +7572,8 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals): """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member never evicts the cached row, so a credential minted off the cached row refused the very first - token exchange as not a member. The loader has to read the row from the database and leave the fresh - row in the cache for the requests the credential makes next.""" + token exchange as not a member. The database source has to read the row from the database and leave + the fresh row in the cache for the requests the credential makes next.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -7589,7 +7589,7 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca proxy_globals.user_api_key_cache = cache proxy_globals.prisma_client = prisma - loaded = await load_active_user_by_id("fresh-jwt-user") + loaded = await load_active_user_by_id("fresh-jwt-user", source="database") assert not isinstance(loaded, str) assert loaded.teams == ["team-a"] @@ -7598,6 +7598,39 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca assert cached.teams == ["team-a"] +@pytest.mark.asyncio +async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_read(proxy_globals): + """Introspection and refresh revalidation run per call, so the loader's default source is the cache: a + cached row answers without a database read, and only a caller that asks for the database row pays for + one.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _reload_active_user_by_id, + load_active_user_by_id, + ) + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="cached-jwt-user", + value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=["team-a"]), + model_type=LiteLLM_UserTable, + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=[]) + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = prisma + + loaded = await load_active_user_by_id("cached-jwt-user") + + assert not isinstance(loaded, str) + assert loaded.teams == ["team-a"] + assert await _reload_active_user_by_id("cached-jwt-user") is None + prisma.db.litellm_usertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 85650c6a05a..04fbbe4a6ce 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -1,6 +1,6 @@ """Tests for minting the ``lite login`` credential from a consented native-client grant.""" -from unittest.mock import ANY, AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock import pytest @@ -10,6 +10,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail _LOAD_USER = "litellm.proxy._experimental.mcp_server.proxy_api_credentials.load_active_user_by_id" @@ -91,7 +92,7 @@ async def test_mint_refuses_a_teamless_grant_for_a_team_member(load_user, fetch_ is refused for a user with teams instead of minting an unscoped credential or drifting onto the first team, on redemption and on every refresh alike.""" assert await mint_proxy_credential("u1", None) == "team_required" - load_user.assert_awaited_once_with("u1") + load_user.assert_awaited_once_with("u1", source="database") fetch_teams.assert_awaited_once_with(ANY, ["team-a", "team-b"]) @@ -126,6 +127,31 @@ async def test_mint_honors_the_consented_team(load_user, fetch_teams): assert decoded.team_model_aliases == {"fast": "gpt-5.4-mini"} +@pytest.mark.asyncio +async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_row(fetch_teams, monkeypatch): + """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member + never evicts the cached row, so a mint off the cached row refused the very first token exchange as not + a member. The mint has to read the database row, whatever the cache holds.""" + from litellm.proxy import proxy_server + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=_user(user_id="stale-cache-user", teams=["team-a"]) + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + minted = await mint_proxy_credential("stale-cache-user", "team-a") + + assert isinstance(minted, MintedProxyCredential) + assert minted.team_id == "team-a" + assert _decoded(minted).team_id == "team-a" + + @pytest.mark.asyncio async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams): assert await mint_proxy_credential("u1", "team-c") == "not_a_member"