From bbb407b65703310692986c14c0f739a526437342 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:40:11 +0000 Subject: [PATCH] fix(responses): bound spend-log session query to prevent OOM --- litellm/constants.py | 1 + .../session_handler.py | 19 ++++++- .../test_session_handler.py | 56 +++++++++++++++++++ 3 files changed, 73 insertions(+), 3 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e104c937a9b..823bfd07d64 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -33,6 +33,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5 DEFAULT_REPLICATE_POLLING_RETRIES = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) +DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION = int(os.getenv("DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION", 1000)) # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 68637ea97b3..af341ed9c09 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -258,7 +258,12 @@ class ResponsesSessionHandler: SQL query SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + + The result set is capped at ``DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION`` most + recent rows (returned in ascending ``endTime`` order) so a large session cannot load + the entire "LiteLLM_SpendLogs" table into memory and OOM the Prisma query engine. """ + from litellm.constants import DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION from litellm.proxy.proxy_server import prisma_client verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) @@ -273,14 +278,22 @@ class ResponsesSessionHandler: SELECT session_id FROM "LiteLLM_SpendLogs" WHERE request_id = $1 + ), + recent_logs AS ( + SELECT * + FROM "LiteLLM_SpendLogs" + WHERE session_id IN (SELECT session_id FROM matching_session) + ORDER BY "endTime" DESC + LIMIT $2 ) SELECT * - FROM "LiteLLM_SpendLogs" - WHERE session_id IN (SELECT session_id FROM matching_session) + FROM recent_logs ORDER BY "endTime" ASC; """ - spend_logs = await prisma_client.db.query_raw(query, previous_response_id) + spend_logs = await prisma_client.db.query_raw( + query, previous_response_id, DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION + ) verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 9c354101e22..9cfc8ac2e4c 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -1,6 +1,8 @@ import json import os import sys +import types +from contextlib import contextmanager from unittest.mock import AsyncMock, patch import pytest @@ -17,6 +19,19 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) +@contextmanager +def _patched_prisma_client(mock_prisma_client): + """ + Expose a stub ``litellm.proxy.proxy_server`` module carrying ``prisma_client`` so the + handler's local ``from litellm.proxy.proxy_server import prisma_client`` resolves + without importing the full (heavy, optional-dependency-laden) proxy server module. + """ + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.prisma_client = mock_prisma_client + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_module}): + yield + + @pytest.mark.asyncio async def test_get_chat_completion_message_history_for_previous_response_id(): """ @@ -435,3 +450,44 @@ async def test_get_chat_completion_message_history_empty_response_dict(): # Verify the session was still created correctly assert result["litellm_session_id"] == "test-session" + + +@pytest.mark.asyncio +async def test_get_all_spend_logs_bounds_query_with_limit(): + """ + Regression test for issue #33666: the session-reconstruction query must not run an + unbounded ``SELECT *`` over "LiteLLM_SpendLogs" (which OOM'd production pods). The + query must carry a LIMIT bound by DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION. + """ + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock(return_value=[]) + + with _patched_prisma_client(mock_prisma_client): + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "resp_previous_id" + ) + + mock_prisma_client.db.query_raw.assert_awaited_once() + call_args = mock_prisma_client.db.query_raw.await_args + executed_query = call_args.args[0] + assert "LIMIT $2" in executed_query + assert call_args.args[2] == litellm.constants.DEFAULT_MAX_SPEND_LOGS_PER_RESPONSES_SESSION + + +@pytest.mark.asyncio +async def test_get_all_spend_logs_returns_most_recent_in_ascending_order(): + """ + When a session has more rows than the cap, the newest rows must be kept (so the tail + of the conversation survives) and returned in ascending endTime order for replay. + """ + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock(return_value=[]) + + with _patched_prisma_client(mock_prisma_client): + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "resp_previous_id" + ) + + executed_query = mock_prisma_client.db.query_raw.await_args.args[0] + assert 'ORDER BY "endTime" DESC' in executed_query + assert executed_query.rstrip().rstrip(";").endswith('ORDER BY "endTime" ASC')