From 93aeabf22a30e60314cecf77559c27c4fbbe16af Mon Sep 17 00:00:00 2001 From: yryzhan-vitech Date: Thu, 23 Jul 2026 09:00:40 +0200 Subject: [PATCH 1/4] fix(spend): disconnect the per-call PrismaClient in global_spend_refresh The REFRESH MATERIALIZED VIEW branch of global_spend_refresh() builds a dedicated PrismaClient (with a long timeout, since the refresh can be slow on large spend tables) and connects it, but never disconnected it. Every /global/spend/refresh call therefore leaked a DB connection until the pool was exhausted, causing 500s on all authenticated endpoints. Wrap the refresh in try/finally so the dedicated client is always disconnected, on both the success and failure paths. The client and its long timeout are kept as-is, so refresh behavior is unchanged. Adds regression tests asserting the client is disconnected on both paths. Fixes #34269 --- .../spend_management_endpoints.py | 9 ++++ .../test_spend_management_endpoints.py | 52 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8fb5570965b..62c08eca3ae 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3126,6 +3126,7 @@ async def global_spend_refresh(): sql_query: Final = """ REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend"; """ + new_client = None try: from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import proxy_logging_obj @@ -3155,6 +3156,14 @@ async def global_spend_refresh(): "message": "Failed to refresh materialized view", "status": "failure", } + finally: + if new_client is not None: + try: + await new_client.db.disconnect() + except Exception: + verbose_proxy_logger.exception( + "Failed to disconnect MonthlyGlobalSpend refresh client" + ) async def global_spend_for_internal_user( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 057193a69db..cae48d33f8b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5299,3 +5299,55 @@ def test_scoped_spend_report_range_at_max_allowed(client, monkeypatch): mock_prisma.db.query_raw.assert_awaited_once() finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_global_spend_refresh_disconnects_client_on_success(monkeypatch): + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + global_spend_refresh, + ) + + fake_singleton = MagicMock() + fake_singleton.db = MagicMock() + fake_singleton.db.query_raw = AsyncMock( + return_value=[{"relname": "MonthlyGlobalSpend", "relkind": "m"}] + ) + monkeypatch.setattr(ps, "prisma_client", fake_singleton) + monkeypatch.setenv("DATABASE_URL", "postgresql://localhost:5432/db") + + fake_client = MagicMock() + fake_client.db = MagicMock() + fake_client.db.connect = AsyncMock() + fake_client.db.query_raw = AsyncMock(return_value=None) + fake_client.db.disconnect = AsyncMock() + + with patch("litellm.proxy.utils.PrismaClient", return_value=fake_client): + result = await global_spend_refresh() + + assert result["status"] == "success" + fake_client.db.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_global_spend_refresh_disconnects_client_on_failure(monkeypatch): + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + global_spend_refresh, + ) + + fake_singleton = MagicMock() + fake_singleton.db = MagicMock() + fake_singleton.db.query_raw = AsyncMock( + return_value=[{"relname": "MonthlyGlobalSpend", "relkind": "m"}] + ) + monkeypatch.setattr(ps, "prisma_client", fake_singleton) + monkeypatch.setenv("DATABASE_URL", "postgresql://localhost:5432/db") + + fake_client = MagicMock() + fake_client.db = MagicMock() + fake_client.db.connect = AsyncMock() + fake_client.db.query_raw = AsyncMock(side_effect=Exception("refresh timed out")) + fake_client.db.disconnect = AsyncMock() + + with patch("litellm.proxy.utils.PrismaClient", return_value=fake_client): + result = await global_spend_refresh() + + assert result["status"] == "failure" + fake_client.db.disconnect.assert_awaited_once() From 8fdb58981538743dce3b90fcbf7e650eac45fd84 Mon Sep 17 00:00:00 2001 From: yryzhan-vitech Date: Thu, 23 Jul 2026 09:36:27 +0200 Subject: [PATCH 2/4] fix(spend): use client-level disconnect for the refresh client Call PrismaClient.disconnect() (which carries the failure-handler/retry wrapper) instead of the raw db.disconnect(), so a transient cleanup failure is handled by the existing disconnect policy rather than only logged. --- .../proxy/spend_tracking/spend_management_endpoints.py | 2 +- .../spend_tracking/test_spend_management_endpoints.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 62c08eca3ae..c39dc45a2b2 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3159,7 +3159,7 @@ async def global_spend_refresh(): finally: if new_client is not None: try: - await new_client.db.disconnect() + await new_client.disconnect() except Exception: verbose_proxy_logger.exception( "Failed to disconnect MonthlyGlobalSpend refresh client" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index cae48d33f8b..b2999065843 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5317,13 +5317,13 @@ async def test_global_spend_refresh_disconnects_client_on_success(monkeypatch): fake_client.db = MagicMock() fake_client.db.connect = AsyncMock() fake_client.db.query_raw = AsyncMock(return_value=None) - fake_client.db.disconnect = AsyncMock() + fake_client.disconnect = AsyncMock() with patch("litellm.proxy.utils.PrismaClient", return_value=fake_client): result = await global_spend_refresh() assert result["status"] == "success" - fake_client.db.disconnect.assert_awaited_once() + fake_client.disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -5344,10 +5344,10 @@ async def test_global_spend_refresh_disconnects_client_on_failure(monkeypatch): fake_client.db = MagicMock() fake_client.db.connect = AsyncMock() fake_client.db.query_raw = AsyncMock(side_effect=Exception("refresh timed out")) - fake_client.db.disconnect = AsyncMock() + fake_client.disconnect = AsyncMock() with patch("litellm.proxy.utils.PrismaClient", return_value=fake_client): result = await global_spend_refresh() assert result["status"] == "failure" - fake_client.db.disconnect.assert_awaited_once() + fake_client.disconnect.assert_awaited_once() From 774df5a0ed7e8826876d351d46abf782ba4da2bf Mon Sep 17 00:00:00 2001 From: yryzhan-vitech Date: Thu, 23 Jul 2026 10:13:10 +0200 Subject: [PATCH 3/4] test(spend): cover disconnect-failure path; fix ruff format - Add a test asserting global_spend_refresh still returns success when the dedicated client's disconnect() raises (the finally block swallows and logs it), covering the cleanup-error branch codecov flagged. - Collapse the disconnect-failure log call to one line to satisfy ruff format. --- .../spend_management_endpoints.py | 4 +-- .../test_spend_management_endpoints.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c39dc45a2b2..79f02d884de 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3161,9 +3161,7 @@ async def global_spend_refresh(): try: await new_client.disconnect() except Exception: - verbose_proxy_logger.exception( - "Failed to disconnect MonthlyGlobalSpend refresh client" - ) + verbose_proxy_logger.exception("Failed to disconnect MonthlyGlobalSpend refresh client") async def global_spend_for_internal_user( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index b2999065843..e94809061b9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5351,3 +5351,30 @@ async def test_global_spend_refresh_disconnects_client_on_failure(monkeypatch): assert result["status"] == "failure" fake_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_global_spend_refresh_swallows_disconnect_error(monkeypatch): + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + global_spend_refresh, + ) + + fake_singleton = MagicMock() + fake_singleton.db = MagicMock() + fake_singleton.db.query_raw = AsyncMock( + return_value=[{"relname": "MonthlyGlobalSpend", "relkind": "m"}] + ) + monkeypatch.setattr(ps, "prisma_client", fake_singleton) + monkeypatch.setenv("DATABASE_URL", "postgresql://localhost:5432/db") + + fake_client = MagicMock() + fake_client.db = MagicMock() + fake_client.db.connect = AsyncMock() + fake_client.db.query_raw = AsyncMock(return_value=None) + fake_client.disconnect = AsyncMock(side_effect=Exception("disconnect failed")) + + with patch("litellm.proxy.utils.PrismaClient", return_value=fake_client): + result = await global_spend_refresh() + + assert result["status"] == "success" + fake_client.disconnect.assert_awaited_once() From a9b565d43175c15d1003a5c3187e34ae42532298 Mon Sep 17 00:00:00 2001 From: yryzhan-vitech Date: Thu, 13 Aug 2026 10:14:23 +0200 Subject: [PATCH 4/4] fix(spend): keep upstream Final on new_client, scope disconnect to nested finally Rebasing onto litellm_internal_staging auto-merged this function without a conflict, but the result reassigned new_client: upstream 2708620d6a5 (feat(lint): enforce Final on locals) annotated the PrismaClient construction as Final, while this branch pre-initialised new_client = None nine lines above. git could not see the collision; basedpyright reports "new_client is declared as Final and cannot be reassigned", which the type_check_gate delta gate in test-linting.yml fails on. Drop the None pre-init and move the disconnect into a nested try/finally that opens after construction. Coverage is unchanged: the db_url-None raise, the imports and any construction failure all happen before a client exists, so the old `if new_client is not None` guard only ever protected paths with nothing to disconnect, and the outer except still catches them. --- .../spend_management_endpoints.py | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 79f02d884de..4397faa94be 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3126,7 +3126,6 @@ async def global_spend_refresh(): sql_query: Final = """ REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend"; """ - new_client = None try: from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import proxy_logging_obj @@ -3142,13 +3141,19 @@ async def global_spend_refresh(): "timeout": 6000, }, ) - await new_client.db.connect() - await _query_raw(new_client, sql_query) - verbose_proxy_logger.info("MonthlyGlobalSpend view refreshed") - return { - "message": "MonthlyGlobalSpend view refreshed", - "status": "success", - } + try: + await new_client.db.connect() + await _query_raw(new_client, sql_query) + verbose_proxy_logger.info("MonthlyGlobalSpend view refreshed") + return { + "message": "MonthlyGlobalSpend view refreshed", + "status": "success", + } + finally: + try: + await new_client.disconnect() + except Exception: + verbose_proxy_logger.exception("Failed to disconnect MonthlyGlobalSpend refresh client") except Exception as e: verbose_proxy_logger.exception("Failed to refresh materialized view - %s", e) @@ -3156,12 +3161,6 @@ async def global_spend_refresh(): "message": "Failed to refresh materialized view", "status": "failure", } - finally: - if new_client is not None: - try: - await new_client.disconnect() - except Exception: - verbose_proxy_logger.exception("Failed to disconnect MonthlyGlobalSpend refresh client") async def global_spend_for_internal_user(