fix(accounting): remove expired delivery receipts as fast as the worker applies them

The worker now applies several batches per run but still removed one cleanup batch
of applied receipts, so the table could outgrow its retention window. Cleanup now
repeats while each pass removes a full batch, for up to five seconds.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Joshua Valluru 2026-11-05 01:07:09 -08:00
parent d69d02191b
commit fed54b70d0
2 changed files with 33 additions and 4 deletions

View file

@ -87,10 +87,15 @@ async def flush_accounting_deliveries(
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
break
await repository.cleanup_applied(
retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")),
batch_size=batch_size * 2,
)
retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1"))
cleanup_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if (
await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2)
< batch_size * 2
or time.monotonic() >= cleanup_deadline
):
break
metrics: Final = await repository.metrics()
record_deliveries(metrics, time.monotonic() - started)
if metrics.batch_count >= int(

View file

@ -103,6 +103,30 @@ async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_trans
await client.db.disconnect()
@pytest.mark.asyncio
async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run(
parity_database: str, monkeypatch: pytest.MonkeyPatch
):
from litellm.accounting.worker import flush_accounting_deliveries
monkeypatch.setenv("DATABASE_URL", parity_database)
monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100")
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database, autocommit=True) as connection:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) '
"SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' "
"FROM generate_series(1, 501) AS n"
)
await client.db.connect()
try:
await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database) as connection:
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,)
finally:
await client.db.disconnect()
@pytest.mark.parametrize("ledger", [False, True])
@pytest.mark.parametrize("exclude_messages", [False, True])
@pytest.mark.asyncio