From 61358572b77b329b3553115939d74b3531389dfb Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:33:09 +0000 Subject: [PATCH] feat(responses): store each bridged turn's path to the root so a fork replays only its ancestors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/_types.py | 6 +- .../handler.py | 17 +- .../session_handler.py | 145 ++++++++++++++---- .../transformation.py | 27 +++- .../test_anthropic_responses_bridge.py | 2 +- .../test_handler.py | 23 +++ .../test_litellm_completion_responses.py | 38 ++++- .../test_session_handler.py | 139 +++++++++++++++-- .../mcp/test_aresponses_api_with_mcp.py | 2 +- 10 files changed, 341 insertions(+), 59 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..3ad789d9c83 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1698,6 +1698,7 @@ MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" +PREVIOUS_RESPONSE_IDS_METADATA_KEY: Final = "litellm_previous_response_ids" COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: Final = "_complexity_router_child_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b7bfc55ade..e313de102f8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2429,7 +2429,9 @@ class ScheduledJobStaggerSettings(LiteLLMPydanticObjectBase): ) -SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset({"status", "cold_storage_object_key"}) +SPEND_LOGS_METADATA_ALWAYS_KEPT_FIELDS: Final = frozenset( + {"status", "cold_storage_object_key", "litellm_previous_response_ids"} +) def _known_spend_logs_metadata_field(name: str) -> str: @@ -3771,6 +3773,8 @@ class SpendLogsMetadata(TypedDict): requester_ip_address: str | None user_agent: ReadOnly[str | None] litellm_call_id: str | None + litellm_previous_response_ids: ReadOnly[NotRequired[list[str] | None]] + """Spend log request ids of the bridged /v1/responses turns this turn continues, root first.""" applied_guardrails: list[str] | None mcp_tool_call_metadata: StandardLoggingMCPToolCall | None vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 802ed6bf9c5..90c8ccff50b 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -22,6 +22,7 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo from litellm.responses.litellm_completion_transformation.tool_search import has_client_tool_search from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, + previous_response_ids_metadata, with_compaction_items_as_params, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator @@ -225,19 +226,23 @@ class LiteLLMCompletionTransformationHandler: **kwargs, ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: previous_response_id: Final[str | None] = responses_api_request.get("previous_response_id") - if previous_response_id: - litellm_completion_request = await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) + completion_request: Final = ( + await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( previous_response_id=previous_response_id, litellm_completion_request=litellm_completion_request, instructions=responses_api_request.get("instructions"), - request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + request_kwargs=request_kwargs, previous_response_client_ids=previous_response_client_ids, stored_turns=stored_turns, input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) + if previous_response_id + else {**litellm_completion_request, **previous_response_ids_metadata((), request_kwargs)} + ) litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **{**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}, + **{**kwargs, **completion_request, "_skip_responses_api_bridge": True}, ) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( @@ -252,11 +257,11 @@ class LiteLLMCompletionTransformationHandler: elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper): return LiteLLMCompletionStreamingIterator( - model=litellm_completion_request.get("model") or "", + model=completion_request.get("model") or "", litellm_custom_stream_wrapper=litellm_completion_response, request_input=request_input, responses_api_request=responses_api_request, - custom_llm_provider=litellm_completion_request.get("custom_llm_provider"), + custom_llm_provider=completion_request.get("custom_llm_provider"), litellm_metadata=kwargs.get("litellm_metadata", {}), tool_call_provider_fields=not embeds_thought_signature_in_tool_call_id( _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7cf163e7ce8..3a272f7fa02 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -5,9 +5,15 @@ from dataclasses import dataclass from itertools import chain from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +from litellm.constants import ( + PREVIOUS_RESPONSE_IDS_METADATA_KEY, + REDACTED_BY_LITELLM, + REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, +) from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, @@ -30,6 +36,7 @@ from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, Messag if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload + from litellm.proxy.utils import PrismaClient from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, SessionMessage, @@ -136,6 +143,40 @@ def _referenced_request_ids(previous_response_id: str, client_response_ids: tupl return tuple(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) +def _stored_previous_response_ids(spend_log: Mapping[str, object]) -> tuple[str, ...] | None: + """The path to the root a turn stored on its row, None for a row written before paths were stored.""" + metadata: Final = spend_log.get("metadata") + try: + parsed: Final = ( + _STORED_METADATA_ADAPTER.validate_json(metadata) + if isinstance(metadata, str) + else _STORED_METADATA_ADAPTER.validate_python(metadata) + ) + except ValidationError: + return None + try: + return tuple(_REQUEST_IDS_ADAPTER.validate_python(parsed.get(PREVIOUS_RESPONSE_IDS_METADATA_KEY))) + except ValidationError: + return None + + +def _path_through_referenced_rows( + referenced_rows: Sequence[Mapping[str, object]], request_ids: Sequence[str] +) -> tuple[str, ...] | None: + """The request ids to replay, root first, ending at the referenced response, or None without a stored path.""" + stored_path: Final = next( + (path for path in map(_stored_previous_response_ids, referenced_rows) if path is not None), None + ) + if stored_path is None: + return None + return tuple(dict.fromkeys((*stored_path, *request_ids))) + + +def _in_path_order(spend_logs: Sequence["SpendLogsPayload"], path: Sequence[str]) -> list["SpendLogsPayload"]: + position: Final = {request_id: index for index, request_id in enumerate(path)} + return sorted(spend_logs, key=lambda spend_log: position.get(spend_log["request_id"], len(path))) + + def _turns_through_referenced_response( spend_logs: Sequence["SpendLogsPayload"], referenced_request_ids: tuple[str, ...] ) -> int: @@ -194,6 +235,55 @@ def _latest_replayable_instructions(proxy_server_requests: Sequence[Mapping[str, return None if latest == REDACTED_BY_LITELLM else latest +_STORED_METADATA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_REQUEST_IDS_ADAPTER: Final = TypeAdapter(list[str]) + +_REFERENCED_ROWS_QUERY: Final = """ + SELECT * + FROM "LiteLLM_SpendLogs" + WHERE request_id = ANY($1::text[]); +""" + +_SAME_OWNER_AS_REFERENCED: Final = """ + EXISTS ( + SELECT 1 + FROM matching_session + WHERE {same_conversation} + AND matching_session."user" IS NOT DISTINCT FROM logs."user" + AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + COALESCE(matching_session."user", '') <> '' + OR COALESCE(matching_session.team_id, '') <> '' + OR matching_session.api_key = logs.api_key + ) + ) +""" + +_OWNER_SCOPED_QUERY: Final = """ + WITH matching_session AS ( + SELECT session_id, api_key, "user", team_id + FROM "LiteLLM_SpendLogs" + WHERE request_id = ANY($1::text[]) + ) + SELECT * + FROM "LiteLLM_SpendLogs" AS logs + WHERE {same_owner} + ORDER BY "endTime" ASC; +""" + +_SESSION_QUERY: Final = _OWNER_SCOPED_QUERY.format( + same_owner=_SAME_OWNER_AS_REFERENCED.format(same_conversation="matching_session.session_id = logs.session_id") +) + +_PATH_QUERY: Final = _OWNER_SCOPED_QUERY.format( + same_owner=_SAME_OWNER_AS_REFERENCED.format(same_conversation="logs.request_id = ANY($2::text[])") +) + + +async def _read_spend_logs(prisma_client: "PrismaClient", query: str, *params: object) -> list["SpendLogsPayload"]: + return cast("list[SpendLogsPayload]", await prisma_client.db.query_raw(query, *params)) + + class ResponsesSessionHandler: @staticmethod async def get_chat_completion_message_history_for_previous_response_id( @@ -264,6 +354,7 @@ class ResponsesSessionHandler: messages=chat_completion_message_history, litellm_session_id=litellm_session_id, instructions=_latest_replayable_instructions(proxy_server_requests), + previous_response_ids=tuple(spend_log["request_id"] for spend_log in all_spend_logs[:replayable_turns]), ) @staticmethod @@ -510,8 +601,11 @@ class ResponsesSessionHandler: """ Get all spend logs for a previous response id - - SQL query + The referenced response's row is read first. A row written with a path to the root + (``litellm_previous_response_ids`` in its metadata) names every earlier turn it continues, so + only those rows plus the referenced one are read, which keeps a fork from replaying its + siblings and the turns sent after the referenced response. A row written before paths were + stored falls back to its whole session: SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then SELECT * FROM spend_logs with that session_id and the same owner @@ -539,40 +633,25 @@ class ResponsesSessionHandler: if prisma_client is None: return [] - query: Final = """ - WITH matching_session AS ( - SELECT session_id, api_key, "user", team_id - FROM "LiteLLM_SpendLogs" - WHERE request_id = ANY($1::text[]) - ) - SELECT * - FROM "LiteLLM_SpendLogs" AS logs - WHERE EXISTS ( - SELECT 1 - FROM matching_session - WHERE matching_session.session_id = logs.session_id - AND matching_session."user" IS NOT DISTINCT FROM logs."user" - AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id - AND ( - COALESCE(matching_session."user", '') <> '' - OR COALESCE(matching_session.team_id, '') <> '' - OR matching_session.api_key = logs.api_key - ) - ) - ORDER BY "endTime" ASC; - """ - max_attempts: Final = 1 if disable_spend_logs else RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS for attempt in range(max_attempts): if attempt: await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) - if spend_logs := await prisma_client.db.query_raw(query, request_ids): - verbose_proxy_logger.debug( - "Found the following spend logs for previous response id %s: %s", - response_id, - json.dumps(spend_logs, indent=4, default=str), - ) - return spend_logs + referenced_rows = await _read_spend_logs(prisma_client, _REFERENCED_ROWS_QUERY, request_ids) + if not referenced_rows: + continue + path = _path_through_referenced_rows(referenced_rows, request_ids) + spend_logs = ( + _in_path_order(await _read_spend_logs(prisma_client, _PATH_QUERY, request_ids, list(path)), path) + if path is not None + else await _read_spend_logs(prisma_client, _SESSION_QUERY, request_ids) + ) + verbose_proxy_logger.debug( + "Found the following spend logs for previous response id %s: %s", + response_id, + json.dumps(spend_logs, indent=4, default=str), + ) + return spend_logs verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) return [] diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0f4ed0611be..d25254a6642 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -42,6 +42,7 @@ from litellm._logging import verbose_logger from litellm.caching import InMemoryCache from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, + PREVIOUS_RESPONSE_IDS_METADATA_KEY, REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, ) @@ -917,6 +918,9 @@ class ChatCompletionSession(TypedDict, total=False): messages: list[SessionMessage] litellm_session_id: str | None instructions: ReadOnly[str | None] + previous_response_ids: ReadOnly[tuple[str, ...]] + """Spend log request ids of the replayed turns, root first, stored on the new turn so a follow-up or a fork + rebuilds exactly this path instead of the whole session.""" ########### End of Initialize Classes used for Responses API ########### @@ -924,10 +928,23 @@ class ChatCompletionSession(TypedDict, total=False): def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]: key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata" + return _with_metadata(key, {CALLER_SESSION_ID_METADATA_KEY: chain_session}, request_kwargs) + + +def previous_response_ids_metadata( + previous_response_ids: Sequence[str], request_kwargs: Mapping[str, object] +) -> Mapping[str, object]: + """The proxy's ``litellm_metadata`` carrying the path to the root, so the spend log row of this turn stores it.""" + return _with_metadata( + "litellm_metadata", {PREVIOUS_RESPONSE_IDS_METADATA_KEY: list(previous_response_ids)}, request_kwargs + ) + + +def _with_metadata(key: str, extra: Mapping[str, object], request_kwargs: Mapping[str, object]) -> Mapping[str, object]: metadata: Final = request_kwargs.get(key) if not isinstance(metadata, dict): return MappingProxyType({}) - return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}} + return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), **extra}} class LiteLLMCompletionResponsesConfig: @@ -1562,7 +1579,13 @@ class LiteLLMCompletionResponsesConfig: if chain_session and litellm.request_correlation_in_logs else MappingProxyType({}) ) - return {**litellm_completion_request, **chain_session_metadata} + replayed_path: Final = chat_completion_session.get("previous_response_ids") + path_metadata: Final[Mapping[str, object]] = ( + previous_response_ids_metadata(replayed_path, {**request_kwargs, **chain_session_metadata}) + if replayed_path is not None + else MappingProxyType({}) + ) + return {**litellm_completion_request, **chain_session_metadata, **path_metadata} @staticmethod def _transform_response_input_param_to_chat_completion_message( diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 5ae985ea7e5..913fed2ba65 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1595,7 +1595,7 @@ async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search [("tool_use", "toolu_standup")], [("tool_result", "toolu_standup"), ("text", None)], ] - assert prisma_client.db.queries == 1 + assert prisma_client.db.queries == 2 @pytest.mark.asyncio diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index 4794f3701c5..0ec88974875 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -22,6 +22,7 @@ import httpx import pytest import litellm +from litellm.constants import PREVIOUS_RESPONSE_IDS_METADATA_KEY from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, @@ -56,6 +57,28 @@ def test_sync_fallback_tags_skip_responses_api_bridge(): assert captured.get("_skip_responses_api_bridge") is True +@pytest.mark.asyncio +async def test_a_first_turn_stores_an_empty_path_to_the_root_so_a_fork_off_it_replays_only_it(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.acompletion", fake_acompletion), pytest.raises(_StopForwarding): + await handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=True, + litellm_metadata={"user_api_key": "hashed"}, + ) + + assert captured["litellm_metadata"] == {"user_api_key": "hashed", PREVIOUS_RESPONSE_IDS_METADATA_KEY: []} + + @pytest.mark.asyncio async def test_async_fallback_tags_skip_responses_api_bridge(): handler = LiteLLMCompletionTransformationHandler() diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b4cd606374a..5bd5d08d4fb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -19,7 +19,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field -from litellm.constants import CALLER_SESSION_ID_METADATA_KEY +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, PREVIOUS_RESPONSE_IDS_METADATA_KEY from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -7109,6 +7109,42 @@ def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_corre assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} +@pytest.mark.parametrize("request_correlation_in_logs", [True, False]) +def test_a_continued_chain_stores_the_replayed_turns_as_its_path_to_the_root( + monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool +): + """The proxy's spend log row for this turn carries the ids it replayed, so a fork off it replays only them.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs) + + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[{"role": "user", "content": "turn 1"}], + litellm_session_id="chain-session", + previous_response_ids=("chatcmpl-root", "chatcmpl-turn-1"), + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_turn_1", + litellm_completion_request={ + "model": _CLAUDE_ON_BEDROCK, + "custom_llm_provider": "bedrock", + "messages": [{"role": "user", "content": "turn 2"}], + }, + instructions=None, + session_lookup=session_lookup, + request_kwargs={"litellm_metadata": {"user_api_key": "hashed"}, "metadata": {"tag": "client"}}, + ) + ) + + assert request["litellm_metadata"] == { + "user_api_key": "hashed", + **({CALLER_SESSION_ID_METADATA_KEY: "chain-session"} if request_correlation_in_logs else {}), + PREVIOUS_RESPONSE_IDS_METADATA_KEY: ["chatcmpl-root", "chatcmpl-turn-1"], + } + assert "metadata" not in request + + @pytest.mark.asyncio async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chunks(): from litellm.responses.litellm_completion_transformation.streaming_iterator import ( 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 f1d841a8e5e..c6c7e4be68b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,8 +1,8 @@ import json import sqlite3 -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from itertools import chain -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -13,6 +13,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient import litellm +from litellm.constants import PREVIOUS_RESPONSE_IDS_METADATA_KEY from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_proxy_server_request_for_spend_logs_payload, ) @@ -445,15 +446,21 @@ def _chat_completion_response(request_id: str, content: str) -> dict: class _FakePrismaDB: + """Serves the next queued result per history read; the referenced-row read peeks at it, dropping a miss.""" + def __init__(self, results): self._results = list(results) self.calls = [] - async def query_raw(self, query, request_ids): + async def query_raw(self, query, request_ids, *path): self.calls.append(tuple(request_ids)) - if not self._results: - return [] - return list(self._results.pop(0)) + if "matching_session" in query: + return list(self._results.pop(0)) if self._results else [] + queued = self._results[0] if self._results else [] + referenced = [row for row in queued if row.get("request_id") in request_ids] + if not referenced and self._results: + self._results.pop(0) + return referenced class _FakePrismaClient: @@ -539,7 +546,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus ("assistant", "OK"), ] assert result["litellm_session_id"] == session_id - assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id] + assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id, request_id] @pytest.mark.asyncio @@ -630,8 +637,8 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id) - assert len(fake_prisma_client.db.calls) == 1 - assert request_id in fake_prisma_client.db.calls[0] + assert len(fake_prisma_client.db.calls) == 2 + assert all(request_id in call for call in fake_prisma_client.db.calls) class _StoredSpendLogs: @@ -1169,18 +1176,27 @@ async def test_a_compacted_session_replays_no_turn_sent_after_the_referenced_res class _SqliteBackedPrismaDB: - def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: + def __init__( + self, + rows: list[tuple[str, str, str, str, str | None, str | None]], + metadata: Mapping[str, Mapping[str, object]] = MappingProxyType({}), + ) -> None: 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, "user" TEXT, team_id TEXT)' + '(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT, metadata TEXT)' + ) + self._connection.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?, ?)', + [(*row, json.dumps(metadata[row[0]]) if row[0] in metadata else None) for row in rows], ) - self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace( - "$1", "?" + sqlite_query: Final = ( + query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?1))") + .replace("= ANY($2::text[])", "IN (SELECT value FROM json_each(?2))") + .replace("IS NOT DISTINCT FROM", "IS") ) sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] @@ -1216,6 +1232,101 @@ async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: st return [row["request_id"] for row in spend_logs] +def _stored_path(*previous_response_ids: str) -> Mapping[str, object]: + return {PREVIOUS_RESPONSE_IDS_METADATA_KEY: list(previous_response_ids)} + + +def _forked_session_db() -> _SqliteBackedPrismaDB: + """One session: root -> a2 -> a3 on branch a, root -> b2 -> b3 on branch b, b2 sent between a2 and a3.""" + return _SqliteBackedPrismaDB( + [ + ("root", "key", "session", "2026-01-01T00:00:01", "", ""), + ("a2", "key", "session", "2026-01-01T00:00:02", "", ""), + ("b2", "key", "session", "2026-01-01T00:00:03", "", ""), + ("a3", "key", "session", "2026-01-01T00:00:04", "", ""), + ("b3", "key", "session", "2026-01-01T00:00:05", "", ""), + ], + metadata={ + "root": _stored_path(), + "a2": _stored_path("root"), + "b2": _stored_path("root"), + "a3": _stored_path("root", "a2"), + "b3": _stored_path("root", "b2"), + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("previous_response_id", "replayed"), + [ + ("root", ["root"]), + ("a2", ["root", "a2"]), + ("b2", ["root", "b2"]), + ("a3", ["root", "a2", "a3"]), + ("b3", ["root", "b2", "b3"]), + ], +) +async def test_a_turn_with_a_stored_path_replays_only_the_path_to_the_root( + previous_response_id: str, replayed: list[str] +): + """A fork shares its session with the other branch, so the stored path, not the session, picks the turns.""" + assert await _history_request_ids(_forked_session_db(), previous_response_id) == replayed + + +@pytest.mark.asyncio +async def test_a_stored_path_is_replayed_in_path_order_whatever_the_clock_says(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("late-root", "key", "session", "2026-01-01T00:00:09", "", ""), + ("child", "key", "session", "2026-01-01T00:00:01", "", ""), + ], + metadata={"late-root": _stored_path(), "child": _stored_path("late-root")}, + ) + + assert await _history_request_ids(sqlite_db, "child") == ["late-root", "child"] + + +@pytest.mark.asyncio +async def test_a_turn_stored_before_paths_existed_still_replays_its_whole_session(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("old-1", "key", "session", "2026-01-01T00:00:01", "", ""), + ("old-2", "key", "session", "2026-01-01T00:00:02", "", ""), + ("old-3", "key", "session", "2026-01-01T00:00:03", "", ""), + ] + ) + + assert await _history_request_ids(sqlite_db, "old-2") == ["old-1", "old-2", "old-3"] + + +@pytest.mark.asyncio +async def test_a_stored_path_never_replays_another_owners_turns(): + """A path naming another key's row (a forged or reused id) replays nothing of that key's history.""" + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("victim-1", "victim-key", "session", "2026-01-01T00:00:01", "", ""), + ("attacker-1", "attacker-key", "session", "2026-01-01T00:00:02", "", ""), + ], + metadata={"victim-1": _stored_path(), "attacker-1": _stored_path("victim-1")}, + ) + + assert await _history_request_ids(sqlite_db, "attacker-1") == ["attacker-1"] + + +@pytest.mark.asyncio +async def test_the_replayed_turns_become_the_next_turns_path_to_the_root(): + fake_prisma_client: Final = _FakePrismaClient(results=[]) + fake_prisma_client.db = _forked_session_db() + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + "b2" + ) + + assert session["previous_response_ids"] == ("root", "b2") + + @pytest.mark.asyncio async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key(): sqlite_db: Final = _SqliteBackedPrismaDB( diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index ec0be896095..8645cde1ab1 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -2366,7 +2366,7 @@ async def test_bridged_auto_execute_follow_up_continues_the_client_conversation_ await _settle(lambda: len(recorder.logged) > 0) follow_up: Final = json.loads(provider.calls[1].request.content)["messages"] - assert spend_logs.asked == (frozenset({"chatcmpl-prior-turn"}),) * 2 + assert spend_logs.asked and set(spend_logs.asked) == {frozenset({"chatcmpl-prior-turn"})} assert [message["role"] for message in follow_up] == ["user", "assistant", "user", "assistant", "user"] assert json.dumps(follow_up).count("My favorite color is teal.") == 1 assert [block["type"] for block in follow_up[3]["content"]] == ["thinking", "tool_use"]