fix(complexity_router): stamp a cache-warming replay warm only when it lands

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/router_strategy/complexity_router/README.md

@@ -177,8 +177,8 @@ Requirements and semantics:
 - **Redis is required.** Session payloads, per-model warmth stamps, and a per-router session index live in Redis so all pods share them and a single pod (via a Redis cron lock) runs the replays. Without Redis, warming logs a warning once and no-ops; requests are unaffected. Sessions are tracked through the index rather than keyspace scans, so Redis Cluster is supported.
 - **Prompt retention consent is a prerequisite.** Warming persists full request payloads (messages, system, tools) in Redis, so capture requires `store_prompts_in_spend_logs: true` and respects message redaction: with the flag off, or with `turn_off_message_logging` (globally or via the per-request redaction header) active, capture warns once and skips.
 - **Only Anthropic and Bedrock models that support prompt caching are warmed.** Other models in the tier pools are left alone. Requests must carry a `metadata.session_id` and exceed the warm set's minimum cacheable token count (`prompt_cache_min_tokens`, default 1024) to be captured.
-- **A replay is admitted like a request.** Each replay is assembled as a request body and put through the proxy's own admission entry points before it is dispatched: the budget reservation (the same call the auth layer makes after `common_checks`, so the key, team, user, end-user, organization and tag counters all apply and are reconciled to the replay's actual cost) and then `ProxyLogging.pre_call_hook` (so RPM, TPM, max-parallel and every configured guardrail apply, on the proxy's own shared counters rather than a private copy). A rejection skips that one replay and is retried on the next tick; the failure hook returns whatever the rejected replay had already reserved. Warming stops for keys that are deleted, blocked, or expired; key state is verified fresh each tick and when the database is unreachable the tick is skipped, so warming pauses until it is reachable again (caches re-warm on the next successful tick).
-- **Warming cost is visible where the customer already looks for cost.** Warming writes no spend logs of its own, so the replay rows are the only record of warming cost that exists, and they carry the same identity block a real request on that key carries: key, team, user and end-user attribution, the key's and team's own `tags` and `spend_logs_metadata`, and the `litellm_cache_warming` tag alongside them so warming is both included in per-tag chargeback and filterable out of it. Replays also fan out to the key-scoped and team-scoped logging callbacks, so a team pointing its traffic at its own Langfuse sees warming there too rather than only in the proxy-wide logs. A `max_tokens=1` replay of a warm prefix bills roughly 10% of the input cost. Key-level cache controls, `disable_fallbacks` and the global-guardrail opt-outs are applied the same way, which means a key that declares its own `cache` controls overrides warming's response-cache bypass exactly as it overrides a caller's.
+- **A replay is admitted like a request.** Each replay is assembled as a request body and put through the proxy's own admission entry points before it is dispatched: the budget reservation (the same call the auth layer makes after `common_checks`, so the key, team, user, end-user, organization and tag counters all apply and are reconciled to the replay's actual cost) and then `ProxyLogging.pre_call_hook` (so RPM, TPM, max-parallel and every configured guardrail apply, on the proxy's own shared counters rather than a private copy). A rejection skips that one replay; the failure hook returns whatever the rejected replay had already reserved, and the attempt is recorded as attempted-but-cold, so the routing pick keeps treating that model as cold while the attempt still paces the retry. Warming stops for keys that are deleted, blocked, or expired; key state is verified fresh each tick and when the database is unreachable the tick is skipped, so warming pauses until it is reachable again (caches re-warm on the next successful tick).
+- **Warming cost is visible where the customer already looks for cost.** Warming writes no spend logs of its own, so the replay rows are the only record of warming cost that exists, and they carry the same identity block a real request on that key carries: key, team, user and end-user attribution, the key's and team's own `tags` and `spend_logs_metadata`, plus a `litellm_cache_warming` marker in `spend_logs_metadata` so warming rows are identifiable in per-tag chargeback without putting warming's own tag on `metadata.tags`, which is an input to deployment selection and which operators can forbid requests from carrying. Replays also fan out to the key-scoped and team-scoped logging callbacks, so a team pointing its traffic at its own Langfuse sees warming there too rather than only in the proxy-wide logs. A `max_tokens=1` replay of a warm prefix bills roughly 10% of the input cost. Key-level cache controls, `disable_fallbacks` and the global-guardrail opt-outs are applied the same way, which means a key that declares its own `cache` controls overrides warming's response-cache bypass exactly as it overrides a caller's.
 - **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.

litellm/router_strategy/complexity_router/cache_warming/refresher.py

@@ -580,7 +580,9 @@ class CacheWarmingRefresher:
         warmth = await store.get_warmth(session_key, warm_models)
         now = time.time()
         due_models = tuple(
-            model for model in warm_models if needs_rewarming(warmth.get(model, 0.0), now, refresh_interval_seconds)
+            model
+            for model in warm_models
+            if needs_rewarming(warmth[model].at if model in warmth else 0.0, now, refresh_interval_seconds)
         )
         if not due_models:
             return
@@ -589,38 +591,68 @@ class CacheWarmingRefresher:
             async with semaphore:
                 if lease_lost.is_set():
                     return
-                admitted = await self._admit_replay(
+                attempted_at = time.time()
+                warmed = await self._replay_once(
                     llm_router=llm_router,
                     payload=payload,
                     record=record,
+                    session_key=session_key,
                     model_group=model_group,
                     key_state=key_state,
                     prisma_client=prisma_client,
                     user_api_key_cache=user_api_key_cache,
                     proxy_logging_obj=proxy_logging_obj,
                 )
-                if admitted is None:
-                    continue
-                data, principal = admitted
-                attempted_at = time.time()
-                try:
-                    await (
-                        llm_router.aanthropic_messages(**data)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # factory-generated router surface is legacy-untyped
-                        if payload.call_surface == "anthropic_messages"
-                        else llm_router.acompletion(**data)  # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportUnknownVariableType, reportArgumentType]  # router overloads are legacy-untyped
-                    )
-                    await proxy_logging_obj.update_request_status(
-                        litellm_call_id=str(data.get("litellm_call_id") or ""), status="success"
-                    )
-                except Exception as exc:  # noqa: BLE001  # one failing replay must not abort the tick
-                    verbose_router_logger.warning(
-                        "cache_warming replay failed for session %s model %s", session_key, model_group, exc_info=True
-                    )
-                    await self._report_rejection(
-                        data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj
-                    )
-                finally:
-                    await store.mark_warm_attempt(session_key, model_group, attempted_at, session_ttl_seconds)
+                await store.mark_warm_attempt(
+                    session_key, model_group, attempted_at, session_ttl_seconds, warmed=warmed
+                )
+
+    async def _replay_once(
+        self,
+        *,
+        llm_router: "Router",
+        payload: CacheWarmingPayload,
+        record: CacheWarmingRecord,
+        session_key: str,
+        model_group: str,
+        key_state: "UserAPIKeyAuth | None",
+        prisma_client: "PrismaClient | None",
+        user_api_key_cache: "DualCache | None",
+        proxy_logging_obj: "ProxyLogging",
+    ) -> bool:
+        """True only when the provider accepted the replay, because a warmth stamp is the claim that the
+        provider now holds this session's prefix on this model. A refused or failed replay leaves it cold, so
+        it is stamped as attempted-but-cold: the pick keeps treating the model as cold, while the attempt
+        still paces the next tick instead of retrying a failing model every 30 seconds."""
+        admitted = await self._admit_replay(
+            llm_router=llm_router,
+            payload=payload,
+            record=record,
+            model_group=model_group,
+            key_state=key_state,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+        if admitted is None:
+            return False
+        data, principal = admitted
+        try:
+            await (
+                llm_router.aanthropic_messages(**data)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # factory-generated router surface is legacy-untyped
+                if payload.call_surface == "anthropic_messages"
+                else llm_router.acompletion(**data)  # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportUnknownVariableType, reportArgumentType]  # router overloads are legacy-untyped
+            )
+        except Exception as exc:  # noqa: BLE001  # one failing replay must not abort the tick
+            verbose_router_logger.warning(
+                "cache_warming replay failed for session %s model %s", session_key, model_group, exc_info=True
+            )
+            await self._report_rejection(data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj)
+            return False
+        await proxy_logging_obj.update_request_status(
+            litellm_call_id=str(data.get("litellm_call_id") or ""), status="success"
+        )
+        return True
 
     async def _admit_replay(
         self,
@@ -675,7 +707,6 @@ class CacheWarmingRefresher:
             await self._report_rejection(data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj)
             return None
 
-    @staticmethod
     @staticmethod
     async def _authorize_and_reserve(
         *,

litellm/router_strategy/complexity_router/cache_warming/store.py

@@ -11,6 +11,7 @@ from litellm.router_strategy.complexity_router.cache_warming.types import (
     CACHE_WARMING_RECORD_SCHEMA_VERSION,
     CacheWarmingAttribution,
     CacheWarmingRecord,
+    WarmthStamp,
 )
 
 _WARMTH_KEY_PREFIX = "complexity_router_cache_warmth:v1"
@@ -82,15 +83,13 @@ def _parse_record(raw: object) -> CacheWarmingRecord | None:
     return record
 
 
-def _parse_warmth(raw: object) -> float | None:
-    if isinstance(raw, (int, float)):
-        return float(raw)
-    if isinstance(raw, (str, bytes)):
-        try:
-            return float(raw)
-        except ValueError:
-            return None
-    return None
+def _parse_warmth(raw: object) -> WarmthStamp | None:
+    try:
+        if isinstance(raw, (str, bytes)):
+            return WarmthStamp.model_validate_json(raw)
+        return WarmthStamp.model_validate(raw)
+    except ValidationError:
+        return None
 
 
 class CacheWarmingStore:
@@ -188,17 +187,21 @@ class CacheWarmingStore:
         if admitted != 1:
             _warn_session_cap_reached(self.auto_router_model_name)
             return
-        await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds)
+        await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds, warmed=True)
 
-    async def mark_warm_attempt(self, key: str, model_group: str, attempted_at: float, ttl_seconds: int) -> None:
+    async def mark_warm_attempt(
+        self, key: str, model_group: str, attempted_at: float, ttl_seconds: int, warmed: bool
+    ) -> None:
         redis_cache = self._require_redis()
         if redis_cache is None:
             return
         await redis_cache.async_set_cache(  # pyright: ignore[reportUnknownMemberType]  # RedisCache is legacy-untyped
-            key=self.warmth_key(key, model_group), value=attempted_at, ttl=ttl_seconds
+            key=self.warmth_key(key, model_group),
+            value=WarmthStamp(at=attempted_at, warmed=warmed).model_dump(),
+            ttl=ttl_seconds,
         )
 
-    async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, float]:
+    async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, WarmthStamp]:
         redis_cache = self._require_redis()
         if redis_cache is None:
             return {}  # mutable-ok: fresh per-call result, not shared state

litellm/router_strategy/complexity_router/cache_warming/types.py

@@ -36,6 +36,17 @@ def needs_rewarming(warmed_at: float, now: float, refresh_interval_seconds: int)
     )
 
 
