fix(proxy): force prisma recreate on postgres cached-plan error - #36428
Conversation
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Greptile SummaryThe PR makes cached-plan recovery force a Prisma engine recreation while preserving singleflight cooldown behavior across writer and read-replica engines.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/routing_prisma_wrapper.py | Centralizes top-level read routing through read_target, allowing recovery logic to identify the engine that served a read. |
| litellm/proxy/utils.py | Adds forced cached-plan recreation, engine-scoped cooldown bookkeeping, and consistent decline handling without an eligible unresolved blocking failure. |
| tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py | Verifies cached-plan fallback forwards forced recreation and observes the pre-query reader engine. |
| tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py | Adds broad regression coverage for cooldown waivers, read-replica identity, failed repairs, generation-guard declines, and escalation cleanup. |
Reviews (6): Last reviewed commit: "fix(proxy): force prisma recreate on pos..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
6c95509 to
48290ca
Compare
|
@greptileai please review the current head 48290ca. The cooldown P1 and the writer_db monkeypatch P2 are both fixed, reasoning in the PR description |
48290ca to
a041a4f
Compare
|
@greptileai please review the current head a041a4f. The read-replica finding was correct and is fixed, along with two related holes it exposed |
a041a4f to
50ea962
Compare
|
@greptileai please review the current head 50ea962. The global failure gate is now scoped per engine, plus two defects it exposed. Reasoning in the description |
50ea962 to
96912fd
Compare
|
@greptileai please review the current head 96912fd. Your 5/5 commit stranded the dead-engine flag on a declined heavy-path recreate, reintroducing #29176. Fixed here |
`_query_first_with_cached_plan_fallback` recovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer with `SELECT 1` first and skips the recreate when it answers, which is right for the IAM token refresh it was added for and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate. Callers now pass `force_recreate` to skip that probe, and only the cached-plan fallback does. Getting past the probe is not enough on its own. Both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing. `force=True` would fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine it observed before the query, and the cooldown is waived only while that engine is still the live one, so the first caller repairs the pool and the rest fall back to the normal cooldown. That engine has to be the one the query actually ran on. `query_first` is a top-level read, so with a read replica configured it is dispatched to the reader and it is the reader's prepared statements that go stale, while `writer_db` names a different engine with its own counter. The observation and the cooldown comparison both go through `read_db`, added alongside `writer_db` and backed by a `read_target` property on the routing wrapper that `__getattr__` now dispatches through so the two cannot drift. The observation carries the wrapper, not just its generation. `read_db` resolves to the reader while it is available and to the writer once it is not, and those counters are independent and both start at zero, so comparing a bare number across that switch pits one engine's counter against another's. Equal by coincidence waives the cooldown for an engine already replaced; unequal gates a caller that needs the recreate. Identity settles it, and is sound because the engine object is never re-pointed without the generation also moving. Three smaller holes on the way out. The waiver is withdrawn once a repair of that same engine has been tried and failed, so a burst collapses onto one attempt instead of each caller running its own recreate serially; the record is keyed per engine rather than counted globally, so an unrelated reconnect failure cannot suppress a stale reader's recovery and a writer failure cannot evict the reader's record. And a forced recreate that the optimistic-lock guard declines is no longer reported as a success on either the direct or the heavy path, since the routing wrapper leaves the reader untouched in that case; a decline is deliberately not counted as a failure, so the caller's own backoff still gets its waiver on the next attempt. A decline on the heavy path clears the dead-engine flag before raising. The clear after the cycle is skipped by any raise, which is right for a failure and wrong here, and the non-forced path already clears it on a decline, so this restores that policy rather than inventing one. Stranding the flag would route the next cycle back down the probe-free heavy branch, where the refreshed generation matches and the recreate kills the healthy engine a refresh just spawned, which is #29176. Clearing that flag is necessary and not sufficient. The escalation check re-arms it whenever the consecutive-failure count sits at the threshold, so a decline that left the count alone sent the very next attempt back down the same path. A decline is raised only at the generation guard, and the generation moves only after a replacement connects, so a decline is proof that a replacement succeeded and the count is reset on it. Fixes #36418
96912fd to
1af9958
Compare
|
@greptileai please review the current head 1af9958. Your escalation P1 was right and is fixed: a decline now resets the failure count. Reasoning in the description |
TLDR
Problem this solves:
How it solves it:
force_recreate, which skips the liveness probeUser Flow
Before: during a rolling upgrade, callers holding perfectly valid virtual keys get 503s from the old pods until those pods are drained
503 {"error":{"message":"Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.","type":"no_db_connection"}}After: the same request recovers on the first retry, so the rollout is invisible to callers
200with the model list200rather than a wait for the retry window to reopenRelevant issues
Fixes #36418
Linear ticket
Pre-Submission checklist
Screenshots / Proof of Fix
Live proxy, re-captured on the current head
Both legs below were captured on a live proxy backed by a real Postgres 16, one leg per tree, with the leg script asserting which side it was on (
grep -c "force_recreate is False"must be 0 on the base and 1 on the head) before it launched anything. Same sequence as the original repro, plus a second migration one second after the first so the reconnect cooldown is in force for it: mint three virtual keys, authenticate with the first so the auth query gets prepared on the pooled connection, apply a migration that changes the result shape ofv.*, authenticate with a fresh key, then do it again immediately.Before, on the base
423b791ee0:After, on the head
1af9958ae8:The control returns 200 on both sides, so the auth path itself is healthy and the 503s are the defect rather than a broken rig. Two distinct engine PIDs are killed a second apart on the head, which is the cooldown-bypass half doing its job: without it the second recreate is the one that gets skipped.
Live proxy, with a read replica configured
Same sequence again with
DATABASE_URL_READ_REPLICAset, soquery_firstis dispatched to a second Prisma engine with its own process, its own prepared-statement cache and its own generation counter. This is the configuration the reader half of the fix is about, and nothing had exercised it live. The leg refuses to report anything unless it first seesread-replica routing enabledin the proxy log, since without that check the run is indistinguishable from the single-engine one above.Before, on the base
423b791ee0:After, on the head
1af9958ae8:Two engine PIDs are killed per recovery rather than one, which is the reader being replaced alongside the writer.
What this capture does NOT show, stated plainly so a green run is not read as more than it is: it does not isolate the
read_dbversuswriter_dbhalf. That only bites when the writer's generation moves independently of the reader's, which needs a concurrent replacement this rig cannot produce, so the mutant that reverts the comparison towriter_dbwould recover here too. That half rests on the mutation checks below and on the code, not on this run.Original repro
Kept for the reader, superseded by the capture above. It predates the rebase onto
423b791ee0and the squash, so it names commits no longer in the PR:b7da09eb2cwas the old base, and21a6457990andfa7c335468are both folded into the current headRepro of the reported production sequence against a live proxy on localhost:4000 backed by local Postgres: generate two virtual keys, authenticate with the first one so the auth query gets prepared on the pooled connection, then apply a migration that changes the result shape of
v.*, then authenticate with the second (uncached) key.Before, at b7da09e:
The engine was never recreated, the retry hit the same stale plan, and the backoff decorator burned its three tries before the request 503'd. That is the reported log line for line.
After, at 21a6457:
No probe-skip line, the engine is recreated, and the single retry of the identical query succeeds.
That left the cooldown gap, still at 21a6457: a second migration one second after the first repair 503s again, and only recovers once the 15s cooldown elapses.
With the cooldown fix, at fa7c335, the same back-to-back migrations both recover on the first retry and no cooldown skip is logged for
postgres_cached_plan_error:Type
🐛 Bug Fix
Changes
_query_first_with_cached_plan_fallbackrecovers from Postgres's "cached plan must not change result type" by recreating the Prisma client, which drops both the server-side plans and the engine's client-side statement-name cache. Since #30183 the shared reconnect path probes the writer withSELECT 1first and skips the recreate when it answers, which is right for the case it was added for, an IAM token refresh that already replaced the engine, and wrong here: the connection is healthy, it is the session's prepared statements that are stale, so the probe always passes and always vetoes the recreate.Callers that only suspect a transport blip keep the probe. Only the cached-plan fallback passes
force_recreate=True.Getting past the probe is not enough on its own: both cooldown checks would still skip the recreate for 15 seconds after any earlier reconnect, which outlives the 10 second auth retry window, so a migration landing in that window kept 503ing.
force=Truewould fix that but would also let every concurrent caller of the same burst kill the engine the first one just built. The caller instead names the engine generation it observed before the query, and the cooldown is waived only while that generation is still live, so the first caller repairs the pool and the rest fall back to the normal cooldown once it has.The generation comes from
read_db, notwriter_db.query_firstis a top-level read, so withDATABASE_URL_READ_REPLICAset the routing wrapper dispatches it to the reader and the stale prepared statements are the reader's.read_dbsits next to the existingwriter_dband resolves through a newread_targetproperty on the wrapper, which__getattr__now uses for its own dispatch so there is one definition of where reads go.Review notes
Greptile scored 4/5 against
21a6457990, which was the first commit only. Its P1, that the reconnect cooldown could still suppress the recreate for longer than the auth retry window, is the defect the second commit fixed, so it is closed on the current head rather than outstandingIts P2 was right and is now fixed. The new probe-bypass test replaced the
writer_dbclass descriptor withmonkeypatch.setattr(PrismaClient, "writer_db", property(...)), which the repo bans. Theprisma_clientfixture already snapspc.dbto a mock, andwriter_dbreturnsself.dbwhen there is no read replica, so the test injects through that seam instead and no longer mutates the class. That also matches production more closely, since writer and db are the same object without a replicaIts second review, on the squashed head, raised the read-replica case and was right.
query_firstis a top-level read, soRoutingPrismaWrapper.__getattr__dispatches it to the reader, and it is the reader's prepared statements that go stale. The waiver was readingwriter_db.engine_generation, a different engine with its own counter, so an unrelated writer reconnect re-armed the cooldown while the poisoned reader kept serving. Both the snapshot and the comparison now readread_db, which is backed by aread_targetproperty on the routing wrapper that__getattr__itself dispatches through, so the routing rule and the recovery cannot drift apartIts third review, on the squashed head, then filed the converse of the failure-path fix and was right about that too. Gating the waiver on
_consecutive_reconnect_failuresmade it global, so a watchdog reconnect failing somewhere unrelated armed the gate and suppressed a poisoned reader's recovery for the length of the cooldown, which is the 503 this PR exists to remove.Those two findings look mutually exclusive and are not. They share a premise, that the gate has to be a failure COUNT, and the premise is wrong. What a burst needs is not "has anything failed recently" but "has a repair of this same engine already been tried and failed", so the record is the engine itself. A caller that names none, which is every watchdog and transport-error caller, writes nothing and leaves a stale reader's waiver intact, while a burst still collapses because the second caller names the engine the first failed to repair. The record is keyed per engine rather than held in one slot, so a writer failure cannot evict the reader's entry and hand the waiver back to a caller whose engine is still broken.
A fourth round found two more, both introduced by earlier rounds of this work rather than by the original commits.
The engine had to be identified, not just counted.
read_dbresolves to the reader while it is available and to the writer once it is not, and those are independent counters that both start at zero and advance on the same reconnect cadence, so comparing a bare generation across that switch pits one engine's counter against the other's. Equal by coincidence waives the cooldown for an engine already replaced, which is the racing-recreates case the singleflight exists to prevent; unequal gates a caller that genuinely needs the recreate, which is the original bug._StaleReadEnginecarries the wrapper with the number and requires both to match. Identity is sound here because the engine object is never re-pointed without the generation also moving:_original_prismais bound at prisma_client.py:204, :604 and :648 only, and the latter two are each followed by a bump, while the routing wrapper's_writerand_readerandPrismaClient.dbare all bound once at construction. Holding the wrapper rather than itsid()is deliberate, since a freed and reallocated object could alias a stored id.The heavy path had no decline check. A forced caller reaches it because the escalation threshold flips
_engine_confirmed_deadafter repeated failures and every cycle after that takes the dead-engine branch, where the recreate's return value was still discarded. Neither pre-existing heavy-path test passesforce_recreate, so nothing covered it.On the decline itself, an earlier version of this section overstated the mechanism and is corrected. The retry happens either way, because the fallback discards the reconnect's return value. What a reported success actually does is reset the consecutive-failure count and log a repair that never happened, and that reset is the part worth avoiding: with failures already accumulated it would re-enable a waiver that should stay withdrawn.
Fifteen mutants, applied one at a time, no survivors. Each fails only the tests named:
..._force_recreate_skips_probe_and_recreates,..._declined_by_the_generation_guard_is_not_reported_as_success_cooldown_appliesalways gates..._bypasses_cooldown_for_still_live_stale_engine,..._reads_generation_from_the_reader_that_served_the_query_cooldown_appliesnever gates a named engine..._honors_cooldown_once_stale_engine_replaced,..._honors_cooldown_once_the_reader_itself_was_replacedforce_recreate..._fallback_reconnects_then_retries_identical_query..._reads_generation_from_the_reader_that_served_the_query..._fallback_reports_the_reader_generation..._withdraws_the_waiver_after_this_generation_failed_to_repair,..._not_evicted_by_a_failure_on_the_other..._declined_by_the_generation_guard_is_not_reported_as_success..._heavy_path_forced_recreate_declined_is_not_reported_as_success..._declined_by_the_generation_guard_is_not_reported_as_success..._keeps_the_waiver_after_an_unrelated_reconnect_failure..._withdraws_the_waiver_after_this_generation_failed_to_repair,..._not_evicted_by_a_failure_on_the_other..._failed_repair_of_one_engine_is_not_evicted_by_a_failure_on_the_other..._gates_when_reads_moved_to_an_engine_of_the_same_generationread_targetignores reader unavailabilitytest_reads_route_to_writer_when_reader_unavailableTwo of those mutants survived an earlier sweep and are worth naming, because both were fixtures chosen to be legible rather than discriminating. Observing the writer generation survived while the routing tests drove
attempt_db_reconnectdirectly and never reached the observation site. Dropping identity survived while every routing fixture gave the reader and writer far-apart generations, so comparing the numbers alone still produced the right answer by arithmetic; the discriminating case is two engines whose generations deliberately coincide. One mutant was retired rather than killed: recording a failure for a caller that named no engine cannot be written once the record is keyed by the engine, since there is no key. That is the shape being better rather than the test being weaker, and the behaviour still has a test.Its fourth review scored 5/5 on the commit before this one. That commit contained a defect this one fixes, so the score was given up deliberately rather than lost. The heavy-path decline raised before the dead-engine flag was cleared, and the clear sits after the cycle's
await, so a decline stranded the flag set. The next cycle would then take the probe-free heavy branch with a now-matching generation and recreate over the healthy engine a refresh had just spawned, which is #29176. The non-forced path never had this, because a decline does not raise for it and it falls through to the clear, so the fix restores an existing policy rather than inventing one. Worth noting the 5/5 summary described the PR as handling declined recreations "consistently", which was the exact property that did not hold.The test for that path made it worse rather than catching it: it asserted the flag stays set, with a comment explaining why that was correct. That assertion restated what the code did, so it could only ever agree with it, and it would have defended the regression against anyone who tried to fix it. It now pins the requirement instead, and a mutant removing the clear fails it. The module docstring carries the general form, next to a note about fixtures whose values are chosen to read clearly rather than to discriminate, which is what let two earlier mutants survive.
Seventeen mutants now, no survivors. Two of this round's fixes are in that count: removing the dead-engine flag clear, and removing the escalation reset.
One limit on that number, stated because it is easy to over-read. It means every defect a single-line mutation can express is pinned, which is narrower than it sounds given that several of the defects found in review were two-cycle interactions no single-line mutation can produce. Concretely: the escalation test drives two attempts, and its second assertion is the only thing in the suite expressing "a later cycle must not reclassify a healthy replacement as dead". Both mutants that kill that test fail on its FIRST assertion and never reach the second, so what the sweep pins there is the post-decline state, while the requirement itself is asserted but not pinned. It would still catch a hand-written regression; the sweep count just cannot vouch for it.
Its fifth review found the converse of that fix, and was right again. Clearing the dead-engine flag is necessary and not sufficient: the escalation check re-arms it whenever the consecutive-failure count is still at the threshold, so a decline that left the count alone sent the very next attempt back down the probe-free heavy path and recreated over the healthy engine. This PR introduced that too, on the same mechanism as the flag, because before declines were distinguished a decline reported success and reset the count. Splitting declines out removed both halves of that bookkeeping and only one was re-established. A decline is raised only at the generation guard, and the generation moves only after a replacement has connected, so a decline is proof that a replacement succeeded and resetting the count on it is correct on its own terms. Note what it proves is that the writer was replaced, not that the caller's engine was repaired, which is exactly the distinction that matters on a read replica.
Known follow-up, not fixed here
Tracked as LIT-5611, which covers all three of the items below.
RoutingPrismaWrapper.recreate_prisma_clientreturns before touching the reader whenever the writer's optimistic-lock guard declines, so a writer-side replacement cancels a reader repair for a reason that has nothing to do with the reader. That coupling is the root of the decline handling in this PR: decoupling the reader recreate would make the declined case unreachable and delete the machinery that handles it. It is left alone deliberately, because changing it alters behaviour for every caller of that method including the IAM token refresh the guard was added for in #29176, which deserves its own change and its own proof rather than riding along here. Two further defects in the same area were found while reviewing this one and are left alone for the same reason. A recreate that raises strands the engine watcher in both branches, since the re-arm never runs and nothing else restores it, leaving a window with no engine-death detection until a later reconnect succeeds. And_recreate_prisma_client_lockedrebinds_original_prismabefore connecting, so a failed connect leaves the wrapper holding an unconnected client with the old engine already killed and the generation unbumped; its sibling_replace_prisma_client_for_token_refresh_lockedconnects into a local first and rebinds only on success, so the correct shape already exists in the same file. That second one is the more serious of the two.Two consequences are worth stating rather than leaving for a reader to find. A token-refresh storm that keeps moving the writer generation can decline repeatedly, and each decline costs one of the lookup's three backoff tries. And there is a narrow window inside a replacement, between the engine being rebound and the generation being bumped, where a caller can observe a stale-but-live engine; it self-corrects through the same decline path on the next backoff pass.
Final Attestation
Link to Devin session: https://app.devin.ai/sessions/56eeed9e4ed24192b9b239e0b46f5774