From 8564c9bd637539dfbe011721d706a7f79fe7b7e7 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:05:39 -0800 Subject: [PATCH] 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 --- .../hooks/model_budget_shared_windows.py | 4 ++ .../proxy/hooks/model_max_budget_limiter.py | 42 +++++++++++++++++++ .../hooks/test_model_max_budget_limiter.py | 23 ++++++++++ 3 files changed, 69 insertions(+) diff --git a/litellm/proxy/hooks/model_budget_shared_windows.py b/litellm/proxy/hooks/model_budget_shared_windows.py index bf7141efdca..b6dfbf39c7d 100644 --- a/litellm/proxy/hooks/model_budget_shared_windows.py +++ b/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) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 29c5dd2d592..c3b382f78de 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/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.""" diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index f1fd02981d6..9384a3f398b 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/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)