+class WarmthStamp(BaseModel):
+    """When a model group was last replayed for a session, and whether that replay actually landed. A failed
+    replay leaves the provider cache cold, so ``warmed`` keeps the pick from preferring it while ``at`` still
+    paces the next attempt."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    at: float
+    warmed: bool
+
+
 class CacheWarmingPayload(BaseModel):
     model_config = ConfigDict(extra="forbid")
 

litellm/router_strategy/complexity_router/complexity_router.py

@@ -671,7 +671,7 @@ class ComplexityRouter(CustomLogger):
             return None
         warmth = await store.get_warmth(record_key, tuple(pool))
         now = time.time()
-        warmed = frozenset(model for model, warmed_at in warmth.items() if is_cache_fresh(warmed_at, now))
+        warmed = frozenset(model for model, stamp in warmth.items() if stamp.warmed and is_cache_fresh(stamp.at, now))
         served = frozenset((record.served_model,)) if is_cache_fresh(record.last_activity, now) else frozenset[str]()
         candidates = tuple(model for model in pool if model in warmed | served)
         if not candidates:

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

@@ -13,6 +13,7 @@ from litellm.router_strategy.complexity_router.cache_warming.capture import (
 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,
+    WarmthStamp,
     decompress_payload,
 )
 from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
@@ -100,13 +101,15 @@ async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth()
     await router._capture_session(_kwargs(), MESSAGES, "claude-sonnet-4-5")
     key = CacheWarmingStore.record_key("smart-router", "hash-1", "sess-1")
     first = json.loads(redis.hashes[SESSIONS_KEY][key])
-    redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(123.0)
+    other_stamp = WarmthStamp(at=123.0, warmed=True).model_dump()
+    redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(other_stamp)
     await router._capture_session(
         _kwargs(), MESSAGES + [{"role": "user", "content": "and rule 8?"}], "claude-sonnet-4-5"
     )
     second = json.loads(redis.hashes[SESSIONS_KEY][key])
-    assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == 123.0
-    assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")]) > 0
+    assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == other_stamp
+    served_stamp = WarmthStamp.model_validate_json(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")])
+    assert served_stamp.at > 0 and served_stamp.warmed is True
     assert second["payload_sha256"] != first["payload_sha256"]
 
 

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

@@ -219,7 +219,42 @@ async def test_a_due_session_is_replayed_on_its_own_surface_and_stamped_warm(sur
     if surface == "anthropic_messages":
         assert call["system"] == "You are a policy assistant"
     stamp = warmth_stamp(redis, record_key, "smart-claude")
-    assert stamp is not None and stamp > 0
+    assert stamp is not None and stamp.at > 0 and stamp.warmed is True
+
+
+@pytest.mark.asyncio
+async def test_a_replay_the_provider_rejected_is_stamped_cold_and_still_paced():
+    """A stamp is the claim that the provider holds this session's prefix on that model, so a replay that
+    failed must not write a warm one: the pick would then send real traffic to a cold model for the provider's
+    whole cache TTL, which is the outcome warming exists to avoid. The attempt is still recorded, so a model
+    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()})
+    await tick(llm_router)
+    assert llm_router.completion_calls == [] and len(llm_router.failed_calls) == 1
+    stamp = warmth_stamp(redis, record_key, "smart-claude")
+    assert stamp is not None and stamp.at > 0 and stamp.warmed is False
+
+    await tick(llm_router)
+    assert len(llm_router.failed_calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_a_replay_refused_by_admission_is_stamped_cold_rather_than_warm():
+    """Same claim, for the replay that never reached the provider at all: a refusal from the request path's
+    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()})
+    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)
+    with registered_callbacks(limiter):
+        await tick(llm_router, active=refresher(keys=keys, limiter=limiter))
+    assert llm_router.completion_calls == []
+    stamp = warmth_stamp(redis, record_key, "smart-claude")
+    assert stamp is not None and stamp.warmed is False
 
 
 @pytest.mark.asyncio

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

@@ -21,6 +21,7 @@ from litellm.router_strategy.complexity_router.cache_warming.types import (
     CacheWarmingAttribution,
     CacheWarmingPayload,
     CacheWarmingRecord,
+    WarmthStamp,
     compress_payload,
 )
 from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
@@ -61,6 +62,7 @@ class ReplayRouter(Router):
         )
         self.cache.redis_cache = redis
         self.completion_calls: list[dict] = []
+        self.failed_calls: list[dict] = []
         self.anthropic_calls: list[dict] = []
         self.replay_delay = replay_delay
         self.failing_message_marker: str | None = None
@@ -76,6 +78,7 @@ class ReplayRouter(Router):
     async def acompletion(self, **kwargs: object):  # pyright: ignore[reportIncompatibleMethodOverride]  # test double narrows the overloads
         marker = self.failing_message_marker
         if marker is not None and marker in json.dumps(kwargs.get("messages"), default=str):
+            self.failed_calls.append(kwargs)
             raise RuntimeError("provider down")
         self._in_flight += 1
         self.max_concurrent = max(self.max_concurrent, self._in_flight)
@@ -247,13 +250,15 @@ def seed_session(
     redis.hashes.setdefault(store.sessions_key(), {})[record_key] = json.dumps(record.model_dump())
     redis.zsets.setdefault(store.index_key(), {})[record_key] = time.time() + 3600
     for model_group, stamp in (warmth or {}).items():
-        redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps(stamp)
+        redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps(
+            WarmthStamp(at=stamp, warmed=True).model_dump()
+        )
     return record_key
 
 
-def warmth_stamp(redis: FakeRedisCache, record_key: str, model_group: str) -> float | None:
+def warmth_stamp(redis: FakeRedisCache, record_key: str, model_group: str) -> WarmthStamp | None:
     raw = redis.data.get(CacheWarmingStore.warmth_key(record_key, model_group))
-    return json.loads(raw) if raw is not None else None
+    return WarmthStamp.model_validate_json(raw) if raw is not None else None
 
 
 def proxy_logging_with_hooks(limiter: object | None = None) -> ProxyLogging:

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -4214,7 +4214,7 @@ class TestWarmAwarePick:
         )
 
     @staticmethod
-    def _seed(redis, warmth, last_activity=None, served_model="fast-claude"):
+    def _seed(redis, warmth, last_activity=None, served_model="fast-claude", warmed=True):
         import json
         import time as time_module
 
@@ -4224,6 +4224,7 @@ class TestWarmAwarePick:
             CacheWarmingAttribution,
             CacheWarmingPayload,
             CacheWarmingRecord,
+            WarmthStamp,
             compress_payload,
         )
 
@@ -4247,7 +4248,9 @@ class TestWarmAwarePick:
         key = store.record_key("warm-router", "hash-w", "warm-sess")
         redis.hashes.setdefault(store.sessions_key(), {})[key] = json.dumps(record.model_dump())
         for model_group, stamp in warmth.items():
-            redis.data[CacheWarmingStore.warmth_key(key, model_group)] = json.dumps(stamp)
+            redis.data[CacheWarmingStore.warmth_key(key, model_group)] = json.dumps(
+                WarmthStamp(at=stamp, warmed=warmed).model_dump()
+            )
 
     @staticmethod
     def _kwargs():
@@ -4272,6 +4275,20 @@ class TestWarmAwarePick:
         assert "cold-model" not in picks
         assert picks <= {"fast-claude", "smart-claude"}
 
+    @pytest.mark.asyncio
+    async def test_a_failed_replay_does_not_make_a_model_look_warm(self, mock_router_instance):
+        """A replay that the provider refused leaves that model's cache cold, so its stamp must not steer the
+        pick onto it: the pick falls back to the served model, whose cache the session's own traffic wrote."""
+        import time as time_module
+
+        redis = self._fresh_redis()
+        self._seed(redis, warmth={"smart-claude": time_module.time()}, served_model="fast-claude", warmed=False)
+        router = self._router(mock_router_instance, redis)
+        picks = {
+            await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, self._kwargs()) for _ in range(20)
+        }
+        assert picks == {"fast-claude"}
+
     @pytest.mark.asyncio
     async def test_stale_warm_entries_fall_back(self, mock_router_instance):
         import time as time_module