feat(cache_warming): warm only the models a session has been served on

Warming resolved its target set from configuration, so every active session was
replayed against a representative of every tier on every interval whether or not it
had ever been routed there. Most sessions never leave their starting tier, so that
spent N replays per interval to keep caches warm that nobody would read, and for a
pooled tier it warmed the wrong member entirely.

The set is now per session and comes from what the session actually did. Capture
records each served model into a per-session Redis set inside the same atomic script
that writes the record, sharing the session's hash tag, so the touched set cannot
disagree with the record it belongs to and expires with it. The refresher replays
exactly that set.

This is the intended shape of the feature: keep a session's own caches alive so
returning to a tier it has already used is a read, rather than pre-warming tiers on
speculation. The first switch to a new tier is a normal cache write, and every visit
after it is warm. warm_models changes meaning accordingly, from a pre-warm list to
an allowlist that narrows what a session may be warmed on, and resolve_warm_models
now returns every model across the tier pools since it bounds eligibility rather
than naming the targets.

Tests that expected a replay on a tier the seeded session had never visited now
declare a session that has been to both, which is the case the feature serves.

litellm/router_strategy/complexity_router/README.md

@@ -144,7 +144,9 @@ Technical code keywords are detected case-insensitively and include:
 
 ## Cache Warming
 
