test(proxy): check stalled failure counters and use the real spend writer

- The Redis stall cell now reads the burst key's live counter and requires at least the full failure spend
- Unit tests drop the charging writer double, bare dicts, list mutation and wall-clock timestamps
- Partial-stream test models are frozen

tests/integration/spend/_partial_stream.py

@@ -9,7 +9,7 @@ import httpx
 from integration._support.client import Gateway, Scenario, eventually, held
 from integration._support.wire import Reply, Request, Wire
 from openai.types.chat import ChatCompletionChunk
-from pydantic import BaseModel, JsonValue, TypeAdapter
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
 
 JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
 BACKEND: Final = "integration-partial-stream"
@@ -20,19 +20,27 @@ ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s OR
 
 
 class _SpendRow(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     spend: float
 
 
 class StatusRow(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     status: str
     spend: float
 
 
 class _BudgetError(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     type: str
 
 
 class _ErrorResponse(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     error: _BudgetError
 
 

tests/integration/spend/test_failure_spend_live_counters_chaos.py

@@ -2,6 +2,7 @@ import json
 import threading
 import uuid
 from concurrent.futures import ThreadPoolExecutor
+from hashlib import sha256
 from pathlib import Path
 from typing import Final
 
@@ -25,6 +26,7 @@ from integration.spend._partial_stream import (
     recovered_cost,
 )
 from pydantic import JsonValue
+from redis import Redis
 
 from litellm.constants import REDIS_SOCKET_TIMEOUT
 
@@ -108,6 +110,17 @@ def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_e
 
         assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT)
         assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST)
+        spent: Final = cost * _BURST
+        with Redis(host=counters.host, port=counters.port) as live_counters:
+            landed: Final = eventually(
+                lambda: float(live_counters.get(f"spend:key:{sha256(burst_key.encode()).hexdigest()}") or 0),
+                lambda counted: counted >= spent - 1e-9,
+                seconds=30,
+                return_last_on_timeout=True,
+            )
+        # The cold counter is seeded from the database, and a spend flush during the stall can land these rows
+        # there first, so it may count them twice, as completed streams held past a flush do
+        assert spent - 1e-9 <= landed <= 2 * spent + 1e-9, f"one broken stream costs {cost}"
 
         budgeted: Final = scenario.key(models=[model], max_budget=cost / 2)
         broken_stream(candidate, model, budgeted)

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -2637,13 +2637,7 @@ def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache:
     return counters
 
 
-def _charging_spend_writer() -> DBSpendUpdateWriter:
-    writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
-    writer.update_database.return_value = True
-    return writer
-
-
-def _failure_with_billed_attempts() -> dict:
+def _failure_with_billed_attempts() -> Mapping[str, object]:
     usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
     return {
         "combined_usage_object": usage,
@@ -2668,7 +2662,7 @@ _LIVE_COUNTERS: Final = {
 }
 
 
-async def _fail_request(logger: ProxyDBLogger, failure: dict) -> None:
+async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None:
     await logger.async_post_call_failure_hook(
         request_data={
             "model": "gpt-4.1",
@@ -2709,14 +2703,14 @@ async def _counter_values(counters: DualCache) -> dict[str, float]:
     ],
 )
 async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
-    live_spend_counters: DualCache, failure: dict, charged: float
+    live_spend_counters: DualCache, failure: Mapping[str, object], charged: float
 ):
     """#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the
     spend log, or an over-budget key keeps passing until the counter is next repaired from the database."""
     for key, spend in _LIVE_COUNTERS.items():
         await live_spend_counters.async_set_cache(key, spend)
 
-    await _fail_request(ProxyDBLogger(spend_writer=_charging_spend_writer), failure)
+    await _fail_request(ProxyDBLogger(), failure)
 
     assert await _counter_values(live_spend_counters) == pytest.approx(
         {key: spend + charged for key, spend in _LIVE_COUNTERS.items()}
@@ -2724,68 +2718,50 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize(("spend_updates_disabled", "writer_charged"), [(True, True), (False, False)])
-async def test_failure_spend_rows_leave_the_live_spend_counters_alone_when_nothing_was_charged(
-    live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch, spend_updates_disabled: bool, writer_charged: bool
+async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled(
+    live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
 ):
-    monkeypatch.setattr(
-        "litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True} if spend_updates_disabled else {}
-    )
-    writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
-    writer.update_database.return_value = writer_charged
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True})
     for key, spend in _LIVE_COUNTERS.items():
         await live_spend_counters.async_set_cache(key, spend)
 
-    await _fail_request(ProxyDBLogger(spend_writer=lambda: writer), _failure_with_billed_attempts())
+    await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts())
 
     assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS)
 
 
 @pytest.mark.asyncio
