fix(accounting): carry a job's model window on every completed poll in ledger mode

The ledger keeps the poll whose delivery claims the job, which need not be the
poll that won the limiter's in-memory charge. Every completed poll now carries
the job's window, so the delivery that lands always has it, while the live
model budget is still charged once.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

litellm/proxy/hooks/model_budget_shared_windows.py

@@ -74,6 +74,10 @@ def captured_model_windows(kwargs: Mapping[str, object] | None) -> tuple[WindowS
     return value.transactions if isinstance(value, CapturedModelWindows) else ()
 
 
+def capturing_model_windows() -> bool:
+    return _MODEL_WINDOWS.get() is not None
+
+
 @contextmanager
 def capture_model_windows(enabled: bool) -> Generator[Callable[[], CapturedModelWindows]]:
     token: Final = _MODEL_WINDOWS.set(() if enabled else None)

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -35,6 +35,7 @@ from litellm.proxy.hooks.model_budget_shared_windows import (
     SharedWindow,
     SharedWindows,
     capture_model_windows,
+    capturing_model_windows,
 )
 from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds
 from litellm.types.llms.openai import AllMessageValues
@@ -1265,6 +1266,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             job=job,
             ttl_seconds=duration_in_seconds(budget_duration),
         ):
+            if capturing_model_windows() and self._shares_windows_in_database(entity_type):
+                await self._capture_job_window(
+                    entity_type=entity_type,
+                    entity_id=entity_id,
+                    resolved=resolved,
+                    spend_key=spend_key,
+                    response_cost=response_cost,
+                )
             if job.repeat_poll_cost > 0:
                 await self._charge_entity(
                     entity_type=entity_type,
@@ -1299,6 +1308,39 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
                 response_cost=response_cost,
             )
 
+    async def _capture_job_window(
+        self,
+        entity_type: Litellm_EntityType,
+        entity_id: str | None,
+        resolved: ResolvedModelBudget,
+        spend_key: str,
+        response_cost: float,
+    ) -> None:
+        """The ledger keeps only the poll whose delivery claims the job, which need not be the poll that won this
+        limiter's charge, so every completed poll carries the job's window and the ledger applies it once."""
+        budget_duration: Final = resolved.budget_config.budget_duration
+        if budget_duration is None:
+            return
+        start_time_key: Final = model_budget_start_time_cache_key(
+            entity_type=entity_type,
+            entity_id=entity_id,
+            budget_model=resolved.budget_model,
+            budget_duration=budget_duration,
+        )
+        if _local_windows(self.dual_cache).get_cache(key=_unshared_spend_key(start_time_key)) is not None:
+            return
+        now: Final = self._clock()
+        await self._shared_windows.add(
+            spend_key=spend_key,
+            budget_duration=budget_duration,
+            started_at=await self._get_or_set_budget_start_time(
+                start_time_key=start_time_key,
+                current_time=now,
+                ttl_seconds=budget_window_seconds(now, budget_duration),
+            ),
+            spend=response_cost,
+        )
+
     def _shares_windows_in_database(self, entity_type: Litellm_EntityType) -> bool:
         """End-user ids are whatever the client sends, so their windows stay on each replica rather than adding a
         database row per id."""

tests/unit/proxy/hooks/test_model_max_budget_limiter.py

@@ -945,3 +945,26 @@ async def test_concurrent_first_charges_on_one_replica_both_count():
         shared_windows=_WindowSpendDatabase().replica(),
     )
     assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4)
+
+
+@pytest.mark.asyncio
+async def test_every_completed_poll_of_a_job_carries_its_window_for_the_ledger():
+    from litellm.proxy.hooks.model_budget_shared_windows import captured_model_windows
+
+    database: Final = _WindowSpendDatabase()
+    limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(
+        dual_cache=DualCache(),
+        shared_windows=DatabaseSharedWindows(
+            prisma_client_source=lambda: SimpleNamespace(db=database, writer_db=database),  # pyright: ignore[reportArgumentType]  # duck-typed prisma double
+        ),
+        defer_window_spend=True,
+    )
+    polls: Final = tuple({**_event("avideo_status", BATCH_COST), POLLED_VIDEO_STATUS_KWARG: "completed"} for _ in "ab")
+    for poll in polls:
+        await _log_success(limiter, poll, response_obj=VideoObject(id="video_first", object="video", status="completed"))
+
+    carried: Final = tuple(
+        sorted((window["entity_id"], window["spend"]) for window in captured_model_windows(poll)) for poll in polls
+    )
+    assert carried[0] and carried[0] == carried[1]
+    assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)