mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(proxy): mint the exchange credential off the database user row, not the cache
JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a team member never evicts the cached user row, so the mint read a row with no teams and refused the very first token exchange for a never-seen user as not a member. The loader now reads the row from the database and leaves the fresh row in the cache for the requests the credential makes next
This commit is contained in:
parent
7cb01cf47f
commit
3a8ac47e99
2 changed files with 44 additions and 5 deletions
|
|
@ -273,7 +273,11 @@ 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."""
|
||||
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
|
||||
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."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
|
|
@ -296,6 +300,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,
|
||||
)
|
||||
except (ProxyException, HTTPException):
|
||||
return "no_active_key"
|
||||
|
|
|
|||
|
|
@ -7568,6 +7568,36 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_
|
|||
assert await _reload_active_user_by_id("sso-user-7") == "faulted"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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."""
|
||||
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
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable
|
||||
)
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"])
|
||||
)
|
||||
proxy_globals.user_api_key_cache = cache
|
||||
proxy_globals.prisma_client = prisma
|
||||
|
||||
loaded = await load_active_user_by_id("fresh-jwt-user")
|
||||
|
||||
assert not isinstance(loaded, str)
|
||||
assert loaded.teams == ["team-a"]
|
||||
cached = await cache.async_get_cache(key="fresh-jwt-user", model_type=LiteLLM_UserTable)
|
||||
assert cached is not None
|
||||
assert cached.teams == ["team-a"]
|
||||
|
||||
|
||||
@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
|
||||
|
|
@ -11866,13 +11896,17 @@ async def test_oauth_refresh_revalidates_the_same_active_user_rule(
|
|||
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"})
|
||||
)
|
||||
user_id: Final = f"jwt-owner-{state}"
|
||||
row: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": state != "inactive"})
|
||||
proxy_server.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row)
|
||||
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
|
||||
assert await _reload_active_user_by_id(user_id) == expected
|
||||
if state != "missing_database":
|
||||
cached: Final = handler.user_api_key_cache.get_cache(user_id, model_type=LiteLLM_UserTable)
|
||||
assert cached is not None
|
||||
assert cached.metadata == row.metadata
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue