From 79bc95d394e45bab07c2d0dc0c4bc1f1a8b6fb97 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:12:57 -0700 Subject: [PATCH] fix(accounting): depend on prior receipts without resubmitting usage --- .../spend_tracking/spend_tracking_utils.py | 8 +++-- litellm/proxy/spend_tracking/usage_capture.py | 16 ++++----- .../test_spend_tracking_utils.py | 26 ++++++++++++++ .../spend_tracking/test_usage_capture.py | 36 ++++++++++++++++++- 4 files changed, 74 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 86d3711754d..7499e377cdf 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -625,8 +625,8 @@ def get_logging_payload( routing_metadata, routing_litellm_params, ) - completion_start_time: Final = TypeAdapter(datetime).validate_python( - kwargs.get("completion_start_time") or end_time + completion_start_time: Final = TypeAdapter[datetime | None](datetime | None).validate_python( + kwargs.get("completion_start_time", end_time) ) call_type: Final = kwargs.get("call_type") cache_hit: Final = kwargs.get("cache_hit", False) @@ -933,7 +933,9 @@ def get_logging_payload( cache_hit=str(cache_hit), startTime=_ensure_datetime_utc(start_time), endTime=_ensure_datetime_utc(end_time), - completionStartTime=_ensure_datetime_utc(completion_start_time), + completionStartTime=( + _ensure_datetime_utc(completion_start_time) if completion_start_time is not None else None + ), model=model_name, user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 795bb131029..84786407c57 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -505,21 +505,21 @@ class _CaptureBuilder: self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance)) -def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: +def _prior_captures(kwargs: Mapping[str, object]) -> tuple[CapturedUsage, ...]: prior: Final = kwargs.get("_accounting_prior_captures", ()) return ( tuple( - chain.from_iterable( - item.writes - for item in TypeAdapter(tuple[object, ...]).validate_python(prior) - if isinstance(item, CapturedUsage) - ) + item for item in TypeAdapter(tuple[object, ...]).validate_python(prior) if isinstance(item, CapturedUsage) ) if isinstance(prior, tuple) else () ) +def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: + return tuple(chain.from_iterable(item.writes for item in _prior_captures(kwargs))) + + def capture_callback( meter: DurableUsageMeter, kwargs: Mapping[str, object], @@ -537,8 +537,8 @@ def capture_callback( payload: Final = get_logging_payload( dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe ) - dependencies: Final = _prior_writes(kwargs) - return CapturedUsage(payload, builder.writes, meter.submit((*dependencies, *builder.writes))) + prerequisites: Final = tuple(item.receipt for item in _prior_captures(kwargs)) + return CapturedUsage(payload, builder.writes, meter.submit(builder.writes, prerequisites=prerequisites)) def captured_callback_kwargs( diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index e2fac424b7d..f52f40ad55d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7561,3 +7561,29 @@ def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadat {**kwargs, "mcp_tool_call_metadata": tool_metadata}, response, now, now ) assert actual == expected + + +@pytest.mark.parametrize("capture_enabled", [False, True]) +@pytest.mark.parametrize("completion_start", [None, "absent", "reported"]) +def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, completion_start: str | None) -> None: + start: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) + end: Final = start + datetime.timedelta(seconds=2) + reported: Final = start + datetime.timedelta(seconds=1) + observations: Final[list[object]] = [] + kwargs: Final = { + "call_type": "acompletion", + "litellm_call_id": "timestamp-call", + "litellm_params": {}, + **({} if completion_start == "absent" else {"completion_start_time": reported if completion_start else None}), + } + payload: Final = spend_tracking_utils.get_logging_payload( + kwargs, + {"id": "timestamp-response"}, + start, + end, + capture_usage=(lambda row, usage: observations.append(row)) if capture_enabled else None, + ) + assert payload["completionStartTime"] == ( + end if completion_start == "absent" else reported if completion_start else None + ) + assert observations == ([payload] if capture_enabled else []) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 85ed9046e8e..993bb8901f0 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -1275,7 +1275,7 @@ async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) captured: Final = other.model_call_details["_accounting_capture"] await runtime.meter.flush() - with pytest.raises(UsageCaptureError, match="identity"): + with pytest.raises(UsageCaptureError, match="dependency"): await captured.receipt finally: await runtime.stop(1) @@ -1302,3 +1302,37 @@ async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usa await runtime.stop(1) assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success") assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prior_committed", [False, True]) +@pytest.mark.parametrize("batch_size", [2, 3]) +async def test_retry_dependencies_do_not_resubmit_usage_or_expand_final_submission( + prior_committed: bool, batch_size: int +) -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = CommitBarrier() + database.release.set() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(batch_size=batch_size, max_pending_records=3)) + logger: Final = _logger() + try: + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) + logger.carry_failed_attempt_usage() + assert runtime.meter is not None + if prior_committed: + await runtime.meter.flush() + logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) + captured: Final = logger.model_call_details.get("_accounting_capture") + assert captured is not None, "Final usage must be admitted independently of prior usage batch size" + await runtime.meter.close() + await captured.receipt + assert runtime.meter.pending_records == 0 + finally: + await runtime.stop(1) + assert database.committed == ("caller-attempt-0", "caller:inbound", "success") + final: Final = json.loads(captured.writes[-1].envelope) + assert final["related_attempt_ids"] == ["caller-attempt-0"] + assert final["usage"]["prompt_tokens"] == 3