From fcb6d2267c09389ae1aa9e80e04a35caa5b8b470 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:51:44 -0700 Subject: [PATCH] fix(spend): keep a batch's claim row out of the logs a proxy was told not to write disable_spend_logs has to keep meaning that no request gets logged, and the row that makes a batch chargeable exactly once is the one row it cannot drop, so with logging off that row now carries only what tells the retrieves apart. SPEND_LOGS_URL deployments get their copy back too: the claim writes straight to this table, so the row is queued as well when an external writer is the one that takes the spend logs. --- litellm/proxy/db/db_spend_update_writer.py | 49 +++++++-- .../proxy/db/test_db_spend_update_writer.py | 99 +++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ee7802a45d3..bdc014d7f13 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -89,6 +89,21 @@ def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" +_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) + + +def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a batch's cost row to what tells the retrieves apart when logging is off. + + A proxy run with spend logs disabled still needs one row per batch to charge it once, + so the row is written either way, but it carries no request of its own: no metadata, + no requester IP, no key, model, or token counts (LIT-7048). + """ + if disable_spend_logs is False: + return payload + return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -336,12 +351,16 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool ) -> bool: if prisma_client is not None and _is_batch_cost_row(payload): - return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + return await self._claim_batch_cost_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) return True - async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + async def _claim_batch_cost_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> bool: """Write the batch's cost row now, or learn that another retrieve already did. Every retrieve of one batch shares this row, so the insert that lands first owns @@ -353,13 +372,17 @@ class DBSpendUpdateWriter: from litellm.repositories.table_repositories import SpendLogsRepository request_id: Final = payload["request_id"] + row: Final = _batch_cost_row_to_write(payload, disable_spend_logs) spend_logs: Final = SpendLogsRepository(prisma_client).table try: claimed: Final = await spend_logs.create_many( - data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list skip_duplicates=True, ) if claimed == 1: + await self._forward_batch_cost_row( + row=row, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) return True existing: Final = await spend_logs.find_unique( where={"request_id": request_id} # mutable-ok: prisma where clause @@ -368,7 +391,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.warning( "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e ) - await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) return True if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": verbose_proxy_logger.warning( @@ -380,10 +403,22 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client) + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + + async def _forward_batch_cost_row( + self, row: Mapping[str, object], prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> None: + """Queue the claimed row for an external spend log writer, which the claim went around. + + With ``SPEND_LOGS_URL`` set the queue posts every spend log to that writer instead of + inserting it, so a batch's cost row reaches it only by being queued here as well. + """ + if disable_spend_logs is True or os.getenv("SPEND_LOGS_URL") is None: + return + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) async def _take_over_uncharged_batch_cost_row( - self, payload: SpendLogsPayload, prisma_client: "PrismaClient" + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: """Take the batch's cost row over from the poll that left it charging nothing. @@ -403,7 +438,7 @@ class DBSpendUpdateWriter: try: taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( data=prisma_client.jsonify_object( - MappingProxyType({field: value for field, value in payload.items() if field != "request_id"}) + MappingProxyType({field: value for field, value in row.items() if field != "request_id"}) ), where={ # mutable-ok: prisma where clause "request_id": request_id, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4efb94b60aa..a8aeaced55f 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import os import re @@ -3169,6 +3170,104 @@ async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disa assert db_writer._batch_database_updates.await_count == 1 +_BATCH_CLAIM_FIELDS = {"request_id", "call_type", "status", "spend", "startTime", "endTime"} + + +def _logged_batch_cost_payload() -> dict: + return { + **_batch_cost_payload(), + "api_key": "0e5b0e9e5f", + "model": "gpt-5.6-luna", + "user": "test-user", + "metadata": '{"batch_models": ["gpt-5.6-luna"]}', + "requester_ip_address": "127.0.0.1", + "proxy_server_request": '{"headers": {"user-agent": "litellm-batch-cost-check"}}', + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("disable_spend_logs", "logs_the_request"), + [(False, True), (True, False)], + ids=["spend_logs_on", "spend_logs_off"], +) +async def test_update_database_claims_a_batch_without_logging_the_request_that_polled_it( + disable_spend_logs: bool, logs_the_request: bool +): + """ + disable_spend_logs has to keep meaning that no request gets logged, and the batch's cost + row is the one row it cannot drop, so with logging off that row carries only what tells + the retrieves apart: no metadata, no requester IP, no key, model, or token counts. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = _logged_batch_cost_payload() + + assert await _update_database_with(db_writer, prisma, payload, disable_spend_logs) is True + + claimed = prisma.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0] + assert set(claimed) == (set(payload) if logs_the_request else _BATCH_CLAIM_FIELDS) + assert claimed["spend"] == 0.25 + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs_url", "forwarded"), + [("http://spend-logs.internal", True), (None, False)], + ids=["an_external_writer_takes_the_rows", "rows_are_written_to_this_db"], +) +async def test_update_database_sends_a_claimed_batch_cost_row_on_to_an_external_spend_log_writer( + monkeypatch, spend_logs_url: str | None, forwarded: bool +): + """ + SPEND_LOGS_URL makes the flush post spend logs to that writer instead of inserting them, + and the claim writes straight to this table, so the batch's row reaches the writer only + by being queued as well. Queueing it with no writer configured would insert it twice. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + if spend_logs_url is None: + monkeypatch.delenv("SPEND_LOGS_URL", raising=False) + else: + monkeypatch.setenv("SPEND_LOGS_URL", spend_logs_url) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + queued = [row["request_id"] for row in prisma.spend_log_transactions] + assert queued == (["batch_abc_batch_cost"] if forwarded else []) + + +@pytest.mark.asyncio +async def test_update_database_forwards_no_batch_cost_row_a_later_retrieve_had_already_claimed(monkeypatch): + """The retrieve that lost the claim charges nothing, so it must not post a row either.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25) + prisma = _spend_logs_prisma(0, existing) + monkeypatch.setenv("SPEND_LOGS_URL", "http://spend-logs.internal") + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False + + assert prisma.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): + """A refused claim is retried through the queue, so what it queues has to stay unlogged too.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _logged_batch_cost_payload(), True) is True + + assert [set(row) for row in prisma.spend_log_transactions] == [_BATCH_CLAIM_FIELDS] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): """An unreachable DB must not drop the batch's only spend row, nor its charge."""