From 70e4273ba1f50fd921c05198fdfc868291dc2d57 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:12:17 -0700 Subject: [PATCH 1/5] fix(responses): make previous_response_id resolve on the bridged path Streaming /v1/responses over the completion bridge minted a fresh resp_{uuid4} for every response, while spend tracking stored the inner chat completion id as request_id. The session lookup queries on request_id, so a follow-up sent with that response id matched no rows and the prior conversation was silently dropped. The iterator now pulls the first upstream chunk before emitting response.created, so created, in_progress and completed all carry the same encoded chat completion id. Two more ways the same history went missing: - The session lookup only read spend logs already written to the DB, so a follow-up sent inside the batch writer's window found nothing. It now also reads the rows still queued in memory. - Input was only accepted as a string or a single dict, so the list shape the Responses API actually sends dropped every user turn from the reconstructed history. --- litellm/proxy/utils.py | 10 + .../session_handler.py | 67 ++++- .../streaming_iterator.py | 82 +++++- .../test_session_handler.py | 235 ++++++++++++++++++ .../test_streaming_iterator_response_id.py | 130 ++++++++++ 5 files changed, 508 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c616d9e8723..9978fa04f40 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6005,6 +6005,16 @@ async def enqueue_spend_logs( ) +async def peek_spend_logs(prisma_client: PrismaClient) -> tuple[SpendLogsPayload, ...]: + """Snapshot the spend logs still waiting for the next flush, leaving the queue intact. + + Reads that need a just-finished request use this, since the batch writer only + reaches the DB every ``PROXY_BATCH_WRITE_AT`` seconds. + """ + async with prisma_client._spend_log_transactions_lock: + return tuple(prisma_client.spend_log_transactions) + + async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index dcff26c5b0c..935c78bc9a1 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -104,10 +105,10 @@ class ResponsesSessionHandler: if proxy_server_request_dict: _response_input_param: Final = proxy_server_request_dict.get("input", None) _messages = proxy_server_request_dict.get("messages", None) - if isinstance(_response_input_param, str): + if isinstance(_response_input_param, (str, list)): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast(ResponseInputParam, _response_input_param) + response_input_param = cast(ResponseInputParam, [_response_input_param]) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -131,8 +132,8 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = spend_log.get("response", "{}") - if isinstance(_response_output, dict) and _response_output and _response_output != {}: + _response_output: Final = ResponsesSessionHandler._get_response_dict_from_spend_log(spend_log) + if _response_output: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -140,6 +141,23 @@ class ResponsesSessionHandler: chat_completion_message_history.append(getattr(choice, "message")) return chat_completion_message_history + @staticmethod + def _get_response_dict_from_spend_log(spend_log: SpendLogsPayload) -> Mapping[str, Any] | None: + """ + Spend logs read from the DB hold `response` as a dict, ones still queued in memory + hold it as a JSON string. + """ + _response_output: Final = spend_log.get("response") + if isinstance(_response_output, dict): + return _response_output or None + if isinstance(_response_output, str): + try: + parsed: Final = json.loads(_response_output) + except json.JSONDecodeError: + return None + return parsed if isinstance(parsed, dict) and parsed else None + return None + @staticmethod async def get_proxy_server_request_from_spend_log( spend_log: SpendLogsPayload, @@ -256,11 +274,12 @@ class ResponsesSessionHandler: SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id """ from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.utils import peek_spend_logs verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) decoded_response_id: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) - previous_response_id = decoded_response_id.get("response_id", previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) if prisma_client is None: return [] @@ -276,12 +295,46 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - spend_logs: Final = await prisma_client.db.query_raw(query, previous_response_id) + written_spend_logs: Final = await prisma_client.db.query_raw(query, response_id) + queued_spend_logs: Final = await peek_spend_logs(prisma_client) + spend_logs: Final = list( + ResponsesSessionHandler._merge_queued_spend_logs( + response_id=response_id, + written_spend_logs=written_spend_logs, + queued_spend_logs=queued_spend_logs, + ) + ) verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", - previous_response_id, + response_id, json.dumps(spend_logs, indent=4, default=str), ) return spend_logs + + @staticmethod + def _merge_queued_spend_logs( + response_id: str, + written_spend_logs: Sequence[SpendLogsPayload], + queued_spend_logs: Sequence[SpendLogsPayload], + ) -> tuple[SpendLogsPayload, ...]: + """ + Append the session's spend logs that the batch writer has not flushed to the DB yet. + + Without this a follow-up sent inside the ``PROXY_BATCH_WRITE_AT`` window sees an + empty session and silently drops the conversation. The queue is FIFO, so anything + still on it is newer than every row already written. + """ + session_ids: Final = frozenset( + session_id + for spend_log in (*written_spend_logs, *queued_spend_logs) + if spend_log.get("request_id") == response_id and (session_id := spend_log.get("session_id")) + ) | frozenset(session_id for spend_log in written_spend_logs if (session_id := spend_log.get("session_id"))) + written_request_ids: Final = frozenset(spend_log.get("request_id") for spend_log in written_spend_logs) + unflushed: Final = tuple( + spend_log + for spend_log in queued_spend_logs + if spend_log.get("session_id") in session_ids and spend_log.get("request_id") not in written_request_ids + ) + return (*written_spend_logs, *unflushed) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index aa5708088b7..a8092edc625 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -86,6 +86,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.final_text: str = "" self._cached_item_id: str | None = None self._cached_response_id: str | None = None + self._buffered_chunk: ModelResponseStream | None = None + self._upstream_exhausted: bool = False + self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} @@ -330,6 +333,59 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) self._pending_tool_events.append(item_done_event) + def _adopt_response_id_from_chunk(self, chunk: ModelResponseStream) -> None: + if self._cached_response_id is not None: + return + chunk_id: Final = getattr(chunk, "id", None) + if chunk_id and isinstance(chunk_id, str): + self._cached_response_id = chunk_id + + async def _aprime_response_id(self) -> None: + """ + Pull the first upstream chunk before `response.created` is emitted so every event + carries the chat completion id that spend tracking stores as `request_id`. + """ + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = await self.litellm_custom_stream_wrapper.__anext__() + except StopAsyncIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _prime_response_id(self) -> None: + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = self.litellm_custom_stream_wrapper.__next__() + except StopIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _take_buffered_chunk(self) -> ModelResponseStream | None: + buffered: Final = self._buffered_chunk + self._buffered_chunk = None + return buffered + + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: + return ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response, + custom_llm_provider=self.custom_llm_provider, + litellm_metadata=self.litellm_metadata, + ) + def _default_response_created_event_data(self) -> dict: # Use cached response ID if available, otherwise generate a new one if self._cached_response_id is None: @@ -388,7 +444,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, - response=ResponsesAPIResponse(**response_created_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -399,7 +455,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, - response=ResponsesAPIResponse(**response_in_progress_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -811,6 +867,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self.finished is True: raise StopAsyncIteration + await self._aprime_response_id() result = self.return_default_initial_events() if result: return result @@ -822,7 +879,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._pending_tool_events.pop(0) try: - chunk = await self.litellm_custom_stream_wrapper.__anext__() + chunk = self._take_buffered_chunk() + if chunk is None: + if self._upstream_exhausted: + raise StopAsyncIteration + chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) self._ensure_output_item_for_chunk(chunk) @@ -912,6 +973,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): while True: if self.finished is True: raise StopIteration + self._prime_response_id() result = self.return_default_initial_events() if result: return result @@ -922,7 +984,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._pending_tool_events: return self._pending_tool_events.pop(0) try: - chunk = self.litellm_custom_stream_wrapper.__next__() + buffered_chunk = self._take_buffered_chunk() + if buffered_chunk is not None: + chunk = buffered_chunk + elif self._upstream_exhausted: + raise StopIteration + else: + chunk = self.litellm_custom_stream_wrapper.__next__() self._ensure_output_item_for_chunk(chunk) # Accumulate provider_specific_fields from chunk and delta for src in ( @@ -1082,11 +1150,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_response.id = self._cached_response_id # Encode the response ID to match non-streaming behavior - encoded_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=responses_api_response, - custom_llm_provider=self.custom_llm_provider, - litellm_metadata=self.litellm_metadata, - ) + encoded_response: Final = self._with_encoded_response_id(responses_api_response) return ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 19f240fa3d4..926e9e0af2a 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -1,3 +1,4 @@ +import asyncio import json from unittest.mock import AsyncMock, patch @@ -10,6 +11,7 @@ from litellm.responses.litellm_completion_transformation import session_handler from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) +from litellm.responses.utils import ResponsesAPIRequestUtils @pytest.mark.asyncio @@ -430,3 +432,236 @@ async def test_get_chat_completion_message_history_empty_response_dict(): # Verify the session was still created correctly assert result["litellm_session_id"] == "test-session" + + +def _chat_completion_response(request_id: str, content: str) -> dict: + return { + "id": request_id, + "object": "chat.completion", + "created": 1748575031, + "model": "claude-haiku-4-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + } + + +class _FakePrismaDB: + def __init__(self, rows): + self._rows = rows + self.calls = [] + + async def query_raw(self, query, *args): + self.calls.append(args) + return list(self._rows) + + +class _FakePrismaClient: + def __init__(self, written_rows, queued_rows): + self.db = _FakePrismaDB(written_rows) + self.spend_log_transactions = list(queued_rows) + self._spend_log_transactions_lock = asyncio.Lock() + + +@pytest.mark.asyncio +async def test_message_history_reconstructs_list_shaped_input(): + """ + The Responses API sends `input` as a list of items, which is what lands in the stored + proxy_server_request. The user turns have to survive session reconstruction. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + mock_spend_logs = [ + { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701", + "proxy_server_request": { + "input": [ + { + "role": "user", + "content": "Remember this: my favorite color is chartreuse.", + } + ], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, "OK"), + } + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs: + mock_get_spend_logs.return_value = mock_spend_logs + + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + assert result["litellm_session_id"] == "a96757c4-c6dc-4c76-b37e-e7dfa526b701" + + +@pytest.mark.asyncio +async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_writer(): + """ + A follow-up sent right after the previous turn arrives before the batch writer has + flushed that turn's spend log, so the row is only in memory. The history has to + include it anyway. + """ + request_id = "chatcmpl-6c1f5f6c-6a2b-4c62-8d1f-0d9d4ce0a1b2" + queued_spend_log = { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3", + "proxy_server_request": json.dumps( + { + "input": [ + { + "role": "user", + "content": "Remember this: my favorite color is chartreuse.", + } + ], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response(request_id, "OK")), + } + fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[queued_spend_log]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + assert result["litellm_session_id"] == "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" + assert fake_prisma_client.spend_log_transactions == [queued_spend_log] + + +@pytest.mark.asyncio +async def test_message_history_merges_written_and_queued_turns_in_order(): + """ + Turn 1 already flushed to the DB, turn 2 still queued: the follow-up sees the whole + conversation, in order, with no row counted twice. + """ + session_id = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1" + first_request_id = "chatcmpl-1111" + second_request_id = "chatcmpl-2222" + written_spend_log = { + "request_id": first_request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": { + "input": [{"role": "user", "content": "My favorite color is chartreuse."}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(first_request_id, "Got it."), + } + queued_spend_log = { + "request_id": second_request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": json.dumps( + { + "input": [{"role": "user", "content": "And my favorite city is Lisbon."}], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response(second_request_id, "Noted.")), + } + fake_prisma_client = _FakePrismaClient( + written_rows=[written_spend_log], + queued_rows=[written_spend_log, queued_spend_log], + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + second_request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "My favorite color is chartreuse."), + ("assistant", "Got it."), + ("user", "And my favorite city is Lisbon."), + ("assistant", "Noted."), + ] + assert result["litellm_session_id"] == session_id + + +@pytest.mark.asyncio +async def test_message_history_ignores_queued_spend_logs_from_other_sessions(): + request_id = "chatcmpl-3333" + written_spend_log = { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "session-a", + "proxy_server_request": { + "input": [{"role": "user", "content": "Hello from session a."}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, "Hi."), + } + other_session_spend_log = { + "request_id": "chatcmpl-4444", + "call_type": "aresponses", + "session_id": "session-b", + "proxy_server_request": json.dumps( + { + "input": [{"role": "user", "content": "Hello from session b."}], + "model": "claude-bridge", + } + ), + "response": json.dumps(_chat_completion_response("chatcmpl-4444", "Hi there.")), + } + fake_prisma_client = _FakePrismaClient( + written_rows=[written_spend_log], + queued_rows=[other_session_spend_log], + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Hello from session a."), + ("assistant", "Hi."), + ] + + +@pytest.mark.asyncio +async def test_message_history_looks_up_the_decoded_chat_completion_id(): + """ + A `previous_response_id` handed back by the proxy is base64 encoded; spend logs store + the bare chat completion id, so that is what the lookup has to query on. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="anthropic", + model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", + response_id=request_id, + ) + fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[]) + + 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 fake_prisma_client.db.calls == [(request_id,)] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py new file mode 100644 index 00000000000..97f35900e9d --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py @@ -0,0 +1,130 @@ +from unittest.mock import MagicMock + +import pytest + +from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + +CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" +RESPONSE_ID_EVENT_TYPES = frozenset( + {"response.created", "response.in_progress", "response.completed"} +) + + +def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +class _FakeStreamWrapper: + def __init__(self, chunks): + self._chunks = list(chunks) + self.logging_obj = MagicMock() + + def __iter__(self): + return self + + def __next__(self): + if not self._chunks: + raise StopIteration + return self._chunks.pop(0) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + +def _response_ids(events) -> list[str]: + return [ + event.response.id + for event in events + if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + +@pytest.mark.asyncio +async def test_streaming_events_share_the_chat_completion_response_id(): + """ + Every event of a bridged stream has to carry the same id, and that id has to decode + to the chat completion id spend tracking stores as `request_id`. Otherwise a + follow-up `previous_response_id` matches no session and the conversation is dropped. + """ + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) + assert decoded["response_id"] == CHAT_COMPLETION_ID + assert decoded["custom_llm_provider"] == "anthropic" + + +def test_sync_streaming_events_share_the_chat_completion_response_id(): + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = list(iterator) + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + assert ( + ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] + == CHAT_COMPLETION_ID + ) + + +@pytest.mark.asyncio +async def test_streaming_emits_every_chunk_after_priming_the_response_id(): + iterator = _build_iterator( + [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] + ) + + events = [event async for event in iterator] + + deltas = "".join( + event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert deltas == "Hello!" + + +@pytest.mark.asyncio +async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): + iterator = _build_iterator([]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert response_ids + assert len(set(response_ids)) == 1 + assert response_ids[0].startswith("resp_") From f89a3693baabf3dba081ba213032ffe7acd39b65 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:46:24 -0700 Subject: [PATCH 2/5] fix(responses): resolve previous_response_id for a just-written turn The session lookup reads spend logs straight out of the database, so a follow-up sent seconds after the turn it chains off found nothing while the row was still queued in the worker that served it, and the conversation was dropped without an error. Responses calls now ask the spend-log writer to flush on its next pass instead of waiting out its poll interval, and the lookup gives a just-finished turn a short second chance. Replaying a session also accepted `input` only as a string or a single dict, so the standard list shape dropped every user turn and left the model with assistant messages alone. --- litellm/constants.py | 2 + litellm/proxy/db/db_spend_update_writer.py | 8 +- litellm/proxy/utils.py | 26 ++- .../session_handler.py | 86 +++------ .../proxy/db/test_db_spend_update_writer.py | 28 +++ .../prisma_and_spend/test_spend_functions.py | 50 ++++- .../test_session_handler.py | 176 +++++++----------- 7 files changed, 192 insertions(+), 184 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c33e5a53b76..aaaddd063e7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1542,6 +1542,8 @@ SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BA SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000"))) SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) +RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3"))) +RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 283194bad7c..0c8c9a853ec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -65,6 +65,7 @@ from litellm.proxy.spend_tracking.savings import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable +from litellm.types.utils import CallTypes if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -73,6 +74,9 @@ else: ProxyLogging = Any +RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -820,9 +824,11 @@ class DBSpendUpdateWriter: ) ) if prisma_client is not None and spend_logs_url is not None or prisma_client is not None: - from litellm.proxy.utils import enqueue_spend_logs + from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) + if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: + request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9978fa04f40..86d954c0913 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3341,6 +3341,7 @@ class _StaleReadEngine: class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() + spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event() spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None tool_usage_transactions: list["ToolUsageTransaction"] = [] @@ -6005,14 +6006,24 @@ async def enqueue_spend_logs( ) -async def peek_spend_logs(prisma_client: PrismaClient) -> tuple[SpendLogsPayload, ...]: - """Snapshot the spend logs still waiting for the next flush, leaving the queue intact. +def request_spend_log_flush() -> None: + """Wake the queue monitor now rather than leaving the rows for its next poll. - Reads that need a just-finished request use this, since the batch writer only - reaches the DB every ``PROXY_BATCH_WRITE_AT`` seconds. + The Responses API hands the client an id it can chain from straight away, and that + lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. + Repeated requests coalesce into the monitor's next pass, so the batching holds. """ - async with prisma_client._spend_log_transactions_lock: - return tuple(prisma_client.spend_log_transactions) + PrismaClient.spend_log_flush_requested.set() + + +async def _wait_for_spend_log_flush_request(interval: float) -> bool: + """Wait out ``interval``, returning early and True when a flush was requested.""" + try: + await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval) + except asyncio.TimeoutError: + return False + PrismaClient.spend_log_flush_requested.clear() + return True async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: @@ -6460,7 +6471,8 @@ async def _monitor_spend_logs_queue( # Exponential backoff when no logs to process current_interval = min(current_interval * backoff_multiplier, max_backoff) - await asyncio.sleep(current_interval) + if await _wait_for_spend_log_flush_request(current_interval): + current_interval = base_interval except Exception as e: spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e) # Continue monitoring even if there's an error, with exponential backoff diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 935c78bc9a1..15267533957 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,5 +1,5 @@ +import asyncio import json -from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -132,8 +132,8 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = ResponsesSessionHandler._get_response_dict_from_spend_log(spend_log) - if _response_output: + _response_output: Final = spend_log.get("response", "{}") + if isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -141,23 +141,6 @@ class ResponsesSessionHandler: chat_completion_message_history.append(getattr(choice, "message")) return chat_completion_message_history - @staticmethod - def _get_response_dict_from_spend_log(spend_log: SpendLogsPayload) -> Mapping[str, Any] | None: - """ - Spend logs read from the DB hold `response` as a dict, ones still queued in memory - hold it as a JSON string. - """ - _response_output: Final = spend_log.get("response") - if isinstance(_response_output, dict): - return _response_output or None - if isinstance(_response_output, str): - try: - parsed: Final = json.loads(_response_output) - except json.JSONDecodeError: - return None - return parsed if isinstance(parsed, dict) and parsed else None - return None - @staticmethod async def get_proxy_server_request_from_spend_log( spend_log: SpendLogsPayload, @@ -272,9 +255,16 @@ class ResponsesSessionHandler: SQL query SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + + A just-finished turn gets a short second chance: the worker that served it may + still be writing its spend log when the follow-up arrives, and an empty result + drops the whole conversation instead of erroring. """ + from litellm.constants import ( + RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS, + RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL, + ) from litellm.proxy.proxy_server import prisma_client - from litellm.proxy.utils import peek_spend_logs verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) @@ -295,46 +285,16 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - written_spend_logs: Final = await prisma_client.db.query_raw(query, response_id) - queued_spend_logs: Final = await peek_spend_logs(prisma_client) - spend_logs: Final = list( - ResponsesSessionHandler._merge_queued_spend_logs( - response_id=response_id, - written_spend_logs=written_spend_logs, - queued_spend_logs=queued_spend_logs, - ) - ) + for attempt in range(RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS): + if attempt: + await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) + if spend_logs := await prisma_client.db.query_raw(query, response_id): + 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 the following spend logs for previous response id %s: %s", - response_id, - json.dumps(spend_logs, indent=4, default=str), - ) - - return spend_logs - - @staticmethod - def _merge_queued_spend_logs( - response_id: str, - written_spend_logs: Sequence[SpendLogsPayload], - queued_spend_logs: Sequence[SpendLogsPayload], - ) -> tuple[SpendLogsPayload, ...]: - """ - Append the session's spend logs that the batch writer has not flushed to the DB yet. - - Without this a follow-up sent inside the ``PROXY_BATCH_WRITE_AT`` window sees an - empty session and silently drops the conversation. The queue is FIFO, so anything - still on it is newer than every row already written. - """ - session_ids: Final = frozenset( - session_id - for spend_log in (*written_spend_logs, *queued_spend_logs) - if spend_log.get("request_id") == response_id and (session_id := spend_log.get("session_id")) - ) | frozenset(session_id for spend_log in written_spend_logs if (session_id := spend_log.get("session_id"))) - written_request_ids: Final = frozenset(spend_log.get("request_id") for spend_log in written_spend_logs) - unflushed: Final = tuple( - spend_log - for spend_log in queued_spend_logs - if spend_log.get("session_id") in session_ids and spend_log.get("request_id") not in written_request_ids - ) - return (*written_spend_logs, *unflushed) + verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) + return [] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 76a80ac2651..ca1827aa38e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2712,3 +2712,31 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey assert mock_prisma_client.db.tx.call_count == 2 proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, expects_flush", + [("aresponses", True), ("responses", True), ("acompletion", False)], +) +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( + call_type: str, expects_flush: bool +): + """ + A `previous_response_id` chained straight off the previous turn reads the DB, so a + Responses row cannot sit in this worker's queue until the monitor's next poll. + """ + from litellm.proxy.utils import PrismaClient + + db_writer = DBSpendUpdateWriter() + prisma = _tool_usage_prisma() + PrismaClient.spend_log_flush_requested.clear() + + await db_writer._insert_spend_log_to_db( + payload={"request_id": "req-1", "call_type": call_type}, + prisma_client=prisma, + ) + + assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] + assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush + PrismaClient.spend_log_flush_requested.clear() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 54d59e690f9..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -11,7 +11,8 @@ Symbols pinned here: from __future__ import annotations import asyncio -from typing import Any, Dict, List +from contextlib import suppress +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock import pytest @@ -526,6 +527,53 @@ async def test_monitor_spend_logs_queue_swallows_errors_and_backs_off( assert sleep_count["n"] == 3 +@pytest.mark.asyncio +async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A requested flush wakes the monitor mid-poll, so a Responses row reaches the DB + before the client can chain a `previous_response_id` off it. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import PrismaClient, request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + PrismaClient.spend_log_flush_requested.clear() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [] + + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.sleep(0.05) + assert not flushed.is_set() + + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) + request_spend_log_flush() + + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + PrismaClient.spend_log_flush_requested.clear() + + def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 926e9e0af2a..4fc288a47cb 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -1,4 +1,3 @@ -import asyncio import json from unittest.mock import AsyncMock, patch @@ -451,20 +450,38 @@ def _chat_completion_response(request_id: str, content: str) -> dict: class _FakePrismaDB: - def __init__(self, rows): - self._rows = rows + def __init__(self, results): + self._results = list(results) self.calls = [] async def query_raw(self, query, *args): self.calls.append(args) - return list(self._rows) + if not self._results: + return [] + return list(self._results.pop(0)) class _FakePrismaClient: - def __init__(self, written_rows, queued_rows): - self.db = _FakePrismaDB(written_rows) - self.spend_log_transactions = list(queued_rows) - self._spend_log_transactions_lock = asyncio.Lock() + def __init__(self, results): + self.db = _FakePrismaDB(results) + + +def _spend_log(request_id: str, session_id: str, prompt: str, answer: str) -> dict: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": { + "input": [{"role": "user", "content": prompt}], + "model": "claude-bridge", + }, + "response": _chat_completion_response(request_id, answer), + } + + +@pytest.fixture +def instant_session_lookup_retries(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.constants, "RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.0) @pytest.mark.asyncio @@ -475,21 +492,12 @@ async def test_message_history_reconstructs_list_shaped_input(): """ request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" mock_spend_logs = [ - { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701", - "proxy_server_request": { - "input": [ - { - "role": "user", - "content": "Remember this: my favorite color is chartreuse.", - } - ], - "model": "claude-bridge", - }, - "response": _chat_completion_response(request_id, "OK"), - } + _spend_log( + request_id, + "a96757c4-c6dc-4c76-b37e-e7dfa526b701", + "Remember this: my favorite color is chartreuse.", + "OK", + ) ] with patch.object( @@ -512,31 +520,22 @@ async def test_message_history_reconstructs_list_shaped_input(): @pytest.mark.asyncio -async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_writer(): +async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flushed_yet( + instant_session_lookup_retries: None, +): """ - A follow-up sent right after the previous turn arrives before the batch writer has - flushed that turn's spend log, so the row is only in memory. The history has to - include it anyway. + A follow-up sent right after the previous turn can beat that turn's spend log to the + DB. The lookup has to try again instead of handing back an empty conversation. """ request_id = "chatcmpl-6c1f5f6c-6a2b-4c62-8d1f-0d9d4ce0a1b2" - queued_spend_log = { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3", - "proxy_server_request": json.dumps( - { - "input": [ - { - "role": "user", - "content": "Remember this: my favorite color is chartreuse.", - } - ], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response(request_id, "OK")), - } - fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[queued_spend_log]) + session_id = "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" + spend_log = _spend_log( + request_id, + session_id, + "Remember this: my favorite color is chartreuse.", + "OK", + ) + fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -548,44 +547,22 @@ async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_wr ("user", "Remember this: my favorite color is chartreuse."), ("assistant", "OK"), ] - assert result["litellm_session_id"] == "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3" - assert fake_prisma_client.spend_log_transactions == [queued_spend_log] + assert result["litellm_session_id"] == session_id + assert fake_prisma_client.db.calls == [(request_id,), (request_id,)] @pytest.mark.asyncio -async def test_message_history_merges_written_and_queued_turns_in_order(): - """ - Turn 1 already flushed to the DB, turn 2 still queued: the follow-up sees the whole - conversation, in order, with no row counted twice. - """ +async def test_message_history_reconstructs_every_turn_of_the_session_in_order(): session_id = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1" first_request_id = "chatcmpl-1111" second_request_id = "chatcmpl-2222" - written_spend_log = { - "request_id": first_request_id, - "call_type": "aresponses", - "session_id": session_id, - "proxy_server_request": { - "input": [{"role": "user", "content": "My favorite color is chartreuse."}], - "model": "claude-bridge", - }, - "response": _chat_completion_response(first_request_id, "Got it."), - } - queued_spend_log = { - "request_id": second_request_id, - "call_type": "aresponses", - "session_id": session_id, - "proxy_server_request": json.dumps( - { - "input": [{"role": "user", "content": "And my favorite city is Lisbon."}], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response(second_request_id, "Noted.")), - } fake_prisma_client = _FakePrismaClient( - written_rows=[written_spend_log], - queued_rows=[written_spend_log, queued_spend_log], + results=[ + [ + _spend_log(first_request_id, session_id, "My favorite color is chartreuse.", "Got it."), + _spend_log(second_request_id, session_id, "And my favorite city is Lisbon.", "Noted."), + ] + ] ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): @@ -604,45 +581,18 @@ async def test_message_history_merges_written_and_queued_turns_in_order(): @pytest.mark.asyncio -async def test_message_history_ignores_queued_spend_logs_from_other_sessions(): - request_id = "chatcmpl-3333" - written_spend_log = { - "request_id": request_id, - "call_type": "aresponses", - "session_id": "session-a", - "proxy_server_request": { - "input": [{"role": "user", "content": "Hello from session a."}], - "model": "claude-bridge", - }, - "response": _chat_completion_response(request_id, "Hi."), - } - other_session_spend_log = { - "request_id": "chatcmpl-4444", - "call_type": "aresponses", - "session_id": "session-b", - "proxy_server_request": json.dumps( - { - "input": [{"role": "user", "content": "Hello from session b."}], - "model": "claude-bridge", - } - ), - "response": json.dumps(_chat_completion_response("chatcmpl-4444", "Hi there.")), - } - fake_prisma_client = _FakePrismaClient( - written_rows=[written_spend_log], - queued_rows=[other_session_spend_log], - ) +async def test_session_lookup_stops_retrying_once_the_budget_is_spent( + instant_session_lookup_retries: None, +): + fake_prisma_client = _FakePrismaClient(results=[]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" ) - messages = result["messages"] - assert [(message.get("role"), message.get("content")) for message in messages] == [ - ("user", "Hello from session a."), - ("assistant", "Hi."), - ] + assert spend_logs == [] + assert len(fake_prisma_client.db.calls) == litellm.constants.RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS @pytest.mark.asyncio @@ -657,7 +607,9 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[]) + fake_prisma_client = _FakePrismaClient( + results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] + ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( From 9d22acab110fb0407059481cc32755b0d3e095b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:22:10 -0700 Subject: [PATCH 3/5] fix(responses): skip the session lookup retry when spend logs are off --- .../session_handler.py | 8 ++++--- .../test_session_handler.py | 21 +++++++++++++++++++ ...ponse_id.py => test_streaming_iterator.py} | 0 3 files changed, 26 insertions(+), 3 deletions(-) rename tests/test_litellm/responses/litellm_completion_transformation/{test_streaming_iterator_response_id.py => test_streaming_iterator.py} (100%) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 15267533957..59ff492a79f 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -258,13 +258,14 @@ class ResponsesSessionHandler: A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result - drops the whole conversation instead of erroring. + drops the whole conversation instead of erroring. Deployments that write no spend + logs at all have nothing to wait for, so they keep the single original query. """ from litellm.constants import ( RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS, RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL, ) - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import disable_spend_logs, prisma_client verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) @@ -285,7 +286,8 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - for attempt in range(RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS): + 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, response_id): diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 4fc288a47cb..df477f6d01e 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -617,3 +617,24 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): ) assert fake_prisma_client.db.calls == [(request_id,)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( + instant_session_lookup_retries: None, +): + """ + A deployment that writes no spend logs has nothing to wait for, so the miss path keeps + the single query it always had. + """ + fake_prisma_client = _FakePrismaClient(results=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( + "litellm.proxy.proxy_server.disable_spend_logs", True + ): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" + ) + + assert spend_logs == [] + assert fake_prisma_client.db.calls == [("chatcmpl-does-not-exist",)] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py similarity index 100% rename from tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_response_id.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py From d2b5034fea9e90e2258b6d91e191371606717f69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:33:23 -0700 Subject: [PATCH 4/5] test(responses): fold the bridged streaming regressions into the mapped test file --- .../test_streaming_iterator.py | 130 ----------------- ...test_streaming_iterator_transformation.py} | 132 +++++++++++++++++- 2 files changed, 129 insertions(+), 133 deletions(-) delete mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py rename tests/test_litellm/responses/litellm_completion_transformation/{test_tool_call_streaming_transformation.py => test_streaming_iterator_transformation.py} (76%) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py deleted file mode 100644 index 97f35900e9d..00000000000 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator.py +++ /dev/null @@ -1,130 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, -) -from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - -CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" -RESPONSE_ID_EVENT_TYPES = frozenset( - {"response.created", "response.in_progress", "response.completed"} -) - - -def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: - return ModelResponseStream( - id=CHAT_COMPLETION_ID, - created=1748575031, - model="claude-haiku-4-5", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - index=0, - delta=Delta(role="assistant", content=content), - finish_reason=finish_reason, - ) - ], - ) - - -class _FakeStreamWrapper: - def __init__(self, chunks): - self._chunks = list(chunks) - self.logging_obj = MagicMock() - - def __iter__(self): - return self - - def __next__(self): - if not self._chunks: - raise StopIteration - return self._chunks.pop(0) - - def __aiter__(self): - return self - - async def __anext__(self): - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - -def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: - return LiteLLMCompletionStreamingIterator( - model="claude-haiku-4-5", - litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), - request_input="What is the weather in San Francisco?", - responses_api_request={}, - custom_llm_provider="anthropic", - litellm_metadata={}, - ) - - -def _response_ids(events) -> list[str]: - return [ - event.response.id - for event in events - if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] - - -@pytest.mark.asyncio -async def test_streaming_events_share_the_chat_completion_response_id(): - """ - Every event of a bridged stream has to carry the same id, and that id has to decode - to the chat completion id spend tracking stores as `request_id`. Otherwise a - follow-up `previous_response_id` matches no session and the conversation is dropped. - """ - iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) - - events = [event async for event in iterator] - - response_ids = _response_ids(events) - assert len(response_ids) == 3 - assert len(set(response_ids)) == 1 - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) - assert decoded["response_id"] == CHAT_COMPLETION_ID - assert decoded["custom_llm_provider"] == "anthropic" - - -def test_sync_streaming_events_share_the_chat_completion_response_id(): - iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) - - events = list(iterator) - - response_ids = _response_ids(events) - assert len(response_ids) == 3 - assert len(set(response_ids)) == 1 - assert ( - ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] - == CHAT_COMPLETION_ID - ) - - -@pytest.mark.asyncio -async def test_streaming_emits_every_chunk_after_priming_the_response_id(): - iterator = _build_iterator( - [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] - ) - - events = [event async for event in iterator] - - deltas = "".join( - event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) - assert deltas == "Hello!" - - -@pytest.mark.asyncio -async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): - iterator = _build_iterator([]) - - events = [event async for event in iterator] - - response_ids = _response_ids(events) - assert response_ids - assert len(set(response_ids)) == 1 - assert response_ids[0].startswith("resp_") diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py similarity index 76% rename from tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index fa6f42609ca..823f656ddc5 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1,18 +1,23 @@ """ -Tests for streaming tool-calls in Responses API transformation. +Tests for the Responses API streaming bridge in +litellm/responses/litellm_completion_transformation/streaming_iterator.py. Ensures that when the underlying chat-completions stream includes tool_calls deltas, LiteLLM emits Responses API streaming events (output_item.added + function_call_arguments.*). Also ensures that tool calls that only appear in the final built response still get emitted -before response.completed. +before response.completed, and that every event of a bridged stream carries the response id +spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock + +import pytest from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import ( Delta, @@ -21,6 +26,68 @@ from litellm.types.utils import ( StreamingChoices, ) +CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" +RESPONSE_ID_EVENT_TYPES = frozenset( + {"response.created", "response.in_progress", "response.completed"} +) + + +def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +class _FakeStreamWrapper: + def __init__(self, chunks): + self._chunks = list(chunks) + self.logging_obj = MagicMock() + + def __iter__(self): + return self + + def __next__(self): + if not self._chunks: + raise StopIteration + return self._chunks.pop(0) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + +def _response_ids(events) -> list[str]: + return [ + event.response.id + for event in events + if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + def test_tool_call_delta_is_emitted_as_responses_events(): iterator = LiteLLMCompletionStreamingIterator( @@ -397,3 +464,62 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): assert arguments_by_call_id["call_b"] == '{"b":' assert arguments_by_call_id["call_a"] != '{"a":1}' assert arguments_by_call_id["call_b"] != '{"b":1}' + + +@pytest.mark.asyncio +async def test_streaming_events_share_the_chat_completion_response_id(): + """ + Every event of a bridged stream has to carry the same id, and that id has to decode + to the chat completion id spend tracking stores as `request_id`. Otherwise a + follow-up `previous_response_id` matches no session and the conversation is dropped. + """ + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) + assert decoded["response_id"] == CHAT_COMPLETION_ID + assert decoded["custom_llm_provider"] == "anthropic" + + +def test_sync_streaming_events_share_the_chat_completion_response_id(): + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = list(iterator) + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + assert ( + ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] + == CHAT_COMPLETION_ID + ) + + +@pytest.mark.asyncio +async def test_streaming_emits_every_chunk_after_priming_the_response_id(): + iterator = _build_iterator( + [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] + ) + + events = [event async for event in iterator] + + deltas = "".join( + event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert deltas == "Hello!" + + +@pytest.mark.asyncio +async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): + iterator = _build_iterator([]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert response_ids + assert len(set(response_ids)) == 1 + assert response_ids[0].startswith("resp_") From 23e64c8b3d89b781642703e4543aa45e56002881 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:25:59 -0700 Subject: [PATCH 5/5] chore(responses): keep the session lookup inside the type-discipline budget --- .../litellm_completion_transformation/session_handler.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 2008006e6cf..1566bb1bdd7 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -108,7 +108,10 @@ class ResponsesSessionHandler: if isinstance(_response_input_param, (str, list)): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast(ResponseInputParam, [_response_input_param]) + response_input_param = cast( + ResponseInputParam, + [_response_input_param], # mutable-ok: a lone input item still has to arrive as a list + ) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -301,4 +304,4 @@ class ResponsesSessionHandler: return spend_logs verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) - return [] + return [] # mutable-ok: an empty result the caller only reads