fix(cache_warming): skip sessions whose tenancy cannot be revalidated

A keyless caller (JWT, and anything else the proxy authenticated without a
virtual key) leaves no key 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 spends
that tenant's budget, consumes its rate limits and reaches the models it grants
for as long as the record lives, with nothing able to notice the association
was withdrawn

A virtual key has the fresh row to answer with, which is why a loaded row wins
wholesale there. This shape has nothing to answer with at all, and authority
that cannot be resolved at use time is not authority warming may spend under,
so the session is skipped instead. A caller with no recorded tenancy is
untouched: there is nothing to go stale and its replay stays unattributed

_replay_principal still restores the recorded fields for that shape rather than
dropping them, so if the skip ever misses a record the replay is
over-constrained by a stale tenant instead of escaping every tenant control on
an empty principal
This commit is contained in:
Tin Chi Lo 2026-07-30 19:07:01 -07:00
parent 4ff2ecd4c8
commit 9a6aa43af9
3 changed files with 57 additions and 16 deletions

View file

@ -15,6 +15,7 @@ from litellm.router_strategy.complexity_router.cache_warming.store import CacheW
from litellm.router_strategy.complexity_router.cache_warming.types import (
CACHE_WARMING_REPLAY_MARKER_KEY,
CACHE_WARMING_REPLAY_TAG,
CacheWarmingAttribution,
CacheWarmingPayload,
CacheWarmingRecord,
decompress_payload,
@ -69,11 +70,11 @@ def _replay_principal(key_state: "UserAPIKeyAuth | None", record: CacheWarmingRe
(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
stamped their tenancy, which the record carries, so every recorded identity field is restored below and
the team, user, organization and project gates resolve by id exactly as for a virtual key. Dropping any
of them is what previously let a JWT caller's replays run against an empty principal and escape every
tenant control
api_key is None) carrying tenancy: never reaches this function, because _unverifiable_keyless_tenancy
skips the session. There is no row to re-read for such a caller, so the tenancy the proxy stamped at
capture is the only copy and no tick can tell whether it still holds. The restore below still covers
that shape rather than dropping the fields, so that if the skip ever misses a record the replay is
over-constrained by a stale tenant instead of escaping every tenant control on an empty principal
(c) direct SDK use with no proxy auth object: nothing was recorded, so there is no tenancy to preserve and
the principal is genuinely limitless, which is the same unattributed warming as before
@ -245,6 +246,20 @@ 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 _excluded_from_warming(key_state: "UserAPIKeyAuth", now: "datetime", proxy_logging_obj: "ProxyLogging") -> bool:
"""Mirrors the canonical auth checks (user_api_key_auth.py:1717 and :2891) because common_checks owns
that policy for real traffic but needs a FastAPI Request. datetime.fromisoformat only accepts a
@ -562,6 +577,15 @@ 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"
)
semaphore = asyncio.Semaphore(self.max_concurrent_replays)
outcomes = await asyncio.gather(
*(
@ -583,7 +607,7 @@ class CacheWarmingRefresher:
lease_lost=lease_lost,
)
for key, record, touched in active
if record.attribution.user_api_key not in excluded_keys
if record.attribution.user_api_key not in excluded_keys and key not in keyless_sessions
),
return_exceptions=True,
)

View file

@ -13,11 +13,15 @@ from datetime import datetime, timedelta, timezone
import pytest
from litellm.caching.dual_cache import DualCache
from litellm.router_strategy.complexity_router.cache_warming.refresher import filter_cache_warmable
from litellm.router_strategy.complexity_router.cache_warming.refresher import (
_unverifiable_keyless_tenancy,
filter_cache_warmable,
)
from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore
from litellm.router_strategy.complexity_router.cache_warming.types import (
CACHE_WARMING_REPLAY_MARKER_KEY,
CACHE_WARMING_REPLAY_TAG,
CacheWarmingAttribution,
)
from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache
@ -356,14 +360,15 @@ async def test_warmth_is_not_shared_between_two_auto_routers_on_one_redis():
@pytest.mark.asyncio
@pytest.mark.parametrize("team_blocked,warmed", [(True, False), (False, True)])
async def test_a_keyless_proxy_caller_is_authorized_through_its_reconstructed_tenancy(team_blocked, warmed):
"""A JWT caller has api_key None, so before this it fell through to the unattributed path and every tenant
control was skipped. Its tenancy is recorded at capture, so the principal is rebuilt from it and the team
gate binds: a blocked team stops the replays, an open team still warms."""
async def test_a_keyless_callers_captured_tenancy_is_never_replayed_on():
"""A JWT caller has api_key None and so leaves no row to re-read, which makes the tenancy the proxy
stamped at capture the only copy of it: once an admin moves that caller out of the team, no tick can tell,
and every later replay still spends that team's budget and rate limits and reaches the models it grants.
A virtual key has the fresh row to answer with; this shape has nothing. The team here is open and fully
resolvable, so nothing but the skip itself can be what stops the replay."""
llm_router, redis = warming_rig(redis=FakeRedisCache())
key_cache = DualCache()
await key_cache.async_set_cache(key="team_id:jwt-team", value=team("jwt-team", blocked=team_blocked, models=[]))
await key_cache.async_set_cache(key="team_id:jwt-team", value=team("jwt-team", blocked=False, models=[]))
seed_session(
redis,
user_api_key=None,
@ -373,9 +378,19 @@ async def test_a_keyless_proxy_caller_is_authorized_through_its_reconstructed_te
touched=_VISITED_BOTH_TIERS,
)
await tick(llm_router, active=refresher(keys=FakeKeyDirectory({})), user_api_key_cache=key_cache)
assert bool(llm_router.completion_calls) is warmed
if warmed:
assert llm_router.completion_calls[0]["metadata"]["user_api_key_team_id"] == "jwt-team"
assert llm_router.completion_calls == []
@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
@pytest.mark.asyncio

View file

@ -219,6 +219,7 @@ def seed_session(
team_id: str | None = None,
user_id: str | None = None,
org_id: str | None = None,
project_id: str | None = None,
touched: tuple[str, ...] | None = None,
) -> str:
payload = CacheWarmingPayload(
@ -243,6 +244,7 @@ def seed_session(
user_api_key_team_id=team_id,
user_api_key_user_id=user_id,
user_api_key_org_id=org_id,
user_api_key_project_id=project_id,
),
auto_router_model_name="smart-router",
)