-Provider prompt caches (Anthropic, Bedrock) are per-model, so a mid-session tier switch pays a fresh cache write on the new model and loses the cache-read discount. `cache_warming` keeps every tier model's prompt cache warm for active sessions: the proxy captures each session's latest payload and a background refresher replays it (`max_tokens=1`) against the other tier models before the provider's ~5 minute cache TTL expires. When the router later switches tiers, the switched-to model already has the session's prefix cached, and the routing pick prefers models whose cache is verifiably warm.
+Provider prompt caches (Anthropic, Bedrock) are per-model, so a session that moves between tiers pays a fresh cache write every time it lands on a model whose cache has expired, and loses the cache-read discount. `cache_warming` keeps a session's own prompt caches alive: the proxy captures each session's latest payload and a background refresher replays it (`max_tokens=1`) against the models that session has been served on, before the provider's ~5 minute cache TTL expires. Returning to a tier the session has already used is then a cache read rather than a fresh write, and the routing pick prefers models whose cache is verifiably warm.
+
+Warming follows the session rather than the configuration: a model is warmed for a session only once that session has actually been routed to it. The first switch to a new tier is therefore a normal cache write, and every subsequent visit to that tier is warm. This is deliberate; pre-warming tiers a session may never reach would spend on caches nobody reads, and most sessions never leave their starting tier.
 
 ```yaml
 model_list:
@@ -162,7 +164,7 @@ model_list:
           session_ttl_seconds: 3600
           idle_timeout_seconds: 600       # stop warming a session this long after its last real request
           max_sessions: 1000
-          # warm_models: [fast-claude, smart-claude]  # default: first member of each tier pool
+          # warm_models: [fast-claude, smart-claude]  # default: every model across the tier pools
 
 general_settings:
   store_prompts_in_spend_logs: true   # consent gate; warming stores full payloads in Redis
@@ -182,7 +184,7 @@ Requirements and semantics:
 - **Guardrails run on replays.** Because a replay goes through `pre_call_hook`, the guardrails configured for the model group (globally and on the deployment) also run on its warming replays. A blocking guardrail costs one skipped warm; a content-rewriting guardrail makes the warm ineffective rather than wrong, because a rewritten prefix simply is not the prefix real traffic sends. Sessions on a key that declares `max_iterations` are skipped instead of warmed, because that limiter counts every request on a `session_id` and cannot be consulted without incrementing it, so warming would consume the caller's own iteration budget.
 - **Multi-deployment groups require deployment affinity.** A replay routes by group name, so with several deployments it warms one member's cache while real traffic spreads across all of them: paying the cache-write premium against 1/N routing odds is worse than not warming, so such a group is skipped (with a one-time warning naming the group and both remedies) unless `DeploymentAffinityCheck` is active for it with the session_id mode, enabled globally via `router_settings.optional_pre_call_checks: ["session_affinity"]` or per group via `router_settings.model_group_affinity_config`. When active, replays carry the session's `session_id` (and the originating key hash), so the affinity check pins warming and real traffic to the same deployment. Single-deployment groups warm regardless. More than one deployment is used as the conservative stand-in for "more than one provider cache domain"
 - **`max_sessions`** caps concurrently warmed sessions per auto-router, enforced atomically at capture; once reached, new sessions are not admitted until existing ones expire.
-- **Interplay with `session_affinity`** (default on): affinity pins a session to its first-turn model, so no tier switch happens and warming buys nothing; with affinity on, captured sessions are still warmed but the pin decides routing. Disable `session_affinity` to let per-turn classification switch tiers and have warming make those switches cache hits.
+- **Interplay with `session_affinity`** (default on): affinity pins a session to its first-turn model, so it only ever touches that one model and warming keeps exactly that cache alive, which is worth having for sessions with long gaps between turns but buys nothing on switching, since no switch happens. Disable `session_affinity` to let per-turn classification move a session between tiers; the first visit to each tier is a normal cache write, and every return to a tier the session has already used is a cache read.
 
 ## Performance
 

litellm/router_strategy/complexity_router/cache_warming/eligibility.py

@@ -7,11 +7,16 @@ if TYPE_CHECKING:
 
 
 def resolve_warm_models(config: "ComplexityRouterConfig") -> tuple[str, ...]:
+    """Every model a session on this auto-router could end up being served by, which is the universe warming
+    draws from rather than the set it warms. A session is warmed only on the members it has actually been
+    served on, so this is used to bound eligibility and the capture-time minimum-token gate. An explicit
+    warm_models narrows the universe rather than widening it: warming never replays against a model the
+    session has not used, so the operator setting is an allowlist over the pool, not a pre-warm list."""
     explicit = config.cache_warming.warm_models
     if explicit:
         return tuple(dict.fromkeys(explicit))
-    first_per_tier = (models if isinstance(models, str) else models[0] for models in config.tiers.values() if models)
-    return tuple(dict.fromkeys(first_per_tier))
+    pooled = (models if isinstance(models, list) else [models] for models in config.tiers.values() if models)
+    return tuple(dict.fromkeys(model for pool in pooled for model in pool))
 
 
 def min_prompt_cache_tokens_for_warm_set(warm_models: tuple[str, ...]) -> int:

litellm/router_strategy/complexity_router/cache_warming/refresher.py

@@ -516,11 +516,12 @@ class CacheWarmingRefresher:
         )
         if not active:
             return
-        tier_models = resolve_warm_models(complexity_router.config)
+        touched = {key: await store.get_touched_models(key) for key, _ in active}
+        allowed = frozenset(resolve_warm_models(complexity_router.config))
         warmable = frozenset(
             filter_cache_warmable(
                 llm_router,
-                tuple(dict.fromkeys((*tier_models, *(record.served_model for _, record in active)))),
+                tuple(dict.fromkeys(model for models in touched.values() for model in models if model in allowed)),
             )
         )
         if not warmable:
@@ -554,7 +555,7 @@ class CacheWarmingRefresher:
                     session_key=key,
                     record=record,
                     warm_models=tuple(
-                        model for model in dict.fromkeys((record.served_model, *tier_models)) if model in warmable
+                        model for model in dict.fromkeys((record.served_model, *touched[key])) if model in warmable
                     ),
                     refresh_interval_seconds=config.refresh_interval_seconds,
                     session_ttl_seconds=config.session_ttl_seconds,

litellm/router_strategy/complexity_router/cache_warming/store.py

@@ -19,11 +19,13 @@ _WARMTH_KEY_PREFIX = "complexity_router_cache_warmth:v1"
 _CAPTURE_SCRIPT = """
 local sessions_key = KEYS[1]
 local index_key = KEYS[2]
+local touched_key = KEYS[3]
 local member = ARGV[1]
 local record_json = ARGV[2]
 local now = tonumber(ARGV[3])
 local expires_at = tonumber(ARGV[4])
 local max_sessions = tonumber(ARGV[5])
+local served_model = ARGV[6]
 local expired = redis.call('ZRANGEBYSCORE', index_key, 0, now)
 if #expired > 0 then
     redis.call('HDEL', sessions_key, unpack(expired))
@@ -34,11 +36,17 @@ if not redis.call('ZSCORE', index_key, member) and redis.call('ZCARD', index_key
 end
 redis.call('HSET', sessions_key, member, record_json)
 redis.call('ZADD', index_key, expires_at, member)
+redis.call('SADD', touched_key, served_model)
 redis.call('EXPIREAT', sessions_key, math.ceil(expires_at))
 redis.call('EXPIREAT', index_key, math.ceil(expires_at))
+redis.call('EXPIREAT', touched_key, math.ceil(expires_at))
 return 1
 """
 
+_TOUCHED_MODELS_SCRIPT = """
+return redis.call('SMEMBERS', KEYS[1])
+"""
+
 _LIST_LIVE_SESSIONS_SCRIPT = """
 local index_key = KEYS[1]
 local now = tonumber(ARGV[1])
@@ -114,6 +122,7 @@ class CacheWarmingStore:
             register(_LIST_LIVE_SESSIONS_SCRIPT) if register else None
         )
         self._get: Callable[..., Awaitable[object]] | None = register(_GET_RECORD_SCRIPT) if register else None
+        self._touched: Callable[..., Awaitable[object]] | None = register(_TOUCHED_MODELS_SCRIPT) if register else None
 
     @staticmethod
     def record_key(auto_router_model_name: str, caller_scope: str, session_id: str) -> str:
@@ -122,6 +131,14 @@ class CacheWarmingStore:
         so without the router in it two warming auto-routers sharing one Redis read each other's warmth."""
         return f"{auto_router_model_name}:{caller_scope}:{session_id}"
 
+    @staticmethod
+    def touched_key(record_key: str) -> str:
+        """The models this session has actually been served on, hash-tagged into the session's own slot family
+        and written inside the capture script so it can never disagree with the record it belongs to. Warming
+        refreshes exactly this set, so a session pays only for caches it has demonstrably used."""
+        auto_router_model_name, _, session_scope = record_key.partition(":")
+        return f"{{cache_warm:v1:{auto_router_model_name}}}:touched:v1:{session_scope}"
+
     @staticmethod
     def warmth_key(record_key: str, model_group: str) -> str:
         """Hash-tagged into the same slot family as the record it belongs to, so a Cluster keeps a session's
@@ -147,6 +164,16 @@ class CacheWarmingStore:
         raw = await self._get(keys=[self.sessions_key()], args=[key])
         return _parse_record(raw)
 
+    async def get_touched_models(self, key: str) -> tuple[str, ...]:
+        if self._require_redis() is None or self._touched is None:
+            return ()
+        raw = await self._touched(keys=[self.touched_key(key)], args=[])
+        try:
+            members = _MEMBERS_ADAPTER.validate_python(raw)
+        except ValidationError:
+            return ()
+        return tuple(member.decode() if isinstance(member, bytes) else member for member in members)
+
     async def upsert_session(
         self,
         *,
@@ -178,8 +205,8 @@ class CacheWarmingStore:
         )
         try:
             admitted = await self._capture(
-                keys=[self.sessions_key(), self.index_key()],
-                args=[key, record.model_dump_json(), now, now + ttl_seconds, max_sessions],
+                keys=[self.sessions_key(), self.index_key(), self.touched_key(key)],
+                args=[key, record.model_dump_json(), now, now + ttl_seconds, max_sessions, served_model],
             )
         except Exception:  # noqa: BLE001  # a capture fault fails closed: no capture beats an uncapped write
             verbose_router_logger.warning("cache_warming capture script failed; skipping capture", exc_info=True)

litellm/router_strategy/complexity_router/config.py

@@ -290,7 +290,9 @@ class CacheWarmingConfig(BaseModel):
     warm_models: tuple[str, ...] | None = Field(
         default=None,
         description=(
-            "Explicit model groups to keep warm; defaults to the first member of each tier pool. "
+            "Restricts warming to these model groups; defaults to every model across the tier pools. "
+            "A session is only ever warmed on models it has actually been served on, so this narrows that "
+            "set rather than pre-warming models the session has not used. "
             "Only Anthropic/Bedrock models that support prompt caching are warmed"
         ),
     )

tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py

@@ -89,30 +89,6 @@ def _stored_records(redis: FakeRedisCache) -> list[dict]:
     return [json.loads(value) for value in redis.hashes.get(SESSIONS_KEY, {}).values()]
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 @pytest.mark.asyncio
 async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth():
     redis = FakeRedisCache()
@@ -132,36 +108,6 @@ async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth()
     assert second["payload_sha256"] != first["payload_sha256"]
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "opt_out",

tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py

@@ -42,6 +42,9 @@ from tests.test_litellm.router_strategy.complexity_router.cache_warming.warming_
 )
 
 _PAST = (datetime.now(timezone.utc) - timedelta(hours=1)).replace(microsecond=0)
+# warming refreshes only models a session has actually been served on, so a test that expects a replay on
+# the other tier has to seed a session that has already been there
+_VISITED_BOTH_TIERS = ("fast-claude", "smart-claude")
 
 
 @pytest.mark.asyncio
@@ -121,7 +124,7 @@ async def test_every_ceiling_the_request_path_enforces_gates_warming(arm, warmed
     counter_key = f"spend:key:{token}"
     if arm == "over-budget":
         await spend_counter_cache.async_set_cache(key=counter_key, value=100.0)
-    seed_session(redis, user_api_key=token, served_model=served, warmth={served: time.time()})
+    seed_session(redis, user_api_key=token, served_model=served, warmth={served: time.time()}, touched=(served, target))
     keys = FakeKeyDirectory({token: key_state(token=token, **fields)})
     try:
         with registered_callbacks(limiter):
@@ -143,7 +146,7 @@ async def test_warming_does_not_double_charge_a_keys_tpm():
 
     limiter, counters = real_limiter()
     llm_router, redis = warming_rig(redis=FakeRedisCache())
-    seed_session(redis, user_api_key="tpm", warmth={"fast-claude": time.time()})
+    seed_session(redis, user_api_key="tpm", warmth={"fast-claude": time.time()}, touched=_VISITED_BOTH_TIERS)
     keys = FakeKeyDirectory({"tpm": key_state(token="tpm", tpm_limit=100_000)})
     with registered_callbacks(limiter):
         await tick(llm_router, active=refresher(keys=keys, limiter=limiter))
@@ -168,7 +171,7 @@ async def test_warming_does_not_reset_a_keys_rate_limit_window():
     limiter, counters = real_limiter(window_size=3600)
     assert limiter.window_size == 3600
     llm_router, redis = warming_rig(redis=FakeRedisCache())
-    seed_session(redis, user_api_key="hourly", warmth={"fast-claude": time.time()})
+    seed_session(redis, user_api_key="hourly", warmth={"fast-claude": time.time()}, touched=_VISITED_BOTH_TIERS)
     keys = FakeKeyDirectory({"hourly": key_state(token="hourly", rpm_limit=100)})
     opened_at = int(time.time()) - 120
     await counters.async_set_cache(key="{api_key:hourly}:window", value=str(opened_at))
@@ -188,7 +191,7 @@ async def test_warming_does_not_collide_hanging_request_tracking():
     proxy_logging = proxy_logging_with_hooks()
     proxy_logging.alerting = ["slack"]
     llm_router, redis = warming_rig(redis=FakeRedisCache())
-    seed_session(redis)
+    seed_session(redis, touched=_VISITED_BOTH_TIERS)
     await tick(llm_router, active=refresher(proxy_logging=proxy_logging))
     call_ids = [call["litellm_call_id"] for call in llm_router.completion_calls]
     assert len(call_ids) == 2 and len(set(call_ids)) == 2
@@ -202,7 +205,9 @@ async def test_warming_does_not_collide_hanging_request_tracking():
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("surface,channel", [("chat_completions", "metadata"), ("anthropic_messages", "litellm_metadata")])
+@pytest.mark.parametrize(
+    "surface,channel", [("chat_completions", "metadata"), ("anthropic_messages", "litellm_metadata")]
+)
 async def test_a_due_session_is_replayed_on_its_own_surface_and_stamped_warm(surface, channel):
     """The replay reaches the surface it was captured from, through that surface's own metadata channel, with
     generation held at the floor and the session_id that lets deployment affinity pin the replay to the same
@@ -210,7 +215,11 @@ async def test_a_due_session_is_replayed_on_its_own_surface_and_stamped_warm(sur
     the system block ride along there. The warmth stamp is what paces the next tick."""
     llm_router, redis = warming_rig(redis=FakeRedisCache())
     record_key = seed_session(
-        redis, call_surface=surface, warmth={"fast-claude": time.time()}, tool_choice={"type": "auto"}
+        redis,
+        call_surface=surface,
+        warmth={"fast-claude": time.time()},
+        touched=_VISITED_BOTH_TIERS,
+        tool_choice={"type": "auto"},
     )
     await tick(llm_router)
     calls = llm_router.anthropic_calls if surface == "anthropic_messages" else llm_router.completion_calls
@@ -239,7 +248,7 @@ async def test_a_replay_the_provider_rejected_is_stamped_cold_and_still_paced():
     failing every replay is retried on the refresh interval rather than on every tick."""
     llm_router, redis = warming_rig(redis=FakeRedisCache())
     llm_router.failing_message_marker = "deployment policy"
-    record_key = seed_session(redis, warmth={"fast-claude": time.time()})
+    record_key = seed_session(redis, warmth={"fast-claude": time.time()}, touched=_VISITED_BOTH_TIERS)
     await tick(llm_router)
     assert llm_router.completion_calls == [] and len(llm_router.failed_calls) == 1
     stamp = warmth_stamp(redis, record_key, "smart-claude")
@@ -255,7 +264,7 @@ async def test_a_replay_refused_by_admission_is_stamped_cold_rather_than_warm():
     own ceilings leaves the cache exactly as cold as a provider failure does."""
     limiter, counters = real_limiter()
     llm_router, redis = warming_rig(redis=FakeRedisCache())
-    record_key = seed_session(redis, user_api_key="k", warmth={"fast-claude": time.time()})
+    record_key = seed_session(redis, user_api_key="k", warmth={"fast-claude": time.time()}, touched=_VISITED_BOTH_TIERS)
     keys = FakeKeyDirectory({"k": key_state(token="k", rpm_limit=1)})
     await counters.async_set_cache(key="{api_key:k}:window", value=str(int(time.time())))
     await counters.async_set_cache(key="{api_key:k}:requests", value=1)
@@ -281,6 +290,7 @@ async def test_session_pacing_bounds_how_often_warming_spends(last_activity_offs
         redis,
         last_activity=now + last_activity_offset if last_activity_offset is not None else None,
         warmth={"fast-claude": now, "smart-claude": now + warmth_offset} if warmth_offset is not None else None,
+        touched=_VISITED_BOTH_TIERS,
     )
     await tick(llm_router)
     assert replayed_models(llm_router) == replayed
@@ -345,8 +355,6 @@ async def test_warmth_is_not_shared_between_two_auto_routers_on_one_redis():
     assert await other.get_warmth(record, ("fast-claude", "smart-claude")) == {}
 
 
-
-
 @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):
@@ -355,15 +363,14 @@ async def test_a_keyless_proxy_caller_is_authorized_through_its_reconstructed_te
     gate binds: a blocked team stops the replays, an open team still warms."""
     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=team_blocked, models=[]))
     seed_session(
         redis,
         user_api_key=None,
         caller_scope="jwt-user",
         team_id="jwt-team",
         warmth={"fast-claude": time.time()},
+        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
@@ -375,7 +382,13 @@ async def test_a_keyless_proxy_caller_is_authorized_through_its_reconstructed_te
 async def test_a_direct_sdk_session_with_no_recorded_identity_still_warms_unattributed():
     """No proxy auth object means no tenancy to preserve, so warming stays unattributed as before."""
     llm_router, redis = warming_rig(redis=FakeRedisCache())
-    seed_session(redis, user_api_key=None, caller_scope="unscoped", warmth={"fast-claude": time.time()})
+    seed_session(
+        redis,
+        user_api_key=None,
+        caller_scope="unscoped",
+        warmth={"fast-claude": time.time()},
+        touched=_VISITED_BOTH_TIERS,
+    )
     await tick(llm_router, active=refresher(keys=FakeKeyDirectory({})))
     assert replayed_models(llm_router) == ["smart-claude"]
     assert llm_router.completion_calls[0]["metadata"]["user_api_key_team_id"] is None
@@ -402,7 +415,13 @@ async def test_the_concurrency_bound_bounds_decompressed_payloads_not_just_repla
 
     llm_router, redis = warming_rig(redis=FakeRedisCache(), replay_delay=0.02)
     for index in range(6):
-        seed_session(redis, session_id=f"sess-{index}", caller_scope=f"hash-{index}", user_api_key=f"hash-{index}")
+        seed_session(
+            redis,
+            session_id=f"sess-{index}",
+            caller_scope=f"hash-{index}",
+            user_api_key=f"hash-{index}",
+            touched=_VISITED_BOTH_TIERS,
+        )
     real_decompress = refresher_module.decompress_payload
     inflated_before_first_replay_completed: list[int] = []
     inflated = 0
@@ -425,6 +444,22 @@ async def test_the_concurrency_bound_bounds_decompressed_payloads_not_just_repla
     assert max(inflated_before_first_replay_completed) <= 2, "payloads inflated must respect the same bound"
 
 
+@pytest.mark.asyncio
+async def test_warming_refreshes_only_the_models_a_session_has_been_served_on():
+    """Warming keeps a session's own caches alive so returning to a tier it has already used is a read; it does
+    not pre-warm tiers on speculation. A session that has only ever been SIMPLE must therefore cost one replay,
+    not one per tier, and the first switch to a new tier is a cold write by design. The second session is the
+    payoff case: having visited both tiers, both stay warm."""
+    llm_router, redis = warming_rig(redis=FakeRedisCache())
+    seed_session(redis, session_id="simple-only", caller_scope="a", served_model="fast-claude")
+    seed_session(redis, session_id="both-tiers", caller_scope="b", touched=_VISITED_BOTH_TIERS)
+    await tick(llm_router)
+    warmed = sorted(replayed_models(llm_router))
+    assert warmed == ["fast-claude", "fast-claude", "smart-claude"], (
+        "the SIMPLE-only session must not pay to warm a tier it has never been routed to"
+    )
+
+
 @pytest.mark.asyncio
 async def test_a_pooled_tier_warms_the_model_the_session_was_actually_served():
     """A tier may be a pool that the router picks from at random, so the member holding this session's cache

tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py

@@ -22,6 +22,7 @@ class FakeRedisCache:
         self.ttls: dict[str, int | None] = {}
         self.hashes: dict[str, dict[str, str]] = {}
         self.zsets: dict[str, dict[str, float]] = {}
+        self.sets: dict[str, set[str]] = {}
         self.expire_calls: list[str] = []
 
     def _namespaced(self, key: str) -> str:
@@ -63,6 +64,12 @@ class FakeRedisCache:
                 return 0
 
             return compare_and_delete
+        if "SMEMBERS" in script:
+
+            async def touched_models(keys: list, args: list) -> list:
+                return [member.encode("utf-8") for member in sorted(self.sets.get(self._namespaced(keys[0]), set()))]
+
+            return touched_models
         if "HGET" in script:
 
             async def get_record(keys: list, args: list) -> str | None:
@@ -74,6 +81,7 @@ class FakeRedisCache:
             async def capture(keys: list, args: list) -> int:
                 sessions = self.hashes.setdefault(self._namespaced(keys[0]), {})
                 index = self.zsets.setdefault(self._namespaced(keys[1]), {})
+                touched = self.sets.setdefault(self._namespaced(keys[2]), set())
                 member, record_json = str(args[0]), str(args[1])
                 now, expires_at, max_sessions = float(args[2]), float(args[3]), int(args[4])
                 for stale in [m for m, score in index.items() if score <= now]:
@@ -83,6 +91,7 @@ class FakeRedisCache:
                     return 0
                 sessions[member] = record_json
                 index[member] = expires_at
+                touched.add(str(args[5]))
                 return 1
 
             return capture
@@ -143,8 +152,6 @@ def test_key_shapes_are_scoped_and_hash_tagged():
     assert CacheWarmingStore.warmth_key(record, "opus").startswith(f"{slot}:")
 
 
-
-
 @pytest.mark.asyncio
 async def test_cap_enforced_atomically_with_the_record_write():
     _warn_session_cap_reached.cache_clear()
@@ -157,14 +164,6 @@ async def test_cap_enforced_atomically_with_the_record_write():
     assert len(await store.list_session_keys(max_sessions=10)) == 2
 
 
-
-
-
-
-
-
-
-
 @pytest.mark.asyncio
 async def test_get_record_returns_none_on_schema_version_mismatch():
     redis = FakeRedisCache()
@@ -172,11 +171,3 @@ async def test_get_record_returns_none_on_schema_version_mismatch():
     key = store.record_key("smart-router", "scope", "s1")
     redis.hashes[store.sessions_key()] = {key: _record_json(schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION + 1)}
     assert await store.get_record(key) is None
-
-
-
-
-
-
-
-

tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py

@@ -11,7 +11,12 @@ from litellm.router_strategy.complexity_router.cache_warming.types import (
 
 @pytest.mark.parametrize(
     "age,fresh",
-    [(0, True), (PROVIDER_PROMPT_CACHE_TTL_SECONDS - 1, True), (PROVIDER_PROMPT_CACHE_TTL_SECONDS, False), (601, False)],
+    [
+        (0, True),
+        (PROVIDER_PROMPT_CACHE_TTL_SECONDS - 1, True),
+        (PROVIDER_PROMPT_CACHE_TTL_SECONDS, False),
+        (601, False),
+    ],
 )
 def test_freshness_is_the_provider_ttl_and_nothing_else(age, fresh):
     """The router's warm-aware pick and the refresher's due-model calculation both read this, so a model can

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,
+    touched: tuple[str, ...] | None = None,
 ) -> str:
     payload = CacheWarmingPayload(
         model=served_model,
@@ -249,6 +250,10 @@ def seed_session(
     record_key = CacheWarmingStore.record_key("smart-router", caller_scope, session_id)
     redis.hashes.setdefault(store.sessions_key(), {})[record_key] = json.dumps(record.model_dump())
     redis.zsets.setdefault(store.index_key(), {})[record_key] = time.time() + 3600
+    # capture SADDs every served model, so a seeded session carries at least the one it was served on
+    redis.sets.setdefault(CacheWarmingStore.touched_key(record_key), set()).update(
+        touched if touched is not None else (served_model,)
+    )
     for model_group, stamp in (warmth or {}).items():
         redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps(
             WarmthStamp(at=stamp, warmed=True).model_dump()

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -3340,9 +3340,7 @@ class TestEscalationKeywords:
         router = ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=mock_router_instance,
-            complexity_router_config={
-                "tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top"}
-            },
+            complexity_router_config={"tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top"}},
         )
         assert router._tier_for_model("shared") == ComplexityTier.COMPLEX
         assert router._tier_for_model("top") == ComplexityTier.REASONING
@@ -3501,19 +3499,20 @@ class TestEscalationKeywords:
     def test_blank_escalation_keywords_are_stripped(self):
         """Blank/whitespace-only phrases are dropped so `"" in message` can't escalate
         every request; surrounding whitespace on real phrases is trimmed."""
-        assert ComplexityRouterConfig(
-            tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
-            escalation_keywords=["", "  "],
-        ).escalation_keywords == []
+        assert (
+            ComplexityRouterConfig(
+                tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
+                escalation_keywords=["", "  "],
+            ).escalation_keywords
+            == []
+        )
         assert ComplexityRouterConfig(
             tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
             escalation_keywords=["  LITELLM ESCALATE  ", ""],
         ).escalation_keywords == ["LITELLM ESCALATE"]
 
     @pytest.mark.asyncio
-    async def test_blank_escalation_keyword_does_not_escalate_everything(
-        self, mock_router_instance, basic_config
-    ):
+    async def test_blank_escalation_keyword_does_not_escalate_everything(self, mock_router_instance, basic_config):
         router = ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=mock_router_instance,
@@ -3533,9 +3532,7 @@ class TestEscalationKeywords:
         router = ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=mock_router_instance,
-            complexity_router_config={
-                "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}
-            },
+            complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}},
         )
         for pinned in ("o1-a", "o1-b", "o1-c"):
             assert router._escalated_pin(pinned) == pinned
@@ -3752,9 +3749,6 @@ class TestWarmAwarePick:
         self._seed(redis, warmth={"smart-claude": time_module.time()}, served_model="fast-claude")
         router = self._router(mock_router_instance, redis, plugins=[ExcludeSmartClaude()])
         picks = {
-            await router._pick_model_for_tier(
-                ComplexityTier.SIMPLE, None, None, self._kwargs()
-            )
-            for _ in range(20)
+            await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, self._kwargs()) for _ in range(20)
         }
         assert picks == {"fast-claude"}