Merge d55efe1e88 into b781d157d7

litellm/proxy/utils.py

@@ -7313,13 +7313,20 @@ class ProxyUpdateSpend:
         start_time: Final = time.time()
         try:
             for i in range(n_retry_times + 1):
+                external_post_attempted = False
                 try:
                     base_url = os.getenv("SPEND_LOGS_URL", None)
                     if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None:
                         if not base_url.endswith("/"):
                             base_url += "/"
                         verbose_proxy_logger.debug("base_url: %s", base_url)
-                        json_data = json.dumps(logs_to_process)
+                        try:
+                            json_data = json.dumps(logs_to_process)
+                        except (TypeError, ValueError):
+                            # No external request has been sent. The batch is safe to replay.
+                            await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process)
+                            raise
+                        external_post_attempted = True
                         response = await db_writer_client.post(
                             url=base_url + "spend/update",
                             data=json_data,
@@ -7332,7 +7339,13 @@ class ProxyUpdateSpend:
                     else:
                         for j in range(0, len(logs_to_process), BATCH_SIZE):
                             batch = logs_to_process[j : j + BATCH_SIZE]
-                            batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch]
+                            try:
+                                batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch]
+                            except (TypeError, ValueError):
+                                # This batch has not reached Prisma. Earlier batches may
+                                # already have committed, so replay only this tail.
+                                await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process[j:])
+                                raise
                             isolation_budget = MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH
                             for statement_rows in spend_log_write_batches(
                                 batch_with_dates,
@@ -7356,6 +7369,10 @@ class ProxyUpdateSpend:
                         )
                     break
                 except Exception as e:
+                    if external_post_attempted:
+                        # Even a transport error can arrive after the remote writer
+                        # committed. Retrying or requeueing could duplicate spend.
+                        raise
                     if not _is_transient_spend_log_write_error(e):
                         if PrismaDBExceptionHandler.is_prisma_error(e):
                             await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process)

tests/unit/proxy/test_update_spend.py

@@ -321,3 +321,72 @@ async def test_update_spend_logs_multiple_batches_with_failure():
 
     # Verify all logs were cleared from transactions
     assert len(prisma_client.spend_log_transactions) == 0
+
+
+# These tests live in the proxy-db-db-and-spend CI shard, unlike the separate
+# prisma_and_spend suite, so their branch coverage contributes to codecov.
+def _spend_log_lifetime_fakes():
+    client = MockPrismaClient()
+    logging = create_mock_proxy_logging()
+    logging.db_spend_update_writer.redis_update_buffer.store_spend_logs_in_redis = AsyncMock(return_value=False)
+    return client, logging
+
+
+@pytest.mark.asyncio
+async def test_spend_log_pre_write_failure_requeues_only_unwritten_tail(monkeypatch):
+    from litellm.proxy.utils import ProxyUpdateSpend
+
+    monkeypatch.delenv("SPEND_LOGS_URL", raising=False)
+    client, logging = _spend_log_lifetime_fakes()
+    rows = [{"request_id": f"committed-{i}"} for i in range(1000)]
+    rows.append({"request_id": "unwritten"})
+
+    def jsonify(row):
+        if row["request_id"] == "unwritten":
+            raise TypeError("local conversion failed")
+        return row
+
+    client.jsonify_object = jsonify
+    with pytest.raises(TypeError, match="local conversion failed"):
+        await ProxyUpdateSpend.update_spend_logs(
+            n_retry_times=0, prisma_client=client, db_writer_client=None,
+            proxy_logging_obj=logging, logs_to_process=rows,
+        )
+    assert client.db.litellm_spendlogs.create_many.await_count >= 1
+    assert client.spend_log_transactions == [{"request_id": "unwritten"}]
+
+
+@pytest.mark.asyncio
+async def test_spend_log_external_preflight_requeues_without_post(monkeypatch):
+    from litellm.proxy.utils import ProxyUpdateSpend
+
+    monkeypatch.setenv("SPEND_LOGS_URL", "http://writer.invalid")
+    client, logging = _spend_log_lifetime_fakes()
+    rows = [{"request_id": "unwritten", "bad": object()}]
+    writer = MagicMock()
+    writer.post = AsyncMock()
+    with pytest.raises(TypeError):
+        await ProxyUpdateSpend.update_spend_logs(
+            n_retry_times=0, prisma_client=client, db_writer_client=writer,
+            proxy_logging_obj=logging, logs_to_process=rows,
+        )
+    writer.post.assert_not_awaited()
+    assert client.spend_log_transactions == rows
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", [ValueError("uncertain delivery"), httpx.ReadError("response lost after send")])
+async def test_spend_log_external_post_failure_never_retries_or_requeues(monkeypatch, failure):
+    from litellm.proxy.utils import ProxyUpdateSpend
+
+    monkeypatch.setenv("SPEND_LOGS_URL", "http://writer.invalid")
+    client, logging = _spend_log_lifetime_fakes()
+    writer = MagicMock()
+    writer.post = AsyncMock(side_effect=failure)
+    with pytest.raises(type(failure)):
+        await ProxyUpdateSpend.update_spend_logs(
+            n_retry_times=2, prisma_client=client, db_writer_client=writer,
+            proxy_logging_obj=logging, logs_to_process=[{"request_id": "maybe-committed"}],
+        )
+    writer.post.assert_awaited_once()
+    assert client.spend_log_transactions == []