diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e7db2d8f276..173805543e8 100644 --- a/litellm/proxy/_types.py +++ b/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" ) diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 3ca62fadb60..af6f4100fce 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/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 diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 05bcbb90679..e465cef2e88 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/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() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 12ffe29f111..65c053371dd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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; /**