fix(accounting): normalize usage metadata before queue admission

litellm/proxy/_types.py

@@ -2480,7 +2480,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
         description="Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. "
         "Other requests remain supported with asynchronous usage capture.",
     )
-    accounting_sync_ack_timeout_seconds: float = Field(5.0, gt=0, description="Maximum response wait for usage commit")
+    accounting_sync_ack_timeout_seconds: float = Field(
+        20.0, gt=0, description="Maximum response wait for usage commit, including the batch flush interval"
+    )
     accounting_sync_ack_failure_policy: Literal["fail_closed", "log"] = Field(
         "fail_closed", description="Return 503 or log the persistence failure after the response wait expires"
     )

litellm/proxy/spend_tracking/usage_capture.py

@@ -131,6 +131,19 @@ class UsageCaptureContext:
     included_attempt_ids: tuple[str, ...] = ()
 
 
+def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite:
+    from sys import maxsize
+
+    from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure
+
+    normalized: Final = _FIELDS.validate_python(safe_json_structure(envelope, max_depth=maxsize))
+    if normalized.get("event") != event.model_dump(mode="json"):
+        raise UsageCaptureError("Usage event text cannot be normalized without changing identity or attribution")
+    serialized: Final = _json(normalized)
+    _JSON_OBJECT.validate_json(serialized)
+    return UsageWrite(event, serialized, usage_delivery_targets(event.dimensions))
+
+
 def usage_writes(
     payload: SpendLogsPayload, usage_json: str | None, context: UsageCaptureContext
 ) -> tuple[UsageWrite, ...]:
@@ -157,27 +170,25 @@ def usage_writes(
         dimensions=dimensions,
         metrics=_usage_metrics(None if context.usage_source == "unknown" else usage),
     )
-    envelope: Final = _json(
-        {
-            "event": outbound.model_dump(mode="json"),
-            "usage": usage,
-            "usage_source": "unknown" if usage is None else context.usage_source,
-            "included_attempt_ids": context.included_attempt_ids,
-            "pricing": _JSON_OBJECT.validate_json(context.pricing),
-            "charge_eligible": context.charge_eligible,
-            "request_count_eligible": context.request_count_eligible,
-            "ownership": context.ownership,
-            "parent_request_id": context.parent_request_id,
-            "job_id": context.job_id,
-            "phase": context.phase,
-            "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps),
-            "finalized_cost": None
-            if context.finalized_cost is None
-            else _JSON_OBJECT.validate_json(context.finalized_cost),
-            "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation),
-        }
-    )
-    outbound_write: Final = UsageWrite(outbound, envelope, usage_delivery_targets(dimensions))
+    envelope: Final = {
+        "event": outbound.model_dump(mode="json"),
+        "usage": usage,
+        "usage_source": "unknown" if usage is None else context.usage_source,
+        "included_attempt_ids": context.included_attempt_ids,
+        "pricing": _JSON_OBJECT.validate_json(context.pricing),
+        "charge_eligible": context.charge_eligible,
+        "request_count_eligible": context.request_count_eligible,
+        "ownership": context.ownership,
+        "parent_request_id": context.parent_request_id,
+        "job_id": context.job_id,
+        "phase": context.phase,
+        "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps),
+        "finalized_cost": None
+        if context.finalized_cost is None
+        else _JSON_OBJECT.validate_json(context.finalized_cost),
+        "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation),
+    }
+    outbound_write: Final = _database_usage_write(outbound, envelope)
     if not context.request_count_eligible or context.caller_id is None:
         return (outbound_write,)
     inbound: Final = InboundUsageEvent(
@@ -195,7 +206,7 @@ def usage_writes(
         ),
     )
     return (
-        UsageWrite(inbound, _json({"event": inbound.model_dump(mode="json")}), usage_delivery_targets(dimensions)),
+        _database_usage_write(inbound, {"event": inbound.model_dump(mode="json")}),
         outbound_write,
     )
 
@@ -529,7 +540,7 @@ async def capture_proxy_response(logging_obj: object, response: object, *, backg
             if not isinstance(captured, CapturedUsage):
                 raise UsageCaptureError("Usage capture was not admitted")
             timeout: Final = TypeAdapter(float).validate_python(
-                settings.get("accounting_sync_ack_timeout_seconds", 5.0)
+                settings.get("accounting_sync_ack_timeout_seconds", 20.0)
             )
             await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout)
     except Exception as error:  # noqa: BLE001  # the configured response policy owns every capture failure

tests/unit/proxy/spend_tracking/test_usage_capture.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from collections.abc import Mapping
+from collections.abc import Awaitable, Mapping
 from datetime import datetime, timezone
 from decimal import Decimal
 from typing import Final
@@ -136,6 +136,86 @@ class StoredUsage:
         self.rows += tuple(row for row in writes if row.event.event_id not in existing)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("text", ["parent\x00value", "pärent-✓-日本", r"\u0000"])
