test(accounting): verify capture races and failure isolation

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -3949,3 +3949,41 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs(
             user_api_key_dict=user_api_key_dict,
         )
     assert len(writer.rows) == rows_written
+
+
+@pytest.mark.asyncio
+async def test_failure_capture_overload_does_not_skip_budget_release(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
+    from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
+    from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _write
+
+    database: Final = CommitBarrier()
+    runtime: Final = UsageCaptureRuntime()
+    runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2))
+    monkeypatch.setattr(proxy_server, "usage_capture", runtime)
+    monkeypatch.setattr(proxy_server, "prisma_client", None)
+    meter: Final = runtime.meter
+    assert meter is not None
+    receipt: Final = meter.submit((_write("occupied-one"), _write("occupied-two")))
+    auth: Final = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False})
+    request: Final = {
+        "call_type": "acompletion",
+        "litellm_call_id": "failed",
+        "model": "capture-model",
+        "start_time": _CALL_START,
+        "litellm_params": {},
+    }
+    try:
+        await database.entered.wait()
+        await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), auth)
+        assert auth.budget_reservation is not None and auth.budget_reservation["finalized"] is True
+        assert "Usage capture failed on the request failure path" in caplog.text
+        assert "_accounting_capture" not in request
+    finally:
+        database.release.set()
+        await runtime.stop(1)
+    await receipt
+    assert database.committed == ("occupied-one", "occupied-two")

tests/unit/proxy/spend_tracking/test_usage_capture.py

@@ -196,6 +196,60 @@ async def test_terminal_retrieval_reuses_accepted_snapshot_but_rejects_changed_u
     assert database.rows == first.writes
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("terminal", [False, True])
+async def test_competing_writers_reuse_only_terminal_pricing_snapshots(terminal: bool) -> None:
+    from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
+    from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback
+    from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
+
+    database: Final = StoredUsage()
+    attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]()
+    competing: Final = asyncio.Event()
+
+    async def insert(writes: tuple[UsageWrite, ...]) -> None:
+        attempts.put_nowait(writes)
+        if attempts.qsize() >= 2:
+            competing.set()
+        await competing.wait()
+        await database.insert(writes)
+
+    meters: Final = tuple(
+        DurableUsageMeter(UsageBatchWriter(insert, database.load).write, UsageCaptureSettings()) for _ in range(2)
+    )
+    kwargs: Final = {
+        "call_type": "aget_responses" if terminal else "acompletion",
+        "litellm_call_id": "caller",
+        "model": "capture-model",
+        "litellm_params": {
+            "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}
+            if terminal
+            else {}
+        },
+    }
+    response: Final = {
+        "id": "job",
+        "status": "completed",
+        "background": terminal,
+        "created_at": NOW.timestamp(),
+        "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7},
+    }
+    captured: Final = tuple(
+        capture_callback(meter, {**kwargs, "response_cost": price}, response, NOW, NOW)
+        for meter, price in zip(meters, (2.0, 9.0), strict=True)
+    )
+    try:
+        await asyncio.gather(*(meter.flush() for meter in meters))
+        results: Final = await asyncio.gather(*(capture.receipt for capture in captured), return_exceptions=True)
+        assert sum(result is None for result in results) == (2 if terminal else 1)
+        assert sum(isinstance(result, UsageCaptureError) for result in results) == (0 if terminal else 1)
+        assert database.rows in tuple(capture.writes for capture in captured)
+        assert len(database.rows) == (1 if terminal else 2)
+        assert attempts.qsize() == (3 if terminal else 2)
+    finally:
+        await asyncio.gather(*(meter.stop(0) for meter in meters))
+
+
 @pytest.mark.asyncio
 async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None:
     from litellm.proxy.spend_tracking.usage_capture import capture_callback

tests/unit/proxy/test_proxy_server.py

@@ -4313,6 +4313,45 @@ async def test_capture_enablement_rejects_unrecoverable_writer_modes(
     assert proxy_server.usage_capture.meter is None
 
 
+@pytest.mark.asyncio
+async def test_capture_configuration_requires_restart_without_replacing_the_active_writer(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
+    from litellm.proxy.utils import PrismaClient
+
+    runtime: Final = UsageCaptureRuntime()
+    monkeypatch.setattr(proxy_server, "usage_capture", runtime)
+    monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True})
+    with pytest.raises(ValueError, match="requires accounting_capture"):
+        proxy_server._configure_usage_capture()
+    client: Final = PrismaClient("postgresql://localhost/unused", ProxyLogging(UserApiKeyCache()))
+    monkeypatch.setattr(proxy_server, "prisma_client", client)
+    monkeypatch.setattr(litellm, "apply_search_free_tier", False)
+    monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false")
+    monkeypatch.delenv("SPEND_LOGS_URL", raising=False)
+    monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 4}})
+    try:
+        proxy_server._configure_usage_capture()
+        active: Final = runtime.meter
+        assert active is not None
+        assert active.settings.batch_size == 4
+        proxy_server._configure_usage_capture()
+        assert runtime.meter is active
+        monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 6}})
+        with pytest.raises(ValueError, match="Changing accounting_capture requires a worker restart"):
+            proxy_server._configure_usage_capture()
+        monkeypatch.setattr(proxy_server, "general_settings", {})
+        with pytest.raises(ValueError, match="Disabling accounting_capture requires a worker restart"):
+            proxy_server._configure_usage_capture()
+        assert runtime.meter is active
+        assert active.settings.batch_size == 4
+    finally:
+        await runtime.stop(1)
+    assert runtime.meter is None
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("path", "payload"),