mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(cache_warming): a loaded key row wins wholesale, never field by field
The replay principal merged the freshly loaded key row with the tenancy captured on the record, taking the record's value whenever the row's was falsy. That cannot distinguish "this field was never captured" from "this association was deliberately removed", and those need opposite outcomes. A key an admin takes out of a team is still valid and still loads; the row correctly reports no team. The merge restored the captured team instead, so every replay for the rest of that session ran as if the key were still a member: reaching models the team grants, spending the team's budget, and consuming the team's rate limits on behalf of a key no longer in it. The same held for organization, project and user. Key deletion and revocation were never affected, since those are dropped a step earlier when the row fails to load at all. The request path has no equivalent seam to get wrong: it re-derives everything from the token on each request and a None simply skips that gate. Warming is the only place holding a snapshot, so it now follows the same rule. A loaded row is used verbatim, Nones included, and the captured tenancy is restored only when there is no row to read, which is the keyless case where the record is the sole source.
This commit is contained in:
parent
049c1abef5
commit
b20f210abc
2 changed files with 29 additions and 4 deletions
|
|
@ -59,6 +59,14 @@ def _replay_principal(key_state: "UserAPIKeyAuth | None", record: CacheWarmingRe
|
|||
"""A replay's principal must be COMPLETE, because the authorization enumerations resolve tenancy by id and
|
||||
a missing field is an absent ceiling rather than a denial. Three shapes of capture, one reconstruction:
|
||||
|
||||
A loaded row wins WHOLESALE, never field by field. A per-field merge cannot tell "this field was never
|
||||
captured" from "this association was deliberately removed", and those need opposite outcomes: a key that
|
||||
an admin took out of a team still loads fine, with team_id correctly None, and falling back to the
|
||||
captured team would let the replay spend that team's budget, consume its rate limits and reach models it
|
||||
grants, on behalf of a key that is no longer in it. The same holds for org, project and user. This is
|
||||
exactly what the request path does, where every gate reads the freshly loaded row and a None simply skips
|
||||
the gate; warming is the only place with a snapshot to be tempted by.
|
||||
|
||||
(a) virtual key: key_state was just read authoritatively from the database, so it is already complete
|
||||
(b) keyless proxy caller (JWT, and anything else the proxy authenticated without a virtual key, where
|
||||
api_key is None): UserAPIKeyAuth is litellm's principal type for those callers too and the proxy
|
||||
|
|
@ -75,13 +83,15 @@ def _replay_principal(key_state: "UserAPIKeyAuth | None", record: CacheWarmingRe
|
|||
|
||||
attribution = record.attribution
|
||||
base = key_state if key_state is not None else UserAPIKeyAuth(api_key=attribution.user_api_key)
|
||||
identity = (
|
||||
{}
|
||||
if key_state is not None
|
||||
else {field: getattr(attribution, f"user_api_key_{field}") for field in _RECONSTRUCTED_IDENTITY_FIELDS}
|
||||
)
|
||||
return base.model_copy(
|
||||
update={ # mutable-ok: pydantic model_copy input, never retained
|
||||
"api_key": base.api_key or base.token,
|
||||
**{
|
||||
field: getattr(base, field) or getattr(attribution, f"user_api_key_{field}")
|
||||
for field in _RECONSTRUCTED_IDENTITY_FIELDS
|
||||
},
|
||||
**identity,
|
||||
"end_user_id": attribution.user_api_key_end_user_id or base.end_user_id,
|
||||
"request_route": route,
|
||||
"budget_reservation": None,
|
||||
|
|
|
|||
|
|
@ -479,3 +479,18 @@ async def test_a_pooled_tier_warms_the_model_the_session_was_actually_served():
|
|||
seed_session(redis, served_model="haiku-b")
|
||||
await tick(llm_router)
|
||||
assert "haiku-b" in replayed_models(llm_router), "the session's own cache must be kept warm"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_key_removed_from_its_team_does_not_keep_replaying_as_that_team():
|
||||
"""The capture stores the tenancy the key had at the time; the tick re-reads the key. When an admin takes
|
||||
a still-valid key out of a team, the fresh row says team_id is None and that is the current truth, so
|
||||
falling back to the captured team would let the replay spend that team's budget, consume its rate limits
|
||||
and reach models it grants, on behalf of a key no longer in it. The request path never merges the two, it
|
||||
reads the loaded row and skips the gate on None, so a loaded row has to win wholesale here too."""
|
||||
llm_router, redis = warming_rig(redis=FakeRedisCache())
|
||||
seed_session(redis, user_api_key="k", team_id="team-A", touched=_VISITED_BOTH_TIERS)
|
||||
keys = FakeKeyDirectory({"k": key_state(token="k")})
|
||||
await tick(llm_router, active=refresher(keys=keys))
|
||||
assert llm_router.completion_calls, "expected a replay"
|
||||
assert all(call["metadata"]["user_api_key_team_id"] is None for call in llm_router.completion_calls)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue