scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(responses): keep reasoning item IDs consistent across streaming events (#51243)

Browse source 
* 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>
This commit is contained in:
joshua-berri 2026-11-11 15:49:46 -08:00 • committed by GitHub
parent cdcece4d32
commit 82471d4603
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 587 additions and 99 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm-rust/crates/inference-responses/src/ids.rs
View file

@ -11,7 +11,7 @@
//! | Id | Applied to | Decoded plaintext |
//! |---|---|---|
//! | `resp_<b64>` | the response `id` | `litellm:custom_llm_provider:{provider};model_id:{model_id};response_id:{id}` |
//! | `encitem_<b64>` | the `id` of an output item with `encrypted_content` | `litellm:model_id:{model_id};item_id:{id}` |
//! | `encitem_<b64>` | 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:<b64>;{content}` | that item's `encrypted_content` | `<b64>` is `model_id:{model_id}`, and the original content follows the `;` unencoded |
//! | `cntr_<b64>` | `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<String, Value>, 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))]

9
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -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,

59
litellm/responses/streaming_iterator.py
View file

@ -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:

163
litellm/responses/utils.py
View file

@ -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

54
tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -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):

71
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -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

192
tests/unit/responses/test_streaming_iterator.py
View file

@ -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()

87
tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -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():
"""