From 174cba307b8c8fbc42bed6a2cdbe7e7d2422fc29 Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:42:07 +0530 Subject: [PATCH 1/5] fix(responses): scope session history lookup to the referenced key --- .../session_handler.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 1d6a47d6365..535ad451092 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -272,7 +272,10 @@ 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 + SELECT session_id, api_key FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id AND api_key = api_key + + Only rows written under the same API key as the referenced response are returned: the session id can be + caller supplied, so matching on it alone would let a caller replay another key's history. A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result @@ -294,13 +297,18 @@ class ResponsesSessionHandler: query: Final = """ WITH matching_session AS ( - SELECT session_id + SELECT session_id, api_key FROM "LiteLLM_SpendLogs" WHERE request_id = $1 ) SELECT * - FROM "LiteLLM_SpendLogs" - WHERE session_id IN (SELECT session_id FROM matching_session) + FROM "LiteLLM_SpendLogs" AS logs + WHERE EXISTS ( + SELECT 1 + FROM matching_session + WHERE matching_session.session_id = logs.session_id + AND matching_session.api_key = logs.api_key + ) ORDER BY "endTime" ASC; """ From 1c9cffe39d840c4876086b06e1e197f47904f136 Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:42:36 +0530 Subject: [PATCH 2/5] Implement SQLite-backed DB for session tests Added a SQLite-backed database class for testing and a new test to ensure session isolation between different API keys. --- .../test_session_handler.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 002a595a5e7..ace4101719c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,4 +1,5 @@ import json +import sqlite3 from typing import Final from unittest.mock import AsyncMock, patch @@ -787,3 +788,38 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: tool_message: Final = result["messages"][2] assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments assert json.loads(tool_message["content"]) == function_output + + +class _SqliteBackedPrismaDB: + def __init__(self, rows: list[tuple[str, str, str, str]]): + self._connection = sqlite3.connect(":memory:") + self._connection.row_factory = sqlite3.Row + self._connection.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" (request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT)' + ) + self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?)', rows) + + async def query_raw(self, query, *args): + return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id(): + """A caller-supplied litellm_session_id can collide with another key's session; history must stay per key.""" + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01"), + ("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02"), + ("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03"), + ("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04"), + ] + ) + fake_prisma_client = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1") + victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1") + + assert [row["request_id"] for row in spend_logs] == ["attacker-1"] + assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"] From e90d01f7374f244357638ce47cab73bc187c6eab Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:43:18 +0530 Subject: [PATCH 3/5] fix(proxy): honor caller litellm_session_id in spend log session_id Updated docstring to clarify session ID resolution logic for spend log rows. --- litellm/proxy/spend_tracking/spend_tracking_utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 94a0f424a09..f24eca38823 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -893,9 +893,8 @@ def _get_session_id_for_spend_log( omit_when_missing: bool, batch_trace_session_id: str | None = None, ) -> str | None: - """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks - the per-request trace ids because those differ between the create call and the cost poller's row.""" + """Resolve the session id for the spend log row: `omit` honors only metadata.session_id, + batch sessions outrank everything, then the caller's litellm_session_id, then trace ids.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None @@ -904,6 +903,9 @@ def _get_session_id_for_spend_log( if batch_trace_session_id is not None: return batch_trace_session_id + caller_session_id: Final = kwargs.get("litellm_session_id") + if caller_session_id: + return str(caller_session_id) if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) if kwargs.get("litellm_trace_id") is not None: From 123b1e3887976285a07dbebaf338a5cb415f07b7 Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:44:08 +0530 Subject: [PATCH 4/5] test(proxy): caller litellm_session_id wins over the trace id --- .../test_spend_tracking_utils.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7a6b933d86e..d558c41bc53 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -237,6 +237,27 @@ def test_batch_session_outranks_the_per_request_trace_id(): assert session_id == "batch-uid-1" + +def test_caller_litellm_session_id_wins_over_the_per_request_trace_id(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_session_id": "sess-1", "litellm_trace_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + ) + assert session_id == "sess-1" + + +def test_batch_session_outranks_a_caller_litellm_session_id(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_session_id": "sess-1", "litellm_trace_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id="batch-uid-1", + ) + assert session_id == "batch-uid-1" + def test_omit_policy_still_suppresses_batch_sessions(): session_id: Final = _get_session_id_for_spend_log( kwargs={}, From 1534adeb3ab88c0ccee9d8514fd2971c1f67ce8c Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+DeepanshuPal@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:44:46 +0530 Subject: [PATCH 5/5] style(proxy): fix blank lines around new tests --- tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index d558c41bc53..7a4d06082ac 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -237,7 +237,6 @@ def test_batch_session_outranks_the_per_request_trace_id(): assert session_id == "batch-uid-1" - def test_caller_litellm_session_id_wins_over_the_per_request_trace_id(): session_id: Final = _get_session_id_for_spend_log( kwargs={"litellm_session_id": "sess-1", "litellm_trace_id": "trace-abc"}, @@ -258,6 +257,7 @@ def test_batch_session_outranks_a_caller_litellm_session_id(): ) assert session_id == "batch-uid-1" + def test_omit_policy_still_suppresses_batch_sessions(): session_id: Final = _get_session_id_for_spend_log( kwargs={},