from __future__ import annotations import asyncio import json import time import traceback import uuid from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable import httpx from openai._streaming import SSEDecoder from typing_extensions import TypeIs import litellm from litellm.constants import ( LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, ) from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.types.responses.streaming_websocket import ( PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, ) class _StreamCachingHandler(Protocol): """The ``_llm_caching_handler`` attached to a logging object, as this module uses it.""" original_function: Callable[..., object] def _should_store_result_in_cache( self, original_function: Callable[..., object], kwargs: Mapping[str, object] ) -> bool: ... class PiiUnmaskingGuardrailCallback(PresidioGuardrailCallback, Protocol): """Guardrail callback that can also reverse its own masking, selected by ``llm_http_handler`` on exactly this attribute.""" def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... class ProjectQuotaCallback(Protocol): async def enforce_project_io_token_quota_for_frame( self, user_api_key_dict: UserAPIKeyAuth | None, requested_model: str | None, estimated_input_tokens: int, estimated_output_tokens: int, ) -> None: ... @lru_cache(maxsize=1) def _get_openai_response_types(): from litellm.types.llms import openai as openai_types return openai_types def _is_json_object(value: object) -> TypeIs[dict[str, object]]: # guard-ok: trivial isinstance; JSON keys are str return isinstance(value, dict) def _is_json_array(value: object) -> TypeIs[list[object]]: # guard-ok: trivial isinstance narrowing return isinstance(value, list) def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verifies every value is str return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) def _load_json_object(payload: str | bytes) -> dict[str, object]: """Parse a JSON payload that the caller consumes as an object.""" return json.loads(payload) def _load_json_value(payload: str | bytes) -> object: """Parse a JSON payload whose top-level shape the caller narrows itself.""" return json.loads(payload) def _model_id_from_metadata(litellm_metadata: dict[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 def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None: if task.cancelled(): return exception: Final = task.exception() if exception is not None: verbose_logger.error("%s failed: %s", task_name, exception) _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( { # mutable-ok: immediately frozen by MappingProxyType "server_error": 500, "rate_limit_exceeded": 429, "insufficient_quota": 429, "vector_store_timeout": 504, "invalid_prompt": 400, "invalid_image": 400, "invalid_image_format": 400, "invalid_base64_image": 400, "invalid_image_url": 400, "image_too_large": 400, "image_too_small": 400, "image_parse_error": 400, "image_content_policy_violation": 400, "invalid_image_mode": 400, "image_file_too_large": 400, "unsupported_image_media_type": 400, "empty_image_file": 400, "failed_to_download_image": 400, "image_file_not_found": 400, "invalid_request_error": 400, "context_length_exceeded": 400, "content_policy_violation": 400, "model_not_found": 400, } ) def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") elif error_obj is not None: raw_message = getattr(error_obj, "message", None) raw_type = getattr(error_obj, "type", None) raw_code = getattr(error_obj, "code", None) else: raw_message = None raw_type = None raw_code = None message: Final = str(raw_message) if raw_message is not None else "Response API in-stream error" error_type: Final = raw_type if isinstance(raw_type, str) else None code: Final = raw_code if isinstance(raw_code, str) else None return message, error_type, code def _status_code_for_error_fields(error_type: str | None, error_code: str | None) -> int: fields: Final = tuple(field for field in (error_code, error_type) if field is not None) if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): return 429 return next( (_ERROR_CODE_HTTP_STATUS[field] for field in fields if field in _ERROR_CODE_HTTP_STATUS), 500, ) def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True status_code: Final = getattr(mapped_exception, "status_code", None) return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. This class contains shared logic for both synchronous and asynchronous iterators. """ def __init__( self, response: httpx.Response, model: str, responses_api_provider_config: BaseResponsesAPIConfig | None, logging_obj: LiteLLMLoggingObj, litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, ): self.response = response self.model = model self.logging_obj = logging_obj self.finished = False self.responses_api_provider_config = responses_api_provider_config self.completed_response: ResponsesAPIStreamingResponse | None = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False self._generated_content = "" self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit: bool | None = None self._persist_completed_response_before_logging = True self._stream_created_time: float = time.time() # track request context for hooks self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider self.request_data: dict[str, object] = request_data or {} self.call_type: str | None = call_type # set hidden params for response headers (e.g., x-litellm-model-id) # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base: Final = get_api_base( model=model or "", optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) self._hidden_params: dict[str, object] = { "model_id": _model_id_from_metadata(litellm_metadata), "api_base": _api_base, "custom_llm_provider": custom_llm_provider, } self._hidden_params["additional_headers"] = process_response_headers( self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: return elapsed: Final = time.time() - self._stream_created_time if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS: raise litellm.Timeout( message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)", model=self.model or "", llm_provider=self.custom_llm_provider or "", ) def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: """Process a single chunk of data from the stream""" if not chunk: return None # NOTE: ``SSEDecoder`` already strips the SSE ``data:`` field prefix, so # the value passed in here is the raw field content. Do not re-run # ``_strip_sse_data_from_chunk`` on it — doing so would incorrectly mangle # payloads whose actual JSON value happens to start with ``data:``. # Handle "[DONE]" marker if chunk == STREAM_SSE_DONE_STRING: self.finished = True return None if self.logging_obj.completion_start_time is None: self.logging_obj._update_completion_start_time(completion_start_time=datetime.now()) try: # Parse the JSON chunk parsed_chunk: Final = _load_json_value(chunk) # Format as ResponsesAPIStreamingResponse if _is_json_object(parsed_chunk): if self.responses_api_provider_config is None: raise ValueError("responses_api_provider_config is required to process live streaming chunks") openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response( model=self.model, parsed_chunk=parsed_chunk, logging_obj=self.logging_obj, ) # Only when the SSE JSON carries a response body (delta events do not). # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: response_object: Final[ResponsesAPIResponse | None] = getattr( openai_responses_api_chunk, "response", None ) if response_object is not None: response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, litellm_metadata=self.litellm_metadata, custom_llm_provider=self.custom_llm_provider, ) setattr(openai_responses_api_chunk, "response", response) # Encode container_id on streaming events so proxy/UI follow-ups route correctly _event_type: Final = getattr(openai_responses_api_chunk, "type", None) if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: _delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta _stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): _item: Final = getattr(openai_responses_api_chunk, "item", None) if _item is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_item, custom_llm_provider=self.custom_llm_provider, model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: _annotation: Final = getattr(openai_responses_api_chunk, "annotation", None) if _annotation is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_annotation, custom_llm_provider=self.custom_llm_provider, model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr( openai_responses_api_chunk, "part", None ) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( _part.get("annotations"), self.custom_llm_provider, _stream_model_id, ) else: ResponsesAPIRequestUtils._encode_container_ids_in_annotations( getattr(_part, "annotations", None), self.custom_llm_provider, _stream_model_id, ) # Wrap encrypted_content in streaming events (output_item.added, output_item.done) if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"): openai_types = _get_openai_response_types() event_type: Final = getattr(openai_responses_api_chunk, "type", None) if event_type in ( openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content: Final = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): model_id: Final = _model_id_from_metadata(self.litellm_metadata) if model_id: wrapped_content: Final = ( ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( encrypted_content, model_id ) ) setattr(item, "encrypted_content", wrapped_content) # Store the completed response (also for incomplete/failed so logging still fires) _chunk_type: Final = getattr(openai_responses_api_chunk, "type", None) openai_types = _get_openai_response_types() if openai_responses_api_chunk and _chunk_type in ( openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: response_obj: Final[ResponsesAPIResponse | None] = getattr( openai_responses_api_chunk, "response", None ) if response_obj: usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: Final[float | None] = self.logging_obj._response_cost_calculator( result=response_obj ) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: # Best-effort usage cost annotation should not break stream replay. pass if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() else: self._handle_logging_completed_response() return openai_responses_api_chunk return None except json.JSONDecodeError: # If we can't parse the chunk, continue return None except Exception as e: # Trigger failure hooks before re-raising # This ensures failures are logged even when _process_chunk is called directly self._handle_failure(e) raise def _log_completed_response(self, *, is_async: bool) -> None: if self._completed_response_logged: return self._completed_response_logged = True if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) # Create a copy for logging to avoid modifying the response object that will be returned to the user # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) # to chat completion format (prompt_tokens/completion_tokens) for internal logging # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) logging_response = self.completed_response if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): try: logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) except Exception: # Fallback to original if serialization fails pass end_time: Final = datetime.now() if is_async: asyncio.create_task( self.logging_obj.dispatch_success_handlers( logging_response, start_time=self.start_time, end_time=end_time, cache_hit=self._completed_response_cache_hit, prefer_async_handlers=True, ) ) else: run_async_function( async_function=self.logging_obj.async_success_handler, result=logging_response, start_time=self.start_time, end_time=end_time, cache_hit=self._completed_response_cache_hit, ) executor.submit( self.logging_obj.success_handler, result=logging_response, cache_hit=self._completed_response_cache_hit, start_time=self.start_time, end_time=end_time, ) self._run_post_success_hooks(end_time=end_time) def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" def _handle_logging_failed_response(self): """ Handle logging for RESPONSE_FAILED events by routing to failure handlers. Unlike _handle_logging_completed_response (which calls success handlers), this constructs an exception from the response error and routes to async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ response_obj: Final[ResponsesAPIResponse | None] = ( getattr(self.completed_response, "response", None) if self.completed_response else None ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None self._record_failed_response_usage(response_obj) self._handle_failure(self._map_error_event_exception(error_info)) def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: self.logging_obj.model_call_details["combined_usage_object"] = ( ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage_obj) ) except (TypeError, ValueError) as usage_error: verbose_logger.debug( "could not record usage for failed responses stream: %s", usage_error, ) return self.logging_obj.model_call_details["response_cost"] = ( self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 ) 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 def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) if chunk_type not in ("error", "response.failed"): return error_obj: Final[object] = ( getattr(getattr(result, "response", None), "error", None) if chunk_type == "response.failed" else getattr(result, "error", None) ) mapped_exception: Final = self._map_error_event_exception(error_obj) if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception raise MidStreamFallbackError( message=str(mapped_exception), model=self.model or "", llm_provider=self.custom_llm_provider or "", original_exception=mapped_exception, generated_content=self._generated_content, is_pre_first_chunk=not self._yielded_first_chunk, ) def _get_completed_response_object(self) -> ResponsesAPIResponse | None: openai_types: Final = _get_openai_response_types() completed_response: Final = self.completed_response if isinstance(completed_response, openai_types.ResponsesAPIResponse): return completed_response response_obj: Final = getattr(completed_response, "response", None) if isinstance(response_obj, openai_types.ResponsesAPIResponse): return response_obj return None def _persist_completed_response_to_cache(self, *, is_async: bool) -> None: if self._completed_response_cached: return completed_response: Final = self.completed_response openai_types: Final = _get_openai_response_types() if getattr(completed_response, "type", None) != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: return response_obj: Final = self._get_completed_response_object() if response_obj is None: return caching_handler: Final[_StreamCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return request_kwargs = getattr(caching_handler, "request_kwargs", None) if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) request_cache_key: Final = request_kwargs.pop("cache_key", None) if preset_cache_key is None: preset_cache_key = request_cache_key if request_kwargs.get("metadata") is None: request_kwargs.pop("metadata", None) request_kwargs.pop("custom_llm_provider", None) if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key if not caching_handler._should_store_result_in_cache( # pyright: ignore[reportPrivateUsage] # no public API original_function=caching_handler.original_function, kwargs=request_kwargs, ): return cache: Final = litellm.cache if cache is None: return cached_response: Final = response_obj.model_dump_json() if is_async: from litellm.caching.caching_handler import create_cache_write_task cache_write_task: Final = create_cache_write_task( lambda: cache.async_add_cache( cached_response, dynamic_cache_object=getattr(caching_handler, "dual_cache", None), **request_kwargs, ) ) cache_write_task.add_done_callback( lambda task: _log_background_task_failure( task, task_name="Responses stream cache write", ) ) else: cache.add_cache( cached_response, dynamic_cache_object=getattr(caching_handler, "dual_cache", None), **request_kwargs, ) self._completed_response_cached = True async def _call_post_streaming_deployment_hook( self, chunk: ResponsesAPIStreamingResponse ) -> ResponsesAPIStreamingResponse: """ Allow callbacks to modify streaming chunks before returning (parity with chat). """ try: # Align with chat pipeline: use logging_obj model_call_details + call_type typed_call_type: CallTypes | None = None if self.call_type is not None: try: typed_call_type = CallTypes(self.call_type) except ValueError: typed_call_type = None if typed_call_type is None: try: typed_call_type = CallTypes(getattr(self.logging_obj, "call_type", None)) except Exception: typed_call_type = None request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {}) callbacks: Final = getattr(litellm, "callbacks", None) or [] hooks_ran = False for callback in callbacks: if hasattr(callback, "async_post_call_streaming_deployment_hook"): hooks_ran = True result = await callback.async_post_call_streaming_deployment_hook( request_data=request_data, response_chunk=chunk, call_type=typed_call_type, ) if result is not None: chunk = result if hooks_ran: setattr(chunk, "_post_streaming_hooks_ran", True) return chunk except Exception: return chunk async def call_post_streaming_hooks_for_testing( self, chunk: ResponsesAPIStreamingResponse ) -> ResponsesAPIStreamingResponse: """ Helper to invoke streaming deployment hooks explicitly (used in tests). """ return await self._call_post_streaming_deployment_hook(chunk) def _run_post_success_hooks(self, end_time: datetime): """ Run post-call deployment hooks and update metadata similar to chat pipeline. """ if self.completed_response is None: return request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: if hasattr(self.logging_obj, "model_call_details"): request_payload.update(self.logging_obj.model_call_details) except Exception: pass if "litellm_params" not in request_payload: try: request_payload["litellm_params"] = getattr(self.logging_obj, "model_call_details", {}).get( "litellm_params", {} ) except Exception: request_payload["litellm_params"] = {} try: update_response_metadata( result=self.completed_response, logging_obj=self.logging_obj, model=self.model, kwargs=request_payload, start_time=self.start_time, end_time=end_time, ) except Exception: # Non-blocking pass try: typed_call_type: CallTypes | None = None if self.call_type is not None: try: typed_call_type = CallTypes(self.call_type) except ValueError: typed_call_type = None except Exception: typed_call_type = None if typed_call_type is None: try: typed_call_type = CallTypes.responses except Exception: typed_call_type = None try: # Call synchronously; async hook will be executed via asyncio.run in a new loop run_async_function( async_function=async_post_call_success_deployment_hook, request_data=request_payload, response=self.completed_response, call_type=typed_call_type, ) except Exception: pass def _handle_failure(self, exception: Exception): """ Trigger failure handlers before bubbling the exception. Only calls handlers once even if called multiple times. """ # Prevent double-calling failure handlers if self._failure_handled: return self._failure_handled = True traceback_exception: Final = traceback.format_exc() try: run_async_function( async_function=self.logging_obj.async_failure_handler, exception=exception, traceback_exception=traceback_exception, start_time=self.start_time, end_time=datetime.now(), ) except Exception: pass try: executor.submit( self.logging_obj.failure_handler, exception, traceback_exception, self.start_time, datetime.now(), ) except Exception: pass async def call_post_streaming_hooks_for_testing( iterator: object, chunk: ResponsesAPIStreamingResponse ) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = ( getattr(iterator, "_call_post_streaming_deployment_hook", None) ) if hook_fn is None: return chunk return await hook_fn(chunk) class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): """ Async iterator for processing streaming responses from the Responses API. """ def __init__( self, response: httpx.Response, model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, ): super().__init__( response, model, responses_api_provider_config, logging_obj, litellm_metadata, custom_llm_provider, request_data, call_type, ) self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes()) def __aiter__(self): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: sse = await self.stream_iterator.__anext__() except StopAsyncIteration: self.finished = True raise StopAsyncIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: raise StopAsyncIteration elif result is not None: self._maybe_raise_for_error_event(result) # Await hook directly instead of run_async_function # (which spawns a thread + event loop per call) result = await self._call_post_streaming_deployment_hook( chunk=result, ) self._yielded_first_chunk = True return result # If result is None, continue the loop to get the next chunk except StopAsyncIteration: # Normal end of stream - don't log as failure raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True if self.completed_response is None: self._handle_failure(e) raise raise StopAsyncIteration from e except httpx.HTTPError as e: # Handle HTTP errors self.finished = True self._handle_failure(e) raise e except Exception as e: self.finished = True self._handle_failure(e) raise e def _handle_logging_completed_response(self): """Handle logging for completed responses in async context""" self._log_completed_response(is_async=True) class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): """ Synchronous iterator for processing streaming responses from the Responses API. """ def __init__( self, response: httpx.Response, model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, ): super().__init__( response, model, responses_api_provider_config, logging_obj, litellm_metadata, custom_llm_provider, request_data, call_type, ) self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes()) def __iter__(self): return self def __next__(self): try: self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: sse = next(self.stream_iterator) except StopIteration: self.finished = True raise StopIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: raise StopIteration elif result is not None: self._maybe_raise_for_error_event(result) # Sync path: use run_async_function for the hook result = run_async_function( async_function=self._call_post_streaming_deployment_hook, chunk=result, ) self._yielded_first_chunk = True return result # If result is None, continue the loop to get the next chunk except StopIteration: # Normal end of stream - don't log as failure raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True if self.completed_response is None: self._handle_failure(e) raise raise StopIteration from e except httpx.HTTPError as e: # Handle HTTP errors self.finished = True self._handle_failure(e) raise e except Exception as e: self.finished = True self._handle_failure(e) raise e def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context""" self._log_completed_response(is_async=False) class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): """ Mock iterator—fake a stream by slicing the full response text into 5 char deltas, then emit a completed event. Models like o1-pro don't support streaming, so we fake it. """ CHUNK_SIZE = 5 def __init__( self, response: httpx.Response, model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, ): transformed: Final = responses_api_provider_config.transform_response_api_response( model=model, raw_response=response, logging_obj=logging_obj, ) super().__init__( response=httpx.Response(200), model=model, responses_api_provider_config=None, logging_obj=logging_obj, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, request_data=request_data, call_type=call_type, ) self._set_events_from_response(transformed=transformed, logging_obj=logging_obj) def _set_events_from_response( self, transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=self.CHUNK_SIZE, ) self._idx = 0 self.completed_response = self._events[-1] def __aiter__(self): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt: Final = self._events[self._idx] self._idx += 1 openai_types: Final = _get_openai_response_types() if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt def __iter__(self): return self def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt: Final = self._events[self._idx] self._idx += 1 openai_types: Final = _get_openai_response_types() if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __init__( self, response: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, request_data: dict[str, object] | None = None, call_type: str | None = None, ): BaseResponsesAPIStreamingIterator.__init__( self, response=httpx.Response(200), model=getattr(response, "model", ""), responses_api_provider_config=None, logging_obj=logging_obj, litellm_metadata=None, custom_llm_provider="cached_response", request_data=request_data, call_type=call_type, ) self._completed_response_cache_hit = True self._persist_completed_response_before_logging = False self._events: list[ResponsesAPIStreamingResponse] = [] self._idx = 0 self._set_events_from_response(transformed=response, logging_obj=logging_obj) def _set_events_from_response( self, transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE, ) self._idx = 0 self.completed_response = self._events[-1] def __aiter__(self): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt: Final = self._events[self._idx] self._idx += 1 openai_types: Final = _get_openai_response_types() if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt def __iter__(self): return self def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt: Final = self._events[self._idx] self._idx += 1 openai_types: Final = _get_openai_response_types() if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt @runtime_checkable class _HasModelDump(Protocol): def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ... @runtime_checkable class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... def _dump_response_object(obj: object) -> dict[str, Any]: if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): return obj return {} def _build_response_status_event( event_type: Literal[ "response.created", "response.in_progress", ], transformed: ResponsesAPIResponse, ) -> ResponsesAPIStreamingResponse: openai_types: Final = _get_openai_response_types() in_progress_response: Final = transformed.model_copy( deep=True, update={"status": "in_progress", "output": []}, ) if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: return openai_types.ResponseCreatedEvent(type=event_type, response=in_progress_response) return openai_types.ResponseInProgressEvent(type=event_type, response=in_progress_response) def _build_content_part_done_event( *, item_id: str, output_index: int, content_index: int, part_payload: dict[str, Any], ) -> ResponsesAPIStreamingResponse | None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") part: PART_UNION_TYPES if part_type == "output_text": annotations: Final = [ openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) for annotation in part_payload.get("annotations", []) or [] ] part = openai_types.ContentPartDonePartOutputText( type="output_text", text=str(part_payload.get("text") or ""), annotations=annotations, logprobs=part_payload.get("logprobs"), ) elif part_type == "refusal": part = openai_types.ContentPartDonePartRefusal( type="refusal", refusal=str(part_payload.get("refusal") or ""), ) elif part_type == "reasoning_text": part = openai_types.ContentPartDonePartReasoningText( type="reasoning_text", reasoning=str(part_payload.get("reasoning") or ""), ) else: return None return openai_types.ContentPartDoneEvent( type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE, item_id=item_id, output_index=output_index, content_index=content_index, part=part, ) def _add_text_like_part_events( *, events: list[ResponsesAPIStreamingResponse], item_id: str, output_index: int, content_index: int, part_payload: dict[str, Any], chunk_size: int, ) -> None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") if part_type == "output_text": text: Final = str(part_payload.get("text") or "") for i in range(0, len(text), chunk_size): events.append( openai_types.OutputTextDeltaEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, output_index=output_index, content_index=content_index, delta=text[i : i + chunk_size], ) ) annotations_payload: Final[Sequence[dict[str, object]]] = part_payload.get("annotations", []) or [] for annotation_index, annotation in enumerate(annotations_payload): events.append( openai_types.OutputTextAnnotationAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, item_id=item_id, output_index=output_index, content_index=content_index, annotation_index=annotation_index, annotation=annotation, ) ) events.append( openai_types.OutputTextDoneEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, item_id=item_id, output_index=output_index, content_index=content_index, text=text, ) ) elif part_type == "refusal": refusal: Final = str(part_payload.get("refusal") or "") for i in range(0, len(refusal), chunk_size): events.append( openai_types.RefusalDeltaEvent( type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA, item_id=item_id, output_index=output_index, content_index=content_index, delta=refusal[i : i + chunk_size], ) ) events.append( openai_types.RefusalDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE, item_id=item_id, output_index=output_index, content_index=content_index, refusal=refusal, ) ) def _build_synthetic_response_events( *, transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = _get_openai_response_types() if litellm.include_cost_in_streaming_usage and logging_obj is not None: usage_obj: Final = transformed.usage if hasattr(transformed, "usage") else None if usage_obj is not None: try: cost: Final[float | None] = logging_obj._response_cost_calculator(result=transformed) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: pass events: Final[list[ResponsesAPIStreamingResponse]] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] sequence_number = 0 output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") events.append( openai_types.OutputItemAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) if item_type == "message": content_parts: Sequence[object] = output_item_payload.get("content", []) or [] for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=item_id, output_index=output_index, content_index=content_index, part=openai_types.BaseLiteLLMOpenAIResponseObject(**part_payload), ) ) _add_text_like_part_events( events=events, item_id=item_id, output_index=output_index, content_index=content_index, part_payload=part_payload, chunk_size=chunk_size, ) done_event = _build_content_part_done_event( item_id=item_id, output_index=output_index, content_index=content_index, part_payload=part_payload, ) if done_event is not None: events.append(done_event) elif item_type == "function_call": arguments = str(output_item_payload.get("arguments") or "") for i in range(0, len(arguments), chunk_size): events.append( openai_types.FunctionCallArgumentsDeltaEvent( type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=item_id, output_index=output_index, delta=arguments[i : i + chunk_size], ) ) events.append( openai_types.FunctionCallArgumentsDoneEvent( type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, item_id=item_id, output_index=output_index, arguments=arguments, ) ) elif item_type == "reasoning": summaries: Sequence[object] = output_item_payload.get("summary", []) or [] for summary_index, summary in enumerate(summaries): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): events.append( openai_types.ReasoningSummaryTextDeltaEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, item_id=item_id, output_index=output_index, summary_index=summary_index, delta=summary_text[i : i + chunk_size], ) ) sequence_number += 1 events.append( openai_types.ReasoningSummaryTextDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, item_id=item_id, output_index=output_index, sequence_number=sequence_number, summary_index=summary_index, text=summary_text, ) ) sequence_number += 1 events.append( openai_types.ReasoningSummaryPartDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, item_id=item_id, output_index=output_index, sequence_number=sequence_number, summary_index=summary_index, part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload), ) ) sequence_number += 1 events.append( openai_types.OutputItemDoneEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, sequence_number=sequence_number, item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) events.append( openai_types.ResponseCompletedEvent( type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=transformed, ) ) return events # --------------------------------------------------------------------------- # WebSocket mode streaming (bidirectional forwarding) # --------------------------------------------------------------------------- from litellm._logging import verbose_logger # Conservative per-frame output-token floor used when a response.create # frame omits max_output_tokens, so a project OTPM quota can't be bypassed # by simply never declaring an output cap. _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR: Final = 1024 # Rough chars-per-token ratio for estimating a frame's input tokens without # resolving a real per-model tokenizer, matching the conservative estimate # the proxy's own rate limiter uses for the same purpose. _FRAME_CHARS_PER_TOKEN_ESTIMATE: Final = 4 def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple[int, int | None]: """Extract a rough input-token count and any explicit max_output_tokens from a ``response.create`` frame, handling both wire shapes: flat: {"type": "response.create", "input": ..., "max_output_tokens": ...} nested: {"type": "response.create", "response": {"input": ..., "max_output_tokens": ...}} """ nested: Final = msg_obj.get("response") params: Final[Mapping[str, object]] = ( nested if _is_json_object(nested) and nested else MappingProxyType( # mutable-ok: immediately frozen filtered frame {k: v for k, v in msg_obj.items() if k != "type"} ) ) text_parts: Final[list[str]] = [] # mutable-ok: local accumulator built in one pass, not shared pending: Final[list[object]] = [ # mutable-ok: explicit worklist avoids recursion params.get("input"), params.get("instructions"), ] while pending: value = pending.pop() if isinstance(value, str): text_parts.append(value) elif _is_json_array(value): for item in value: if isinstance(item, str): text_parts.append(item) elif _is_json_object(item): pending.append(item.get("content")) pending.append(item.get("text")) total_chars: Final = sum(len(part) for part in text_parts) estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0 max_output_tokens: Final = params.get("max_output_tokens") return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None async def _enforce_frame_project_quota( quota_callbacks: Sequence[ProjectQuotaCallback], user_api_key_dict: UserAPIKeyAuth | None, model: str | None, raw_message: str, ) -> None: """Charge one response.create frame's estimated tokens against every registered project ITPM/OTPM quota callback, in isolation from PII masking / logging so a malformed frame still reaches those callbacks.""" if not quota_callbacks: return try: msg_obj: Final = _load_json_value(raw_message) except (json.JSONDecodeError, TypeError): return if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create": return estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj) estimated_output_tokens: Final = ( explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR ) for callback in quota_callbacks: await callback.enforce_project_io_token_quota_for_frame( user_api_key_dict=user_api_key_dict, requested_model=model, estimated_input_tokens=estimated_input_tokens, estimated_output_tokens=estimated_output_tokens, ) RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", "response.failed", "response.incomplete", "error", ] RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"}) class ResponsesWebSocketStreaming: """ Manages bidirectional WebSocket forwarding for the Responses API WebSocket mode (wss://.../v1/responses). Unlike the Realtime API, the Responses API WebSocket mode: - Uses response.create as the client-to-server event - Streams back the same events as the HTTP streaming Responses API - Supports previous_response_id for incremental continuation - Supports generate: false for warmup - One response at a time per connection (sequential, no multiplexing) """ def __init__( self, websocket: ResponsesClientWebSocket, backend_ws: ResponsesBackendWebSocket, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, first_message: str | None = None, guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.request_data: dict[str, object] = request_data or {} self.messages: list[dict[str, object]] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message self.guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model def _should_store_event(self, event_obj: Mapping[str, object]) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: str | bytes | dict[str, object]) -> None: if isinstance(event, bytes): event = event.decode("utf-8") if isinstance(event, str): try: event_obj = _load_json_object(event) except (json.JSONDecodeError, TypeError): return else: event_obj = event if self._should_store_event(event_obj): self.messages.append(event_obj) def _collect_input_from_client_event(self, message: object) -> None: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): msg_obj = _load_json_object(message) elif _is_json_object(message): msg_obj = message else: return if msg_obj.get("type") != "response.create": return input_items: Final = msg_obj.get("input", []) if isinstance(input_items, str): self.input_messages.append({"role": "user", "content": input_items}) return if _is_json_array(input_items): for item in input_items: if not _is_json_object(item): continue if item.get("type") == "message" and item.get("role") == "user": content = item.get("content", []) if isinstance(content, str): self.input_messages.append({"role": "user", "content": content}) elif _is_json_array(content): for c in content: if _is_json_object(c) and c.get("type") == "input_text": text = c.get("text", "") if text: self.input_messages.append({"role": "user", "content": text}) except (json.JSONDecodeError, AttributeError, TypeError): pass def _store_input(self, message: object) -> None: self._collect_input_from_client_event(message) if self.logging_obj: self.logging_obj.pre_call(input=message, api_key="") 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: asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) async def backend_to_client(self) -> None: """Forward events from backend WebSocket to the client.""" import websockets try: while True: try: raw_response = await self.backend_ws.recv(decode=False) except TypeError: raw_response = await self.backend_ws.recv() if isinstance(raw_response, bytes): response_str = raw_response.decode("utf-8") else: response_str = raw_response # When apply_to_output masking is active, suppress delta events # and the text-bearing "done" events. Per-fragment Presidio # cannot reliably catch PII spanning multiple delta chunks (e.g. # "alice@" + "example.com"), and the done events carry the full # output text that response.completed already delivers in # fully-masked form; forwarding them would leak unmasked PII # before response.completed arrives. The client receives only the # masked response.completed. if self.output_guardrail_callbacks: try: _evt_payload: Mapping[str, object] = _load_json_object(response_str) _evt_type = _evt_payload.get("type") except (json.JSONDecodeError, TypeError): _evt_type = None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: continue unmasked_str = self._unmask_response_event(response_str) output_masked_str = await self._mask_response_completed(unmasked_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) await self.websocket.send_text(output_masked_str) except websockets.exceptions.ConnectionClosed as e: verbose_logger.debug("Responses WS backend connection closed: %s", e) except Exception as e: verbose_logger.exception("Error in responses WS backend_to_client: %s", e) finally: await self._log_messages() def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. Handles both shapes: flat: ``{"type": "response.create", "model": "...", ...}`` nested: ``{"type": "response.create", "response": {"model": "...", ...}}`` Returns True if the object was modified. """ if not self.authorized_model: return False modified = False nested: Final = msg_obj.get("response") if _is_json_object(nested): if nested.get("model") != self.authorized_model: nested["model"] = self.authorized_model modified = True if "model" in msg_obj and msg_obj["model"] != self.authorized_model: msg_obj["model"] = self.authorized_model modified = True elif msg_obj.get("model") != self.authorized_model: msg_obj["model"] = self.authorized_model modified = True return modified async def _mask_response_create(self, message: str) -> str: """ Enforce the authorized model and apply Presidio PII masking to a ``response.create`` message before it is forwarded to the upstream provider. - Overwrites any ``model`` field with the connection-authorized model to prevent deployment-substitution attacks (always applied). - Walks the ``input`` and ``instructions`` fields, calls ``check_pii`` on every text block, and stores the resulting ``pii_tokens`` map in ``self.request_data["metadata"]`` for later unmasking. Non-``response.create`` messages are returned unchanged. """ try: msg_obj: Final = _load_json_object(message) except (json.JSONDecodeError, TypeError): return message if msg_obj.get("type") != "response.create": return message # Always enforce the authorized model, even when PII masking is off. model_modified: Final = self._enforce_authorized_model(msg_obj) if not self.guardrail_callbacks: return json.dumps(msg_obj) if model_modified else message if "metadata" not in self.request_data: self.request_data["metadata"] = {} modified = model_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) # response.create carries client text in two shapes: # flat: {"type": "response.create", "input": ..., "instructions": ...} # nested: {"type": "response.create", "response": {"input": ..., "instructions": ...}} # Mask "input" and "instructions" in both shapes so PII is never # forwarded unmasked regardless of where the client places it. nested_candidate = msg_obj.get("response") nested_response = nested_candidate if _is_json_object(nested_candidate) else None text_containers: list[tuple[dict[str, object], str]] = [] for container in (msg_obj, nested_response): if container is None: continue if "input" in container: text_containers.append((container, "input")) if isinstance(container.get("instructions"), str): text_containers.append((container, "instructions")) for container, key in text_containers: field_value = container[key] if isinstance(field_value, str): container[key] = await cb.check_pii( text=field_value, output_parse_pii=True, presidio_config=presidio_config, request_data=self.request_data, ) modified = True elif _is_json_array(field_value): for item in field_value: if not _is_json_object(item): continue for item_field in ("content", "output"): value = item.get(item_field) if isinstance(value, str): item[item_field] = await cb.check_pii( text=value, output_parse_pii=True, presidio_config=presidio_config, request_data=self.request_data, ) modified = True elif _is_json_array(value): for block in value: if not _is_json_object(block): continue block_text = block.get("text") if block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES and isinstance( block_text, str ): block["text"] = await cb.check_pii( text=block_text, output_parse_pii=True, presidio_config=presidio_config, request_data=self.request_data, ) modified = True return json.dumps(msg_obj) if modified else message # Delta event types whose ``delta`` field may contain PII tokens. _DELTA_EVENT_TYPES = frozenset( { "response.output_text.delta", "response.reasoning_summary_text.delta", "response.refusal.delta", "response.function_call_arguments.delta", } ) # Terminal events that carry the full output text or tool-call arguments # already delivered by ``response.completed``. Suppressed when output masking # is active so the unmasked copy never reaches the client before the masked # completed event. _OUTPUT_DONE_EVENT_TYPES = frozenset( { "response.output_text.done", "response.content_part.done", "response.output_item.done", "response.function_call_arguments.done", "response.reasoning_summary_text.done", "response.reasoning_summary_part.done", } ) def _unmask_response_event(self, response_str: str) -> str: """ Apply Presidio PII unmasking to backend events before forwarding to the client. Handles two shapes: - ``response.completed``: walks ``response.output[*].content[*].text`` - streaming delta events (``response.output_text.delta``, etc.): replaces tokens in the ``delta`` field Uses the ``pii_tokens`` map stored during ``_mask_response_create`` to replace every token (e.g. ````) with the original value. Events with no stored tokens are returned unchanged. """ if not self.guardrail_callbacks: return response_str metadata: Final = self.request_data.get("metadata") raw_pii_tokens: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None pii_tokens: Final[Mapping[str, str]] = raw_pii_tokens if _is_str_mapping(raw_pii_tokens) else {} if not pii_tokens: return response_str try: evt_obj: Final = _load_json_object(response_str) except (json.JSONDecodeError, TypeError): return response_str cb: Final = self.guardrail_callbacks[0] event_type: Final = evt_obj.get("type") if event_type == "response.completed": modified = False response_obj: Final = evt_obj.get("response") if not _is_json_object(response_obj): return response_str output_items: Final = response_obj.get("output") for output_item in output_items if _is_json_array(output_items) else []: if not _is_json_object(output_item): continue content = output_item.get("content") if not _is_json_array(content): continue for content_block in content: if not _is_json_object(content_block): continue text = content_block.get("text") if isinstance(text, str): unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker text, pii_tokens ) if unmasked != text: content_block["text"] = unmasked modified = True return json.dumps(evt_obj) if modified else response_str if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker delta, pii_tokens ) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) return response_str async def _mask_response_completed(self, response_str: str) -> str: """ Apply Presidio output masking (apply_to_output=True) to the ``response.completed`` event before it is forwarded to the client. Walks ``response.output[*].content[*].text`` and masks every text block, as well as ``response.output[*].arguments`` on function-call items and ``response.output[*].summary[*].text`` on reasoning items. Delta and ``*.done`` events are suppressed upstream in ``backend_to_client`` when output masking is active, so only the authoritative full-output view reaches this method; events of other types are returned unchanged. """ if not self.output_guardrail_callbacks: return response_str try: evt_obj: Final[Mapping[str, object]] = _load_json_object(response_str) except (json.JSONDecodeError, TypeError): return response_str if evt_obj.get("type") != "response.completed": return response_str modified = False for cb in self.output_guardrail_callbacks: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) response_obj = evt_obj.get("response") if not _is_json_object(response_obj): continue output_items = response_obj.get("output") for output_item in output_items if _is_json_array(output_items) else []: if not _is_json_object(output_item): continue arguments = output_item.get("arguments") if isinstance(arguments, str): masked_args = await cb.check_pii( text=arguments, output_parse_pii=False, presidio_config=presidio_config, request_data=self.request_data, ) if masked_args != arguments: output_item["arguments"] = masked_args modified = True summary = output_item.get("summary") if _is_json_array(summary): for summary_block in summary: if not _is_json_object(summary_block): continue summary_text = summary_block.get("text") if isinstance(summary_text, str): masked_summary = await cb.check_pii( text=summary_text, output_parse_pii=False, presidio_config=presidio_config, request_data=self.request_data, ) if masked_summary != summary_text: summary_block["text"] = masked_summary modified = True content = output_item.get("content") if not _is_json_array(content): continue for content_block in content: if not _is_json_object(content_block): continue text = content_block.get("text") if isinstance(text, str): masked = await cb.check_pii( text=text, output_parse_pii=False, presidio_config=presidio_config, request_data=self.request_data, ) if masked != text: content_block["text"] = masked modified = True return json.dumps(evt_obj) if modified else response_str async def _enforce_or_reject_frame(self, message: str) -> bool: """Run the per-frame project quota check. On rejection, sends an ``error`` event to the client and reports that the frame must be dropped instead of forwarded, so the connection stays open for the client to retry once the window resets. """ try: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) except RateLimitError as e: try: await self.websocket.send_text( json.dumps( # mutable-ok: WebSocket wire payload requires JSON objects { # mutable-ok: WebSocket wire payload requires JSON objects "type": "error", "error": { # mutable-ok: nested WebSocket error object "type": "rate_limit_exceeded", "message": str(e), }, } ) ) except Exception: # noqa: BLE001, S110 # client may already be gone pass return False return True async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) self._store_input(masked_first) self._store_event(masked_first) await self.backend_ws.send(masked_first) while True: message = await self.websocket.receive_text() if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) self._store_input(masked) self._store_event(masked) await self.backend_ws.send(masked) 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.""" forward_task: Final = asyncio.create_task(self.backend_to_client()) try: await self.client_to_backend() except Exception: pass finally: if not forward_task.done(): forward_task.cancel() try: await forward_task except asyncio.CancelledError: pass try: await self.backend_ws.close() except Exception: pass # --------------------------------------------------------------------------- # Managed WebSocket mode (HTTP-backed, provider-agnostic) # --------------------------------------------------------------------------- _RESPONSE_CREATE_PARAMS: Final[frozenset[str]] = ( _get_openai_response_types().ResponsesAPIRequestParams.__required_keys__ | _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__ ) _MANAGED_WS_SKIP_KWARGS: Final[frozenset[str]] = frozenset( { "litellm_logging_obj", "litellm_call_id", "aresponses", "_aresponses_websocket", "user_api_key_dict", } ) _WARMUP_RESPONSE_ID_PREFIX: Final = "resp_warmup_" class ManagedResponsesWebSocketHandler: """ Handles Responses API WebSocket mode for providers that do not expose a native ``wss://`` responses endpoint. Instead of proxying to a provider WebSocket, this handler: - Listens for ``response.create`` events from the client - Makes HTTP streaming calls via ``litellm.aresponses(stream=True)`` - Serialises and forwards every streaming event back over the WebSocket - Supports ``previous_response_id`` for multi-turn conversations via in-memory session tracking (avoids async DB-write timing issues) - Supports sequential requests over a single persistent connection This makes every provider that LiteLLM can reach over HTTP available on the WebSocket transport without any provider-specific changes. """ def __init__( self, websocket: ResponsesClientWebSocket, model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, litellm_metadata: dict[str, Any] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.litellm_metadata: dict[str, Any] = litellm_metadata or {} self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( "deployment_model_name" ) self.api_key = api_key self.api_base = api_base self.timeout = timeout self.custom_llm_provider = custom_llm_provider self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. # Keyed by the DECODED (pre-encoding) response ID from response.completed. # This avoids the async DB-write race condition where spend logs haven't # been committed yet when the next response.create arrives. self._session_history: dict[str, list[dict[str, object]]] = {} # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @staticmethod def _serialize_chunk(chunk: Any) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if isinstance(chunk, _HasModelDumpJson): return chunk.model_dump_json(exclude_none=True) if isinstance(chunk, _HasModelDump): return json.dumps(chunk.model_dump(exclude_none=True), default=str) if _is_json_object(chunk): return json.dumps(chunk, default=str) return json.dumps(str(chunk)) except Exception as exc: verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc) return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: try: await self.websocket.send_text( json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) ) except Exception: pass def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]: """ Return accumulated message history for *previous_response_id*. The key is the *decoded* response ID (the raw provider response ID before LiteLLM base64-encodes it into the ``resp_...`` format). """ decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) raw_id: Final = decoded.get("response_id", previous_response_id) return list(self._session_history.get(raw_id, [])) def _store_history(self, response_id: str, messages: list[dict[str, object]]) -> None: """ Store the complete accumulated message history for *response_id*. Replaces any prior value — callers are responsible for passing the full history (prior turns + current input + new output). """ self._session_history[response_id] = messages @staticmethod def _extract_response_id(completed_event: dict[str, object]) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. """ resp_obj: Final = completed_event.get("response", {}) raw_id: Final = resp_obj.get("id") if _is_json_object(resp_obj) else None encoded_id: Final[str | None] = raw_id if isinstance(raw_id, str) else None if not encoded_id: return None decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) return decoded.get("response_id", encoded_id) @staticmethod def _extract_output_messages( completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into Responses API message dicts suitable for the next turn's ``input``. """ resp_obj: Final = completed_event.get("response", {}) if not isinstance(resp_obj, dict): return [] messages: Final[list[dict[str, object]]] = [] for item in resp_obj.get("output", []) or []: if not isinstance(item, dict): continue item_type = item.get("type") role = item.get("role", "assistant") if item_type == "message": content_parts = item.get("content") or [] text_parts = [ p.get("text", "") for p in content_parts if isinstance(p, dict) and p.get("type") in ("output_text", "text") ] text = "".join(text_parts) if text: messages.append( { "type": "message", "role": role, "content": [{"type": "output_text", "text": text}], } ) elif item_type == "function_call": messages.append(item) return messages @staticmethod def _input_to_messages(input_val: object) -> list[dict[str, object]]: """ Normalise the ``input`` field of a ``response.create`` event to a list of Responses API message dicts. """ if isinstance(input_val, str): return [ { "type": "message", "role": "user", "content": [{"type": "input_text", "text": input_val}], } ] if _is_json_array(input_val): return [item for item in input_val if _is_json_object(item)] return [] # ------------------------------------------------------------------ # _process_response_create sub-methods # ------------------------------------------------------------------ async def _parse_message(self, raw_message: str) -> dict[str, object] | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" try: msg_obj: Final = _load_json_object(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None if msg_obj.get("type") != "response.create": # Silently ignore non-response.create messages (e.g. warmup pings) return None return msg_obj @staticmethod def _is_warmup_frame(msg_obj: dict[str, object]) -> bool: """Return True for a response.create whose generate flag is false.""" nested: Final = msg_obj.get("response") source: Final = nested if _is_json_object(nested) and nested else msg_obj return source.get("generate") is False @staticmethod def _is_warmup_response_id(response_id: str | None) -> bool: """Return True for synthetic warmup IDs that only exist on this connection.""" if not response_id: return False decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id) raw_id: Final = decoded.get("response_id", response_id) return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]: nested: Final = msg_obj.get("response") if _is_json_object(nested) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]: """Build a minimal completed Responses API object for a warmup ack.""" source: Final = self._warmup_source_params(msg_obj) wire_model: Final = source.get("model") or self.model_group or self.model return { "id": f"{_WARMUP_RESPONSE_ID_PREFIX}{uuid.uuid4().hex}", "object": "response", "created_at": int(time.time()), "status": "completed", "model": wire_model, "output": [], "usage": { "input_tokens": 0, "output_tokens": 0, "total_tokens": 0, }, } async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None: """ Acknowledge a generate=false prewarm without calling the provider. Codex blocks on the warmup turn until it receives response.created and response.completed over the WebSocket. Managed HTTP providers cannot honor an empty-input warmup, so we synthesize the completion locally. """ response: Final = self._build_warmup_response(msg_obj) for event_type, status in ( ("response.created", "in_progress"), ("response.completed", "completed"), ): event = { "type": event_type, "response": {**response, "status": status}, } serialized = self._serialize_chunk(event) if serialized is None: continue await self.websocket.send_text(serialized) @staticmethod def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested: Final = msg_obj.get("response") response_params: Final[dict[str, object]] = ( nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { param: response_params[param] for param in _RESPONSE_CREATE_PARAMS if param in response_params and response_params[param] is not None } def _apply_history( self, call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], ) -> None: """Prepend in-memory turn history, or fall back to DB-based reconstruction.""" if not previous_response_id: return if self._is_warmup_response_id(previous_response_id): verbose_logger.debug( "ManagedResponsesWS: ignoring synthetic warmup previous_response_id=%s", previous_response_id, ) return if prior_history: call_kwargs["input"] = prior_history + current_messages verbose_logger.debug( "ManagedResponsesWS: prepended %d history messages for previous_response_id=%s", len(prior_history), previous_response_id, ) else: verbose_logger.debug( "ManagedResponsesWS: no in-memory history for previous_response_id=%s; " "falling back to DB-based session reconstruction", previous_response_id, ) # Fall back to DB-based session reconstruction (may work for # cross-connection multi-turn when spend logs are committed) call_kwargs["previous_response_id"] = previous_response_id @staticmethod def _resolve_provider(model: str | None) -> str | None: """Resolve the LLM provider for a model string, or None if unresolvable.""" if not model: return None try: from litellm import get_llm_provider _, provider, _, _ = get_llm_provider(model=model) return provider except Exception: return None def _same_provider(self, model: str | None) -> bool: """Return True if model uses the same LLM provider as the connection model.""" if model is None or model == self.model: return True event_provider: Final = self._resolve_provider(model) if event_provider is None: return False return event_provider == self._connection_provider def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key if self.api_base is not None: call_kwargs["api_base"] = self.api_base if self.timeout is not None: call_kwargs["timeout"] = self.timeout # Only force connection-level custom_llm_provider when the per-event model # uses the same provider as the connection model. If the provider differs # (e.g., connection is vertex_ai but event says openai/gpt-4), let litellm # re-resolve from the model string. Same-provider model variants (e.g., # vertex_ai/gemini-2.0 -> vertex_ai/gemini-1.5) still inherit the provider. if self.custom_llm_provider is not None and self._same_provider(model): call_kwargs["custom_llm_provider"] = self.custom_llm_provider if self.litellm_metadata: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) @staticmethod def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} if not isinstance(proxy_server_request, dict): return body: Final = dict(proxy_server_request.get("body") or {}) body["input"] = call_kwargs.get("input") body["store"] = call_kwargs.get("store") body["model"] = model for k in ("tools", "tool_choice", "instructions", "metadata"): if k in call_kwargs and call_kwargs[k] is not None: body[k] = call_kwargs[k] proxy_server_request = {**proxy_server_request, "body": body} if "litellm_metadata" not in call_kwargs: call_kwargs["litellm_metadata"] = {} call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. Captures the ``response.completed`` event type from the chunk object directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ completed_event: dict[str, object] | None = ( None # rebind-ok: captures the completed event once the stream yields it ) stream_response: Final = await litellm.aresponses(model=model, **call_kwargs) async for chunk in stream_response: if chunk is None: continue # Read type from the object before serializing to avoid double JSON parse chunk_type = getattr(chunk, "type", None) or (chunk.get("type") if isinstance(chunk, dict) else None) serialized = self._serialize_chunk(chunk) if serialized is None: continue if chunk_type == "response.completed" and completed_event is None: try: completed_event = _load_json_object(serialized) except Exception: pass try: await self.websocket.send_text(serialized) except Exception as send_exc: verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc) return completed_event # Client disconnected return completed_event def _save_turn_history( self, completed_event: dict[str, object] | None, prior_history: list[dict[str, object]], current_messages: list[dict[str, object]], ) -> None: """Store this turn in in-memory history for future previous_response_id lookups.""" if completed_event is None: return new_response_id: Final = self._extract_response_id(completed_event) if not new_response_id: return output_msgs: Final = self._extract_output_messages(completed_event) all_messages: Final = prior_history + current_messages + output_msgs self._store_history(new_response_id, all_messages) verbose_logger.debug( "ManagedResponsesWS: stored %d messages for response_id=%s", len(all_messages), new_response_id, ) # ------------------------------------------------------------------ # Core request handler # ------------------------------------------------------------------ async def _process_response_create(self, raw_message: str) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, and forward every streaming event to the client. Multi-turn support via in-memory session history ------------------------------------------------ When ``previous_response_id`` is present in the event: 1. Look up the accumulated message history in ``self._session_history`` (keyed by the decoded provider response ID). 2. Prepend those messages to the current ``input`` so the model has full conversation context. 3. After the stream completes, extract the new response ID and output messages from ``response.completed`` and store them in ``self._session_history`` for the next turn. This in-memory approach avoids the async DB-write race condition that occurs when spend logs haven't been committed by the time the second ``response.create`` arrives over the same WebSocket connection. """ msg_obj: Final = await self._parse_message(raw_message) if msg_obj is None: return # generate=false is a prompt-cache warmup hint (sent by codex prewarm). # Native provider sockets handle it server-side, but there is no HTTP # equivalent and the frame carries empty input. Managed providers must # synthesize a completion so clients like Codex can proceed. if self._is_warmup_frame(msg_obj): try: await self._send_warmup_ack(msg_obj) except Exception as exc: verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return try: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message ) except RateLimitError as e: await self._send_error(str(e), error_type="rate_limit_exceeded") return call_kwargs: Final = self._build_base_call_kwargs(msg_obj) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. requested_model: Final[str | None] = call_kwargs.pop("model", None) model: Final[str] = ( self.model if requested_model is None or requested_model == self.model_group else requested_model ) previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history prior_history: Final = self._get_history_messages(previous_response_id) if previous_response_id else [] self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) self._inject_credentials(call_kwargs, model=model) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) try: completed_event: Final = await self._stream_and_forward(model, call_kwargs) except Exception as exc: verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc) await self._send_error(str(exc)) return self._save_turn_history(completed_event, prior_history, current_messages) # ------------------------------------------------------------------ # Main entry point # ------------------------------------------------------------------ async def run(self) -> None: """ Main loop: accept ``response.create`` events sequentially and handle each one before waiting for the next message. """ try: if self.first_message is not None: await self._process_response_create(self.first_message) while True: try: message = await self.websocket.receive_text() except Exception as exc: verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc) break await self._process_response_create(message) except Exception as exc: verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) await self._send_error(f"Internal server error: {exc}")