fix(responses): bound spend-log session query to prevent OOM

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.

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",

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')