fix(proxy): prevent Prisma engine orphan leak on shutdown disconnect failure (#26619)

litellm/proxy/proxy_cli.py

@@ -45,6 +45,7 @@ def _build_db_connection_url_params(
     connect_timeout: Optional[Union[int, float]] = None,
     socket_timeout: Optional[Union[int, float]] = None,
     extra_params: Optional[dict] = None,
+    existing_url: Optional[str] = None,
 ) -> dict:
     """Build the Prisma DATABASE_URL query params controlling connection pool behavior.
 
@@ -53,10 +54,39 @@ def _build_db_connection_url_params(
     omitted when None so Prisma's defaults apply. `extra_params` is an
     untyped passthrough — keys it provides win over the named arguments above,
     so it can be used to override any default we set here.
+
+    The four `keepalives*` defaults are libpq query-string params that get
+    forwarded to the underlying Postgres driver. Without them, dead TCP
+    sessions from killed Prisma engine subprocesses sit on the DB side for
+    the OS keepalive default (~7200s) and accumulate toward max_connections
+    under worker churn. The defaults here cause the kernel to detect dead
+    peers in ~90s instead. See BerriAI/litellm#26619.
+
+    `existing_url` is the DATABASE_URL these params will be appended to. Any
+    keepalive key already present in that URL's query string is left as-is,
+    so operators who have set explicit values (e.g. `keepalives=0`) keep
+    them. `extra_params` still wins last for both keepalives and
+    everything else.
     """
-    params: dict = {
-        "connection_limit": connection_limit,
+    keepalive_defaults: dict = {
+        "keepalives": "1",
+        "keepalives_idle": "60",
+        "keepalives_interval": "10",
+        "keepalives_count": "3",
     }
+    existing_keys: set = set()
+    if existing_url:
+        try:
+            existing_keys = set(
+                urlparse.parse_qs(urlparse.urlparse(existing_url).query).keys()
+            )
+        except Exception:
+            existing_keys = set()
+
+    params: dict = {"connection_limit": connection_limit}
+    for k, v in keepalive_defaults.items():
+        if k not in existing_keys:
+            params[k] = v
     if pool_timeout is not None:
         params["pool_timeout"] = pool_timeout
     if connect_timeout is not None:
@@ -1087,24 +1117,33 @@ def run_server(  # noqa: PLR0915
             try:
                 from litellm.secret_managers.main import get_secret
 
-                connection_url_params = _build_db_connection_url_params(
-                    connection_limit=db_connection_pool_limit,
-                    pool_timeout=db_connection_timeout,
-                    connect_timeout=db_connect_timeout,
-                    socket_timeout=db_socket_timeout,
-                    extra_params=db_extra_connection_params,
-                )
                 if os.getenv("DATABASE_URL", None) is not None:
                     database_url = get_secret("DATABASE_URL", default_value=None)
+                    url_str = str(database_url) if database_url else None
                     modified_url = append_query_params(
-                        str(database_url) if database_url else None,
-                        connection_url_params,
+                        url_str,
+                        _build_db_connection_url_params(
+                            connection_limit=db_connection_pool_limit,
+                            pool_timeout=db_connection_timeout,
+                            connect_timeout=db_connect_timeout,
+                            socket_timeout=db_socket_timeout,
+                            extra_params=db_extra_connection_params,
+                            existing_url=url_str,
+                        ),
                     )
                     os.environ["DATABASE_URL"] = modified_url
                 if os.getenv("DIRECT_URL", None) is not None:
-                    database_url = os.getenv("DIRECT_URL")
+                    direct_url = os.getenv("DIRECT_URL")
                     modified_url = append_query_params(
-                        database_url, connection_url_params
+                        direct_url,
+                        _build_db_connection_url_params(
+                            connection_limit=db_connection_pool_limit,
+                            pool_timeout=db_connection_timeout,
+                            connect_timeout=db_connect_timeout,
+                            socket_timeout=db_socket_timeout,
+                            extra_params=db_extra_connection_params,
+                            existing_url=direct_url,
+                        ),
                     )
                     os.environ["DIRECT_URL"] = modified_url
                 subprocess.run(["prisma"], capture_output=True)

litellm/proxy/proxy_server.py

@@ -966,7 +966,17 @@ async def _run_pw_migration():
         except Exception as e:
             verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}")
 
-    await proxy_shutdown_event()  # type: ignore[reportGeneralTypeIssues]
+    # Shutdown disconnect failures must not abort the uvicorn lifespan — if the
+    # lifespan exits with an exception, uvicorn skips the Python-level
+    # `atexit`/SIGTERM cleanup and the Prisma query-engine subprocess is left
+    # as an orphan whose Postgres connections sit on the DB side for the OS
+    # keepalive default (~2h). See BerriAI/litellm#26619.
+    try:
+        await proxy_shutdown_event()  # type: ignore[reportGeneralTypeIssues]
+    except Exception:
+        verbose_proxy_logger.exception(
+            "proxy_shutdown_event failed (continuing shutdown to avoid orphan subprocesses)"
+        )
 
 
 def _generate_stable_operation_id(route: Any) -> str:

litellm/proxy/utils.py

@@ -4147,14 +4147,6 @@ async def connect(self):
             )
             raise e
 
-    # Define a retrying strategy with exponential backoff
-    @backoff.on_exception(
-        backoff.expo,
-        Exception,  # base exception to catch for the backoff
-        max_tries=3,  # maximum number of retries
-        max_time=10,  # maximum total time to retry for
-        on_backoff=on_backoff,  # specifying the function to call on backoff
-    )
     async def disconnect(self):
         start_time = time.time()
         try:
@@ -4175,7 +4167,12 @@ async def disconnect(self):
                     traceback_str=error_traceback,
                 )
             )
