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"