-async def test_a_spend_counter_outage_still_writes_every_failure_spend_row():
-    writer: Final = _charging_spend_writer()
+async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row():
     counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable"))
 
-    await _fail_request(
-        ProxyDBLogger(spend_writer=lambda: writer, spend_counters=lambda: counters), _failure_with_billed_attempts()
-    )
+    await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts())
 
-    assert [call.kwargs["response_cost"] for call in writer.update_database.call_args_list] == [0.5, 0.25, 0.5]
     assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5]
 
 
 @pytest.mark.asyncio
+@pytest.mark.timeout(30)
 async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers():
     store_answers: Final = asyncio.Event()
-    landed: Final = asyncio.Event()
-    raised_costs: Final[list[object]] = []
+    raised_cost: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future()
 
     async def slow_counter_store(**counter_args: object) -> None:
         await store_answers.wait()
-        raised_costs.append(counter_args["response_cost"])
-        landed.set()
+        raised_cost.set_result(counter_args["response_cost"])
 
-    logger: Final = ProxyDBLogger(
-        spend_writer=_charging_spend_writer,
-        spend_counters=lambda: slow_counter_store,
-        failure_counter_wait_seconds=0.05,
-    )
+    logger: Final = ProxyDBLogger(spend_counters=lambda: slow_counter_store, failure_counter_wait_seconds=0)
 
-    await asyncio.wait_for(
-        _fail_request(
-            logger,
-            {
-                "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
-                "response_cost": 0.125,
-            },
-        ),
-        timeout=1,
+    await _fail_request(
+        logger,
+        {
+            "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
+            "response_cost": 0.125,
+        },
     )
-    assert raised_costs == []
+    assert not raised_cost.done()
 
     store_answers.set()
-    await asyncio.wait_for(landed.wait(), timeout=1)
-    assert raised_costs == [0.125]
+    assert await raised_cost == 0.125
 
 
 @pytest.mark.asyncio
@@ -2800,7 +2776,8 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ
     await window_queue.flush_and_get_aggregated_window_spend_transactions()
     key_cache: Final = DualCache()
     monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache)
-    reset_at: Final = datetime.now(timezone.utc) + timedelta(days=10)
+    request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc)
+    reset_at: Final = request_start + timedelta(days=10)
     await key_cache.async_set_cache(
         "hashed-key",
         UserAPIKeyAuth(
@@ -2808,10 +2785,9 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ
             budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}],
         ),
     )
-    request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc)
 
     await _fail_request(
-        ProxyDBLogger(spend_writer=_charging_spend_writer),
+        ProxyDBLogger(),
         {
             "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
             "response_cost": 0.125,
@@ -2845,11 +2821,10 @@ async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_
     )
 
     await _fail_request(
-        ProxyDBLogger(spend_writer=_charging_spend_writer),
+        ProxyDBLogger(),
         {
             **_failure_with_billed_attempts(),
-            "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"]},
-            "standard_logging_object": {"id": "rid", "model_id": "dep-b"},
+            "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}},
         },
     )