From 5bfd6432fddbda6f8b45a082fc88d1e676d1d3f2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 30 Jul 2026 19:31:36 -0700 Subject: [PATCH] fix(cache_warming): drop a session once its captured tenancy stops holding A record binds a payload to a tenant, and a replay spends that tenant's budget, consumes its rate limits, reaches the models it grants and fans out to its logging callbacks. So the only safe question at replay time is whether the tenancy stamped at capture still holds, and until now only one way of failing that question was handled. Keyless callers were skipped because their tenancy cannot be re-read. A virtual key whose tenancy CHANGED was not: the row was read, the principal took the new tenant wholesale, and the old payload was replayed under it. A key moved between teams therefore attributed prompts captured under the previous team to the new one and fanned them out to the new team's logging callbacks, where its members could read them. Neither available answer is right on its own, which is the tell that the record itself is no longer valid: keeping the captured tenant spends a budget for a key that has left it, and adopting the current tenant discloses the old tenant's prompts to the new one. _unverifiable_keyless_tenancy is generalized rather than joined by a second predicate, since both cases are the same question with different answers. It now returns why the tenancy failed, or None when it is confirmed to still hold, and the tick skips anything that is not a confirmed match. The session re-captures under whatever is true on its next real turn, so the cost of being wrong is one skipped warm. --- .../cache_warming/refresher.py | 76 +++++++++++++------ .../cache_warming/test_refresher.py | 47 +++++++----- 2 files changed, 82 insertions(+), 41 deletions(-) diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index 2df0119676f..9549873c724 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -4,7 +4,7 @@ import time import uuid from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import litellm from litellm._logging import verbose_router_logger @@ -47,6 +47,19 @@ _RECONSTRUCTED_IDENTITY_FIELDS = ("user_id", "team_id", "org_id", "project_id") _REPLAY_MAX_OUTPUT_TOKENS = 1 +_STALE_TENANCY_WARNINGS = { + "unverifiable": ( + "cache_warming: a keyless caller's tenancy is stamped once at capture and has no key row to " + "re-read, so it cannot be revalidated when a replay runs; sessions carrying it are skipped " + "rather than replayed on a tenant the caller may no longer belong to" + ), + "reassigned": ( + "cache_warming: a key's tenancy changed after its sessions were captured, so those payloads belong " + "to a tenant the key has left; they are skipped rather than replayed under the new tenant, and " + "re-capture on the session's next real turn" + ), +} + def _replay_surface(payload: CacheWarmingPayload) -> "tuple[str, CallTypesLiteral]": """Route and call type as the owning endpoints declare them, so hooks branching on either see the @@ -246,18 +259,34 @@ def _stamp_identity(data: "dict[str, object]", principal: "UserAPIKeyAuth") -> N ) -def _unverifiable_keyless_tenancy(attribution: CacheWarmingAttribution) -> bool: - """A keyless caller (JWT, and anything else the proxy authenticated without a virtual key) leaves no row - to re-read, so the tenancy the proxy stamped at capture is the only copy of it and no tick can tell - whether the caller still belongs to that team, organization, project or user. Replaying on it would - spend that tenant's budget, consume its rate limits and reach the models it grants for as long as the - record lives, with nothing able to notice the association was withdrawn. Authority that cannot be - resolved at use time is not authority warming may spend under, so those sessions are skipped. A caller - with no recorded tenancy at all is untouched: there is nothing to go stale and its replay stays - unattributed exactly as before.""" - return attribution.user_api_key is None and any( - getattr(attribution, f"user_api_key_{field}") is not None for field in _RECONSTRUCTED_IDENTITY_FIELDS - ) +def _tenancy_no_longer_holds( + attribution: CacheWarmingAttribution, key_state: "UserAPIKeyAuth | None" +) -> "Literal['unverifiable', 'reassigned'] | None": + """One question asked of every captured session: does the tenancy stamped at capture still hold? + + A record binds a payload to a tenant, and a replay spends that tenant's budget, consumes its rate limits, + reaches the models it grants and fans out to its logging callbacks. Authority that cannot be confirmed at + use time is not authority warming may spend under, so anything but a confirmed match is skipped, and the + session re-captures under whatever is true on its next real turn. + + ``unverifiable``: a keyless caller (JWT, and anything the proxy authenticated without a virtual key) + leaves no row to re-read, so the captured tenancy is the only copy and no tick can tell whether the + caller still belongs to it. + + ``reassigned``: the row exists and disagrees. A key moved between teams keeps its sessions, and replaying + them attributes a payload captured under the old tenant to the new one, whose members then read those + prompts through their own logging callbacks. Neither answer is right on its own here: keeping the + captured tenant spends a tenant's budget for a key no longer in it, and adopting the current one hands + the old tenant's prompts to the new one. The record simply is not valid any more. + + A caller with no recorded tenancy at all is untouched: there is nothing to go stale, and its replay stays + unattributed exactly as before. end_user is deliberately not compared, being request-scoped rather than a + property of the key; callers differing there already land on separate partitions.""" + captured = tuple(getattr(attribution, f"user_api_key_{field}") for field in _RECONSTRUCTED_IDENTITY_FIELDS) + if key_state is None: + return "unverifiable" if any(value is not None for value in captured) else None + current = tuple(getattr(key_state, field) for field in _RECONSTRUCTED_IDENTITY_FIELDS) + return "reassigned" if captured != current else None def _excluded_from_warming(key_state: "UserAPIKeyAuth", now: "datetime", proxy_logging_obj: "ProxyLogging") -> bool: @@ -577,15 +606,18 @@ class CacheWarmingRefresher: for key in attributed if (row := key_states.get(key)) is None or _excluded_from_warming(row, checked_at, proxy_logging_obj) ) - keyless_sessions = frozenset( - key for key, record, _ in active if _unverifiable_keyless_tenancy(record.attribution) - ) - if keyless_sessions: - warn_once( - "cache_warming: a keyless caller's tenancy is stamped once at capture and has no key row to " - "re-read, so it cannot be revalidated when a replay runs; sessions carrying it are skipped " - "rather than replayed on a tenant the caller may no longer belong to" + stale_tenancy = { + key: reason + for key, record, _ in active + if ( + reason := _tenancy_no_longer_holds( + record.attribution, key_states.get(record.attribution.user_api_key or "") + ) ) + is not None + } + for reason in frozenset(stale_tenancy.values()): + warn_once(_STALE_TENANCY_WARNINGS[reason]) semaphore = asyncio.Semaphore(self.max_concurrent_replays) outcomes = await asyncio.gather( *( @@ -607,7 +639,7 @@ class CacheWarmingRefresher: lease_lost=lease_lost, ) for key, record, touched in active - if record.attribution.user_api_key not in excluded_keys and key not in keyless_sessions + if record.attribution.user_api_key not in excluded_keys and key not in stale_tenancy ), return_exceptions=True, ) diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py index 421aec6ce0b..c4eb6159b66 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py @@ -14,7 +14,7 @@ import pytest from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.cache_warming.refresher import ( - _unverifiable_keyless_tenancy, + _tenancy_no_longer_holds, filter_cache_warmable, ) from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore @@ -382,15 +382,23 @@ async def test_a_keyless_callers_captured_tenancy_is_never_replayed_on(): @pytest.mark.parametrize("field", ["user_id", "team_id", "org_id", "project_id"]) -def test_every_recorded_tenancy_field_makes_a_keyless_record_unverifiable(field): - """Each of the four ids selects tenant objects that carry their own budgets, rate limits and model grants, - so any one of them going stale is enough to make a replay spend under an association the caller may have - lost. A record with no tenancy at all keeps warming unattributed, which is the shape a direct SDK caller - has and nothing about it can go stale.""" - keyless = CacheWarmingAttribution(**{f"user_api_key_{field}": "stale-tenant"}) - assert _unverifiable_keyless_tenancy(keyless) is True - assert _unverifiable_keyless_tenancy(CacheWarmingAttribution()) is False - assert _unverifiable_keyless_tenancy(keyless.model_copy(update={"user_api_key": "hash-1"})) is False +def test_a_captured_tenancy_is_warmed_only_while_it_is_confirmed_to_still_hold(field): + """Each of the four ids selects tenant objects carrying their own budgets, rate limits, model grants and + logging callbacks, so any one of them going stale is enough to make a replay spend and disclose under an + association that no longer applies. Only a confirmed match warms: an unreadable tenancy and a contradicted + one are both skipped, because keeping the captured tenant spends for a key that left it while adopting the + current one hands the old tenant's prompt to the new one. A record with no tenancy at all is untouched, + which is the direct-SDK shape where nothing can go stale.""" + from litellm.proxy._types import UserAPIKeyAuth + + captured = CacheWarmingAttribution(user_api_key="hash-1", **{f"user_api_key_{field}": "tenant-a"}) + keyless = captured.model_copy(update={"user_api_key": None}) + + assert _tenancy_no_longer_holds(keyless, None) == "unverifiable" + assert _tenancy_no_longer_holds(CacheWarmingAttribution(), None) is None + assert _tenancy_no_longer_holds(captured, UserAPIKeyAuth(**{field: "tenant-a"})) is None + assert _tenancy_no_longer_holds(captured, UserAPIKeyAuth(**{field: "tenant-b"})) == "reassigned" + assert _tenancy_no_longer_holds(captured, UserAPIKeyAuth()) == "reassigned" @pytest.mark.asyncio @@ -497,15 +505,16 @@ async def test_a_pooled_tier_warms_the_model_the_session_was_actually_served(): @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.""" +@pytest.mark.parametrize("current_team", [None, "team-B"]) +async def test_a_session_is_dropped_once_its_key_leaves_the_tenant_it_was_captured_under(current_team): + """Capture stores the tenancy the key had; the tick re-reads it. Once those disagree there is no correct + way to replay the record, which is why it is dropped rather than reconciled. Keeping the captured tenant + spends a budget, consumes rate limits and reaches models on behalf of a key no longer in it. Adopting the + current one attributes a payload captured under the old tenant to the new one and fans it out to the new + tenant's logging callbacks, so its members read prompts captured before the key ever reached them. Both + arms are covered: removed from a team, and moved to another one.""" 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")}) + keys = FakeKeyDirectory({"k": key_state(token="k", team_id=current_team)}) 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) + assert llm_router.completion_calls == []