mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
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.
This commit is contained in:
parent
3a8ac47e99
commit
167edf2769
4 changed files with 78 additions and 13 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue