scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

Merge pull request #41893 from BerriAI/litellm_fix_responses_ws_encrypted_content_affinity

Browse source 
fix(responses): restore encrypted_content and apply affinity on the native WebSocket relay
This commit is contained in:
Mateo Wang 2026-09-18 21:09:17 -07:00 • committed by GitHub
commit b46612cfeb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1019 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -6604,7 +6604,7 @@ class BaseLLMHTTPHandler:
first_message: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
**kwargs: Any,
):
) -> Exception | None:
"""
Handles Responses API WebSocket mode.
@ -6638,7 +6638,7 @@ class BaseLLMHTTPHandler:
**kwargs,
)
await handler.run()
return
return None
import websockets
from websockets.asyncio.client import ClientConnection
@ -6757,9 +6757,10 @@ class BaseLLMHTTPHandler:
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
quota_callbacks=_ws_quota_callbacks,
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
)
await streaming.bidirectional_forward()
return await streaming.bidirectional_forward()
except websockets.exceptions.InvalidStatusCode as e:
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
@ -6773,6 +6774,7 @@ class BaseLLMHTTPHandler:
pass
else:
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
return None
def image_edit_handler(
self,

88
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -1,7 +1,8 @@
import asyncio
import contextlib
import json
import time
from collections.abc import AsyncIterator, Awaitable, Mapping
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
from enum import Enum
from functools import partial
from types import MappingProxyType
@ -12,10 +13,12 @@ import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from openai.types.responses.response_create_params import ResponseInputParam
from pydantic import BaseModel, ConfigDict, ValidationError
from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_api_usage as _blocked_responses_api_usage,
@ -1291,7 +1294,8 @@ async def cancel_response(
async def _read_ws_model_from_first_frame(
websocket: WebSocket,
) -> tuple | None:
query_model: str | None = None,
) -> tuple[str, str] | None:
"""Read the first WS frame and return (model, raw_message), or None on error.
Sends an appropriate error frame and closes the socket before returning None.
@ -1340,7 +1344,7 @@ async def _read_ws_model_from_first_frame(
await websocket.close(code=1008, reason="Invalid first message")
return None
model: Final = _extract_model_from_first_ws_event(first_event)
model: Final = query_model or _extract_model_from_first_ws_event(first_event)
if not model:
await websocket.send_text(
json.dumps(
@ -1371,6 +1375,38 @@ def _extract_model_from_first_ws_event(first_event: Any) -> str | None:
return (nested.get("model") if isinstance(nested, dict) else None) or first_event.get("model")
class _ResponseCreateRoutingHints(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
input: str | Sequence[object] | None = None
previous_response_id: str | None = None
response: "_ResponseCreateRoutingHints | None" = None
def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, object]:
try:
frame: Final = _ResponseCreateRoutingHints.model_validate_json(first_message)
except ValidationError:
return EMPTY_MAPPING
nested: Final = frame.response or frame
hints: Final = {
"input": frame.input if nested.input is None else nested.input,
"previous_response_id": (
frame.previous_response_id if nested.previous_response_id is None else nested.previous_response_id
),
}
return MappingProxyType({key: value for key, value in hints.items() if value is not None})
def _responses_ws_failure_frame(failure: Exception) -> str:
raw_status: Final = getattr(failure, "status_code", None)
status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
error_type: Final = (
"rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error"
)
return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}})
async def _enforce_responses_ws_first_frame_model_auth(
request: Request,
model: str,
@ -1457,19 +1493,16 @@ async def responses_websocket_endpoint(
accept_kwargs["subprotocol"] = requested_protocols[0]
await websocket.accept(**accept_kwargs)
first_message: str | None = None
if not model:
result: Final = await _read_ws_model_from_first_frame(websocket)
if result is None:
return
model, first_message = result
result: Final = await _read_ws_model_from_first_frame(websocket, query_model=model)
if result is None:
return
resolved_model, first_message = result
data: dict[str, object] = {
"model": model,
"model": resolved_model,
"websocket": websocket,
"first_message": first_message,
}
if first_message is not None:
data["first_message"] = first_message
# Construct a synthetic Request for pre-call processing
headers_list: Final = list(websocket.scope.get("headers") or [])
@ -1482,7 +1515,7 @@ async def responses_websocket_endpoint(
request: Final = Request(scope=scope)
request._url = websocket.url
_body_bytes: Final = json.dumps({"model": model}).encode()
_body_bytes: Final = json.dumps({"model": resolved_model}).encode()
async def return_body():
return _body_bytes
@ -1492,10 +1525,10 @@ async def responses_websocket_endpoint(
# Phase 1: pre-call processing (auth, guardrails, rate limits)
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
if first_message is not None:
if not model:
await _enforce_responses_ws_first_frame_model_auth(
request=request,
model=model,
model=resolved_model,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
@ -1514,7 +1547,7 @@ async def responses_websocket_endpoint(
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
model=resolved_model,
route_type="_aresponses_websocket",
)
except Exception as e:
@ -1536,16 +1569,31 @@ async def responses_websocket_endpoint(
await websocket.close(code=1008, reason="Pre-call error")
return
routed_data: Final = dict(
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
)
# Phase 2: route to upstream provider
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call: Final = await route_request(
data=data,
data=routed_data,
route_type="_aresponses_websocket",
llm_router=llm_router,
user_model=user_model,
)
await llm_call
except Exception:
failure: Final = await llm_call
if isinstance(failure, Exception):
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=failure,
request_data=routed_data,
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket error")
with contextlib.suppress(Exception):
await websocket.send_text(_responses_ws_failure_frame(e))
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=routed_data,
)
await websocket.close(code=1011, reason="Internal server error")

41
litellm/responses/main.py
View file

@ -1,5 +1,6 @@
import asyncio
import contextvars
import json
from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
@ -8,7 +9,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast
import httpx
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
@ -2274,6 +2275,27 @@ def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | d
return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None
_RESPONSES_WS_ROUTING_HINT_KEYS: Final = frozenset({"input", "previous_response_id"})
def _first_ws_frame_with_routed_input(first_message: str, routed_input: object) -> str:
try:
frame: Final = _JSON_OBJECT_ADAPTER.validate_json(first_message)
except ValidationError:
return first_message
if frame is None or routed_input is None:
return first_message
raw_nested: Final = frame.get("response")
nested: Final = _JSON_OBJECT_ADAPTER.validate_python(raw_nested) if isinstance(raw_nested, Mapping) else None
if nested is not None and nested.get("input") is not None:
if nested["input"] == routed_input:
return first_message
return json.dumps({**frame, "response": {**nested, "input": routed_input}})
if frame.get("input") == routed_input:
return first_message
return json.dumps({**frame, "input": routed_input})
def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults:
default_reasoning: Final = _deployment_reasoning_default(kwargs)
candidate_params: Final[dict[str, object]] = {
@ -2295,11 +2317,11 @@ async def _aresponses_websocket(
api_key: str | None = None,
timeout: float | None = None,
**kwargs,
):
) -> Exception | None:
"""
Private function to handle the Responses API WebSocket mode.
For PROXY use only.
For PROXY use only. Returns the provider failure that ended the connection, if any.
Resolves the LLM provider from ``model``, looks up the matching
``BaseResponsesAPIConfig``, and hands off to
@ -2364,10 +2386,14 @@ async def _aresponses_websocket(
"api_base",
"api_key",
"timeout",
"first_message",
*_RESPONSES_WS_ROUTING_HINT_KEYS,
}
remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys}
deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
first_message: Final = kwargs.get("first_message")
await base_llm_http_handler.async_responses_websocket(
return await base_llm_http_handler.async_responses_websocket(
model=resolved_model,
websocket=websocket,
logging_obj=litellm_logging_obj,
@ -2375,9 +2401,14 @@ async def _aresponses_websocket(
api_base=resolved_api_base,
api_key=resolved_api_key,
timeout=timeout,
first_message=(
_first_ws_frame_with_routed_input(first_message, kwargs.get("input"))
if isinstance(first_message, str)
else None
),
user_api_key_dict=kwargs.get("user_api_key_dict"),
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
request_defaults=_build_responses_websocket_request_defaults(kwargs),
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
**remaining_kwargs,
)

186
litellm/responses/streaming_iterator.py
View file

@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import copy
import json
import time
import traceback
@ -154,7 +155,7 @@ def _load_json_value(payload: str | bytes) -> object:
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
return model_id if isinstance(model_id, str) else None
@ -229,6 +230,29 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None
return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500)
def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=model,
custom_llm_provider=custom_llm_provider,
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
return True
@ -592,26 +616,7 @@ class BaseResponsesAPIStreamingIterator:
)
def _map_error_event_exception(self, error_obj: object) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=self.model or "",
custom_llm_provider=self.custom_llm_provider or "",
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "")
def _maybe_raise_for_error_event(self, result: object) -> None:
chunk_type: Final = getattr(result, "type", None)
@ -1695,6 +1700,65 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
_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 event.get("error")
response: Final = event.get("response")
return response.get("error") if _is_json_object(response) else None
def _restore_input_item_ids(items: Sequence[object]) -> Sequence[object]:
return ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(copy.deepcopy(list(items))) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs
def _restored_container_fields(container: Mapping[str, object]) -> Mapping[str, object]:
input_items: Final = container.get("input")
previous_response_id: Final = container.get("previous_response_id")
restored: Final = {
"input": _restore_input_item_ids(input_items) if _is_json_array(input_items) else input_items,
"previous_response_id": (
ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id)
if isinstance(previous_response_id, str)
else previous_response_id
),
}
return MappingProxyType({key: value for key, value in restored.items() if value != container.get(key)})
def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> dict[str, object] | None:
nested: Final = msg_obj.get("response")
nested_fields: Final = _restored_container_fields(nested) if _is_json_object(nested) else EMPTY_MAPPING
top_fields: Final = _restored_container_fields(msg_obj)
if not nested_fields and not top_fields:
return None
restored_nested: Final = (
{"response": {**nested, **nested_fields}} if _is_json_object(nested) and nested_fields else EMPTY_MAPPING
)
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( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
encrypted_content=encrypted_content, model_id=model_id
)
return {**event_obj, "item": {**item, "encrypted_content": wrapped_content}}
class ResponsesWebSocketStreaming:
"""
@ -1721,6 +1785,7 @@ class ResponsesWebSocketStreaming:
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
custom_llm_provider: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
):
self.websocket = websocket
@ -1728,6 +1793,9 @@ class ResponsesWebSocketStreaming:
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
litellm_metadata: Final = self.request_data.get("litellm_metadata")
self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {}
self.custom_llm_provider: str | None = custom_llm_provider
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
@ -1796,13 +1864,65 @@ class ResponsesWebSocketStreaming:
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
def _failure_exception(self) -> Exception | None:
failed_event: Final = next(
(event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
)
if failed_event is None:
return None
return _map_stream_error_to_exception(
_ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or ""
)
async def _log_messages(self) -> None:
if not self.logging_obj:
return
if self.input_messages:
self.logging_obj.model_call_details["messages"] = self.input_messages
if self.messages:
if not self.messages:
return
exception: Final = self._failure_exception()
if exception is None:
asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True))
return
self._record_usage_for_failure()
traceback_exception: Final = "".join(traceback.format_exception(exception))
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
)
def _record_usage_for_failure(self) -> None:
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject
usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
self.messages
)
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages)
service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None
logging_result: Final = LiteLLMRealtimeStreamLoggingObject(
usage=usage, results=self.messages, service_tier=service_tier
)
response_cost: Final = self.logging_obj._response_cost_calculator(result=logging_result) or 0.0 # pyright: ignore[reportPrivateUsage] # as the HTTP streaming iterator does
self.logging_obj.record_partial_usage_for_failure(usage, response_cost)
def _wrap_response_event(self, response_str: str) -> str:
try:
event_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
response: Final = event_obj.get("response")
if _is_json_object(response):
wrapped_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
responses_api_response=response,
custom_llm_provider=self.custom_llm_provider,
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)
return response_str if wrapped_event is None else json.dumps(wrapped_event)
async def backend_to_client(self) -> None:
"""Forward events from backend WebSocket to the client."""
@ -1839,12 +1959,13 @@ class ResponsesWebSocketStreaming:
unmasked_str = self._unmask_response_event(response_str)
output_masked_str = await self._mask_response_completed(unmasked_str)
wrapped_str = self._wrap_response_event(output_masked_str)
# Log the output-masked form so PII redacted by apply_to_output
# guardrails does not appear in success logs.
self._store_event(output_masked_str)
self._store_event(wrapped_str)
await self.websocket.send_text(output_masked_str)
await self.websocket.send_text(wrapped_str)
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.debug("Responses WS backend connection closed: %s", e)
@ -1913,19 +2034,22 @@ class ResponsesWebSocketStreaming:
if parsed.get("type") != "response.create":
return message
msg_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = msg_obj != parsed
authorized_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = authorized_obj != parsed
# Always enforce the authorized model, even when PII masking is off.
model_modified: Final = self._enforce_authorized_model(msg_obj)
model_modified: Final = self._enforce_authorized_model(authorized_obj)
restored_obj: Final = _restore_wrapped_ids_in_response_create(authorized_obj)
msg_obj: Final = authorized_obj if restored_obj is None else restored_obj
frame_modified: Final = model_modified or restored_obj is not None or defaults_applied
if not self.guardrail_callbacks:
return json.dumps(msg_obj) if model_modified or defaults_applied else message
return json.dumps(msg_obj) if frame_modified else message
if "metadata" not in self.request_data:
self.request_data["metadata"] = {}
modified = model_modified or defaults_applied
modified = frame_modified
guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks)
for cb in guardrail_cbs:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
@ -2209,8 +2333,7 @@ class ResponsesWebSocketStreaming:
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
"""Run both forwarding directions concurrently."""
async def bidirectional_forward(self) -> Exception | None:
forward_task: Final = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
@ -2227,6 +2350,7 @@ class ResponsesWebSocketStreaming:
await self.backend_ws.close()
except Exception:
pass
return self._failure_exception()
# ---------------------------------------------------------------------------

2
litellm/utils.py
View file

@ -2008,7 +2008,7 @@ def client(original_function):
result=result,
call_type=call_type,
)
elif call_type == CallTypes.arealtime.value:
elif call_type in (CallTypes.arealtime.value, CallTypes.aresponses_websocket.value):
return result
### POST-CALL RULES ###
post_call_processing(

29
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -1068,6 +1068,35 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch):
assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False
@pytest.mark.asyncio
async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch: pytest.MonkeyPatch):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.responses.main import base_llm_http_handler
success_events = []
class CaptureLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
success_events.append(response_obj)
monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai")
with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure)
):
outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test")
await asyncio.sleep(0)
with contextlib.suppress(asyncio.TimeoutError):
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
assert outcome is failure
assert success_events == []
@pytest.mark.asyncio
async def test_agenerate_content_marks_litellm_params_async():
"""LIT-4475: the async ``agenerate_content`` entrypoint must plant

230
tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
View file

@ -617,6 +617,201 @@ class TestResponsesWSFirstFrameModelAuth:
mock_model_auth.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
@pytest.mark.parametrize("query_model", [None, "gpt-4o-mini"])
async def test_endpoint_routes_on_first_frame_input_and_previous_response_id(
self, nested: bool, query_model: str | None
):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
replayed_input = [{"type": "reasoning", "id": "encitem_abc", "encrypted_content": "litellm_enc:abc;blob"}]
payload = {"model": "gpt-4o-mini", "input": replayed_input, "previous_response_id": "resp_prev"}
first_frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload}
raw_first_frame = json.dumps(first_frame)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(return_value=raw_first_frame)
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
async def fake_llm_call():
return None
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests below
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; the payload it hands to routing is what is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam where the first frame's input and previous_response_id become observable
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
return_value=fake_llm_call(),
) as mock_route_request,
):
await responses_websocket_endpoint(
websocket=ws,
model=query_model,
user_api_key_dict=MagicMock(),
)
ws.receive_text.assert_awaited_once()
routed = mock_route_request.await_args.kwargs["data"]
assert routed["model"] == "gpt-4o-mini"
assert routed["input"] == replayed_input
assert routed["previous_response_id"] == "resp_prev"
assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini"
assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket"
ws.close.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_rejected", [True, False])
async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
)
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
failure = litellm.BadRequestError(
message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai"
)
async def fake_llm_call():
return failure if provider_rejected else None
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
user_api_key_dict = MagicMock()
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam that hands back the relay's outcome
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
return_value=fake_llm_call(),
),
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
),
):
await responses_websocket_endpoint(
websocket=ws,
model=None,
user_api_key_dict=user_api_key_dict,
)
ws.close.assert_not_awaited()
if not provider_rejected:
proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
return
proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
assert booked["original_exception"] is failure
assert booked["user_api_key_dict"] is user_api_key_dict
assert booked["request_data"]["model"] == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
rejection = litellm.RateLimitError(
message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai"
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
user_api_key_dict = MagicMock()
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint tells the client is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam that raises the affinity rejection
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
side_effect=rejection,
),
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
),
):
await responses_websocket_endpoint(
websocket=ws,
model=None,
user_api_key_dict=user_api_key_dict,
)
frame = json.loads(ws.send_text.await_args.args[0])
assert frame["type"] == "error"
assert frame["status"] == 429
assert frame["error"]["type"] == "rate_limit_exceeded"
assert "cooling down" in frame["error"]["message"]
ws.close.assert_awaited_once_with(code=1011, reason="Internal server error")
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
assert booked["original_exception"] is rejection
assert booked["user_api_key_dict"] is user_api_key_dict
assert booked["request_data"]["model"] == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_reruns_model_auth_for_first_frame_model(self):
from starlette.requests import Request
@ -743,6 +938,41 @@ class TestReadWSModelFromFirstFrameErrors:
ws.send_text.assert_not_awaited()
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_query_model_wins_over_first_frame_model(self):
from litellm.proxy.response_api_endpoints.endpoints import (
_read_ws_model_from_first_frame,
)
raw = json.dumps({"type": "response.create", "model": "gpt-4o", "input": []})
ws = MagicMock()
ws.receive_text = AsyncMock(return_value=raw)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group")
assert result == ("reasoning-group", raw)
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_query_model_satisfies_a_first_frame_without_model(self):
from litellm.proxy.response_api_endpoints.endpoints import (
_read_ws_model_from_first_frame,
)
raw = json.dumps({"type": "response.create", "input": []})
ws = MagicMock()
ws.receive_text = AsyncMock(return_value=raw)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group")
assert result == ("reasoning-group", raw)
ws.send_text.assert_not_awaited()
ws.close.assert_not_awaited()
class TestManagedResponsesSameProvider:
def _handler(self, model, custom_llm_provider=None):

88
tests/test_litellm/responses/test_responses_api_request_body.py
View file

@ -424,6 +424,94 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_
assert mock_ws.call_args.kwargs["custom_llm_provider"] == "openai"
@pytest.mark.asyncio
async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs(): # test-quality-ok: the relay kwargs are the only place a dropped key is observable; the provider socket behind them is the boundary
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=[{"type": "message", "role": "user", "content": "hi"}],
previous_response_id="resp_prev",
)
mock_ws.assert_awaited_once()
assert "input" not in mock_ws.call_args.kwargs
assert "previous_response_id" not in mock_ws.call_args.kwargs
_STRIPPED_WS_INPUT = [{"role": "user", "content": "hi"}]
_ORIGINAL_WS_INPUT = [
{"type": "reasoning", "id": "rs_1", "encrypted_content": "blob-from-a-removed-deployment", "summary": []},
*_STRIPPED_WS_INPUT,
]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested: bool): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
body = {"model": "gpt-5.6", "input": _ORIGINAL_WS_INPUT, "store": False}
first_message = json.dumps(
{"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
)
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=list(_STRIPPED_WS_INPUT),
first_message=first_message,
)
forwarded = json.loads(mock_ws.call_args.kwargs["first_message"])
container = forwarded["response"] if nested else forwarded
assert container["input"] == _STRIPPED_WS_INPUT
assert container["store"] is False
assert container["model"] == "gpt-5.6"
assert forwarded["type"] == "response.create"
@pytest.mark.asyncio
async def test_aresponses_websocket_forwards_the_first_frame_verbatim_when_routing_left_the_input_alone(): # test-quality-ok: the relay kwargs are the boundary; byte-identical passthrough is only observable there
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
first_message = '{"type": "response.create", "model": "gpt-5.6", "input": [{"role": "user", "content": "hi"}]}'
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=list(_STRIPPED_WS_INPUT),
first_message=first_message,
)
assert mock_ws.call_args.kwargs["first_message"] == first_message
_INJECTION_POINT_INPUT = [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "hi"}]
_SYSTEM_POINT = {"location": "message", "role": "system"}
_USER_POINT = {"location": "message", "role": "user"}

407
tests/test_litellm/responses/test_responses_websocket_all_providers.py
View file

@ -1502,6 +1502,34 @@ class TestNativeWebSocketDeploymentDefaults:
assert dict(request_defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"}
assert dict(request_defaults.overrides) == {"provider_default": "configured"}
@pytest.mark.asyncio
async def test_aresponses_websocket_keeps_first_frame_routing_hints_out_of_the_defaults(
self, monkeypatch: pytest.MonkeyPatch
):
import importlib
from unittest.mock import AsyncMock
responses_main = importlib.import_module("litellm.responses.main")
stub = MagicMock()
stub.async_responses_websocket = AsyncMock()
monkeypatch.setattr(responses_main, "base_llm_http_handler", stub)
await responses_main._aresponses_websocket.__wrapped__(
model="openai/gpt-5-pro",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
reasoning_effort="high",
input=[{"id": "encitem_abc", "type": "reasoning", "encrypted_content": "litellm_enc:abc"}],
previous_response_id="resp_first_turn",
)
call_kwargs = stub.async_responses_websocket.call_args.kwargs
assert dict(call_kwargs["request_defaults"].fill_missing) == {"reasoning": {"effort": "high"}}
assert "input" not in call_kwargs
assert "previous_response_id" not in call_kwargs
class TestNativeWebSocketGuardrails:
@pytest.mark.asyncio
@ -2927,3 +2955,382 @@ class TestNativeWebSocketUrlConstruction:
mock_config.get_websocket_url.assert_called_once()
_, call_kwargs = mock_config.get_websocket_url.call_args
assert call_kwargs["litellm_params"]["api_version"] == "2025-04-01-preview"
_AFFINITY_METADATA = {
"model_info": {"id": "dep-1"},
"encrypted_content_affinity_enabled": True,
}
def _wrapped_reasoning_item():
from litellm.responses.utils import ResponsesAPIRequestUtils
return {
"type": "reasoning",
"id": ResponsesAPIRequestUtils._build_encrypted_item_id("dep-1", "rs_orig"),
"encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1"),
"summary": [],
}
class TestNativeWebSocketEncryptedContentAffinity:
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_client_to_backend_restores_wrapped_ids(self, nested: bool):
from unittest.mock import AsyncMock
from litellm.responses.utils import ResponsesAPIRequestUtils
wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="dep-1", response_id="resp_orig"
)
payload = {
"input": [_wrapped_reasoning_item(), {"type": "message", "role": "user", "content": "hi"}],
"previous_response_id": wrapped_previous,
}
frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload}
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
websocket = MagicMock()
websocket.receive_text = AsyncMock(side_effect=[json.dumps(frame), Exception("stop")])
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={})
await handler.client_to_backend()
sent = json.loads(backend_ws.send.await_args_list[0][0][0])
body = sent["response"] if nested else sent
assert body["input"][0]["id"] == "rs_orig"
assert body["input"][0]["encrypted_content"] == "gAAAA-blob"
assert body["input"][1] == {"type": "message", "role": "user", "content": "hi"}
assert body["previous_response_id"] == "resp_orig"
@pytest.mark.asyncio
async def test_client_to_backend_leaves_unwrapped_frames_untouched(self):
from unittest.mock import AsyncMock
frame = json.dumps({"type": "response.create", "input": "hello", "previous_response_id": "resp_raw"})
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
websocket = MagicMock()
websocket.receive_text = AsyncMock(side_effect=[frame, Exception("stop")])
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={})
await handler.client_to_backend()
assert backend_ws.send.await_args_list[0][0][0] == frame
@pytest.mark.asyncio
async def test_backend_to_client_wraps_ids_when_affinity_is_enabled(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []}
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}),
json.dumps(
{
"type": "response.completed",
"response": {"id": "resp_1", "output": [dict(reasoning_item)], "usage": {"total_tokens": 3}},
}
),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={"litellm_metadata": dict(_AFFINITY_METADATA)},
custom_llm_provider="openai",
)
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])
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 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
async def test_backend_to_client_wraps_only_response_id_without_affinity(self):
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []}
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}),
json.dumps({"type": "response.completed", "response": {"id": "resp_1", "output": [dict(reasoning_item)]}}),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={"litellm_metadata": {"model_info": {"id": "dep-1"}}},
custom_llm_provider="openai",
)
await handler.backend_to_client()
item_done = json.loads(websocket.send_text.await_args_list[0][0][0])
assert item_done["item"] == reasoning_item
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] == reasoning_item
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failure_frame, expected_status",
[
(
{
"type": "error",
"error": {
"type": "invalid_request_error",
"code": "invalid_encrypted_content",
"message": "The encrypted content for item rs_1 could not be verified.",
},
},
400,
),
(
{
"type": "response.failed",
"response": {
"id": "resp_1",
"status": "failed",
"error": {"code": "server_error", "message": "upstream blew up"},
},
},
500,
),
],
)
async def test_backend_to_client_books_failure_frames_as_failures(
self, failure_frame: dict[str, object], expected_status: int
):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
json.dumps(failure_frame),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.0)
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
await handler.backend_to_client()
await asyncio.sleep(0)
logging_obj.dispatch_success_handlers.assert_not_awaited()
logging_obj.dispatch_failure_handlers.assert_awaited_once()
exception = logging_obj.dispatch_failure_handlers.await_args[0][0]
assert exception.status_code == expected_status
assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception)
@pytest.mark.asyncio
async def test_backend_to_client_bills_completed_turns_before_a_failure(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
}
),
json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.01)
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={})
await handler.backend_to_client()
await asyncio.sleep(0)
logging_obj.record_partial_usage_for_failure.assert_called_once()
usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0]
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert response_cost == 0.01
logging_obj.dispatch_success_handlers.assert_not_awaited()
logging_obj.dispatch_failure_handlers.assert_awaited_once()
@pytest.mark.asyncio
async def test_bidirectional_forward_returns_the_provider_failure(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
backend_drained = asyncio.Event()
backend_events = [
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
json.dumps(
{
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "invalid_encrypted_content",
"message": "could not be verified",
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
backend_ws.close = AsyncMock()
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.0)
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
failure = await handler.bidirectional_forward()
assert isinstance(failure, Exception)
assert failure.status_code == 400
assert "could not be verified" in str(failure)
@pytest.mark.asyncio
async def test_bidirectional_forward_returns_none_after_a_completed_turn(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
backend_drained = asyncio.Event()
backend_events = [
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
backend_ws.close = AsyncMock()
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
assert await handler.bidirectional_forward() is None