From 5485c90ffcae2ca2bf4375cf312faddb6cf81f2d Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:25:18 -0700 Subject: [PATCH] fix(proxy): raise live budget counters for failure spend rows A stream that breaks mid-flight writes its recovered cost to the spend log, but the failure hook had already released the reservation and never touched the live counters, so an over-budget key kept passing until the counter was next repaired from the database Every charged failure row with a cost now raises the same counters the success path raises. The error waits at most a second for them, so a hung Redis cannot hold it back Fixes #35568 --- .../proxy/hooks/proxy_track_cost_callback.py | 117 ++++++++- litellm/proxy/proxy_server.py | 2 +- .../spend/test_partial_stream_live_budget.py | 195 +++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 235 +++++++++++++++++- 4 files changed, 540 insertions(+), 9 deletions(-) create mode 100644 tests/integration/spend/test_partial_stream_live_budget.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 994310c6619..2062d154a75 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final +from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -126,6 +127,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) ) +FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 + def _proxy_stamped_used_client_oauth_token( request_data: Mapping[str, object], request_route: str | None @@ -189,6 +192,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _proxy_spend_counters() -> "_IncrementSpendCounters": + from litellm.proxy.proxy_server import increment_spend_counters + + return increment_spend_counters + + DeploymentLookup: TypeAlias = Callable[[str], Deployment | None] @@ -315,11 +324,16 @@ class _ProxyDBLogger(CustomLogger): message_logging: bool = True, spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, deployment_lookup: DeploymentLookup = _proxy_deployment, + spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters, + failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS, ) -> None: super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) self.spend_event_producer = spend_event_producer self._spend_writer: Final = spend_writer self._deployment_lookup: Final = deployment_lookup + self._spend_counters: Final = spend_counters + self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds + self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime @@ -504,13 +518,16 @@ class _ProxyDBLogger(CustomLogger): start_time: datetime, spend_rows: Sequence[tuple[float, Mapping[str, object]]], ) -> None: - """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair.""" - for row_cost, row_kwargs in spend_rows: + """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the + live spend counters raised by every charged row that cost something.""" + token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + charged: Final = [ await self._spend_writer().update_database( - token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + token=token, response_cost=row_cost, user_id=user_api_key_dict.user_id, - end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, + end_user_id=end_user_id, team_id=user_api_key_dict.team_id, kwargs=dict(row_kwargs), completion_response=original_exception, @@ -519,6 +536,91 @@ class _ProxyDBLogger(CustomLogger): org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) + for row_cost, row_kwargs in spend_rows + ] + counted_rows: Final = tuple( + row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + ) + if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): + await self._increment_failure_spend_counters( + user_api_key_dict, token, end_user_id, counted_rows, start_time + ) + + async def _increment_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + """The counters the success path raises for the same spend, sent before the client sees the error. A + counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of + holding the error back. The reservation was already released, so every counter takes the full row cost.""" + increment: Final = asyncio.create_task( + self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time) + ) + self._pending_increments.add(increment) + increment.add_done_callback(self._pending_increments.discard) + finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds) + if not finished: + verbose_proxy_logger.warning( + "Spend counters for failure spend rows took over %ss, finishing them in the background", + self._failure_counter_wait_seconds, + ) + + async def _raise_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + for row_cost, row_kwargs in counted_rows: + await self._raise_row_spend_counters( + user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time + ) + await flush_post_call_redis_batches() + + async def _raise_row_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + row_cost: float, + row_kwargs: Mapping[str, object], + start_time: datetime, + ) -> None: + standard_logging_object: Final = row_kwargs.get("standard_logging_object") + sl_object: Final = ( + cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict + if isinstance(standard_logging_object, dict) + else None + ) + metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata")) + served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields( + standard_logging_object + ).get("model_id") + try: + await self._spend_counters()( + token=token, + team_id=user_api_key_dict.team_id, + user_id=user_api_key_dict.user_id, + response_cost=row_cost, + org_id=user_api_key_dict.org_id, + end_user_id=end_user_id, + tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)), + request_started_at=start_time, + model_access_groups=debitable_model_access_groups( + attributed=get_request_model_access_groups(row_kwargs), + served_model_id=served_model_id if isinstance(served_model_id, str) else None, + router=get_llm_router(), + ), + project_id=user_api_key_dict.project_id, + ) + except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing + verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") async def _PROXY_track_cost_callback( self, @@ -915,7 +1017,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: return tuple(sorted(str(key) for key in metadata)) -def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: +def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None: return budget_reservation_from_metadata(metadata) @@ -950,6 +1052,7 @@ class _IncrementSpendCounters(Protocol): tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> None: ... @@ -966,7 +1069,7 @@ async def _update_database_and_spend_counters( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, @@ -1027,7 +1130,7 @@ async def _update_database_and_spend_counters_in_batch( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None, model_access_groups: Sequence[str] | None, project_id: str | None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8d6c9dd925b..3fe36f03207 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3414,7 +3414,7 @@ async def increment_spend_counters( user_id: str | None, response_cost: float | None, org_id: str | None = None, - budget_reservation: dict | None = None, + budget_reservation: dict[str, object] | None = None, end_user_id: str | None = None, tags: list[str] | None = None, request_started_at: datetime | None = None, diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py new file mode 100644 index 00000000000..fd7ca2de522 --- /dev/null +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -0,0 +1,195 @@ +import json +import uuid +from collections.abc import Callable, Iterable, Iterator +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_BACKEND: Final = "integration-partial-stream" +_PIECE: Final = "partial stream billed tokens " +_DELIVERED_PIECES: Final = 12 +_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' + + +class _SpendRow(BaseModel): + spend: float + + +class _BudgetError(BaseModel): + type: str + + +class _ErrorResponse(BaseModel): + error: _BudgetError + + +def _chunk(delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: + body: Final = { + "id": "chatcmpl-partial", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": _BACKEND, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def _completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": _BACKEND, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() + + +def _peer(delivered: int) -> Callable[[Request], Reply]: + frames: Final = ( + *(_chunk({"role": "assistant", "content": _PIECE}) for _ in range(_DELIVERED_PIECES)), + _chunk({}, "stop"), + b"data: [DONE]\n\n", + ) + + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply( + content_type="text/event-stream", chunks=frames, abort_after=delivered, pause_between_chunks=0.02 + ) + return Reply(body=_completion()) + + return respond + + +def _priced_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"openai/{_BACKEND}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def _stream(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "stream": True, + "messages": [{"role": "user", "content": "partial stream"}], + }, + key=key, + ) + + +def _broken_stream(gateway: Gateway, model: str, key: str) -> str: + response: Final = _stream(gateway, model, key) + assert response.status_code == 200, response.text + events: Final = tuple( + line.removeprefix("data: ") + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert "error" in _JSON_OBJECT.validate_json(events[-1]), response.text + return "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + + +def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: + for chunk in chunks: + yield from (choice.delta.content or "" for choice in chunk.choices) + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) + found: Final = held( + lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + ) + return tuple(_SpendRow.model_validate(row).spend for row in found) + + +def _chats_sent(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def _recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: + key: Final = scenario.key(models=[model], max_budget=100) + assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES + (spend,) = _failure_spend(gateway, key) + assert spend > 0 + return spend + + +def _assert_budget_refusal(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text + + +@pytest.mark.timeout(240) +def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: + with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire) + cost: Final = _recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + _chats_sent(wire) + + assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES + after: Final = _chat(gateway, model, key) + + _assert_budget_refusal(after) + assert _failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) + assert _chats_sent(wire) == 1 + + +@pytest.mark.timeout(240) +def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire) + cost: Final = _recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES + assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES + after: Final = _chat(gateway, model, key) + + _assert_budget_refusal(after) + assert _failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) + + +@pytest.mark.timeout(240) +def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: + with wire_server(_peer(0)) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=0.05) + + refused: Final = _stream(gateway, model, key) + after: Final = _chat(gateway, model, key) + + assert refused.status_code == 500, refused.text + + assert after.status_code == 200, after.text + assert _failure_spend(gateway, key) == (0.0,) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index bd997f10322..36072d7eb79 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -11,6 +11,7 @@ import pytest import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer @@ -26,7 +27,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError -from litellm.proxy.utils import ProxyUpdateSpend +from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -2550,6 +2551,238 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 +@pytest.fixture() +def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache: + counters: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + 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: + usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + return { + "combined_usage_object": usage, + "response_cost": 0.5, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25), + BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5), + ), + } + + +_LIVE_COUNTERS: Final = { + "spend:key:hashed-key": 0.2, + "spend:team:team-1": 0.3, + "spend:team_member:user-1:team-1": 0.35, + "spend:user:user-1": 0.4, + "spend:end_user:end-user-1": 0.45, + "spend:org:org-1": 0.5, + "spend:project:project-1": 0.55, + "spend:tag:batch-jobs": 0.6, + "spend:model_access_group:premium": 0.65, +} + + +async def _fail_request(logger: ProxyDBLogger, failure: dict) -> None: + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]}, + "proxy_server_request": {"request_id": "rid"}, + **failure, + }, + original_exception=Exception("stream broke mid-flight"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", + team_id="team-1", + user_id="user-1", + end_user_id="end-user-1", + org_id="org-1", + project_id="project-1", + ), + ) + + +async def _counter_values(counters: DualCache) -> dict[str, float]: + return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ( + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + 0.125, + ), + (_failure_with_billed_attempts(), 1.25), + ({}, 0.0), + ], +) +async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( + live_spend_counters: DualCache, failure: dict, 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) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + charged for key, spend in _LIVE_COUNTERS.items()} + ) + + +@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 +): + 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 + 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()) + + 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() + counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable")) + + await _fail_request( + ProxyDBLogger(spend_writer=lambda: writer, 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 +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]] = [] + + async def slow_counter_store(**counter_args: object) -> None: + await store_answers.wait() + raised_costs.append(counter_args["response_cost"]) + landed.set() + + logger: Final = ProxyDBLogger( + spend_writer=_charging_spend_writer, + spend_counters=lambda: slow_counter_store, + failure_counter_wait_seconds=0.05, + ) + + 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, + ) + assert raised_costs == [] + + store_answers.set() + await asyncio.wait_for(landed.wait(), timeout=1) + assert raised_costs == [0.125] + + +@pytest.mark.asyncio +async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row + it already logged from being counted twice.""" + from litellm.proxy.proxy_server import proxy_logging_obj + + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue + 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) + await key_cache.async_set_cache( + "hashed-key", + UserAPIKeyAuth( + token="hashed-key", + 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), + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + "start_time": request_start, + }, + ) + + assert [ + (row["entity_type"], row["entity_id"], row["spend"], row["started_at"]) + for row in await window_queue.flush_and_get_aggregated_window_spend_transactions() + ] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")] + + +@pytest.mark.asyncio +async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the + final attempt's groups either.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + "model_info": {"id": "dep-b", "access_groups": ["pool"]}, + } + ] + ), + ) + + await _fail_request( + ProxyDBLogger(spend_writer=_charging_spend_writer), + { + **_failure_with_billed_attempts(), + "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"]}, + "standard_logging_object": {"id": "rid", "model_id": "dep-b"}, + }, + ) + + assert { + group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}") + for group in ("premium", "pool") + } == pytest.approx({"premium": 0.25, "pool": 1.25}) + + @pytest.mark.asyncio async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): """MCP tool calls may only carry user_api_key; user/team rollups still need user_id."""