From 9a6aa43af937ff196a785743d22ec25d5dfdd661 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 30 Jul 2026 19:07:01 -0700 Subject: [PATCH] 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 --- .../cache_warming/refresher.py | 36 +++++++++++++++---- .../cache_warming/test_refresher.py | 35 ++++++++++++------ .../cache_warming/warming_rig.py | 2 ++ 3 files changed, 57 insertions(+), 16 deletions(-) diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index 07c60f21bc4..2df0119676f 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -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, ) 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 7910ebcbb28..421aec6ce0b 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 @@ -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 diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py index c8fa0e5ac66..9d4050073e7 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py @@ -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", )