From 61358572b77b329b3553115939d74b3531389dfb Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:33:09 +0000 Subject: [PATCH 1/7] 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"] From 779d7e2f5f919b1b3f353b0094cf0e02938f1150 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:41:22 +0000 Subject: [PATCH 2/7] fix(responses): keep the spend log read-only typed so the type discipline gate holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 +- .../session_handler.py | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e313de102f8..d4d94240ae3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3773,7 +3773,7 @@ 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]] + litellm_previous_response_ids: ReadOnly[NotRequired[Sequence[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 diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 3a272f7fa02..2b732315e4c 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -172,7 +172,7 @@ def _path_through_referenced_rows( return tuple(dict.fromkeys((*stored_path, *request_ids))) -def _in_path_order(spend_logs: Sequence["SpendLogsPayload"], path: Sequence[str]) -> list["SpendLogsPayload"]: +def _in_path_order(spend_logs: Sequence["SpendLogsPayload"], path: Sequence[str]) -> Sequence["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))) @@ -280,8 +280,11 @@ _PATH_QUERY: Final = _OWNER_SCOPED_QUERY.format( ) -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)) +async def _read_spend_logs(prisma_client: "PrismaClient", query: str, *params: object) -> Sequence["SpendLogsPayload"]: + return cast( # cast-ok: query_raw is untyped; spend log rows are read as is + "Sequence[SpendLogsPayload]", + await prisma_client.db.query_raw(query, *params), + ) class ResponsesSessionHandler: @@ -597,7 +600,7 @@ class ResponsesSessionHandler: async def get_all_spend_logs_for_previous_response_id( previous_response_id: str, client_response_ids: tuple[str, ...] = (), - ) -> "list[SpendLogsPayload]": + ) -> "Sequence[SpendLogsPayload]": """ Get all spend logs for a previous response id From 6ec0fe1f7101e1ae9662de65e6f3ead624bcd660 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:42:50 +0000 Subject: [PATCH 3/7] fix(responses): type the first turn's completion request so the basedpyright gate holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/handler.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 90c8ccff50b..52333fef4bd 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -227,7 +227,7 @@ class LiteLLMCompletionTransformationHandler: ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: previous_response_id: Final[str | None] = responses_api_request.get("previous_response_id") request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) - completion_request: Final = ( + completion_request: Final[Mapping[str, object]] = ( await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( previous_response_id=previous_response_id, litellm_completion_request=litellm_completion_request, @@ -238,7 +238,10 @@ class LiteLLMCompletionTransformationHandler: 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)} + else { + **_REQUEST_KWARGS_ADAPTER.validate_python(litellm_completion_request), + **previous_response_ids_metadata((), request_kwargs), + } ) litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( @@ -257,11 +260,11 @@ class LiteLLMCompletionTransformationHandler: elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper): return LiteLLMCompletionStreamingIterator( - model=completion_request.get("model") or "", + model=litellm_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=completion_request.get("custom_llm_provider"), + custom_llm_provider=litellm_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) From ba55a860429dbf380883e60dc02839a62b7fce46 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:45:01 +0000 Subject: [PATCH 4/7] fix(responses): record every replayed turn as the path when continuing a session stored without one Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../session_handler.py | 2 +- .../test_session_handler.py | 20 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 2b732315e4c..6aab3ec5c17 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -357,7 +357,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]), + previous_response_ids=tuple(spend_log["request_id"] for spend_log in all_spend_logs[:replayed_turns]), ) @staticmethod 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 c6c7e4be68b..4fefab5f84a 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1327,6 +1327,26 @@ async def test_the_replayed_turns_become_the_next_turns_path_to_the_root(): assert session["previous_response_ids"] == ("root", "b2") +@pytest.mark.asyncio +async def test_a_turn_continuing_a_session_stored_before_paths_existed_records_every_turn_it_replayed(): + """The legacy replay feeds the whole session to the model, so the new turn's path must name all of it.""" + fake_prisma_client: Final = _FakePrismaClient( + results=[ + [ + _spend_log("old-1", "session", "My favorite number is 7.", "ok"), + _spend_log("old-2", "session", "My second number is 42.", "ok"), + ] + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + "old-1" + ) + + assert session["previous_response_ids"] == ("old-1", "old-2") + + @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( From 022d0ee29e26a96e0394a57a65ce30db07faa218 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:47:53 +0000 Subject: [PATCH 5/7] fix(responses): type the bridged completion kwargs so the basedpyright gate holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/litellm_completion_transformation/handler.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 52333fef4bd..59362e6dbe4 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -244,8 +244,13 @@ class LiteLLMCompletionTransformationHandler: } ) + completion_kwargs: Final[Mapping[str, object]] = { + **request_kwargs, + **completion_request, + "_skip_responses_api_bridge": True, + } litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **{**kwargs, **completion_request, "_skip_responses_api_bridge": True}, + **completion_kwargs ) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( From d26f3dcbb62838d24ba75646bd648ca2279f1a51 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 22:59:04 +0000 Subject: [PATCH 6/7] test(integrations): add the previous response ids key to the pub sub spend log fixture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/integrations/gcs_pubsub/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json index 6e5461c3b60..4b9dd300b71 100644 --- a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json +++ b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"billed_retry_attempt\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", + "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"billed_retry_attempt\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"litellm_previous_response_ids\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From 540f3342b96fc87e58961e16ac2343d51d2887ac Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 23:15:42 +0000 Subject: [PATCH 7/7] fix(responses): stamp the path into the request's own metadata dict so router io token refunds stay deduped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 7 +++++- .../test_spend_management_endpoints.py | 2 +- .../test_handler.py | 8 +++---- .../test_litellm_completion_responses.py | 24 +++++++++++++------ 4 files changed, 28 insertions(+), 13 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d25254a6642..65d7418298d 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -944,7 +944,12 @@ def _with_metadata(key: str, extra: Mapping[str, object], request_kwargs: Mappin metadata: Final = request_kwargs.get(key) if not isinstance(metadata, dict): return MappingProxyType({}) - return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), **extra}} + _STR_KEY_DICT_ADAPTER.validate_python(metadata) + channel: Final = cast( # cast-ok: validated on the line above, and the router dedupes io token refunds by popping sentinels from this exact dict, so a copy would refund twice + "dict[str, object]", metadata + ) + channel.update(extra) + return {key: channel} class LiteLLMCompletionResponsesConfig: diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 6598b94e60f..f043e168d43 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3814,7 +3814,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "billed_retry_attempt": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "litellm_previous_response_ids": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "billed_retry_attempt": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index 0ec88974875..2129195e880 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -59,6 +59,7 @@ def test_sync_fallback_tags_skip_responses_api_bridge(): @pytest.mark.asyncio async def test_a_first_turn_stores_an_empty_path_to_the_root_so_a_fork_off_it_replays_only_it(): + litellm_metadata: Final[dict[str, object]] = {"user_api_key": "hashed"} handler = LiteLLMCompletionTransformationHandler() captured: dict = {} @@ -73,10 +74,11 @@ async def test_a_first_turn_stores_an_empty_path_to_the_root_so_a_fork_off_it_re responses_api_request={}, custom_llm_provider="openai", _is_async=True, - litellm_metadata={"user_api_key": "hashed"}, + litellm_metadata=litellm_metadata, ) assert captured["litellm_metadata"] == {"user_api_key": "hashed", PREVIOUS_RESPONSE_IDS_METADATA_KEY: []} + assert captured["litellm_metadata"] is litellm_metadata @pytest.mark.asyncio @@ -346,9 +348,7 @@ async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_h ) assert provider.request_body is not None, "the bridged turn never reached the provider" - assert [ - (message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"] - ] == [ + assert [(message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"]] == [ ("user", json.dumps([{"type": "text", "text": "My favorite number is 7."}])), ("assistant", json.dumps([{"type": "text", "text": "Noted."}])), ("user", json.dumps([{"type": "text", "text": "Double it"}])), 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 5bd5d08d4fb..6000f6ab1ad 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 @@ -18,7 +18,6 @@ from openai.types.responses.response_function_web_search import ( 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, 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 @@ -26,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, anthropic_messages_pt, ) +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks from litellm.responses.litellm_completion_transformation.tool_search import tool_search_routing from litellm.responses.litellm_completion_transformation.transformation import ( @@ -6876,7 +6876,11 @@ def test_a_new_message_after_a_paused_search_turn_replays_only_that_turn_without [ {"role": "user", "content": "Plan a Tokyo trip."}, _reasoning_item("rs_tokyo", tokyo_thinking), - {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Tokyo is sunny."}]}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Tokyo is sunny."}], + }, _hosted_search("ws_srvtoolu_01Tokyo", "completed"), {"role": "user", "content": "Now Osaka."}, _reasoning_item("rs_osaka", _ANSWERED_THINKING), @@ -7083,8 +7087,6 @@ def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_c monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str ): monkeypatch.setattr(litellm, "request_correlation_in_logs", True) - sent_kwargs: Final = deepcopy(request_kwargs) - request: Final = _continued_chain_request(request_kwargs) assert request["litellm_trace_id"] == "chain-session" @@ -7093,7 +7095,7 @@ def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_c **TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]), CALLER_SESSION_ID_METADATA_KEY: "chain-session", } - assert request_kwargs == sent_kwargs + assert request[logged_metadata_key] is request_kwargs[logged_metadata_key] def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation( @@ -7252,7 +7254,11 @@ def test_a_tool_output_with_an_unhashable_call_id_is_a_400_naming_its_input_item ) -> None: history: Final = cast( ResponseInputParam, - [_WEATHER_ASK, _weather_call(call_id="call_weather"), {"type": output_type, "call_id": call_id, "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id="call_weather"), + {"type": output_type, "call_id": call_id, "output": "sunny"}, + ], ) with pytest.raises(litellm.BadRequestError) as raised: @@ -7356,7 +7362,11 @@ def test_a_non_string_role_still_converts_for_a_target_that_judges_it_itself() - id="dict_function_call_id", ), pytest.param( - [_WEATHER_ASK, _weather_call(call_id=7), {"type": "function_call_output", "call_id": "7", "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id=7), + {"type": "function_call_output", "call_id": "7", "output": "sunny"}, + ], "input[1].call_id", "int", id="int_function_call_id",