-            raise e
+            # Do NOT re-raise: propagating from disconnect() aborts the
+            # uvicorn lifespan shutdown and leaves the Prisma query-engine
+            # subprocess as an orphan with leaked Postgres connections.
+            # disconnect() is only called from shutdown paths; no caller
+            # benefits from the re-raise, and the lifespan is actively
+            # harmed by it. See BerriAI/litellm#26619.
 
     def _get_engine_pid(self) -> int:
         try:

tests/proxy_unit_tests/test_aproxy_startup.py

@@ -13,6 +13,7 @@
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import pytest, logging, asyncio
+from unittest.mock import AsyncMock
 import litellm
 from litellm.proxy.proxy_server import (
     router,
@@ -94,3 +95,50 @@ async def test_proxy_gunicorn_startup_config_dict():
 
 
 # test_proxy_gunicorn_startup()
+
+
+@pytest.mark.asyncio
+async def test_proxy_lifespan_swallows_shutdown_event_exception(monkeypatch, caplog):
+    """
+    Regression test for BerriAI/litellm#26619.
+
+    `proxy_shutdown_event()` is invoked inside the uvicorn lifespan
+    `proxy_startup_event` context manager. If it raises and the call site
+    has no try/except, the lifespan exits with an exception, uvicorn skips
+    the Python-level `atexit`/SIGTERM cleanup, and the Prisma query-engine
+    subprocess is left as an orphan whose Postgres connections sit on the
+    DB side for ~2h. This test pins the wrapped-in-try/except behavior.
+    """
+    from litellm._logging import verbose_proxy_logger
+
+    setattr(litellm.proxy.proxy_server, "prisma_client", None)
+    database_url = os.environ.pop("DATABASE_URL", None)
+
+    filepath = os.path.dirname(os.path.abspath(__file__))
+    config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
+    os.environ["WORKER_CONFIG"] = config_fp
+
+    raising_shutdown = AsyncMock(
+        side_effect=RuntimeError("simulated disconnect failure")
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.proxy_shutdown_event", raising_shutdown
+    )
+
+    try:
+        verbose_proxy_logger.setLevel(logging.DEBUG)
+        caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name)
+
+        # Driving the full lifespan: if the patch is missing, RuntimeError
+        # propagates out of the __aexit__ here and the test fails.
+        async with proxy_startup_event(app=None) as _:
+            pass
+
+        raising_shutdown.assert_awaited_once()
+        # The exception should have been logged, not swallowed silently.
+        assert any(
+            "proxy_shutdown_event failed" in rec.getMessage() for rec in caplog.records
+        ), "expected proxy_shutdown_event failure to be logged"
+    finally:
+        if database_url is not None:
+            os.environ["DATABASE_URL"] = database_url

tests/test_litellm/proxy/test_proxy_cli.py

@@ -581,7 +581,55 @@ def test_build_db_connection_url_params_defaults(self):
         from litellm.proxy.proxy_cli import _build_db_connection_url_params
 
         params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60)
-        assert params == {"connection_limit": 10, "pool_timeout": 60}
+        assert params == {
+            "connection_limit": 10,
+            "pool_timeout": 60,
+            "keepalives": "1",
+            "keepalives_idle": "60",
+            "keepalives_interval": "10",
+            "keepalives_count": "3",
+        }
+
+    def test_build_db_connection_url_params_includes_keepalives(self):
+        """Default libpq keepalive params are emitted so the kernel reaps dead
+        Postgres sessions from killed Prisma engine subprocesses in ~90s
+        instead of the OS default (~7200s). See BerriAI/litellm#26619."""
+        from litellm.proxy.proxy_cli import _build_db_connection_url_params
+
+        params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60)
+        assert params["keepalives"] == "1"
+        assert params["keepalives_idle"] == "60"
+        assert params["keepalives_interval"] == "10"
+        assert params["keepalives_count"] == "3"
+
+    def test_build_db_connection_url_params_keepalives_overridable_via_extras(self):
+        """Operators can opt out of the keepalive defaults by passing
+        `extra_params` because `extra_params.update` runs last."""
+        from litellm.proxy.proxy_cli import _build_db_connection_url_params
+
+        params = _build_db_connection_url_params(
+            connection_limit=10,
+            pool_timeout=60,
+            extra_params={"keepalives": "0"},
+        )
+        assert params["keepalives"] == "0"
+
+    def test_build_db_connection_url_params_respects_existing_url_keepalives(self):
+        """If the DATABASE_URL already has a keepalive value, the default
+        for that key is not emitted — operators who set explicit
+        keepalive values in their URL keep them."""
+        from litellm.proxy.proxy_cli import _build_db_connection_url_params
+
+        params = _build_db_connection_url_params(
+            connection_limit=10,
+            pool_timeout=60,
+            existing_url=("postgresql://u:p@h/db?keepalives=0&keepalives_idle=120"),
+        )
+        assert "keepalives" not in params
+        assert "keepalives_idle" not in params
+        # Keys NOT in the URL still get the default.
+        assert params["keepalives_interval"] == "10"
+        assert params["keepalives_count"] == "3"
 
     def test_build_db_connection_url_params_omits_none_timeouts(self):
         from litellm.proxy.proxy_cli import _build_db_connection_url_params

