fix(proxy): isolate poison spend-log rows so one bad record can't drop the whole batch - #31705
Conversation
|
|
Greptile SummaryThis PR makes spend-log batch writes tolerate one bad row without dropping the full batch. The main changes are:
Confidence Score: 4/5The change is narrowly scoped to spend-log persistence and preserves the existing retry path for database availability failures. The implementation includes targeted tests for poison-row isolation, connection-error propagation, and the isolation attempt budget, with no remaining review comments.
 T-Rex LogsWhat T-Rex did
Reviews (5): Last reviewed commit: "fix(proxy): isolate poison spend-log row..." | Re-trigger Greptile |
Greptile SummaryThis PR makes spend-log batch writes survive one bad row. The main changes are:
Confidence Score: 5/5The change is narrowly scoped to spend-log batch insertion and preserves the existing retry path for database connectivity failures. The implementation is covered by targeted tests for both poison-row isolation and connection-error propagation, and no code issues were identified.
T-Rex LogsWhat T-Rex did
Reviews (1): Last reviewed commit: "fix(proxy): isolate poison spend-log row..." | Re-trigger Greptile |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
eed97c8 to
3a28122
Compare
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
3a28122 to
cb614be
Compare
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…p the whole batch update_spend_logs flushes the queue with a single create_many per batch, so one row carrying bytes Postgres refuses (a residual NUL byte is the canonical case) fails the entire insert and drops every good spend log alongside it. PR #29515 strips NUL bytes from the JSON columns, but the scalar string columns (end_user, model, session_id, ...) still flow through unsanitized, so a poisoned row can still reach the write and take a batch of up to 1000 good rows down with it. On a genuine data-layer rejection the batch is now bisected so the good rows still persist and only the offending row is dropped and logged with its request_id. The classification lives in PrismaDBExceptionHandler.is_prisma_data_error (matched by exact type so systemic subclasses like a missing table are not mistaken for a single poison row), which keeps prisma an in-function import and litellm.proxy.utils importable without the proxy extra. Transport failures, including the "can't reach database server" outage that prisma mislabels as a DataError, are re-raised unchanged so the existing connection-retry path still runs and a transient outage never turns into silent per-row data loss. The bisection carries a per-batch attempt budget that hard-caps how many create_many calls the isolation may issue (checked before any insert, decremented per call, threaded through the recursion), so an authenticated caller flooding poisoned rows cannot amplify one failed bulk insert into unbounded failed inserts and log lines; once the budget is spent the still-failing remainder is dropped wholesale under a single log line. Resolves LIT-4103
cb614be to
bc3cdd4
Compare
|
Addressed both review points in the latest push. The isolation budget now caps create_many attempts directly rather than single-row drops. MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH is a hard ceiling on the number of inserts the isolation may issue for one batch; it is checked before any insert (so an exhausted budget never even attempts a write, which also fixes the ordering point about the single-row branch), decremented once per create_many call, and threaded through the recursion, so the total inserts are bounded by the budget regardless of how many rows are poisoned. The new test asserts the attempt count stays at or below the cap and below the input row count for a single-batch poison flood, and it is a mutation-killer (removing the budget branch makes it fail). Also removed the unused real_spend_log_error local the inline comment flagged. |
…p the whole batch (BerriAI#31705) update_spend_logs flushes the queue with a single create_many per batch, so one row carrying bytes Postgres refuses (a residual NUL byte is the canonical case) fails the entire insert and drops every good spend log alongside it. PR BerriAI#29515 strips NUL bytes from the JSON columns, but the scalar string columns (end_user, model, session_id, ...) still flow through unsanitized, so a poisoned row can still reach the write and take a batch of up to 1000 good rows down with it. On a genuine data-layer rejection the batch is now bisected so the good rows still persist and only the offending row is dropped and logged with its request_id. The classification lives in PrismaDBExceptionHandler.is_prisma_data_error (matched by exact type so systemic subclasses like a missing table are not mistaken for a single poison row), which keeps prisma an in-function import and litellm.proxy.utils importable without the proxy extra. Transport failures, including the "can't reach database server" outage that prisma mislabels as a DataError, are re-raised unchanged so the existing connection-retry path still runs and a transient outage never turns into silent per-row data loss. The bisection carries a per-batch isolation budget so an authenticated caller flooding poisoned rows cannot amplify one failed bulk insert into ~2N failed inserts and N log lines; once the budget is spent the still-failing remainder is dropped wholesale under a single log line. Resolves LIT-4103
3 participants
Relevant issues
Linear ticket
Resolves LIT-4103
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewScreenshots / Proof of Fix
update_spend_logsflushes the queue with a singlecreate_manyper batch, which Postgres executes as one atomic statement. If a single row carries bytes Postgres refuses, the whole insert fails and every good row in that batch is dropped with it. PR #29515 strips NUL bytes from the JSON columns (messages,response,request_tags,proxy_server_request,metadata), but the scalar string columns are still written unsanitized, so a NUL byte in the requestuserfield reaches theend_usercolumn and reproduces the original 22xxx encoding rejection.Live proxy against real Postgres and a real Anthropic model. A NUL byte cannot travel in a shell argv, so the three requests are sent from a tiny client that injects
chr(0)into theuserfield of one of them; everything else is observed with psql. Two requests carry a cleanuser, the third carries a NUL byte, and all three land in the same flush batch (proxy_batch_write_at: 5).Before the fix, the whole batch is lost. All three responses are 200, but none of the spend logs reach the DB:
proxy log:
After the fix, the two good rows persist and only the poisoned row is dropped and dead-lettered with its
request_id:proxy log:
The "Error in spend logs queue monitor" line that fires on the unfixed batch is gone after the fix, since the batch no longer crashes the flush job.
A transient outage must never be treated as a poison row. The "can't reach database server" failure that prisma mislabels as a
DataErroris re-raised unchanged throughis_database_service_unavailable_error, so it keeps flowing into the existing connection-retry path instead of being bisected into silent per-row drops. This is pinned bytest_update_spend_logs_reraises_connection_masquerade_dataerror, and the salvage behavior bytest_update_spend_logs_isolates_poison_row_and_persists_good_rows; both fail on the pre-fix code.Type
🐛 Bug Fix
Changes
update_spend_logsnow routes each batch through_create_spend_logs_with_poison_isolation. On a genuine data-layer rejection it bisects the batch so the good rows still persist and only the offending row is dropped and logged with itsrequest_id; the disjoint-halves recursion writes each row at most once andskip_duplicateskeeps any cross-retry re-attempt idempotent against therequest_idprimary key. Transport failures, including the connection outage prisma surfaces as aDataError, are re-raised so the caller's retry/backoff path is unchanged and a blip never becomes data loss.The classification lives in a new
PrismaDBExceptionHandler.is_prisma_data_errorhelper that matches the baseDataErrorby exact type, so the specific subclasses (UniqueViolationError,TableNotFoundError,MissingRequiredValueError, ...) are excluded and a systemic failure like a missing table surfaces loudly instead of being bisected into silent per-row drops. Keeping the check there preserves the in-functionprismaimport that module already uses, solitellm.proxy.utilsstays importable without the proxy extra.The bisection carries a per-batch attempt budget (
MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH) that hard-caps how manycreate_manycalls the isolation may issue, checked before any insert and decremented per call, so an authenticated caller flooding poisoned rows cannot amplify one failed bulk insert into unbounded failed inserts and log lines. Once the budget is spent the still-failing remainder is dropped wholesale under one log line, which is the pre-existing drop-the-batch behavior, so the isolation salvages the common sparse-poison case while the DB work under abuse stays bounded by the budget regardless of how many rows are poisoned.This is scoped to the batch-isolation gap that PR #29515 left open; it does not change the NUL-byte stripping or add sanitization for the scalar columns, so whatever still slips through to the write is now contained to its own row instead of taking the batch down.
The isolation covers the in-process Prisma write only. When
SPEND_LOGS_URLis set the batch is POSTed to an external writer that owns its own atomicity, and the separateLiteLLM_SpendLogToolIndexwrite is already fault-isolated behind its own non-fatal try/except, so neither is in scope here.