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 == []