fix(proxy): self-heal Prisma connection for auth and runtime - #21706
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR adds Prisma DB self-healing to address an outage pattern where DB connectivity drops and recovers, but the Prisma client remains in a degraded state until process restart. The implementation includes three key components:
All new functionality is well-tested with mock-only unit tests covering success/failure paths, cooldown behavior, lock timeout races, timeout budgets, and watchdog lifecycle. E2E proof is provided in the PR description. Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/db/exception_handler.py | Narrows is_database_connection_error to only match actual connection-related Prisma errors (keyword-based matching for generic PrismaError, plus explicit checks for ClientNotConnectedError and HTTPClientClosedError). This is a correctness improvement - previously all PrismaError subtypes (including DataError, UniqueViolationError, etc.) were incorrectly classified as connection errors. |
| litellm/proxy/utils.py | Adds Prisma self-healing infrastructure to PrismaClient: reconnect primitive with cooldown + singleflight lock, background DB health watchdog loop, and extensive configuration via env vars. Lock acquisition uses asyncio.wait with defensive race-condition cleanup for Python 3.9-3.11 compatibility. All reconnect paths have explicit timeout bounds. |
| litellm/proxy/auth/auth_checks.py | Adds _fetch_key_object_from_db_with_reconnect helper that wraps the DB lookup in get_key_object with a single reconnect-on-connection-error retry. Uses short lock/reconnect timeouts (0.1s/2.0s) to minimize auth path latency impact. |
| litellm/proxy/proxy_server.py | Adds watchdog lifecycle hooks: starts the health watchdog task after Prisma client initialization and stops it during shutdown. Changes are minimal and defensive (hasattr guards). Import sorting is cosmetic. |
| tests/test_litellm/proxy/auth/test_auth_checks.py | Adds two mock-only tests for the auth reconnect path: one verifying successful reconnect-and-retry, one verifying error propagation when reconnect fails. No real network calls. |
| tests/test_litellm/proxy/db/test_exception_handler.py | Updates tests to reflect the narrowed is_database_connection_error behavior: connection-specific Prisma errors are now tested separately from non-connection Prisma errors, with the latter now expected to return False. |
| tests/test_litellm/proxy/db/test_prisma_self_heal.py | New comprehensive test file covering reconnect success, cooldown skip, lock timeout, lock-leak race condition, cooldown timestamp placement, direct DB ops usage, timeout budgets, watchdog reconnect triggers (DB error and probe timeout), and start/stop lifecycle. All mock-only. |
Sequence Diagram
sequenceDiagram
participant Client
participant AuthChecks as auth_checks.py
participant PrismaClient as PrismaClient
participant Watchdog as Health Watchdog
participant DB as Prisma DB
Note over Watchdog,DB: Background health probe (every 30s)
loop Every _db_health_watchdog_interval_seconds
Watchdog->>DB: query_raw("SELECT 1")
alt DB healthy
DB-->>Watchdog: OK
else DB connection error or timeout
Watchdog->>PrismaClient: attempt_db_reconnect(reason=watchdog)
PrismaClient->>DB: disconnect()
PrismaClient->>DB: connect()
PrismaClient->>DB: query_raw("SELECT 1")
DB-->>PrismaClient: OK / error
end
end
Note over Client,DB: Auth request path
Client->>AuthChecks: get_key_object(hashed_token)
AuthChecks->>PrismaClient: get_data(token, combined_view)
alt DB lookup succeeds
PrismaClient-->>AuthChecks: key object
else DB connection error
PrismaClient-->>AuthChecks: connection error
AuthChecks->>PrismaClient: attempt_db_reconnect(reason=auth, timeout=2s, lock_timeout=0.1s)
alt Reconnect succeeds
PrismaClient-->>AuthChecks: True
AuthChecks->>PrismaClient: get_data(token, combined_view) [retry]
PrismaClient-->>AuthChecks: key object
else Reconnect fails or skipped (cooldown/lock)
PrismaClient-->>AuthChecks: False
AuthChecks-->>Client: raise original error
end
end
Last reviewed commit: 8b8c58a
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| try: | ||
| await self.connect() | ||
| await self.health_check() | ||
| verbose_proxy_logger.info( | ||
| "Prisma DB reconnect succeeded. reason=%s", reason | ||
| ) | ||
| return True | ||
| except Exception as reconnect_err: | ||
| verbose_proxy_logger.error( | ||
| "Prisma DB reconnect failed. reason=%s error=%s", | ||
| reason, | ||
| reconnect_err, | ||
| ) | ||
| return False |
There was a problem hiding this comment.
Reconnect on auth path can block ~30s
attempt_db_reconnect calls disconnect(), connect(), and health_check() — each decorated with @backoff.on_exception(max_tries=3, max_time=10). In the worst case (DB still unreachable), a single auth request triggering this path could block for up to ~30 seconds before returning a failure. Meanwhile, other concurrent auth requests will queue on _db_reconnect_lock.
Consider either:
- Using shorter timeouts/fewer retries when called from the auth path (e.g. a
timeoutparameter), or - Letting only the watchdog perform full reconnects and having the auth path just surface the error immediately (relying on the watchdog to heal in the background).
Context Used: Rule from dashboard - What: Avoid creating new database requests or Router objects in the critical request path.
Why: Cre... (source)
| ) | ||
| return False | ||
|
|
||
| self._db_last_reconnect_attempt_ts = now |
There was a problem hiding this comment.
Cooldown starts before reconnect completes
_db_last_reconnect_attempt_ts is set at the start of the reconnect attempt. If the reconnect takes, say, 20 seconds and fails, the 15-second default cooldown has already elapsed by the time the failure returns — meaning the next request will immediately trigger another reconnect attempt. Consider updating the timestamp after the reconnect attempt completes (on both success and failure paths) to ensure the full cooldown elapses between attempts.
| self._db_last_reconnect_attempt_ts = now | |
| self._db_last_reconnect_attempt_ts = now # will be updated again after attempt |
| proxy_logging_obj=proxy_logging_obj, | ||
| ) | ||
| except Exception as e: | ||
| from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler |
There was a problem hiding this comment.
Inline import in hot auth path
The PrismaDBExceptionHandler import is placed inside the except block. Per the project's code style guidelines, imports should be at the top of the file. There's no circular import risk here since litellm.proxy.db.exception_handler doesn't import from auth_checks. Moving this to the top-level imports would be consistent with how auth_exception_handler.py imports the same class.
| from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler | |
| from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler # TODO: move to top-level imports |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
Addressed review feedback:\n- Removed inline import in \n- Added module-level import\n- Extracted DB reconnect/retry logic into helper \n- Re-ran reconnect auth tests: |
|
Addressed review feedback:
|
|
Addressed Greptile feedback on this PR:
Validation run:
|
|
Validated the Greptile 3/5 items against current PR head ():\n\n1. Auth-path reconnect latency: addressed via bounded auth reconnect timeout () + fast reconnect cycle path () with default .\n2. Cooldown timing: addressed by updating in after reconnect attempt completes.\n3. Inline import in auth hot path: addressed with top-level import and helper extraction ().\n\nRe-ran targeted tests on current branch:\n- \n\nCommand:\n....... [100%] -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html |
|
Validated the Greptile 3/5 items against current PR head (
Re-ran targeted tests on current branch:
Command: Greptile summary appears based on older reviewed commit ( |
|
@greptile-apps please re-review this PR on latest head The previously flagged points (auth reconnect latency bound, cooldown timestamp timing, and inline import in auth path) have been addressed and validated with targeted tests. |
| if timeout_seconds is None: | ||
| try: | ||
| await self.disconnect() | ||
| except Exception as disconnect_err: | ||
| verbose_proxy_logger.debug( | ||
| "Prisma DB disconnect before reconnect failed (ignored): %s", | ||
| disconnect_err, | ||
| ) | ||
| await self.connect() | ||
| await self.health_check() | ||
| return |
There was a problem hiding this comment.
Watchdog reconnect may silently skip connect() after failed disconnect()
When timeout_seconds is None (the watchdog path), if self.disconnect() raises an exception (caught and ignored on line 3561-3565), the Prisma client's internal is_connected() flag may still be True. Since self.connect() at line 3566 guards with if self.db.is_connected() is False, it would silently skip the actual reconnection attempt. The subsequent health_check() would then run against the same stale connection.
The timeout path (lines 3570-3579) correctly bypasses this by calling self.db.disconnect() / self.db.connect() directly. Consider applying the same approach for the watchdog path — calling self.db.disconnect() and self.db.connect() directly rather than going through the backoff-decorated wrappers that have the is_connected() guard.
| await asyncio.wait_for(self.db.connect(), timeout=timeout_seconds) | ||
| await asyncio.wait_for(self.db.query_raw("SELECT 1"), timeout=timeout_seconds) |
There was a problem hiding this comment.
Auth reconnect timeout applies per-step, actual max is 2x
timeout_seconds is applied to both self.db.connect() and self.db.query_raw("SELECT 1") independently. With the default auth reconnect timeout of 2.0s, the total reconnect cycle can block the auth request for up to 4 seconds (2s connect + 2s query), not 2s as the config name suggests. Consider either documenting this or wrapping both steps in a single overall timeout.
|
@greptile-apps please re-review latest head Addressed the two latest comments on
Validation: |
| if timeout_seconds is None: | ||
| await _do_direct_reconnect() | ||
| return |
There was a problem hiding this comment.
Watchdog reconnect path has no timeout bound
When timeout_seconds is None (the watchdog path), _do_direct_reconnect() is called without any asyncio.wait_for wrapper. If self.db.connect() or self.db.query_raw("SELECT 1") hangs (e.g., TCP-level timeout waiting for an unreachable host), attempt_db_reconnect will hold _db_reconnect_lock indefinitely. While the lock is held, auth-path reconnect attempts will also block waiting on it, effectively stalling auth lookups.
Consider applying a generous upper-bound timeout for the watchdog path as well, e.g.:
| if timeout_seconds is None: | |
| await _do_direct_reconnect() | |
| return | |
| if timeout_seconds is None: | |
| await asyncio.wait_for(_do_direct_reconnect(), timeout=30.0) | |
| return |
Context Used: Rule from dashboard - What: Avoid creating new database requests or Router objects in the critical request path.
Why: Cre... (source)
| async def _db_health_watchdog_loop(self) -> None: | ||
| from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler | ||
|
|
||
| while True: | ||
| try: | ||
| await asyncio.sleep(self._db_health_watchdog_interval_seconds) | ||
| await self.health_check() | ||
| except asyncio.CancelledError: | ||
| break | ||
| except Exception as e: | ||
| if PrismaDBExceptionHandler.is_database_connection_error(e): | ||
| await self.attempt_db_reconnect( | ||
| reason="db_health_watchdog_connection_error" | ||
| ) | ||
| else: | ||
| verbose_proxy_logger.debug( | ||
| "Prisma DB health watchdog observed non-DB error: %s", e | ||
| ) |
There was a problem hiding this comment.
Watchdog health_check uses backoff retries before reaching reconnect
health_check() is decorated with @backoff.on_exception(max_tries=3, max_time=10), so each watchdog iteration will retry the health check with exponential backoff for up to 10 seconds before the exception propagates here. Combined with the unbounded attempt_db_reconnect call, a single watchdog cycle when the DB is down could take 10s (health_check retries) + unbounded time (reconnect with no timeout) before the loop can sleep and try again.
Consider calling self.db.query_raw("SELECT 1") directly instead of self.health_check() to avoid the backoff retries in the watchdog loop, since the watchdog already has its own retry-via-sleep cadence.
|
@greptile-apps review |
|
Superseded by the next comment with full details (formatting issue in this comment body). |
|
@greptile-apps please re-review latest head Addressed the remaining 4/5 summary deductions:
Validation:
|
| try: | ||
| await asyncio.wait_for( | ||
| self._db_reconnect_lock.acquire(), | ||
| timeout=lock_timeout_seconds, | ||
| ) |
There was a problem hiding this comment.
wait_for(lock.acquire()) can leak the lock on timeout
On Python 3.9–3.11, asyncio.wait_for(lock.acquire(), timeout=...) has a known race condition: if the lock becomes available at the exact moment the timeout fires, the lock is acquired but TimeoutError is still raised. Since the lock is never released, all subsequent callers will deadlock on it permanently.
This was fixed in CPython 3.12 (bpo-98048), but since litellm supports Python ≥3.9, this path is vulnerable.
A safer pattern is to use asyncio.wait_for with a wrapper that releases on cancellation, or use a manual approach:
| try: | |
| await asyncio.wait_for( | |
| self._db_reconnect_lock.acquire(), | |
| timeout=lock_timeout_seconds, | |
| ) | |
| result = asyncio.get_event_loop().create_future() | |
| async def _acquire_with_timeout(): | |
| await self._db_reconnect_lock.acquire() | |
| if not result.done(): | |
| result.set_result(True) | |
| task = asyncio.ensure_future(_acquire_with_timeout()) | |
| try: | |
| await asyncio.wait_for(asyncio.shield(task), timeout=lock_timeout_seconds) | |
| except asyncio.TimeoutError: | |
| task.cancel() | |
| # If the lock was acquired between timeout and cancel, release it | |
| if self._db_reconnect_lock.locked(): | |
| try: | |
| self._db_reconnect_lock.release() | |
| except RuntimeError: | |
| pass |
Alternatively, a simpler approach: use a asyncio.wait_for only on Python 3.12+ and fall back to a manual timer on older versions, or just use asyncio.timeout() (3.11+).
Context Used: Rule from dashboard - What: Avoid creating new database requests or Router objects in the critical request path.
Why: Cre... (source)
|
@greptile-apps please re-review latest head Addressed latest comment about
Tests:
|
…ore broad is_database_connection_error #21706 narrowed is_database_connection_error to only match PrismaErrors with connectivity keywords. That was correct for the reconnect use case but broke the allow_requests_on_db_unavailable behavior as a side effect. Fix: add is_database_transport_error with the narrow keyword-gated check and use it in the two reconnect call sites (auth_checks key lookup retry, utils.py watchdog). Restore is_database_connection_error to its original broad behavior (any PrismaError) so allow_requests_on_db_unavailable works correctly. Update the tests added in #21706 to assert against is_database_transport_error instead of is_database_connection_error.
…se_transport_error for reconnect Any PrismaError should be treated as a DB connection error for the allow_requests_on_db_unavailable feature and 503 responses. The narrow keyword-based check is now in is_database_transport_error, which is what the reconnect logic in auth_checks.py should use. Fixes test_delete_access_group_503_on_db_connection_error and test_handle_authentication_error_db_unavailable failures caused by PR #21706 narrowing is_database_connection_error.
…se_transport_error for reconnect Any PrismaError should be treated as a DB connection error for the allow_requests_on_db_unavailable feature and 503 responses. The narrow keyword-based check is now in is_database_transport_error, which is what the reconnect logic in auth_checks.py should use. Fixes test_delete_access_group_503_on_db_connection_error and test_handle_authentication_error_db_unavailable failures caused by PR #21706 narrowing is_database_connection_error.
…se_transport_error for reconnect (#21796) Any PrismaError should be treated as a DB connection error for the allow_requests_on_db_unavailable feature and 503 responses. The narrow keyword-based check is now in is_database_transport_error, which is what the reconnect logic in auth_checks.py should use. Fixes test_delete_access_group_503_on_db_connection_error and test_handle_authentication_error_db_unavailable failures caused by PR #21706 narrowing is_database_connection_error.
…#21706) * fix(proxy): add prisma reconnect primitive and db watchdog * fix(proxy): start and stop prisma watchdog in lifecycle * fix(auth): retry key lookup once after prisma reconnect * test(proxy): add prisma self-heal watchdog coverage * test(auth): cover reconnect-once behavior for key lookup * refactor(auth): extract db reconnect helper and remove inline import * fix(proxy): apply reconnect cooldown after attempt and add auth timeout path * fix(auth): bound reconnect latency on key lookup path * test(auth): assert reconnect timeout argument in key lookup * test(proxy): verify reconnect cooldown timestamp set after attempt * fix(proxy): harden prisma reconnect cycle semantics * test(proxy): cover watchdog reconnect + timeout budget * fix(proxy): bound watchdog probe and reconnect paths * test(proxy): cover watchdog timeout and probe behavior * fix(proxy): narrow prisma db connection error classification * fix(proxy): add auth reconnect lock timeout budget * fix(auth): pass lock timeout for db reconnect retries * test(proxy): cover narrow prisma connection error detection * test(proxy): add reconnect lock-timeout behavior coverage * test(auth): assert reconnect lock timeout argument * fix(proxy): avoid lock leak race in reconnect lock timeout path * test(proxy): cover reconnect lock-timeout race cleanup
…se_transport_error for reconnect (BerriAI#21796) Any PrismaError should be treated as a DB connection error for the allow_requests_on_db_unavailable feature and 503 responses. The narrow keyword-based check is now in is_database_transport_error, which is what the reconnect logic in auth_checks.py should use. Fixes test_delete_access_group_503_on_db_connection_error and test_handle_authentication_error_db_unavailable failures caused by PR BerriAI#21706 narrowing is_database_connection_error.
1 participant
Summary
This PR adds Prisma self-healing for the outage pattern where DB connectivity drops, network recovers, but LiteLLM continues failing auth lookups until process restart.
Changes
PrismaClientget_key_objectfor DB-connection-class errorsWhy
During transient network failures, Prisma can remain in a degraded connection state. Before this, there was no proactive runtime reconnect loop, and auth DB lookups could continue failing after network recovery.
E2E: Live Proxy + Postgres
GET /key/info200GET /key/info500GET /key/info20014sGET /v1/models200GET /v1/models401GET /v1/models20015sGET /health/readiness200db=connectedin all sampled statesLive log proof captured on this branch
Source:
.context/e2e_compare/after_pr_proof_lines.txt