tests/test_litellm/proxy/test_proxy_server.py

@@ -906,6 +906,90 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
         assert master_key == test_resolved_key
 
 
+@pytest.mark.asyncio
+async def test_proxy_lifespan_runs_shutdown_event_on_clean_exit(monkeypatch):
+    """
+    Happy-path coverage for the try/except around `await proxy_shutdown_event()`
+    inside the uvicorn lifespan (BerriAI/litellm#26619). With shutdown
+    mocked to a clean no-op, the lifespan exits without entering the
+    except branch.
+    """
+    import litellm
+    from litellm.proxy.proxy_server import proxy_startup_event
+
+    setattr(litellm.proxy.proxy_server, "prisma_client", None)
+    database_url = os.environ.pop("DATABASE_URL", None)
+
+    filepath = os.path.dirname(os.path.abspath(__file__))
+    # Use the existing repo test config that other proxy_startup_event tests use.
+    config_fp = os.path.normpath(
+        os.path.join(
+            filepath,
+            "..",
+            "..",
+            "proxy_unit_tests",
+            "test_configs",
+            "test_config_no_auth.yaml",
+        )
+    )
+    os.environ["WORKER_CONFIG"] = config_fp
+
+    clean_shutdown = AsyncMock(return_value=None)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.proxy_shutdown_event", clean_shutdown
+    )
+
+    try:
+        async with proxy_startup_event(app=None) as _:
+            pass
+        clean_shutdown.assert_awaited_once()
+    finally:
+        if database_url is not None:
+            os.environ["DATABASE_URL"] = database_url
+
+
+@pytest.mark.asyncio
+async def test_prisma_disconnect_does_not_reraise_on_failure():
+    """
+    Regression test for BerriAI/litellm#26619.
+
+    `PrismaClient.disconnect()` is only ever called from shutdown paths.
+    If it re-raises on failure, the uvicorn lifespan exits with an
+    exception, uvicorn skips the Python-level `atexit`/SIGTERM cleanup,
+    and the Prisma query-engine subprocess is left as an orphan whose
+    Postgres connections sit on the DB side for ~2h. This test pins the
+    "swallow + log + schedule failure_handler" behavior.
+    """
+    from litellm.proxy.utils import PrismaClient
+
+    # Bypass __init__ (which imports prisma binaries) — we're only testing
+    # the disconnect() failure path, which needs `self.db` and
+    # `self.proxy_logging_obj` and nothing else.
+    client = PrismaClient.__new__(PrismaClient)
+
+    boom = AsyncMock(side_effect=RuntimeError("simulated engine disconnect failure"))
+    fake_db = MagicMock()
+    fake_db.disconnect = boom
+    client.db = fake_db
+
+    failure_handler = AsyncMock()
+    fake_proxy_logging = MagicMock()
+    fake_proxy_logging.failure_handler = failure_handler
+    client.proxy_logging_obj = fake_proxy_logging
+
+    # If the fix regresses, this call propagates RuntimeError and the test fails.
+    await client.disconnect()
+
+    boom.assert_awaited()
+    # The failure_handler is scheduled via asyncio.create_task — yield once
+    # to let the scheduled coroutine run before asserting.
+    await asyncio.sleep(0)
+    failure_handler.assert_awaited_once()
+    call_kwargs = failure_handler.await_args.kwargs
+    assert call_kwargs["call_type"] == "disconnect"
+    assert isinstance(call_kwargs["original_exception"], RuntimeError)
+
+
 def test_team_info_masking():
     """
     Test that sensitive team information is properly masked
@@ -5859,15 +5943,16 @@ async def fresh_lock(_counter_key):
         if call.kwargs.get("nx") is True
     ]
     assert len(nx_writes) == 2
-    assert sorted(set_results) == [False, True], (
-        f"expected exactly one SET NX winner and one loser, got {set_results}"
-    )
+    assert sorted(set_results) == [
+        False,
+        True,
+    ], f"expected exactly one SET NX winner and one loser, got {set_results}"
     # Loser path executed: after the winner's SET NX returned True, the
     # losing coalesced() call falls back to async_get_cache to read the
     # winner's value rather than re-seeding.
-    assert get_after_set_count >= 1, (
-        "loser branch (else: read back winner's value) was never exercised"
-    )
+    assert (
+        get_after_set_count >= 1
+    ), "loser branch (else: read back winner's value) was never exercised"
 
 
 @pytest.mark.asyncio