From e12a4243b658e31684983884c1f5cd4cce2d539e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:13:28 +0000 Subject: [PATCH] fix(proxy): keep operator-pinned max_idle_connection_lifetime URL values over the default Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_cli.py | 28 +++++---- tests/test_litellm/proxy/test_proxy_cli.py | 68 +++++++++------------- 2 files changed, 44 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index eb40d7c2d0e..d42bf9715e7 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -7,8 +7,9 @@ import re import subprocess import sys import urllib.parse as urlparse -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import click @@ -64,7 +65,6 @@ def _build_db_connection_url_params( pool_timeout: float | None, connect_timeout: float | None = None, socket_timeout: float | None = None, - idle_connection_lifetime: float | None = None, disable_prepared_statements: bool = False, extra_params: dict | None = None, ) -> dict: @@ -88,8 +88,6 @@ def _build_db_connection_url_params( params["connect_timeout"] = connect_timeout if socket_timeout is not None: params["socket_timeout"] = socket_timeout - if idle_connection_lifetime is not None: - params["max_idle_connection_lifetime"] = idle_connection_lifetime if disable_prepared_statements: params["pgbouncer"] = "true" if extra_params: @@ -1085,7 +1083,9 @@ def run_server( db_connection_timeout: int | float | None = 60 db_connect_timeout: int | float | None = None db_socket_timeout: int | float | None = None - db_connection_idle_lifetime: int | float | None = ( # rebind-ok: overwritten from general_settings when a config is provided + db_connection_idle_lifetime: ( + int | float | None + ) = ( # rebind-ok: overwritten from general_settings when a config is provided LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value ) db_disable_prepared_statements: bool = False @@ -1190,7 +1190,7 @@ def run_server( db_connection_timeout = LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value db_connect_timeout = general_settings.get("database_connect_timeout") db_socket_timeout = general_settings.get("database_socket_timeout") - db_connection_idle_lifetime = general_settings.get( # rebind-ok: default declared above for the no-config path + db_connection_idle_lifetime = general_settings.get( # rebind-ok: default set for the no-config path "database_connection_idle_lifetime", LiteLLMDatabaseConnectionPool.database_connection_idle_lifetime.value, ) @@ -1261,10 +1261,18 @@ def run_server( pool_timeout=db_connection_timeout, connect_timeout=db_connect_timeout, socket_timeout=db_socket_timeout, - idle_connection_lifetime=db_connection_idle_lifetime, disable_prepared_statements=db_disable_prepared_statements, extra_params=db_extra_connection_params, ) + # The idle lifetime is applied add-if-missing so a value the operator + # pinned on the URL itself keeps winning over the built-in default. + idle_lifetime_params: Final[Mapping[str, int | float]] = MappingProxyType( + { + "max_idle_connection_lifetime": lifetime + for lifetime in (db_connection_idle_lifetime,) + if lifetime is not None + } + ) if os.getenv("DATABASE_URL", None) is not None: database_url = get_secret("DATABASE_URL", default_value=None) resolved_url: Final[str | None] = str(database_url) if database_url else None @@ -1282,11 +1290,11 @@ def run_server( writer_url, connection_url_params, ) - os.environ["DATABASE_URL"] = modified_url + os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, idle_lifetime_params) if os.getenv("DIRECT_URL", None) is not None: database_url = os.getenv("DIRECT_URL") modified_url = append_query_params(database_url, connection_url_params) - os.environ["DIRECT_URL"] = modified_url + os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, idle_lifetime_params) # The reader pool is a real pool against the same configured cap, so it # gets the allowlisted pool params. Schema-affecting ones, including any # the operator smuggled in through database_extra_connection_params, stay @@ -1303,7 +1311,7 @@ def run_server( _with_query_value(read_replica_url, "options", reader_options) if reader_options else read_replica_url, - reader_shareable_params(connection_url_params), + reader_shareable_params(MappingProxyType({**idle_lifetime_params, **connection_url_params})), ) subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 9f4ec20ce16..0ad9049b228 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -879,42 +879,25 @@ class TestProxyInitializationHelpers: assert appended_params["pgbouncer"] == "true" assert appended_params["statement_cache_size"] == 0 - def test_build_db_connection_url_params_includes_idle_lifetime(self): - from litellm.proxy.proxy_cli import _build_db_connection_url_params - - params = _build_db_connection_url_params( - connection_limit=10, - pool_timeout=60, - idle_connection_lifetime=60, - ) - assert params["max_idle_connection_lifetime"] == 60 - - def test_build_db_connection_url_params_omits_none_idle_lifetime(self): - from litellm.proxy.proxy_cli import _build_db_connection_url_params - - params = _build_db_connection_url_params( - connection_limit=10, - pool_timeout=60, - idle_connection_lifetime=None, - ) - assert "max_idle_connection_lifetime" not in params - - def test_build_db_connection_url_params_extra_overrides_idle_lifetime(self): - from litellm.proxy.proxy_cli import _build_db_connection_url_params - - params = _build_db_connection_url_params( - connection_limit=10, - pool_timeout=60, - idle_connection_lifetime=60, - extra_params={"max_idle_connection_lifetime": 300}, - ) - assert params["max_idle_connection_lifetime"] == 300 - @pytest.mark.parametrize( - "general_settings, expected_idle_lifetime", + "general_settings, database_url, expected_idle_lifetime", [ - ({}, 60), - ({"database_connection_idle_lifetime": 30}, 30), + ({}, "postgresql://test:test@localhost:5432/test", "60"), + ( + {"database_connection_idle_lifetime": 30}, + "postgresql://test:test@localhost:5432/test", + "30", + ), + ( + {"database_connection_idle_lifetime": None}, + "postgresql://test:test@localhost:5432/test", + None, + ), + ( + {}, + "postgresql://test:test@localhost:5432/test?max_idle_connection_lifetime=300", + "300", + ), ], ) @patch("subprocess.run") @@ -930,6 +913,7 @@ class TestProxyInitializationHelpers: mock_atexit_register, mock_subprocess_run, general_settings, + database_url, expected_idle_lifetime, ): from click.testing import CliRunner @@ -948,7 +932,7 @@ class TestProxyInitializationHelpers: mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( return_value={ "general_settings": { - "database_url": "postgresql://test:test@localhost:5432/test", + "database_url": database_url, **general_settings, } } @@ -972,10 +956,6 @@ class TestProxyInitializationHelpers: patch( # test-quality-ok: keeps the boot test from binding a real port, same as sibling boot tests "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" ) as mock_get_args, - patch( # test-quality-ok: capture point for the assembled URL params, same as sibling boot tests - "litellm.proxy.proxy_cli.append_query_params", - side_effect=lambda url, params: str(url), - ) as mock_append_query_params, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -991,9 +971,13 @@ class TestProxyInitializationHelpers: assert ( result.exit_code == 0 ), f"exit_code={result.exit_code}, output={result.output}" - mock_append_query_params.assert_called() - appended_params = mock_append_query_params.call_args.args[1] - assert appended_params["max_idle_connection_lifetime"] == expected_idle_lifetime + final_query = dict( + urlparse.parse_qsl(urlparse.urlparse(os.environ["DATABASE_URL"]).query) + ) + if expected_idle_lifetime is None: + assert "max_idle_connection_lifetime" not in final_query + else: + assert final_query["max_idle_connection_lifetime"] == expected_idle_lifetime def test_build_db_connection_url_params_disable_prepared_statements(self): from litellm.proxy.proxy_cli import _build_db_connection_url_params