+async def test_capture_normalizes_database_text_without_poisoning_neighbors(text: str) -> None:
+    from litellm.proxy.spend_tracking.usage_capture import PricingCapture, capture_callback
+
+    database: Final = StoredUsage()
+    batches: Final = asyncio.Queue[tuple[UsageWrite, ...]]()
+
+    async def persist(writes: tuple[UsageWrite, ...]) -> None:
+        assert all("\x00" not in (json.loads(write.envelope).get("parent_request_id") or "") for write in writes)
+        batches.put_nowait(writes)
+        await database.insert(writes)
+
+    meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(), wait_for_batch=_manual_flush)
+    nested: Final = json.loads("[" * 100 + "3" + "]" * 100)
+    kwargs: Final = {
+        "call_type": "acompletion",
+        "litellm_call_id": "caller",
+        "model": "capture-model",
+        "search_context_size": text,
+        "_accounting_pricing_snapshot": PricingCapture(
+            "{}", "{}", json.dumps({"request": {"search_context_size": text}, "nested": nested})
+        ),
+        "litellm_params": {"metadata": {"parent_request_id": text}},
+    }
+    response: Final = {"id": "provider", "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}}
+    captured: Final = capture_callback(meter, kwargs, response, NOW, NOW)
+    neighbor: Final = capture_callback(
+        meter,
+        {**kwargs, "litellm_call_id": "neighbor", "litellm_params": {}},
+        {**response, "id": "neighbor-provider"},
+        NOW,
+        NOW,
+    )
+    try:
+        await meter.flush()
+        await asyncio.gather(captured.receipt, neighbor.receipt)
+        assert batches.qsize() == 1
+        assert len(batches.get_nowait()) == 4
+        envelope: Final = json.loads(captured.writes[-1].envelope)
+        assert envelope["parent_request_id"] == text.replace("\x00", "")
+        assert envelope["pricing"]["request"]["search_context_size"] == text.replace("\x00", "")
+        assert envelope["pricing"]["nested"] == nested
+        assert envelope["event"] == captured.writes[-1].event.model_dump(mode="json")
+        assert envelope["usage"] == response["usage"]
+        assert kwargs["litellm_params"] == {"metadata": {"parent_request_id": text}}
+        repeated: Final = capture_callback(meter, kwargs, response, NOW, NOW)
+        await meter.flush()
+        await repeated.receipt
+        assert repeated.writes == captured.writes
+        assert len(database.rows) == 4
+    finally:
+        await meter.stop(0)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("invalid_field", ["identity", "context"])
+async def test_invalid_capture_text_is_rejected_before_admission(invalid_field: str) -> None:
+    from litellm.proxy.spend_tracking.usage_capture import capture_callback
+
+    database: Final = StoredUsage()
+    meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush)
+    kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "caller", "litellm_params": {}}
+    try:
+        accepted: Final = capture_callback(meter, kwargs, {"id": "provider"}, NOW, NOW)
+        invalid: Final = (
+            {**kwargs, "litellm_call_id": "caller\x00"}
+            if invalid_field == "identity"
+            else {**kwargs, "litellm_params": {"metadata": {"parent_request_id": "parent\ud800"}}}
+        )
+        with pytest.raises((UsageCaptureError, ValueError)):
+            capture_callback(meter, invalid, {"id": "other"}, NOW, NOW)
+        assert meter.pending_records == 2
+        await meter.flush()
+        await accepted.receipt
+        assert tuple(write.event.event_id for write in database.rows) == ("caller:inbound", "provider")
+    finally:
+        await meter.stop(0)
+
+
 def _logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging:
     from litellm.litellm_core_utils.litellm_logging import Logging
 
@@ -348,6 +428,7 @@ async def test_stream_capture_ignores_chunks_and_preserves_terminal_usage_extras
 @pytest.mark.parametrize("call_type", ["acompletion", "atext_completion", "aembedding", "aresponses"])
 async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch: pytest.MonkeyPatch) -> None:
     from litellm.proxy import proxy_server
+    from litellm.proxy._types import ConfigGeneralSettings
     from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response
     from litellm.types.utils import ModelResponse, Usage
 
@@ -356,6 +437,14 @@ async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch:
     runtime.start(database.write, UsageCaptureSettings(batch_size=2))
     monkeypatch.setattr(proxy_server, "usage_capture", runtime)
     monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True})
+    deadlines: Final = asyncio.Queue[float | None]()
+    original_wait: Final = asyncio.wait_for
+
+    async def wait_with_deadline(awaitable: Awaitable[object], timeout: float | None) -> object:
+        deadlines.put_nowait(timeout)
+        return await original_wait(awaitable, timeout)
+
+    monkeypatch.setattr(asyncio, "wait_for", wait_with_deadline)
     response: Final = asyncio.create_task(
         capture_proxy_response(
             _logger(call_type),
@@ -364,6 +453,9 @@ async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch:
     )
     try:
         await database.entered.wait()
+        deadline: Final = deadlines.get_nowait()
+        assert deadline == ConfigGeneralSettings().accounting_sync_ack_timeout_seconds
+        assert deadline is not None and deadline > UsageCaptureSettings().flush_interval_seconds + 10
         assert not response.done()
         assert database.committed == ()
         database.release.set()

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -29959,8 +29959,8 @@ export interface components {
             accounting_sync_ack_failure_policy: "fail_closed" | "log";
             /**
              * Accounting Sync Ack Timeout Seconds
-             * @description Maximum response wait for usage commit
-             * @default 5
+             * @description Maximum response wait for usage commit, including the batch flush interval
+             * @default 20
              */
             accounting_sync_ack_timeout_seconds: number;
             /**
@@ -49531,7 +49531,7 @@ export interface components {
             batch_size: number;
             /**
              * Flush Interval Seconds
-             * @default 0.2
+             * @default 5
              */
             flush_interval_seconds: number;
             /**