diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8ceb9a0fea9..0dc06f9f183 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3202fabe74e..3b6afc34063 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index a5912bb42b1..5a4a08b760c 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 08f5ec236a7..195214b077c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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() # --------------------------------------------------------------------------- diff --git a/litellm/utils.py b/litellm/utils.py index e30e8cde86d..430052b7935 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 8ce5357dc94..836ac42e1f5 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 63faef2366f..f53fbde6b51 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -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): diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 5fced458208..6b5aab932ec 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -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"} diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 50fbfb592a5..2fe9f231f14 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -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