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:
Tin Chi Lo 2026-07-30 11:48:02 -07:00
parent 049c1abef5
commit b20f210abc
2 changed files with 29 additions and 4 deletions

View file

@ -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,

View file

@ -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)