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", )