diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index ca9293325ff..92ce5b49285 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -185,6 +185,40 @@ def get_complete_url( default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION, ) + def supports_native_websocket(self) -> bool: + return True + + def get_websocket_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + """ + Azure Responses WebSocket endpoint is at /openai/v1/responses with no + api-version query param. Auth is via Authorization header, model is sent + in the response.create body — not the URL. + """ + if api_base is None: + raise ValueError("api_base is required for Azure WebSocket") + + parsed_url = httpx.URL(api_base) + path = parsed_url.path.rstrip("/") + # Strip existing /openai/responses path if the api_base already contains it + for suffix in ("/openai/v1/responses", "/openai/responses"): + if path.endswith(suffix): + path = path[: -len(suffix)] + break + scheme = "wss" if parsed_url.scheme == "https" else "ws" + return str( + parsed_url.copy_with( + scheme=scheme, path=f"{path}/openai/v1/responses", query=None + ) + ) + + def model_in_websocket_url(self) -> bool: + # Azure sends the model in the response.create body, not the URL + return False + ######################################################### ########## DELETE RESPONSE API TRANSFORMATION ############## ######################################################### diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 407d5ad8146..c61ce52b530 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -258,6 +258,31 @@ def supports_native_websocket(self) -> bool: """ return False + def get_websocket_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + """ + Return the wss:// URL for the provider's native Responses WebSocket endpoint. + + Defaults to converting the HTTP URL from get_complete_url. Providers whose + WebSocket path differs from their HTTP path (e.g. Azure uses + /openai/v1/responses without api-version) should override this. + """ + http_url = self.get_complete_url( + api_base=api_base, litellm_params=litellm_params + ) + return http_url.replace("https://", "wss://").replace("http://", "ws://") + + def model_in_websocket_url(self) -> bool: + """ + Return True if the model should be appended as a ?model= query param to + the WebSocket URL. Providers that identify the model via the request body + (e.g. Azure Responses API) should override this to return False. + """ + return True + ######################################################### ########## CANCEL RESPONSE API TRANSFORMATION ########## ######################################################### diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 25424feaeb4..ca61a430abe 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5575,7 +5575,7 @@ async def async_realtime_calls_handler( ) raise - async def async_responses_websocket( + async def async_responses_websocket( # noqa: PLR0915 self, model: str, websocket: Any, @@ -5628,7 +5628,11 @@ async def async_responses_websocket( import websockets from websockets.asyncio.client import ClientConnection - litellm_params = GenericLiteLLMParams() + litellm_params = GenericLiteLLMParams( + api_base=api_base, + api_key=api_key, + **kwargs, + ) headers = responses_api_provider_config.validate_environment( headers={}, model=model, @@ -5637,21 +5641,21 @@ async def async_responses_websocket( if api_key: headers["Authorization"] = f"Bearer {api_key}" - http_url = responses_api_provider_config.get_complete_url( + ws_url = responses_api_provider_config.get_websocket_url( api_base=api_base, - litellm_params={}, + litellm_params=dict(litellm_params), ) - ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") - # OpenAI's WebSocket responses endpoint requires ?model= in the URL, - # matching the Realtime API convention (wss://.../v1/realtime?model=...). - # Use urllib.parse so existing query params (e.g. api-version) are preserved. - _parsed = urlparse(ws_url) - _qs = parse_qs(_parsed.query) - if "model" not in _qs: - _qs["model"] = [model] - ws_url = urlunparse( - _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) - ) + # Some providers (e.g. OpenAI) require ?model= in the WebSocket URL. + # Providers that send the model in the request body (e.g. Azure) set + # model_in_websocket_url() to False to suppress this append. + if responses_api_provider_config.model_in_websocket_url(): + _parsed = urlparse(ws_url) + _qs = parse_qs(_parsed.query) + if "model" not in _qs: + _qs["model"] = [model] + ws_url = urlunparse( + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) + ) try: ssl_context = get_shared_realtime_ssl_context() @@ -5679,6 +5683,41 @@ async def async_responses_websocket( _request_data: Dict[str, Any] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata + + _ws_guardrail_callbacks: list = [] + _ws_output_guardrail_callbacks: list = [] + try: + import litellm as _litellm + + # Use duck-typing so any guardrail that exposes the PII + # masking interface works, not just _OPTIONAL_PresidioPIIMasking. + # This avoids a layering violation (SDK importing from proxy). + _ws_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable( + getattr(cb, "get_presidio_settings_from_request_data", None) + ) + and callable(getattr(cb, "_unmask_pii_text", None)) + and getattr(cb, "output_parse_pii", False) + ] + _ws_output_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable( + getattr(cb, "get_presidio_settings_from_request_data", None) + ) + and getattr(cb, "apply_to_output", False) + ] + except Exception as _guardrail_exc: + verbose_logger.warning( + "Responses WebSocket: failed to collect guardrail " + "callbacks — PII masking will be skipped. Error: %s", + _guardrail_exc, + ) + streaming = ResponsesWebSocketStreaming( websocket=websocket, backend_ws=cast(ClientConnection, backend_ws), @@ -5686,6 +5725,9 @@ async def async_responses_websocket( user_api_key_dict=user_api_key_dict, request_data=_request_data, first_message=first_message, + guardrail_callbacks=_ws_guardrail_callbacks, + output_guardrail_callbacks=_ws_output_guardrail_callbacks, + authorized_model=model, ) await streaming.bidirectional_forward() diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index f7dd68aec55..b5319797cc6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -35,7 +35,6 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( - LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -479,90 +478,137 @@ async def process_output_streaming_response( ) -> List[Any]: """ Process output streaming response by applying guardrails to text content. + + Mirrors the Chat Completions handler pattern: extract text from the final + chunk, apply the guardrail, then write the result back in-place so the + caller sees the modified content (e.g. PII tokens replaced). + + For ``response.completed`` events (the normal end-of-stream signal) we + use the same per-item extraction + task-mapping approach as + ``process_output_response`` so that unmasking / blocking works correctly + for every output item. """ + if not responses_so_far: + return responses_so_far final_chunk = responses_so_far[-1] + # Accept both plain dicts and Pydantic models (BaseLiteLLMOpenAIResponseObject + # exposes a .get() shim, so all the .get() calls below work for both). + if not (isinstance(final_chunk, dict) or hasattr(final_chunk, "get")): + return responses_so_far + + # ------------------------------------------------------------------ # + # Case 1: response.completed — full response is available in the # + # final chunk; iterate output items, apply guardrail, write back. # + # ------------------------------------------------------------------ # + if final_chunk.get("type") == "response.completed": + response_obj = final_chunk.get("response") or {} + if not hasattr(response_obj, "get"): + return responses_so_far + outputs: List[Any] = response_obj.get("output") or [] + + texts_to_check: List[str] = [] + tool_calls_to_check: List[ChatCompletionToolCallChunk] = [] + task_mappings: List[Tuple[int, int]] = [] + + for output_idx, output_item in enumerate(outputs): + self._extract_output_text_and_images( + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts_to_check, + images_to_check=[], + task_mappings=task_mappings, + tool_calls_to_check=tool_calls_to_check, + ) + + if texts_to_check or tool_calls_to_check: + if request_data is None: + request_data = {} + if "response" not in request_data: + request_data["response"] = response_obj + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if tool_calls_to_check: + inputs["tool_calls"] = cast( + List[ChatCompletionToolCallChunk], tool_calls_to_check + ) + response_model = response_obj.get("model") + if response_model: + inputs["model"] = response_model + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Write guardrailed texts back into the output items in-place. + # final_chunk is a reference into responses_so_far so this + # mutates the list that the caller holds. + await self._apply_guardrail_responses_to_output( + response=response_obj, + responses=guardrailed_texts, + task_mappings=task_mappings, + ) + + return responses_so_far + # ------------------------------------------------------------------ # + # Case 2: response.output_item.done — extract tool calls only. # + # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": - # convert openai response to model response model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( final_chunk ) - tool_calls = model_response_stream.choices[0].delta.tool_calls if tool_calls: inputs = GenericGuardrailAPIInputs() inputs["tool_calls"] = cast( List[ChatCompletionToolCallChunk], tool_calls ) - # Include model information if available if ( hasattr(model_response_stream, "model") and model_response_stream.model ): inputs["model"] = model_response_stream.model - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) - return responses_so_far - elif final_chunk.get("type") == "response.completed": - # convert openai response to model response - outputs = final_chunk.get("response", {}).get("output", []) + return responses_so_far - model_response_choices = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( - output_items=outputs, - handle_raw_dict_callback=None, - ) - - if model_response_choices: - tool_calls = model_response_choices[0].message.tool_calls - text = model_response_choices[0].message.content - guardrail_inputs = GenericGuardrailAPIInputs() - if text: - guardrail_inputs["texts"] = [text] - if tool_calls: - guardrail_inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - # Include model information from the response if available - response_model = final_chunk.get("response", {}).get("model") - if response_model: - guardrail_inputs["model"] = response_model - if tool_calls or text: - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=guardrail_inputs, - request_data=request_data if request_data is not None else {}, - input_type="response", - logging_obj=litellm_logging_obj, - ) - return responses_so_far - else: - verbose_proxy_logger.debug( - "Skipping output guardrail - model response has no choices" - ) - # model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) - # tool_calls = model_response_stream.choices[0].tool_calls - # convert openai response to model response + # ------------------------------------------------------------------ # + # Fallback: apply guardrail to the accumulated text string. # + # No structured write-back is possible here; guardrails that only # + # need to block/flag (not rewrite) still work correctly. # + # ------------------------------------------------------------------ # string_so_far = self.get_streaming_string_so_far(responses_so_far) - inputs = GenericGuardrailAPIInputs(texts=[string_so_far]) - # Try to get model from the final chunk if available - if isinstance(final_chunk, dict): + if string_so_far: + fallback_inputs = GenericGuardrailAPIInputs(texts=[string_so_far]) response_model = ( final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None ) if response_model: - inputs["model"] = response_model - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=request_data if request_data is not None else {}, - input_type="response", - logging_obj=litellm_logging_obj, - ) + fallback_inputs["model"] = response_model + await guardrail_to_apply.apply_guardrail( + inputs=fallback_inputs, + request_data=request_data if request_data is not None else {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) return responses_so_far def _check_streaming_has_ended(self, responses_so_far: List[Any]) -> bool: @@ -721,7 +767,7 @@ def _extract_output_text_and_images( async def _apply_guardrail_responses_to_output( self, - response: "ResponsesAPIResponse", + response: Union["ResponsesAPIResponse", Dict[Any, Any]], responses: List[str], task_mappings: List[Tuple[int, int]], ) -> None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 033e1d0b8e7..e723c07e3c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1194,9 +1194,10 @@ async def _stream_apply_output_masking( return if not all_chunks: verbose_proxy_logger.warning( - "Presidio apply_to_output: streaming response contained only " - "bytes chunks (Anthropic native SSE). Output PII masking was " - "skipped for this response." + "Presidio apply_to_output: streaming response contained no " + "ModelResponseStream chunks (e.g. raw SSE bytes or an empty " + "upstream stream). Output PII masking was skipped for this " + "response." ) return @@ -1258,6 +1259,37 @@ def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: Dict[str, str]) -> bytes: return "\n".join(result_lines).encode("utf-8") + def _unmask_responses_api_completed_chunk( + self, chunk: Any, pii_tokens: Dict[str, str] + ) -> None: + """ + Unmask PII tokens in-place for a ``response.completed`` Responses API event. + + The chunk carries a ``response`` attribute (ResponsesAPIResponse) whose + ``output`` list holds message items. Each item has a ``content`` list of + blocks; text blocks expose a ``.text`` string attribute. We walk the tree + and replace every PII token with its original value. + """ + response_obj = getattr(chunk, "response", None) + if response_obj is None: + return + + output = getattr(response_obj, "output", None) or [] + for output_item in output: + content = getattr(output_item, "content", None) or [] + for content_block in content: + if isinstance(content_block, dict): + if isinstance(content_block.get("text"), str): + content_block["text"] = self._unmask_pii_text( + content_block["text"], pii_tokens + ) + elif hasattr(content_block, "text") and isinstance( + content_block.text, str + ): + content_block.text = self._unmask_pii_text( + content_block.text, pii_tokens + ) + async def _stream_pii_unmasking( self, response: Any, @@ -1274,16 +1306,36 @@ async def _stream_pii_unmasking( pii_tokens: Dict[str, str] = metadata.get("pii_tokens", {}) remaining_chunks: List[ModelResponseStream] = [] + saw_non_chat_chunk = False try: async for chunk in response: if isinstance(chunk, ModelResponseStream): - remaining_chunks.append(chunk) + if saw_non_chat_chunk: + yield chunk + else: + remaining_chunks.append(chunk) elif isinstance(chunk, bytes): if pii_tokens: yield self._unmask_sse_bytes_chunk(chunk, pii_tokens) # type: ignore[misc] else: yield chunk # type: ignore[misc] continue + else: + # /v1/responses events: unmask response.completed text in-place. + # A mixed stream can't be reassembled, so flush buffered chat + # chunks in order before passthrough instead of dropping them. + if remaining_chunks and not saw_non_chat_chunk: + for buffered_chunk in remaining_chunks: + yield buffered_chunk + remaining_chunks = [] + chunk_type = getattr(chunk, "type", None) + if chunk_type == "response.completed" and pii_tokens: + self._unmask_responses_api_completed_chunk(chunk, pii_tokens) + saw_non_chat_chunk = True + yield chunk + + if saw_non_chat_chunk: + return if not remaining_chunks: return diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index dfc43bc29b5..27ecba83be7 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1230,6 +1230,10 @@ def _build_synthetic_response_events( "error", ] +RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES = frozenset( + {"input_text", "output_text", "text"} +) + class ResponsesWebSocketStreaming: """ @@ -1252,6 +1256,9 @@ def __init__( user_api_key_dict: Optional[Any] = None, request_data: Optional[Dict] = None, first_message: Optional[str] = None, + guardrail_callbacks: Optional[List[Any]] = None, + output_guardrail_callbacks: Optional[List[Any]] = None, + authorized_model: Optional[str] = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1261,6 +1268,11 @@ def __init__( self.messages: list[Dict] = [] self.input_messages: list[Dict[str, str]] = [] self.first_message = first_message + self.guardrail_callbacks: List[Any] = guardrail_callbacks or [] + self.output_guardrail_callbacks: List[Any] = output_guardrail_callbacks or [] + # Model name authorized at connection time; enforced on every + # response.create frame to prevent deployment-substitution attacks. + self.authorized_model: Optional[str] = authorized_model def _should_store_event(self, event_obj: dict) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1351,8 +1363,33 @@ async def backend_to_client(self) -> None: else: response_str = raw_response - self._store_event(response_str) - await self.websocket.send_text(response_str) + # 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_type = json.loads(response_str).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: # type: ignore verbose_logger.debug("Responses WS backend connection closed: %s", e) @@ -1361,20 +1398,316 @@ async def backend_to_client(self) -> None: finally: await self._log_messages() + def _enforce_authorized_model(self, msg_obj: dict) -> 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 = msg_obj.get("response") + if isinstance(nested, dict): + 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 = json.loads(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 = 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 + for cb in self.guardrail_callbacks: + 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_response = ( + msg_obj.get("response") + if isinstance(msg_obj.get("response"), dict) + else None + ) + text_containers: list[tuple[dict, 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 isinstance(field_value, list): + for item in field_value: + if not isinstance(item, dict): + 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 isinstance(value, list): + for block in value: + if ( + isinstance(block, dict) + and block.get("type") + in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES + and isinstance(block.get("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 + + pii_tokens: Dict[str, str] = (self.request_data.get("metadata") or {}).get( + "pii_tokens", {} + ) + if not pii_tokens: + return response_str + + try: + evt_obj = json.loads(response_str) + except (json.JSONDecodeError, TypeError): + return response_str + + cb = self.guardrail_callbacks[0] + event_type = evt_obj.get("type") + + if event_type == "response.completed": + modified = False + response_obj = evt_obj.get("response") or {} + if not isinstance(response_obj, dict): + return response_str + for output_item in response_obj.get("output") or []: + if not isinstance(output_item, dict): + continue + content = output_item.get("content") or [] + if not isinstance(content, list): + continue + for content_block in content: + if not isinstance(content_block, dict): + continue + text = content_block.get("text") + if isinstance(text, str): + unmasked = cb._unmask_pii_text(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 = evt_obj.get("delta") + if isinstance(delta, str): + unmasked = cb._unmask_pii_text(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 = json.loads(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") or {} + if not isinstance(response_obj, dict): + continue + for output_item in response_obj.get("output") or []: + if not isinstance(output_item, dict): + 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") or [] + if isinstance(summary, list): + for summary_block in summary: + if not isinstance(summary_block, dict): + 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") or [] + if not isinstance(content, list): + continue + for content_block in content: + if not isinstance(content_block, dict): + 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 client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None: - self._store_input(self.first_message) - self._store_event(self.first_message) - await self.backend_ws.send(self.first_message) # type: ignore[union-attr] + masked_first = 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) # type: ignore[union-attr] while True: message = await self.websocket.receive_text() - - self._store_input(message) - self._store_event(message) - await self.backend_ws.send(message) # type: ignore[union-attr] + masked = await self._mask_response_create(message) + self._store_input(masked) + self._store_event(masked) + await self.backend_ws.send(masked) # type: ignore[union-attr] except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index aee6ccc2e76..49cd1b71ef2 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -900,6 +900,20 @@ async def test_process_output_streaming_response_missing_output_key(self): # Should return the responses unchanged assert result == responses_so_far + @pytest.mark.asyncio + async def test_process_output_streaming_response_null_response(self): + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + responses_so_far = [{"type": "response.completed", "response": None}] + + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result == responses_so_far + @pytest.mark.asyncio async def test_process_output_streaming_response_unrecognized_output_type(self): """Test that streaming response with unrecognized output types doesn't raise IndexError @@ -996,6 +1010,105 @@ async def test_process_output_streaming_response_with_valid_output(self): # Should return the responses assert result == responses_so_far + @pytest.mark.asyncio + async def test_process_output_streaming_response_writes_back_guardrailed_text(self): + """Guardrailed text must be written back into the response.completed chunk in-place.""" + + class RewriteGuardrail(CustomGuardrail): + """Replaces '' with 'john@example.com' to simulate PII unmasking.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + inputs["texts"] = [ + t.replace("", "john@example.com") for t in texts + ] + return inputs + + handler = OpenAIResponsesHandler() + guardrail = RewriteGuardrail(guardrail_name="test-rewrite") + + responses_so_far = [ + {"type": "response.output_text.delta", "delta": "send to "}, + {"type": "response.output_text.delta", "delta": ""}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "send to "}, + ], + } + ], + "status": "completed", + }, + }, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + completed_chunk = next( + c + for c in result + if isinstance(c, dict) and c.get("type") == "response.completed" + ) + output_text = completed_chunk["response"]["output"][0]["content"][0]["text"] + assert ( + output_text == "send to john@example.com" + ), f"Expected PII token to be unmasked in response.completed output, got: {output_text!r}" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_pass_through_unchanged(self): + """A pass-through guardrail must not modify the output text.""" + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="pass-through") + + original_text = "No PII here, just normal text." + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_456", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "id": "msg_456", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": original_text}], + } + ], + "status": "completed", + }, + } + ] + + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + output_text = result[-1]["response"]["output"][0]["content"][0]["text"] + assert output_text == original_text + class TestGetStructuredMessages: """Test the get_structured_messages method for Responses API handler.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 8a5eeeff367..3efc42523f1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2329,6 +2329,164 @@ async def mock_stream(): assert "Output PII masking was skipped" in warning_msg +@pytest.mark.asyncio +async def test_output_parse_pii_streaming_responses_events_passthrough( + mock_user_api_key, +): + """ + Regression test: when output_parse_pii=True and pii_tokens exist, /v1/responses + streaming events must pass through instead of being dropped. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + response_events = [ + {"type": "response.created", "response": {"id": "resp_1"}}, + {"type": "response.output_text.delta", "delta": "Hello"}, + { + "type": "response.completed", + "response": {"id": "resp_1", "status": "completed"}, + }, + ] + + async def mock_stream(): + for event in response_events: + yield event + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={ + "metadata": { + "pii_tokens": {"": "john@example.com"}, + } + }, + ): + collected.append(chunk) + + assert collected == response_events + + +@pytest.mark.asyncio +async def test_output_parse_pii_streaming_responses_completed_event_unmasked( + mock_user_api_key, +): + """ + When output_parse_pii=True, a /v1/responses ``response.completed`` event + (a Pydantic ResponseCompletedEvent, as produced in production) must have its + output text unmasked in-place before being forwarded to the client. + """ + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + completed_event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=1, + output=[ + GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText( + type="output_text", + text="Reach me at today.", + annotations=[], + ) + ], + ) + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + async def mock_stream(): + yield completed_event + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={ + "metadata": { + "pii_tokens": {"": "john@example.com"}, + } + }, + ): + collected.append(chunk) + + assert collected == [completed_event] + assert ( + collected[0].response.output[0].content[0].text + == "Reach me at john@example.com today." + ) + + +@pytest.mark.asyncio +async def test_output_parse_pii_streaming_mixed_chunks_flushes_buffered( + mock_user_api_key, +): + """ + Regression test: when output_parse_pii=True and a stream mixes buffered + ModelResponseStream chunks with a /v1/responses event, the buffered chat + chunks must still be forwarded (in order) instead of being dropped at the + saw_non_chat_chunk early return. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + class FakeResponsesEvent: + def __init__(self, event_type: str): + self.type = event_type + + model_chunk = ModelResponseStream( + id="chatcmpl-mixed-unmask-1", + choices=[], + created=1, + model="gpt-4", + object="chat.completion.chunk", + system_fingerprint=None, + ) + response_completed = FakeResponsesEvent("response.completed") + + async def mock_stream(): + yield model_chunk + yield response_completed + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={ + "metadata": { + "pii_tokens": {"": "john@example.com"}, + } + }, + ): + collected.append(chunk) + + assert collected == [model_chunk, response_completed] + + @pytest.mark.asyncio async def test_anonymize_text_uses_correct_positions_no_parse_pii(): """ 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 1981651797d..dc83ebd40d5 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -7,6 +7,9 @@ 3. Providers without native websocket support use ManagedResponsesWebSocketHandler """ +import json +from unittest.mock import MagicMock + import pytest from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig @@ -46,12 +49,59 @@ def test_openai_supports_native_websocket(self): ), "OpenAI should support native websocket" def test_azure_supports_native_websocket(self): - """Azure should support native websocket (inherits from OpenAI)""" + """Azure should support native websocket""" config = AzureOpenAIResponsesAPIConfig() assert ( config.supports_native_websocket() is True ), "Azure should support native websocket" + def test_azure_websocket_url_uses_v1_path(self): + """Azure WebSocket URL must use /openai/v1/responses (no api-version)""" + config = AzureOpenAIResponsesAPIConfig() + url = config.get_websocket_url( + api_base="https://myresource.cognitiveservices.azure.com", + litellm_params={"api_version": "2025-04-01-preview"}, + ) + assert url == "wss://myresource.cognitiveservices.azure.com/openai/v1/responses" + assert "api-version" not in url + + def test_azure_websocket_url_strips_existing_path(self): + """api_base that already contains /openai/responses must be cleaned""" + config = AzureOpenAIResponsesAPIConfig() + url = config.get_websocket_url( + api_base="https://myresource.cognitiveservices.azure.com/openai/responses", + litellm_params={}, + ) + assert url == "wss://myresource.cognitiveservices.azure.com/openai/v1/responses" + + def test_azure_websocket_url_strips_query_params(self): + config = AzureOpenAIResponsesAPIConfig() + url = config.get_websocket_url( + api_base="https://myresource.cognitiveservices.azure.com/openai/responses?api-version=2024-05-01-preview", + litellm_params={}, + ) + assert url == "wss://myresource.cognitiveservices.azure.com/openai/v1/responses" + + def test_azure_websocket_url_requires_api_base(self): + config = AzureOpenAIResponsesAPIConfig() + with pytest.raises(ValueError): + config.get_websocket_url(api_base=None, litellm_params={}) + + def test_azure_model_not_in_websocket_url(self): + """Azure sends the model in the body, so it must not be appended to the URL""" + assert AzureOpenAIResponsesAPIConfig().model_in_websocket_url() is False + + def test_openai_default_websocket_url_converts_scheme(self): + """The base get_websocket_url default converts the HTTP endpoint to wss://""" + config = OpenAIResponsesAPIConfig() + url = config.get_websocket_url( + api_base="https://api.openai.com/v1", litellm_params={} + ) + assert url == "wss://api.openai.com/v1/responses" + + def test_openai_model_in_websocket_url_default(self): + assert OpenAIResponsesAPIConfig().model_in_websocket_url() is True + def test_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() @@ -777,6 +827,1079 @@ async def test_managed_handler_handles_invalid_json(self): assert "Invalid JSON" in error_event +class TestNativeWebSocketGuardrails: + @pytest.mark.asyncio + async def test_response_create_injects_authorized_model(self): + import json + from unittest.mock import MagicMock + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + handler = ResponsesWebSocketStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=MagicMock(), + authorized_model="authorized-deployment", + ) + + flat_message = await handler._mask_response_create( + json.dumps({"type": "response.create", "input": "hi"}) + ) + nested_message = await handler._mask_response_create( + json.dumps({"type": "response.create", "response": {"input": "hi"}}) + ) + + assert json.loads(flat_message)["model"] == "authorized-deployment" + assert ( + json.loads(nested_message)["response"]["model"] == "authorized-deployment" + ) + + @pytest.mark.asyncio + async def test_completed_event_with_null_response_passes_through(self): + from unittest.mock import MagicMock + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + class Guardrail: + def get_presidio_settings_from_request_data(self, request_data): + return None + + def _unmask_pii_text(self, text, pii_tokens): + return text + + event = '{"type":"response.completed","response":null}' + guardrail = Guardrail() + handler = ResponsesWebSocketStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=MagicMock(), + request_data={"metadata": {"pii_tokens": {"": "secret"}}}, + guardrail_callbacks=[guardrail], + output_guardrail_callbacks=[guardrail], + ) + + assert handler._unmask_response_event(event) == event + assert await handler._mask_response_completed(event) == event + + @pytest.mark.asyncio + async def test_output_masking_suppresses_delta_without_calling_presidio(self): + import json + from unittest.mock import AsyncMock, MagicMock + + import websockets.exceptions + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + class RecordingGuardrail: + def __init__(self): + self.check_pii_calls = [] + + def get_presidio_settings_from_request_data(self, request_data): + return None + + def _unmask_pii_text(self, text, pii_tokens): + return text + + async def check_pii( + self, text, output_parse_pii, presidio_config, request_data + ): + self.check_pii_calls.append(text) + return text + + class FakeBackendWS: + def __init__(self, events): + self._events = list(events) + + async def recv(self, decode=False): + if self._events: + return self._events.pop(0) + raise websockets.exceptions.ConnectionClosed(None, None) + + guardrail = RecordingGuardrail() + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + delta_event = json.dumps( + {"type": "response.output_text.delta", "delta": "alice@example.com"} + ) + completed_event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + {"type": "output_text", "text": "alice@example.com"} + ] + } + ] + }, + } + ) + + handler = ResponsesWebSocketStreaming( + websocket=client_ws, + backend_ws=FakeBackendWS([delta_event, completed_event]), + logging_obj=logging_obj, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + # The delta event must be suppressed without ever invoking Presidio, + # so check_pii is called exactly once (for the completed event only). + assert guardrail.check_pii_calls == ["alice@example.com"] + client_ws.send_text.assert_called_once() + sent_payload = client_ws.send_text.call_args[0][0] + assert json.loads(sent_payload)["type"] == "response.completed" + + @pytest.mark.asyncio + async def test_output_masking_suppresses_text_bearing_done_events(self): + import json + from unittest.mock import AsyncMock, MagicMock + + import websockets.exceptions + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + class MaskingGuardrail: + def __init__(self): + self.check_pii_calls = [] + + def get_presidio_settings_from_request_data(self, request_data): + return None + + def _unmask_pii_text(self, text, pii_tokens): + return text + + async def check_pii( + self, text, output_parse_pii, presidio_config, request_data + ): + self.check_pii_calls.append(text) + return text.replace("alice@example.com", "") + + class FakeBackendWS: + def __init__(self, events): + self._events = list(events) + + async def recv(self, decode=False): + if self._events: + return self._events.pop(0) + raise websockets.exceptions.ConnectionClosed(None, None) + + guardrail = MaskingGuardrail() + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + done_events = [ + json.dumps( + {"type": "response.output_text.done", "text": "alice@example.com"} + ), + json.dumps( + { + "type": "response.content_part.done", + "part": {"type": "output_text", "text": "alice@example.com"}, + } + ), + json.dumps( + { + "type": "response.output_item.done", + "item": { + "type": "message", + "content": [ + {"type": "output_text", "text": "alice@example.com"} + ], + }, + } + ), + ] + completed_event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + {"type": "output_text", "text": "alice@example.com"} + ] + } + ] + }, + } + ) + + handler = ResponsesWebSocketStreaming( + websocket=client_ws, + backend_ws=FakeBackendWS(done_events + [completed_event]), + logging_obj=logging_obj, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + # Text-bearing done events carry the full output before response.completed + # arrives; they must be suppressed so unmasked PII never reaches the + # client, and Presidio is only invoked for response.completed. + assert guardrail.check_pii_calls == ["alice@example.com"] + client_ws.send_text.assert_called_once() + sent_payload = client_ws.send_text.call_args[0][0] + assert json.loads(sent_payload)["type"] == "response.completed" + assert "alice@example.com" not in sent_payload + assert "" in sent_payload + + +class _FakeWSGuardrail: + """Presidio-like guardrail double for the WebSocket masking hooks. + + ``check_pii`` replaces each known PII string with its token. When + ``output_parse_pii`` is True (input masking) the token->original map is + persisted into ``request_data["metadata"]["pii_tokens"]`` so the response + path can reverse it. ``_unmask_pii_text`` performs that reversal. + """ + + def __init__(self, mask_map=None): + self.mask_map = mask_map or {"alice@example.com": ""} + self.output_parse_pii = True + self.apply_to_output = True + + def get_presidio_settings_from_request_data(self, request_data): + return None + + async def check_pii(self, text, output_parse_pii, presidio_config, request_data): + masked = text + tokens = {} + for original, token in self.mask_map.items(): + if original in masked: + masked = masked.replace(original, token) + tokens[token] = original + if output_parse_pii and tokens: + metadata = request_data.setdefault("metadata", {}) + metadata.setdefault("pii_tokens", {}).update(tokens) + return masked + + def _unmask_pii_text(self, text, pii_tokens): + for token, original in pii_tokens.items(): + text = text.replace(token, original) + return text + + +def _make_streaming(**kwargs): + from unittest.mock import MagicMock + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + kwargs.setdefault("websocket", MagicMock()) + kwargs.setdefault("backend_ws", MagicMock()) + kwargs.setdefault("logging_obj", MagicMock()) + return ResponsesWebSocketStreaming(**kwargs) + + +class TestNativeWebSocketGuardrailMasking: + """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" + + @pytest.mark.asyncio + async def test_mask_response_create_flat_string_input(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, + guardrail_callbacks=[guardrail], + authorized_model="auth-model", + ) + + masked = await handler._mask_response_create( + json.dumps( + {"type": "response.create", "input": "email alice@example.com now"} + ) + ) + obj = json.loads(masked) + + assert obj["model"] == "auth-model" + assert obj["input"] == "email now" + assert handler.request_data["metadata"]["pii_tokens"] == { + "": "alice@example.com" + } + + @pytest.mark.asyncio + async def test_mask_response_create_list_content_string(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": [ + { + "type": "message", + "role": "user", + "content": "ping alice@example.com", + } + ], + } + ) + ) + obj = json.loads(masked) + + assert obj["input"][0]["content"] == "ping " + + @pytest.mark.asyncio + async def test_mask_response_create_input_text_blocks(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "alice@example.com"}, + {"type": "input_image", "image_url": "http://x"}, + ], + } + ], + } + ) + ) + obj = json.loads(masked) + blocks = obj["input"][0]["content"] + + assert blocks[0]["text"] == "" + assert blocks[1]["image_url"] == "http://x" + + @pytest.mark.asyncio + async def test_mask_response_create_function_call_output_string(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "tool returned alice@example.com", + } + ], + } + ) + ) + obj = json.loads(masked) + + assert obj["input"][0]["output"] == "tool returned " + assert handler.request_data["metadata"]["pii_tokens"] == { + "": "alice@example.com" + } + + @pytest.mark.asyncio + async def test_mask_response_create_function_call_output_blocks(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [ + {"type": "output_text", "text": "alice@example.com"}, + {"type": "input_image", "image_url": "http://x"}, + ], + } + ], + } + ) + ) + obj = json.loads(masked) + blocks = obj["input"][0]["output"] + + assert blocks[0]["text"] == "" + assert blocks[1]["image_url"] == "http://x" + + @pytest.mark.asyncio + async def test_mask_response_create_nested_shape(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, + guardrail_callbacks=[guardrail], + authorized_model="auth-model", + ) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "response": {"input": "alice@example.com", "model": "spoofed"}, + } + ) + ) + obj = json.loads(masked) + + assert obj["response"]["model"] == "auth-model" + assert obj["response"]["input"] == "" + + @pytest.mark.asyncio + async def test_mask_response_create_flat_instructions(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": "hi", + "instructions": "reply to alice@example.com", + } + ) + ) + obj = json.loads(masked) + + assert obj["instructions"] == "reply to " + assert handler.request_data["metadata"]["pii_tokens"] == { + "": "alice@example.com" + } + + @pytest.mark.asyncio + async def test_mask_response_create_nested_instructions(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "response": { + "input": "hi", + "instructions": "email alice@example.com", + }, + } + ) + ) + obj = json.loads(masked) + + assert obj["response"]["instructions"] == "email " + assert handler.request_data["metadata"]["pii_tokens"] == { + "": "alice@example.com" + } + + @pytest.mark.asyncio + async def test_mask_response_create_non_create_unchanged(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, + guardrail_callbacks=[guardrail], + authorized_model="auth-model", + ) + + message = json.dumps({"type": "response.cancel", "input": "alice@example.com"}) + assert await handler._mask_response_create(message) == message + + @pytest.mark.asyncio + async def test_mask_response_create_invalid_json_unchanged(self): + handler = _make_streaming( + request_data={}, guardrail_callbacks=[_FakeWSGuardrail()] + ) + assert await handler._mask_response_create("not json {{{") == "not json {{{" + + @pytest.mark.asyncio + async def test_mask_response_create_model_only_without_guardrails(self): + handler = _make_streaming(request_data={}, authorized_model="auth-model") + + masked = await handler._mask_response_create( + json.dumps({"type": "response.create", "input": "alice@example.com"}) + ) + obj = json.loads(masked) + + assert obj["model"] == "auth-model" + assert obj["input"] == "alice@example.com" + + @pytest.mark.asyncio + async def test_mask_response_create_no_op_without_model_or_guardrails(self): + handler = _make_streaming(request_data={}) + message = json.dumps({"type": "response.create", "input": "alice@example.com"}) + assert await handler._mask_response_create(message) == message + + @pytest.mark.asyncio + async def test_mask_response_create_list_with_non_dict_item(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + + masked = await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "input": [ + "not-a-dict", + { + "type": "message", + "role": "user", + "content": "alice@example.com", + }, + ], + } + ) + ) + obj = json.loads(masked) + assert obj["input"][0] == "not-a-dict" + assert obj["input"][1]["content"] == "" + + def test_enforce_authorized_model_no_authorized_model(self): + handler = _make_streaming(request_data={}) + assert handler._enforce_authorized_model({"model": "anything"}) is False + + def test_enforce_authorized_model_nested_with_top_level_model(self): + handler = _make_streaming(request_data={}, authorized_model="auth-model") + msg = {"response": {"model": "spoofed"}, "model": "also-spoofed"} + assert handler._enforce_authorized_model(msg) is True + assert msg["response"]["model"] == "auth-model" + assert msg["model"] == "auth-model" + + @pytest.mark.asyncio + async def test_unmask_response_event_completed(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={ + "metadata": {"pii_tokens": {"": "alice@example.com"}} + }, + guardrail_callbacks=[guardrail], + ) + + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + {"type": "output_text", "text": "to "} + ] + } + ] + }, + } + ) + unmasked = json.loads(handler._unmask_response_event(event)) + assert ( + unmasked["response"]["output"][0]["content"][0]["text"] + == "to alice@example.com" + ) + + @pytest.mark.asyncio + async def test_unmask_response_event_delta(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={ + "metadata": {"pii_tokens": {"": "alice@example.com"}} + }, + guardrail_callbacks=[guardrail], + ) + + event = json.dumps( + {"type": "response.output_text.delta", "delta": ""} + ) + unmasked = json.loads(handler._unmask_response_event(event)) + assert unmasked["delta"] == "alice@example.com" + + def test_unmask_response_event_no_tokens_unchanged(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming(request_data={}, guardrail_callbacks=[guardrail]) + event = json.dumps( + {"type": "response.output_text.delta", "delta": ""} + ) + assert handler._unmask_response_event(event) == event + + def test_unmask_response_event_no_guardrails_unchanged(self): + handler = _make_streaming( + request_data={"metadata": {"pii_tokens": {"": "x"}}} + ) + event = json.dumps({"type": "response.completed", "response": {}}) + assert handler._unmask_response_event(event) == event + + def test_unmask_response_event_invalid_json_unchanged(self): + handler = _make_streaming( + request_data={"metadata": {"pii_tokens": {"": "x"}}}, + guardrail_callbacks=[_FakeWSGuardrail()], + ) + assert handler._unmask_response_event("not json {{{") == "not json {{{" + + def test_unmask_response_event_non_dict_response_unchanged(self): + handler = _make_streaming( + request_data={"metadata": {"pii_tokens": {"": "x"}}}, + guardrail_callbacks=[_FakeWSGuardrail()], + ) + event = json.dumps({"type": "response.completed", "response": ["bad-shape"]}) + assert handler._unmask_response_event(event) == event + + def test_unmask_response_event_malformed_output_items_unchanged(self): + handler = _make_streaming( + request_data={ + "metadata": {"pii_tokens": {"": "alice@example.com"}} + }, + guardrail_callbacks=[_FakeWSGuardrail()], + ) + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + "not-a-dict", + {"content": "not-a-list"}, + {"content": ["not-a-dict-block"]}, + ] + }, + } + ) + assert handler._unmask_response_event(event) == event + + def test_unmask_response_event_other_event_type_unchanged(self): + handler = _make_streaming( + request_data={"metadata": {"pii_tokens": {"": "x"}}}, + guardrail_callbacks=[_FakeWSGuardrail()], + ) + event = json.dumps( + {"type": "response.in_progress", "delta": ""} + ) + assert handler._unmask_response_event(event) == event + + @pytest.mark.asyncio + async def test_mask_response_completed_event(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[guardrail] + ) + + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + { + "type": "output_text", + "text": "contact alice@example.com", + } + ] + } + ] + }, + } + ) + masked = json.loads(await handler._mask_response_completed(event)) + assert ( + masked["response"]["output"][0]["content"][0]["text"] + == "contact " + ) + + @pytest.mark.asyncio + async def test_mask_response_completed_masks_function_call_arguments(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[guardrail] + ) + + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "type": "function_call", + "name": "send_email", + "arguments": '{"to": "alice@example.com"}', + } + ] + }, + } + ) + masked = json.loads(await handler._mask_response_completed(event)) + assert ( + masked["response"]["output"][0]["arguments"] + == '{"to": ""}' + ) + + @pytest.mark.asyncio + async def test_mask_response_completed_masks_reasoning_summary(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[guardrail] + ) + + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": "user is alice@example.com", + } + ], + } + ] + }, + } + ) + masked = json.loads(await handler._mask_response_completed(event)) + assert ( + masked["response"]["output"][0]["summary"][0]["text"] + == "user is " + ) + + @pytest.mark.asyncio + async def test_mask_response_completed_delta_unchanged(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[guardrail] + ) + + event = json.dumps( + {"type": "response.output_text.delta", "delta": "alice@example.com"} + ) + assert await handler._mask_response_completed(event) == event + + @pytest.mark.asyncio + async def test_mask_response_completed_no_guardrails_unchanged(self): + handler = _make_streaming(request_data={}) + event = json.dumps( + {"type": "response.output_text.delta", "delta": "alice@example.com"} + ) + assert await handler._mask_response_completed(event) == event + + @pytest.mark.asyncio + async def test_mask_response_completed_invalid_json_unchanged(self): + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[_FakeWSGuardrail()] + ) + assert await handler._mask_response_completed("not json {{{") == "not json {{{" + + @pytest.mark.asyncio + async def test_mask_response_completed_malformed_unchanged(self): + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[_FakeWSGuardrail()] + ) + event = json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + "not-a-dict", + {"content": "not-a-list"}, + {"content": ["not-a-dict-block"]}, + ] + }, + } + ) + assert await handler._mask_response_completed(event) == event + + @pytest.mark.asyncio + async def test_mask_response_completed_non_dict_response_unchanged(self): + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[_FakeWSGuardrail()] + ) + event = json.dumps({"type": "response.completed", "response": ["bad"]}) + assert await handler._mask_response_completed(event) == event + + @pytest.mark.asyncio + async def test_client_to_backend_masks_and_enforces_model(self): + from unittest.mock import AsyncMock + + guardrail = _FakeWSGuardrail() + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + websocket = MagicMock() + websocket.receive_text = AsyncMock( + side_effect=[ + json.dumps( + {"type": "response.create", "input": "ping alice@example.com"} + ), + Exception("stop"), + ] + ) + + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + request_data={}, + first_message=json.dumps( + {"type": "response.create", "input": "alice@example.com"} + ), + guardrail_callbacks=[guardrail], + authorized_model="auth-model", + ) + + await handler.client_to_backend() + + assert backend_ws.send.await_count == 2 + first_sent = json.loads(backend_ws.send.await_args_list[0][0][0]) + assert first_sent["model"] == "auth-model" + assert first_sent["input"] == "" + second_sent = json.loads(backend_ws.send.await_args_list[1][0][0]) + assert second_sent["model"] == "auth-model" + assert second_sent["input"] == "ping " + assert handler.request_data["metadata"]["pii_tokens"] == { + "": "alice@example.com" + } + + @pytest.mark.asyncio + async def test_backend_to_client_suppresses_deltas_and_masks_completed(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + guardrail = _FakeWSGuardrail() + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + {"type": "response.output_text.delta", "delta": "alice@example.com"} + ), + json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + { + "type": "output_text", + "text": "contact alice@example.com", + } + ] + } + ] + }, + } + ), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + websocket.send_text.assert_awaited_once() + forwarded = json.loads(websocket.send_text.await_args[0][0]) + assert forwarded["type"] == "response.completed" + assert ( + forwarded["response"]["output"][0]["content"][0]["text"] + == "contact " + ) + + @pytest.mark.asyncio + async def test_backend_to_client_suppresses_function_call_arguments_done(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + guardrail = _FakeWSGuardrail() + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "response.function_call_arguments.done", + "arguments": '{"to": "alice@example.com"}', + } + ), + json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "type": "function_call", + "name": "send_email", + "arguments": '{"to": "alice@example.com"}', + } + ] + }, + } + ), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + # The unmasked function-call arguments must never reach the client; only + # the masked response.completed is forwarded. + websocket.send_text.assert_awaited_once() + sent_payload = websocket.send_text.await_args[0][0] + forwarded = json.loads(sent_payload) + assert forwarded["type"] == "response.completed" + assert ( + forwarded["response"]["output"][0]["arguments"] + == '{"to": ""}' + ) + assert "alice@example.com" not in sent_payload + + @pytest.mark.asyncio + async def test_backend_to_client_suppresses_reasoning_summary_text_done(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + guardrail = _FakeWSGuardrail() + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "response.reasoning_summary_text.done", + "text": "contact alice@example.com", + } + ), + json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "content": [ + { + "type": "output_text", + "text": "done", + } + ] + } + ] + }, + } + ), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + # The reasoning-summary done event carries the full reasoning text before + # response.completed arrives; it must be suppressed so unmasked PII never + # reaches the client. + websocket.send_text.assert_awaited_once() + sent_payload = websocket.send_text.await_args[0][0] + assert json.loads(sent_payload)["type"] == "response.completed" + assert "alice@example.com" not in sent_payload + + @pytest.mark.asyncio + async def test_backend_to_client_suppresses_reasoning_summary_part_done(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + guardrail = _FakeWSGuardrail() + websocket = MagicMock() + websocket.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "response.reasoning_summary_part.done", + "part": { + "type": "summary_text", + "text": "user is alice@example.com", + }, + } + ), + json.dumps( + { + "type": "response.completed", + "response": { + "output": [ + { + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": "user is alice@example.com", + } + ], + } + ] + }, + } + ), + Exception("stop"), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + + handler = _make_streaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data={}, + output_guardrail_callbacks=[guardrail], + ) + + await handler.backend_to_client() + + # The reasoning-summary part-done event carries the full reasoning text + # before response.completed arrives; it must be suppressed, and the + # reasoning summary in response.completed must itself be masked. + websocket.send_text.assert_awaited_once() + sent_payload = websocket.send_text.await_args[0][0] + forwarded = json.loads(sent_payload) + assert forwarded["type"] == "response.completed" + assert ( + forwarded["response"]["output"][0]["summary"][0]["text"] + == "user is " + ) + assert "alice@example.com" not in sent_payload + + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" @@ -1000,9 +2123,7 @@ async def __aexit__(self, *args): mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) mock_config.supports_native_websocket.return_value = True - mock_config.get_complete_url.return_value = ( - "https://api.openai.com/v1/responses" - ) + mock_config.get_websocket_url.return_value = "wss://api.openai.com/v1/responses" mock_config.validate_environment.return_value = {} mock_logging = MagicMock() @@ -1051,8 +2172,8 @@ async def __aexit__(self, *args): mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) mock_config.supports_native_websocket.return_value = True - mock_config.get_complete_url.return_value = ( - "https://custom.example.com/v1/responses?api-version=2024-05-01" + mock_config.get_websocket_url.return_value = ( + "wss://custom.example.com/v1/responses?api-version=2024-05-01" ) mock_config.validate_environment.return_value = {} @@ -1084,3 +2205,49 @@ async def __aexit__(self, *args): assert qs.get("api-version") == [ "2024-05-01" ], f"existing param lost: {captured_urls[0]}" + + @pytest.mark.asyncio + async def test_ws_passes_litellm_params_to_get_websocket_url(self): + """Deployment api_version must reach get_websocket_url (Azure WS URL).""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_websocket_url.return_value = ( + "wss://example.openai.azure.com/openai/v1/responses" + ) + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + class FakeConnect: + def __init__(self, url, **kwargs): + pass + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-5.3-codex", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + api_base="https://example.openai.azure.com", + api_version="2025-04-01-preview", + ) + + 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"