From 82471d4603029571e28719181a6bbeff28ff2706 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 11 Nov 2026 15:49:46 -0800 Subject: [PATCH] fix(responses): keep reasoning item IDs consistent across streaming events (#51243) * fix(responses): keep encrypted reasoning IDs consistent across stream events * fix(responses): preserve immutable stream events during affinity encoding * fix(responses): type affinity encoding input boundaries * fix(responses): use immutable affinity item updates --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../crates/inference-responses/src/ids.rs | 51 ++++- .../streaming_iterator.py | 9 +- litellm/responses/streaming_iterator.py | 59 ++---- litellm/responses/utils.py | 163 ++++++++++++--- .../test_streaming_iterator_transformation.py | 54 ++++- .../test_responses_websocket_all_providers.py | 71 ++++++- .../unit/responses/test_streaming_iterator.py | 192 +++++++++++++++++- .../test_encrypted_content_affinity_check.py | 87 +++++++- 8 files changed, 587 insertions(+), 99 deletions(-) diff --git a/litellm-rust/crates/inference-responses/src/ids.rs b/litellm-rust/crates/inference-responses/src/ids.rs index b51aea1559e..fbba8f949a0 100644 --- a/litellm-rust/crates/inference-responses/src/ids.rs +++ b/litellm-rust/crates/inference-responses/src/ids.rs @@ -11,7 +11,7 @@ //! | Id | Applied to | Decoded plaintext | //! |---|---|---| //! | `resp_` | the response `id` | `litellm:custom_llm_provider:{provider};model_id:{model_id};response_id:{id}` | -//! | `encitem_` | the `id` of an output item with `encrypted_content` | `litellm:model_id:{model_id};item_id:{id}` | +//! | `encitem_` | the `id` of a reasoning or compaction output item, or of any item with `encrypted_content` | `litellm:model_id:{model_id};item_id:{id}` | //! | `litellm_enc:;{content}` | that item's `encrypted_content` | `` is `model_id:{model_id}`, and the original content follows the `;` unencoded | //! | `cntr_` | `container_id` on output items, `code_interpreter_call`, and message annotations | `litellm:custom_llm_provider:{provider};model_id:{model_id};container_id:{id}` | //! @@ -25,8 +25,9 @@ //! - `encitem_` and `litellm_enc:` carry encrypted-content affinity: encrypted reasoning only //! decrypts under the API key that produced it, so the router pins replays to that deployment. //! They are stamped only when the router sets `encrypted_content_affinity_enabled` and a -//! deployment `model_id` is known (see [`IdScope`]). Items without an `id`, such as Codex's, -//! carry the deployment in the wrapped content alone. +//! deployment `model_id` is known (see [`IdScope`]). Reasoning and compaction items get an +//! `encitem_` id by type, so it matches the one Python streams before their content arrives. +//! Items without an `id`, such as Codex's, carry the deployment in the wrapped content alone. //! - `cntr_` lets container calls find the provider and deployment that own a container. A //! container id that already names a provider is left as it is. //! - On the way in, [`restore_input`] turns `encitem_` ids and `litellm_enc:` content in the @@ -51,6 +52,7 @@ const RESPONSE_PREFIX: &str = "resp_"; const ITEM_PREFIX: &str = "encitem_"; const CONTAINER_PREFIX: &str = "cntr_"; const WRAPPED_CONTENT_PREFIX: &str = "litellm_enc:"; +const ENCRYPTED_ITEM_TYPES: [&str; 2] = ["reasoning", "compaction"]; /// Python's non-strict `base64.b64decode`: characters outside the alphabet are skipped and /// non-zero trailing bits are accepted @@ -211,15 +213,23 @@ pub fn decode_container_id(id: &str) -> DecodedContainerId { } fn encode_encrypted_item(item: &mut Map, model_id: &str) { - let Some(Value::String(content)) = item.get_mut("encrypted_content") else { - return; + let wrapped = match item.get_mut("encrypted_content") { + Some(Value::String(content)) if !content.is_empty() => { + *content = wrap_encrypted_content(content, model_id); + true + } + _ => false, }; - if content.is_empty() { + let encrypted_type = item + .get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| ENCRYPTED_ITEM_TYPES.contains(&item_type)); + if !wrapped && !encrypted_type { return; } - *content = wrap_encrypted_content(content, model_id); if let Some(Value::String(id)) = item.get_mut("id") && !id.is_empty() + && decode_encrypted_item_id(id).is_none() { *id = encrypted_item_id(model_id, id); } @@ -456,6 +466,33 @@ mod tests { ); } + #[rstest] + #[case::reasoning_without_ciphertext( + json!({"type": "reasoning", "id": "rs_0a1b", "summary": []}), + json!({"type": "reasoning", "id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=", "summary": []}) + )] + #[case::compaction( + json!({"type": "compaction", "id": "cmp_0a1b", "encrypted_content": "gAAAA"}), + json!({ + "type": "compaction", + "id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOmNtcF8wYTFi", + "encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA" + }) + )] + #[case::already_encoded( + json!({"type": "reasoning", "id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=", "summary": []}), + json!({"type": "reasoning", "id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=", "summary": []}) + )] + #[case::message( + json!({"type": "message", "id": "msg_1", "content": []}), + json!({"type": "message", "id": "msg_1", "content": []}) + )] + fn encodes_item_ids_by_type_under_affinity(#[case] item: Value, #[case] expected: Value) { + let mut encoded = response("resp_1", json!([item])); + encode_response(&mut encoded, &scope(Some("openai"), Some("deploy-1"), true)); + assert_eq!(Value::Array(encoded.output), json!([expected])); + } + #[rstest] #[case::affinity_disabled(scope(Some("openai"), Some("deploy-1"), false))] #[case::no_model_id(scope(Some("openai"), None, true))] diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 489ab8a317d..ea02e4830d2 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -187,7 +187,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.request_input: str | ResponseInputParam = request_input self.responses_api_request: ResponsesAPIOptionalRequestParams = responses_api_request self.custom_llm_provider: str | None = custom_llm_provider - self.litellm_metadata: dict | None = litellm_metadata or {} + self.litellm_metadata: dict[str, object] | None = litellm_metadata or {} self.tool_call_provider_fields: Final = tool_call_provider_fields # Store lightweight dict snapshots for stream_chunk_builder to reduce # repeated Pydantic attribute access in end-of-stream assembly. @@ -1331,7 +1331,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _stamp_sequence_number(self, event: BaseLiteLLMOpenAIResponseObject) -> BaseLiteLLMOpenAIResponseObject: stamped: Final = event.model_copy(update={"sequence_number": self._sequence_number}) self._sequence_number += 1 - return stamped + affinity_model_id: Final = ResponsesAPIRequestUtils.affinity_model_id(self.litellm_metadata) + return ( + stamped + if affinity_model_id is None + else ResponsesAPIRequestUtils.encode_stream_event_for_affinity(stamped, affinity_model_id) + ) async def __anext__( self, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 347537277df..dca93436efa 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -498,11 +498,17 @@ class BaseResponsesAPIStreamingIterator: "type": ResponsesAPIStreamEvents.ERROR.value, } ) - openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response( + transformed_chunk: Final = self.responses_api_provider_config.transform_streaming_response( model=self.model, parsed_chunk=typed_chunk, logging_obj=self.logging_obj, ) + affinity_model_id: Final = ResponsesAPIRequestUtils.affinity_model_id(self.litellm_metadata) + openai_responses_api_chunk: Final = ( + transformed_chunk + if affinity_model_id is None + else ResponsesAPIRequestUtils.encode_stream_event_for_affinity(transformed_chunk, affinity_model_id) + ) # Only when the SSE JSON carries a response body (delta events do not). # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a @@ -571,27 +577,6 @@ class BaseResponsesAPIStreamingIterator: _stream_model_id, ) - # Wrap encrypted_content in streaming events (output_item.added, output_item.done) - if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"): - openai_types = get_openai_response_types() - event_type: Final = getattr(openai_responses_api_chunk, "type", None) - if event_type in ( - openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - ): - item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) - if item: - encrypted_content: Final = getattr(item, "encrypted_content", None) - if encrypted_content and isinstance(encrypted_content, str): - model_id: Final = _model_id_from_metadata(self.litellm_metadata) - if model_id: - wrapped_content: Final = ( - ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) - ) - setattr(item, "encrypted_content", wrapped_content) - # Store the completed response (also for incomplete/failed so logging still fires) _chunk_type: Final = getattr(openai_responses_api_chunk, "type", None) openai_types = get_openai_response_types() @@ -2493,9 +2478,6 @@ class _ResponsesWsTurn: ) -_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) - - def _ws_event_error(event: Mapping[str, object]) -> object: if event.get("type") == "error": return _error_event_error_object(event) @@ -2535,24 +2517,6 @@ def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> di return {**msg_obj, **top_fields, **restored_nested} -def _wrap_output_item_encrypted_content( - event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object] -) -> dict[str, object] | None: - if not litellm_metadata.get("encrypted_content_affinity_enabled"): - return None - model_id: Final = _model_id_from_metadata(litellm_metadata) - item: Final = event_obj.get("item") - if model_id is None or not _is_json_object(item): - return None - encrypted_content: Final = item.get("encrypted_content") - if not isinstance(encrypted_content, str) or not encrypted_content: - return None - wrapped_content: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( - encrypted_content=encrypted_content, model_id=model_id - ) - return {**event_obj, "item": {**item, "encrypted_content": wrapped_content}} - - class ResponsesWebSocketStreaming: """ Manages bidirectional WebSocket forwarding for the Responses API @@ -2758,9 +2722,12 @@ class ResponsesWebSocketStreaming: litellm_metadata=self.litellm_metadata, ) return json.dumps({**event_obj, "response": wrapped_response}) - if event_obj.get("type") not in _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: - return response_str - wrapped_event: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata) + affinity_model_id: Final = ResponsesAPIRequestUtils.affinity_model_id(self.litellm_metadata) + wrapped_event: Final = ( + None + if affinity_model_id is None + else ResponsesAPIRequestUtils.affinity_encoded_stream_event(event_obj, affinity_model_id) + ) return response_str if wrapped_event is None else json.dumps(wrapped_event) async def backend_to_client(self) -> None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 991cf6213c6..dd3bd1c7e35 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -13,7 +13,13 @@ from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire pay import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, decoded_json_object, iter_items +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, + attribute_or_none, + decoded_json_object, + iter_items, +) from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import ( @@ -51,6 +57,19 @@ _PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[ TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None] ] = TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None) _REASONING_FIELDS_ONLY_ITS_ORIGIN_RESOLVES: Final = frozenset({"encrypted_content", "id", "status"}) +_ENCRYPTED_OUTPUT_ITEM_TYPES: Final = frozenset({"reasoning", "compaction"}) +_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) +_REASONING_ITEM_EVENT_TYPES: Final = frozenset( + { + "response.reasoning_summary_part.added", + "response.reasoning_summary_part.done", + "response.reasoning_summary_text.delta", + "response.reasoning_summary_text.done", + "response.reasoning_text.delta", + "response.reasoning_text.done", + } +) +_CONTENT_PART_EVENT_TYPES: Final = frozenset({"response.content_part.added", "response.content_part.done"}) class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object): @@ -133,6 +152,7 @@ def _as_input_text_part(part: object) -> object: _RequestInputT: Final = TypeVar("_RequestInputT") +_StreamEventT: Final = TypeVar("_StreamEventT", bound=BaseModel) def _provider_mapped_params( @@ -521,6 +541,95 @@ class ResponsesAPIRequestUtils: _decode_encrypted_item_id = decode_encrypted_item_id + @staticmethod + def affinity_model_id(litellm_metadata: Mapping[str, object] | None) -> str | None: + if not litellm_metadata or not litellm_metadata.get("encrypted_content_affinity_enabled"): + return None + model_info: Final = litellm_metadata.get("model_info") + model_id: Final = model_info.get("id") if _is_object_dict(model_info) else None + return model_id if isinstance(model_id, str) and model_id else None + + @staticmethod + def affinity_item_id(item_type: object, item_id: object, encrypted_content: object, model_id: str) -> str | None: + """The ``encitem_`` id an output item carries under encrypted-content affinity, or ``None`` to keep its id. + + Reasoning and compaction items qualify by type, so a stream encodes them before their + ``encrypted_content`` arrives and every event agrees with the final response. + """ + if not isinstance(item_id, str) or not item_id: + return None + if ResponsesAPIRequestUtils.decode_encrypted_item_id(item_id) is not None: + return None + encrypted_type: Final = isinstance(item_type, str) and item_type in _ENCRYPTED_OUTPUT_ITEM_TYPES + if not encrypted_type and not (isinstance(encrypted_content, str) and encrypted_content): + return None + return ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id) + + @staticmethod + def _reasoning_event_item_id(event_type: object, part_type: object, item_id: object, model_id: str) -> str | None: + if not isinstance(event_type, str): + return None + scoped_to_reasoning: Final = event_type in _REASONING_ITEM_EVENT_TYPES or ( + event_type in _CONTENT_PART_EVENT_TYPES and part_type == "reasoning_text" + ) + if not scoped_to_reasoning: + return None + return ResponsesAPIRequestUtils.affinity_item_id("reasoning", item_id, None, model_id) + + @staticmethod + def _affinity_item_updates( + item_type: object, item_id: object, encrypted_content: object, model_id: str + ) -> Mapping[str, str]: + encoded_id: Final = ResponsesAPIRequestUtils.affinity_item_id(item_type, item_id, encrypted_content, model_id) + wrapped: Final = ( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(encrypted_content, model_id) + if isinstance(encrypted_content, str) and encrypted_content + else None + ) + return MappingProxyType( + {key: value for key, value in (("id", encoded_id), ("encrypted_content", wrapped)) if value is not None} + ) + + @staticmethod + def encode_stream_event_for_affinity(event: _StreamEventT, model_id: str) -> _StreamEventT: + """The stream event with the item id and wrapped ``encrypted_content`` the final response carries.""" + event_type: Final = attribute_or_none(event, "type") + if isinstance(event_type, str) and event_type in _OUTPUT_ITEM_EVENT_TYPES: + item: Final = attribute_or_none(event, "item") + if not isinstance(item, BaseModel): + return event + item_updates: Final = ResponsesAPIRequestUtils._affinity_item_updates( + attribute_or_none(item, "type"), + attribute_or_none(item, "id"), + attribute_or_none(item, "encrypted_content"), + model_id, + ) + return event.model_copy(update={"item": item.model_copy(update=item_updates)}) if item_updates else event + part: Final = attribute_or_none(event, "part") + encoded_item_id: Final = ResponsesAPIRequestUtils._reasoning_event_item_id( + event_type, + part.get("type") if _is_object_dict(part) else attribute_or_none(part, "type"), + attribute_or_none(event, "item_id"), + model_id, + ) + return event if encoded_item_id is None else event.model_copy(update={"item_id": encoded_item_id}) + + @staticmethod + def affinity_encoded_stream_event(event: Mapping[str, object], model_id: str) -> dict[str, object] | None: + """The JSON stream event with the ids and wrapped content the final response carries, or ``None``.""" + event_type: Final = event.get("type") + item: Final = event.get("item") + if isinstance(event_type, str) and event_type in _OUTPUT_ITEM_EVENT_TYPES and _is_object_dict(item): + item_updates: Final = ResponsesAPIRequestUtils._affinity_item_updates( + item.get("type"), item.get("id"), item.get("encrypted_content"), model_id + ) + return {**event, "item": {**item, **item_updates}} if item_updates else None + part: Final = event.get("part") + encoded_item_id: Final = ResponsesAPIRequestUtils._reasoning_event_item_id( + event_type, part.get("type") if _is_object_dict(part) else None, event.get("item_id"), model_id + ) + return None if encoded_item_id is None else {**event, "item_id": encoded_item_id} + @staticmethod def wrap_encrypted_content_with_model_id(encrypted_content: str, model_id: str) -> str: """Wrap encrypted_content with model_id metadata for affinity routing. @@ -576,7 +685,7 @@ class ResponsesAPIRequestUtils: response: Union["ResponsesAPIResponse", dict[str, object]], model_id: str | None, ) -> Union["ResponsesAPIResponse", dict[str, object]]: - """Rewrite item IDs for output items that contain ``encrypted_content``. + """Rewrite item IDs for reasoning, compaction and other output items that contain ``encrypted_content``. Encodes ``model_id`` into the item ID so that follow-up requests can be routed back to the originating deployment without any cache lookup. @@ -593,40 +702,28 @@ class ResponsesAPIRequestUtils: else: output = getattr(response, "output", None) - if not isinstance(output, list): + items: Final = as_object_list(output) + if items is None: return response - for item in output: - if isinstance(item, dict): - item_id = item.get("id") - encrypted_content = item.get("encrypted_content") - - if encrypted_content and isinstance(encrypted_content, str): - # Always wrap encrypted_content with model_id for redundancy - item["encrypted_content"] = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( - encrypted_content, model_id + for item in items: + if _is_object_dict(item): + item.update( + ResponsesAPIRequestUtils._affinity_item_updates( + item.get("type"), item.get("id"), item.get("encrypted_content"), model_id ) - # Also encode the ID if present - if item_id and isinstance(item_id, str): - item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id) - else: - item_id = getattr(item, "id", None) - encrypted_content = getattr(item, "encrypted_content", None) - - if encrypted_content and isinstance(encrypted_content, str): - # Always wrap encrypted_content with model_id for redundancy - try: - item.encrypted_content = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) - except AttributeError: - pass - # Also encode the ID if present - if item_id and isinstance(item_id, str): - try: - item.id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id) - except AttributeError: - pass + ) + continue + for key, value in ResponsesAPIRequestUtils._affinity_item_updates( + attribute_or_none(item, "type"), + attribute_or_none(item, "id"), + attribute_or_none(item, "encrypted_content"), + model_id, + ).items(): + try: + setattr(item, key, value) + except AttributeError: + pass return response diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a6e3cc7832..0dce041cce8 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -81,14 +81,14 @@ class _FakeStreamWrapper: return self._chunks.pop(0) -def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: +def _build_iterator(chunks, litellm_metadata: dict[str, object] | None = None) -> 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={}, + litellm_metadata=litellm_metadata or {}, ) @@ -1774,6 +1774,56 @@ async def test_completed_output_follows_the_streamed_output_indexes(sync_mode: b assert [item.id for item in completed.response.output] == [event.item.id for event in added] +_REASONING_ITEM_EVENT_TYPES: Final = frozenset( + { + ResponsesAPIStreamEvents.RESPONSE_PART_ADDED, + ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + } +) + + +def _reasoning_item_ids(events: list[BaseLiteLLMOpenAIResponseObject]) -> set[str]: + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + streamed_items: Final = { + event.item.id for event in events if getattr(getattr(event, "item", None), "type", None) == "reasoning" + } + scoped_events: Final = { + event.item_id for event in events if getattr(event, "type", None) in _REASONING_ITEM_EVENT_TYPES + } + completed_items: Final = {item.id for item in completed.response.output if item.type == "reasoning"} + return streamed_items | scoped_events | completed_items + + +@pytest.mark.parametrize( + "stream_chunks", + [_reasoning_then_text_chunks, _signed_thinking_and_text_in_one_chunk], + ids=["reasoning_then_text_chunks", "signed_thinking_and_text_in_one_chunk"], +) +@pytest.mark.parametrize( + "litellm_metadata, origin", + [({"encrypted_content_affinity_enabled": True, "model_info": {"id": "dep-1"}}, "dep-1"), ({}, None)], + ids=["affinity-on", "affinity-off"], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_reasoning_item_keeps_one_id_across_bridged_events_and_completed_output( + stream_chunks: StreamChunks, litellm_metadata: dict[str, object], origin: str | None, sync_mode: bool +): + events: Final = await _collect_events(_build_iterator(stream_chunks(), litellm_metadata), sync_mode) + + reasoning_ids: Final = _reasoning_item_ids(events) + assert len(reasoning_ids) == 1 + decoded: Final = ResponsesAPIRequestUtils.decode_encrypted_item_id(next(iter(reasoning_ids))) + assert (decoded or {}).get("model_id") == origin + message_ids: Final = {event.item.id for event in events if _is_message_item(event)} + assert message_ids + assert [ResponsesAPIRequestUtils.decode_encrypted_item_id(item_id) for item_id in message_ids] == [None] + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_reasoning_only_stream_lists_no_message_item_it_never_announced(sync_mode: bool): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 9a73054b2e3..fe4ee0c4bbe 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -4692,6 +4692,22 @@ class TestNativeWebSocketEncryptedContentAffinity: backend_ws = MagicMock() backend_ws.recv = AsyncMock( side_effect=[ + json.dumps( + { + "type": "response.output_item.added", + "output_index": 0, + "item": {"type": "reasoning", "id": "rs_1", "summary": []}, + } + ), + json.dumps( + { + "type": "response.reasoning_summary_text.delta", + "item_id": "rs_1", + "output_index": 0, + "summary_index": 0, + "delta": "thinking", + } + ), json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}), json.dumps( { @@ -4716,20 +4732,65 @@ class TestNativeWebSocketEncryptedContentAffinity: await handler.backend_to_client() wrapped_content = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1") - item_done = json.loads(websocket.send_text.await_args_list[0][0][0]) + added, summary_delta, item_done, completed = ( + json.loads(call[0][0]) for call in websocket.send_text.await_args_list[:4] + ) assert item_done["item"]["encrypted_content"] == wrapped_content - completed = json.loads(websocket.send_text.await_args_list[1][0][0]) assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id( custom_llm_provider="openai", model_id="dep-1", response_id="resp_1" ) - assert completed["response"]["output"][0]["id"] == ResponsesAPIRequestUtils._build_encrypted_item_id( - "dep-1", "rs_1" - ) + assert { + added["item"]["id"], + summary_delta["item_id"], + item_done["item"]["id"], + completed["response"]["output"][0]["id"], + } == {ResponsesAPIRequestUtils._build_encrypted_item_id("dep-1", "rs_1")} assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content await asyncio.sleep(0) logged = logging_obj.dispatch_success_handlers.await_args[0][0] assert logged[0]["response"]["id"] == completed["response"]["id"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"note": "a frame without a type"}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 1, "delta": "ok"}, + { + "type": "response.content_part.added", + "item_id": "msg_1", + "output_index": 1, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + }, + {"type": "response.output_item.done", "output_index": 1, "item": {"type": "message", "id": "msg_1"}}, + { + "type": "response.output_item.done", + "output_index": 0, + "item": {"type": "reasoning", "id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXAtMTtpdGVtX2lkOnJzXzE="}, + }, + ], + ids=["no-type", "message-delta", "message-content-part", "message-item", "already-encoded-reasoning"], + ) + async def test_backend_to_client_forwards_events_affinity_does_not_own_unchanged(self, frame: dict[str, object]): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + raw_frame: Final = json.dumps(frame) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=[raw_frame, Exception("stop")]) + handler: Final = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={"litellm_metadata": dict(_AFFINITY_METADATA)}, + custom_llm_provider="openai", + ) + + await handler.backend_to_client() + + assert websocket.send_text.await_args_list[0][0][0] == raw_frame + @pytest.mark.asyncio async def test_backend_to_client_wraps_only_response_id_without_affinity(self): from unittest.mock import AsyncMock diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index ec14f67ce17..7fa5d6bfa80 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest import websockets.exceptions +from pydantic import BaseModel, TypeAdapter from pydantic_core import PydanticSerializationError import litellm @@ -43,6 +44,7 @@ from litellm.types.llms.openai import ( ) from contextlib import suppress from litellm.responses import streaming_iterator as streaming_module +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.utils import CallTypes from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from types import SimpleNamespace @@ -86,6 +88,7 @@ def _make_iterator( trailing_error: Optional[Exception] = None, config: Mock | None = None, request_data: dict | None = None, + litellm_metadata: dict[str, object] | None = None, ) -> ResponsesAPIStreamingIterator: async def aiter_bytes(): for evt in sse_events: @@ -102,7 +105,7 @@ def _make_iterator( model="gpt-4o-mini", responses_api_provider_config=config or _mock_config(), logging_obj=logging_obj, - litellm_metadata={}, + litellm_metadata=litellm_metadata or {}, custom_llm_provider="openai", request_data=request_data, ) @@ -561,6 +564,193 @@ async def test_event_without_a_type_reaches_the_caller_instead_of_ending_the_str assert untyped.model_extra == {"note": "hook output without a type"} +_REASONING_ITEM_ID: Final = "rs_affinity" +_MESSAGE_ITEM_ID: Final = "msg_affinity" +_STREAMED_REASONING_ITEM: Final = { + "id": _REASONING_ITEM_ID, + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thinking"}], + "encrypted_content": "ciphertext", +} +_STREAMED_MESSAGE_ITEM: Final = { + "id": _MESSAGE_ITEM_ID, + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], +} +_REASONING_SCOPED_EVENT_TYPES: Final = ( + ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + "response.reasoning_text.delta", + ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, +) + + +def _reasoning_response_body(status: str, output: list[dict[str, object]]) -> dict[str, object]: + return { + "id": "resp_affinity", + "object": "response", + "created_at": 1, + "status": status, + "model": "gpt-5-mini", + "output": output, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + + +_REASONING_STREAM_EVENTS: Final = [ + _sse_event( + {"type": "response.created", "sequence_number": 0, "response": _reasoning_response_body("in_progress", [])} + ), + _sse_event( + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {"id": _REASONING_ITEM_ID, "type": "reasoning", "summary": []}, + } + ), + _sse_event( + { + "type": "response.content_part.added", + "sequence_number": 2, + "item_id": _REASONING_ITEM_ID, + "output_index": 0, + "content_index": 0, + "part": {"type": "reasoning_text", "text": ""}, + } + ), + _sse_event( + { + "type": "response.reasoning_text.delta", + "sequence_number": 3, + "item_id": _REASONING_ITEM_ID, + "output_index": 0, + "content_index": 0, + "delta": "thinking", + } + ), + _sse_event( + { + "type": "response.reasoning_summary_text.delta", + "sequence_number": 4, + "item_id": _REASONING_ITEM_ID, + "output_index": 0, + "summary_index": 0, + "delta": "thinking", + } + ), + _sse_event( + {"type": "response.output_item.done", "sequence_number": 5, "output_index": 0, "item": _STREAMED_REASONING_ITEM} + ), + _sse_event( + { + "type": "response.output_text.delta", + "sequence_number": 6, + "item_id": _MESSAGE_ITEM_ID, + "output_index": 1, + "content_index": 0, + "delta": "ok", + } + ), + _sse_event( + {"type": "response.output_item.done", "sequence_number": 7, "output_index": 1, "item": _STREAMED_MESSAGE_ITEM} + ), + _sse_event( + { + "type": "response.completed", + "sequence_number": 8, + "response": _reasoning_response_body("completed", [_STREAMED_REASONING_ITEM, _STREAMED_MESSAGE_ITEM]), + } + ), +] + + +_ITEM_RECORD: Final = TypeAdapter(dict[str, object]) + + +def _item_record(item: object) -> dict[str, object]: + return _ITEM_RECORD.validate_python(item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_metadata, origin", + [ + ({"encrypted_content_affinity_enabled": True, "model_info": {"id": "dep-1"}}, "dep-1"), + ({"encrypted_content_affinity_enabled": False, "model_info": {"id": "dep-1"}}, None), + ({"encrypted_content_affinity_enabled": True}, None), + ], + ids=["affinity-on", "affinity-off", "no-model-id"], +) +async def test_streamed_reasoning_item_keeps_one_id_and_replays_as_one_input_item( + litellm_metadata: dict[str, object], origin: str | None +): + iterator: Final = _make_iterator( + sse_events=_REASONING_STREAM_EVENTS, + logging_obj=_logging_obj_stub(), + config=OpenAIResponsesAPIConfig(), + litellm_metadata=litellm_metadata, + ) + + events: Final = [event async for event in iterator] + + item_events: Final = [ + _item_record(event.item) + for event in events + if event.type in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ] + completed: Final = next( + event.response for event in events if event.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + output_items: Final = [_item_record(item) for item in completed.output] + scoped_item_ids: Final = {event.type: event.item_id for event in events if getattr(event, "item_id", None)} + expected_reasoning_id: Final = ( + _REASONING_ITEM_ID + if origin is None + else ResponsesAPIRequestUtils._build_encrypted_item_id(origin, _REASONING_ITEM_ID) + ) + assert {item["id"] for item in [*item_events, *output_items] if item["type"] == "reasoning"} | { + scoped_item_ids[event_type] for event_type in _REASONING_SCOPED_EVENT_TYPES + } == {expected_reasoning_id} + assert {item["id"] for item in [*item_events, *output_items] if item["type"] == "message"} | { + scoped_item_ids[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA] + } == {_MESSAGE_ITEM_ID} + + done_reasoning: Final = next(item for item in reversed(item_events) if item["type"] == "reasoning") + completed_reasoning: Final = next(item for item in output_items if item["type"] == "reasoning") + assert done_reasoning["encrypted_content"] == completed_reasoning["encrypted_content"] + assert ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id( + str(done_reasoning["encrypted_content"]) + ) == (origin, "ciphertext") + + replayed: Final = list( + { + record["id"]: {key: record[key] for key in ("id", "type", "encrypted_content")} + for record in (done_reasoning, completed_reasoning) + }.values() + ) + assert ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(replayed) == [ + {"id": _REASONING_ITEM_ID, "type": "reasoning", "encrypted_content": "ciphertext"} + ] + + +@pytest.mark.asyncio +async def test_output_item_added_without_an_item_streams_through_under_affinity(): + iterator: Final = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_item.added", "sequence_number": 0, "output_index": 0, "item": None}) + ], + logging_obj=_logging_obj_stub(), + config=OpenAIResponsesAPIConfig(), + litellm_metadata={"encrypted_content_affinity_enabled": True, "model_info": {"id": "dep-1"}}, + ) + + added: Final = await iterator.__anext__() + + assert (added.type, added.output_index, added.item) == (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, None) + + @pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) def test_sync_transport_error_before_any_output_raises_fallback_error(trailing_error): logging_obj: Final = _failure_tracking_logging_obj() diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index e908481ca2b..ae5acffab31 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -3,7 +3,7 @@ Tests for encrypted_content_affinity pre-call check. The mechanism works without any cache and supports two encoding strategies: -1. **Items with IDs**: item IDs for output items with `encrypted_content` are rewritten to +1. **Items with IDs**: IDs of reasoning, compaction and other `encrypted_content` output items are rewritten to `encitem_{base64("litellm:model_id:{model_id};item_id:{original_id}")}`. 2. **Items without IDs** (Codex): encrypted_content itself is wrapped with model_id metadata: @@ -18,10 +18,12 @@ The mechanism works without any cache and supports two encoding strategies: import datetime import time from collections.abc import Mapping +from dataclasses import dataclass from typing import Final from unittest.mock import AsyncMock, patch import pytest +from pydantic import TypeAdapter import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager @@ -132,6 +134,13 @@ class TestEncryptedItemIdCodec: assert decoded["item_id"] == original_item_id +@dataclass(frozen=True, slots=True) +class _ReadOnlyReasoningItem: + id: str + type: str + encrypted_content: str + + class TestUpdateEncryptedContentItemIds: def test_rewrites_encrypted_items_in_dict_response(self): model_id = "deployment-1" @@ -157,6 +166,52 @@ class TestUpdateEncryptedContentItemIds: result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, None) assert result["output"][0]["id"] == "rs_xyz" + @pytest.mark.parametrize( + "response", + [ + {"id": "resp_failed", "status": "failed", "error": {"code": "server_error", "message": "boom"}}, + { + "id": "resp_failed", + "status": "failed", + "error": {"code": "server_error", "message": "boom"}, + "output": None, + }, + ], + ids=["output-absent", "output-none"], + ) + def test_passes_a_response_without_output_through_unchanged(self, response: dict[str, object]): + expected: Final = {**response} + result: Final = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response, "deployment-1" + ) + assert result == expected + + def test_leaves_read_only_output_items_as_returned(self): + result: Final = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + {"id": "resp_123", "output": [_ReadOnlyReasoningItem("rs_xyz", "reasoning", "secret")]}, "deployment-1" + ) + assert result["output"] == [_ReadOnlyReasoningItem("rs_xyz", "reasoning", "secret")] + + @pytest.mark.parametrize( + "item, encoded", + [ + ({"id": "rs_xyz", "type": "reasoning", "summary": []}, True), + ({"id": "cmp_xyz", "type": "compaction", "encrypted_content": "secret"}, True), + ({"id": "ws_xyz", "type": "web_search_call"}, False), + ({"id": ResponsesAPIRequestUtils._build_encrypted_item_id("dep-0", "rs_xyz"), "type": "reasoning"}, False), + ], + ids=["reasoning-before-ciphertext", "compaction", "non-reasoning", "already-encoded"], + ) + def test_encodes_reasoning_and_compaction_ids_by_type(self, item: dict[str, object], encoded: bool): + result: Final = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + {"id": "resp_123", "output": [dict(item)]}, "deployment-1" + ) + assert result["output"][0]["id"] == ( + ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-1", str(item["id"])) + if encoded + else item["id"] + ) + class TestEncryptedContentWrapping: def test_wrap_and_unwrap_encrypted_content(self): @@ -219,10 +274,10 @@ class TestRestoreEncryptedContentItemIds: request_input = [ {"type": "message", "id": "msg_abc123", "role": "assistant"}, {"type": "reasoning", "id": encoded_id, "encrypted_content": "secret"}, + {"type": "item_reference", "id": encoded_id}, ] restored = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(request_input) - assert restored[0]["id"] == "msg_abc123" - assert restored[1]["id"] == original_id + assert [item["id"] for item in restored] == ["msg_abc123", original_id, original_id] def test_unwraps_encrypted_content(self): """Test that wrapped encrypted_content is unwrapped before forwarding.""" @@ -715,6 +770,32 @@ async def test_encrypted_content_affinity_does_not_create_litellm_metadata_for_c assert len(result) == 1 +@pytest.mark.asyncio +async def test_item_reference_to_a_reasoning_item_returned_without_ciphertext_pins_its_deployment(): + returned: Final = ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + {"id": "resp_stored", "output": [{"type": "reasoning", "id": "rs_stored", "summary": []}]}, + "openai", + {"model_info": {"id": "dep-2"}, "encrypted_content_affinity_enabled": True}, + ) + returned_output: Final = TypeAdapter(list[dict[str, object]]).validate_python(returned["output"]) + deployments: Final = [ + {"model_info": {"id": "dep-1"}, "litellm_params": {"model": "openai/gpt-5.1-codex"}}, + {"model_info": {"id": "dep-2"}, "litellm_params": {"model": "openai/gpt-5.1-codex"}}, + ] + + pinned: Final = await EncryptedContentAffinityCheck().async_filter_deployments( + model="gpt-5.1-codex", + healthy_deployments=deployments, + messages=None, + request_kwargs={ + "input": [{"type": "item_reference", "id": returned_output[0]["id"]}], + "litellm_metadata": {}, + }, + ) + + assert pinned == [deployments[1]] + + @pytest.mark.asyncio async def test_encrypted_content_affinity_preserves_litellm_metadata_for_responses(): """