From 5869dd85aa03589203ad65c2e27165bce1bbf358 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:11:09 +0530 Subject: [PATCH 01/26] fix(responses): Presidio PII masking for Azure WebSocket and streaming Wire Presidio into native Responses WebSocket forwarding and fix streaming output unmasking so masked tokens are restored for HTTP and WS clients. Co-authored-by: Cursor --- .../llms/azure/responses/transformation.py | 26 +++ .../llms/base_llm/responses/transformation.py | 17 ++ litellm/llms/custom_httpx/llm_http_handler.py | 37 ++++- .../guardrail_translation/handler.py | 148 ++++++++++++------ .../guardrails/guardrail_hooks/presidio.py | 49 +++++- litellm/responses/streaming_iterator.py | 130 ++++++++++++++- ...test_openai_responses_guardrail_handler.py | 99 ++++++++++++ .../guardrail_hooks/test_presidio.py | 41 +++++ .../test_responses_websocket_all_providers.py | 75 ++++++++- 9 files changed, 548 insertions(+), 74 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index ca9293325ff..1dfc618c100 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -185,6 +185,32 @@ 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") + import httpx + + base = api_base.rstrip("/") + # Strip existing /openai/responses path if the api_base already contains it + for suffix in ("/openai/v1/responses", "/openai/responses"): + if base.endswith(suffix): + base = base[: -len(suffix)] + break + ws_url = base.replace("https://", "wss://").replace("http://", "ws://") + return f"{ws_url}/openai/v1/responses" + ######################################################### ########## DELETE RESPONSE API TRANSFORMATION ############## ######################################################### diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 407d5ad8146..bae7bc54f1d 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -258,6 +258,23 @@ 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://") + ######################################################### ########## 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..c7ec5cb20d1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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,14 +5641,14 @@ 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. + # Some providers (e.g. OpenAI) require ?model= in the WebSocket URL. + # Providers that encode the model differently (e.g. Azure sends it in the + # response.create body) return a URL that already has no model param — + # we only add it when not already present. _parsed = urlparse(ws_url) _qs = parse_qs(_parsed.query) if "model" not in _qs: @@ -5679,6 +5683,24 @@ async def async_responses_websocket( _request_data: Dict[str, Any] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata + + _ws_guardrail_callbacks: list = [] + try: + import litellm as _litellm + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + _ws_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if isinstance(cb, _OPTIONAL_PresidioPIIMasking) + and getattr(cb, "output_parse_pii", False) + and not getattr(cb, "apply_to_output", False) + ] + except Exception: + pass + streaming = ResponsesWebSocketStreaming( websocket=websocket, backend_ws=cast(ClientConnection, backend_ws), @@ -5686,6 +5708,7 @@ 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, ) 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..f870a218076 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -479,90 +479,134 @@ 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": + outputs: List[Any] = final_chunk.get("response", {}).get("output", []) + + 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"] = final_chunk.get("response", {}) + 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 = final_chunk.get("response", {}).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=final_chunk.get("response", {}), + 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", []) - - model_response_choices = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( - output_items=outputs, - handle_raw_dict_callback=None, - ) + return responses_so_far - 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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 033e1d0b8e7..4b2e8fc59c8 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,6 +1306,7 @@ 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): @@ -1284,6 +1317,16 @@ async def _stream_pii_unmasking( else: yield chunk # type: ignore[misc] continue + else: + # /v1/responses events: unmask response.completed text in-place. + 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..08bfdfb33b5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1252,6 +1252,7 @@ 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, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1261,6 +1262,7 @@ 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 [] def _should_store_event(self, event_obj: dict) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1352,7 +1354,8 @@ async def backend_to_client(self) -> None: response_str = raw_response self._store_event(response_str) - await self.websocket.send_text(response_str) + unmasked_str = self._unmask_response_completed(response_str) + await self.websocket.send_text(unmasked_str) except websockets.exceptions.ConnectionClosed as e: # type: ignore verbose_logger.debug("Responses WS backend connection closed: %s", e) @@ -1361,20 +1364,135 @@ async def backend_to_client(self) -> None: finally: await self._log_messages() + async def _mask_response_create(self, message: str) -> str: + """ + Apply Presidio PII masking to a ``response.create`` message before it + is forwarded to the upstream provider. + + Walks the ``input`` field of the message (string or list of message + items), 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. + """ + if not self.guardrail_callbacks: + return message + + try: + msg_obj = json.loads(message) + except (json.JSONDecodeError, TypeError): + return message + + if msg_obj.get("type") != "response.create": + return message + + if "metadata" not in self.request_data: + self.request_data["metadata"] = {} + + modified = False + for cb in self.guardrail_callbacks: + presidio_config = cb.get_presidio_settings_from_request_data( + self.request_data + ) + input_data = msg_obj.get("input", []) + + if isinstance(input_data, str): + msg_obj["input"] = await cb.check_pii( + text=input_data, + output_parse_pii=True, + presidio_config=presidio_config, + request_data=self.request_data, + ) + modified = True + + elif isinstance(input_data, list): + for item in input_data: + if not isinstance(item, dict): + continue + content = item.get("content", []) + if isinstance(content, str): + item["content"] = await cb.check_pii( + text=content, + output_parse_pii=True, + presidio_config=presidio_config, + request_data=self.request_data, + ) + modified = True + elif isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "input_text" + 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 + + def _unmask_response_completed(self, response_str: str) -> str: + """ + Apply Presidio PII unmasking to a ``response.completed`` event before + it is forwarded back to the client. + + Uses the ``pii_tokens`` map stored during ``_mask_response_create`` to + replace every token (e.g. ````) with the original + value. Events that are not ``response.completed``, or where no tokens + were stored, 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 + + if evt_obj.get("type") != "response.completed": + return response_str + + cb = self.guardrail_callbacks[0] + modified = False + response_obj = evt_obj.get("response", {}) + for output_item in response_obj.get("output", []): + for content_block in output_item.get("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 + async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None: + masked_first = await self._mask_response_create(self.first_message) self._store_input(self.first_message) - self._store_event(self.first_message) - await self.backend_ws.send(self.first_message) # type: ignore[union-attr] + self._store_event(masked_first) + await self.backend_ws.send(masked_first) # type: ignore[union-attr] while True: message = await self.websocket.receive_text() - + masked = await self._mask_response_create(message) self._store_input(message) - self._store_event(message) - await self.backend_ws.send(message) # type: ignore[union-attr] + 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..5a4457c7ebd 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 @@ -996,6 +996,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..9f6084fe9fc 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,47 @@ 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_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..966f6688926 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -46,12 +46,31 @@ 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_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() @@ -1000,9 +1019,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 +1068,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 +1101,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" From 3d26750f4a6f5e83d9b76244acd7dffe24f36c24 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:18:17 +0530 Subject: [PATCH 02/26] Fix unused imports in responses handlers --- litellm/llms/azure/responses/transformation.py | 1 - litellm/llms/openai/responses/guardrail_translation/handler.py | 1 - 2 files changed, 2 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 1dfc618c100..643962073d9 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -200,7 +200,6 @@ def get_websocket_url( """ if api_base is None: raise ValueError("api_base is required for Azure WebSocket") - import httpx base = api_base.rstrip("/") # Strip existing /openai/responses path if the api_base already contains it diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index f870a218076..b09d4ef2da7 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 From b5a4c07d7676680e83940dd9cc605d3b289c09e3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:26:29 +0530 Subject: [PATCH 03/26] fix(responses): address Greptile review - Azure WebSocket model URL and PII logging - Add model_in_websocket_url() to BaseResponsesAPIConfig (default True) so providers can opt out of ?model= being appended to WebSocket URLs. - Override model_in_websocket_url() to return False for Azure, since Azure sends the model in the response.create body, not the URL query string. - Use this flag in llm_http_handler to conditionally append ?model=. - Pass masked message to _store_input() instead of the original PII-containing message so logging destinations do not receive unmasked PII. Co-Authored-By: Claude Sonnet 4.6 --- .../llms/azure/responses/transformation.py | 4 ++++ .../llms/base_llm/responses/transformation.py | 8 +++++++ litellm/llms/custom_httpx/llm_http_handler.py | 22 ++++++++++--------- litellm/responses/streaming_iterator.py | 4 ++-- 4 files changed, 26 insertions(+), 12 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 643962073d9..1f6ad804107 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -210,6 +210,10 @@ def get_websocket_url( ws_url = base.replace("https://", "wss://").replace("http://", "ws://") return f"{ws_url}/openai/v1/responses" + 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 bae7bc54f1d..c61ce52b530 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -275,6 +275,14 @@ def get_websocket_url( ) 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 c7ec5cb20d1..f02664f5374 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5646,16 +5646,18 @@ async def async_responses_websocket( litellm_params=dict(litellm_params), ) # Some providers (e.g. OpenAI) require ?model= in the WebSocket URL. - # Providers that encode the model differently (e.g. Azure sends it in the - # response.create body) return a URL that already has no model param — - # we only add it when not already present. - _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()})) - ) + # 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() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 08bfdfb33b5..7cd24dd0a26 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1483,14 +1483,14 @@ async def client_to_backend(self) -> None: try: if self.first_message is not None: masked_first = await self._mask_response_create(self.first_message) - self._store_input(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() masked = await self._mask_response_create(message) - self._store_input(message) + self._store_input(masked) self._store_event(masked) await self.backend_ws.send(masked) # type: ignore[union-attr] From 4a392c6feecce97e844308d2414cd6dff406b913 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:28:38 +0530 Subject: [PATCH 04/26] fix(responses): mask nested response.create input format for Presidio PII Handle the nested {"type":"response.create","response":{"input":[...]}} format in _mask_response_create. Previously only the flat top-level input was masked; the nested shape bypassed Presidio and forwarded raw PII upstream. Now both shapes are normalized and masked before forwarding. Co-Authored-By: Claude Sonnet 4.6 --- litellm/responses/streaming_iterator.py | 86 ++++++++++++++----------- 1 file changed, 49 insertions(+), 37 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cd24dd0a26..6b9f4e59887 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1394,44 +1394,56 @@ async def _mask_response_create(self, message: str) -> str: presidio_config = cb.get_presidio_settings_from_request_data( self.request_data ) - input_data = msg_obj.get("input", []) - - if isinstance(input_data, str): - msg_obj["input"] = await cb.check_pii( - text=input_data, - output_parse_pii=True, - presidio_config=presidio_config, - request_data=self.request_data, - ) - modified = True + # response.create supports two shapes: + # flat: {"type": "response.create", "input": [...], ...} + # nested: {"type": "response.create", "response": {"input": [...], ...}} + # Mask both so PII is never forwarded unmasked regardless of shape. + nested_response = msg_obj.get("response") if isinstance(msg_obj.get("response"), dict) else None + input_containers: list[tuple[dict, str]] = [] + if "input" in msg_obj: + input_containers.append((msg_obj, "input")) + if nested_response is not None and "input" in nested_response: + input_containers.append((nested_response, "input")) + + for container, key in input_containers: + input_data = container[key] + + if isinstance(input_data, str): + container[key] = await cb.check_pii( + text=input_data, + output_parse_pii=True, + presidio_config=presidio_config, + request_data=self.request_data, + ) + modified = True - elif isinstance(input_data, list): - for item in input_data: - if not isinstance(item, dict): - continue - content = item.get("content", []) - if isinstance(content, str): - item["content"] = await cb.check_pii( - text=content, - output_parse_pii=True, - presidio_config=presidio_config, - request_data=self.request_data, - ) - modified = True - elif isinstance(content, list): - for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "input_text" - 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 + elif isinstance(input_data, list): + for item in input_data: + if not isinstance(item, dict): + continue + content = item.get("content", []) + if isinstance(content, str): + item["content"] = await cb.check_pii( + text=content, + output_parse_pii=True, + presidio_config=presidio_config, + request_data=self.request_data, + ) + modified = True + elif isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "input_text" + 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 From 6150fac4954cd76f5342b6fa3e8e227e73887e1e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:35:24 +0530 Subject: [PATCH 05/26] style: apply black formatting to llm_http_handler and streaming_iterator Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +--- litellm/responses/streaming_iterator.py | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f02664f5374..2d501a76a54 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5654,9 +5654,7 @@ async def async_responses_websocket( if "model" not in _qs: _qs["model"] = [model] ws_url = urlunparse( - _parsed._replace( - query=urlencode({k: v[0] for k, v in _qs.items()}) - ) + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) ) try: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6b9f4e59887..83f5ad65544 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1398,7 +1398,11 @@ async def _mask_response_create(self, message: str) -> str: # flat: {"type": "response.create", "input": [...], ...} # nested: {"type": "response.create", "response": {"input": [...], ...}} # Mask both so PII is never forwarded unmasked regardless of shape. - nested_response = msg_obj.get("response") if isinstance(msg_obj.get("response"), dict) else None + nested_response = ( + msg_obj.get("response") + if isinstance(msg_obj.get("response"), dict) + else None + ) input_containers: list[tuple[dict, str]] = [] if "input" in msg_obj: input_containers.append((msg_obj, "input")) From ea49d3640d331a648b38c381fef2815fc00904f8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:41:15 +0530 Subject: [PATCH 06/26] style: suppress PLR0915 on async_responses_websocket Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2d501a76a54..1ba12243d3a 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, From fbde3736329fa583c0c994f115b8d2a7d42cfe1e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 11:57:00 +0530 Subject: [PATCH 07/26] fix(responses): add apply_to_output masking on Responses API WebSocket path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the WebSocket guardrail filter excluded callbacks with apply_to_output=True, leaving model-generated PII unmasked before returning to the client. - Collect apply_to_output callbacks separately in llm_http_handler and pass them to ResponsesWebSocketStreaming as output_guardrail_callbacks. - Add _mask_response_completed method that calls check_pii(output_parse_pii=False) on text blocks in response.completed events, masking model output PII. - backend_to_client now chains unmask (pii_tokens) → mask (apply_to_output) before forwarding each event to the client. Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 9 +++- litellm/responses/streaming_iterator.py | 48 ++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1ba12243d3a..df99e32f53a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5685,6 +5685,7 @@ async def async_responses_websocket( # noqa: PLR0915 _request_data["litellm_metadata"] = litellm_metadata _ws_guardrail_callbacks: list = [] + _ws_output_guardrail_callbacks: list = [] try: import litellm as _litellm from litellm.proxy.guardrails.guardrail_hooks.presidio import ( @@ -5696,7 +5697,12 @@ async def async_responses_websocket( # noqa: PLR0915 for cb in _litellm.callbacks if isinstance(cb, _OPTIONAL_PresidioPIIMasking) and getattr(cb, "output_parse_pii", False) - and not getattr(cb, "apply_to_output", False) + ] + _ws_output_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if isinstance(cb, _OPTIONAL_PresidioPIIMasking) + and getattr(cb, "apply_to_output", False) ] except Exception: pass @@ -5709,6 +5715,7 @@ async def async_responses_websocket( # noqa: PLR0915 request_data=_request_data, first_message=first_message, guardrail_callbacks=_ws_guardrail_callbacks, + output_guardrail_callbacks=_ws_output_guardrail_callbacks, ) await streaming.bidirectional_forward() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 83f5ad65544..dc9fa7b6ae0 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1253,6 +1253,7 @@ def __init__( request_data: Optional[Dict] = None, first_message: Optional[str] = None, guardrail_callbacks: Optional[List[Any]] = None, + output_guardrail_callbacks: Optional[List[Any]] = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1263,6 +1264,7 @@ def __init__( 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 [] def _should_store_event(self, event_obj: dict) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1355,7 +1357,8 @@ async def backend_to_client(self) -> None: self._store_event(response_str) unmasked_str = self._unmask_response_completed(response_str) - await self.websocket.send_text(unmasked_str) + output_masked_str = await self._mask_response_completed(unmasked_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) @@ -1494,6 +1497,49 @@ def _unmask_response_completed(self, response_str: str) -> str: return json.dumps(evt_obj) if modified else response_str + async def _mask_response_completed(self, response_str: str) -> str: + """ + Apply Presidio output masking (apply_to_output=True) to a + ``response.completed`` event before it is forwarded to the client. + + Masks model-generated PII in the output text blocks. Events that are + not ``response.completed`` 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", {}) + for output_item in response_obj.get("output", []): + for content_block in output_item.get("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: From 112d5cea1c2ee2834eceea2527633abc1dec90a8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 12:41:00 +0530 Subject: [PATCH 08/26] fix(responses): unmask PII tokens in streaming delta events and warn on guardrail init failure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename _unmask_response_completed -> _unmask_response_event and extend it to also unmask response.output_text.delta (and other delta types) so real-time streaming clients receive original values, not PII tokens. - Split the broad except-and-swallow into ImportError (expected in SDK-only environments) vs Exception (unexpected — now logs a warning so operators know masking is disabled). Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 8 ++- litellm/responses/streaming_iterator.py | 67 +++++++++++++------ 2 files changed, 52 insertions(+), 23 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index df99e32f53a..9f5f97590cd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5704,8 +5704,14 @@ async def async_responses_websocket( # noqa: PLR0915 if isinstance(cb, _OPTIONAL_PresidioPIIMasking) and getattr(cb, "apply_to_output", False) ] - except Exception: + except ImportError: pass + except Exception as _guardrail_exc: + verbose_logger.warning( + "Responses WebSocket: failed to collect Presidio guardrail " + "callbacks — PII masking will be skipped. Error: %s", + _guardrail_exc, + ) streaming = ResponsesWebSocketStreaming( websocket=websocket, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index dc9fa7b6ae0..3f5b17f33f8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1356,7 +1356,7 @@ async def backend_to_client(self) -> None: response_str = raw_response self._store_event(response_str) - unmasked_str = self._unmask_response_completed(response_str) + unmasked_str = self._unmask_response_event(response_str) output_masked_str = await self._mask_response_completed(unmasked_str) await self.websocket.send_text(output_masked_str) @@ -1454,15 +1454,29 @@ async def _mask_response_create(self, message: str) -> str: return json.dumps(msg_obj) if modified else message - def _unmask_response_completed(self, response_str: str) -> str: + # 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", + } + ) + + def _unmask_response_event(self, response_str: str) -> str: """ - Apply Presidio PII unmasking to a ``response.completed`` event before - it is forwarded back to the client. + 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 that are not ``response.completed``, or where no tokens - were stored, are returned unchanged. + value. Events with no stored tokens are returned unchanged. """ if not self.guardrail_callbacks: return response_str @@ -1478,24 +1492,33 @@ def _unmask_response_completed(self, response_str: str) -> str: except (json.JSONDecodeError, TypeError): return response_str - if evt_obj.get("type") != "response.completed": - return response_str - cb = self.guardrail_callbacks[0] - modified = False - response_obj = evt_obj.get("response", {}) - for output_item in response_obj.get("output", []): - for content_block in output_item.get("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 + event_type = evt_obj.get("type") - return json.dumps(evt_obj) if modified else response_str + if event_type == "response.completed": + modified = False + response_obj = evt_obj.get("response", {}) + for output_item in response_obj.get("output", []): + for content_block in output_item.get("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: """ From 9e8ffb3934495208308c37b9d957440a63dfc1f5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 13:03:12 +0530 Subject: [PATCH 09/26] fix(responses): enforce authorized model on WebSocket frames and remove proxy import Security: add _enforce_authorized_model to ResponsesWebSocketStreaming that overwrites both flat and nested model fields in every response.create frame with the connection-authorized model, preventing deployment-substitution attacks where an authenticated user sends a different model name in the frame body after connecting with an allowed model. Layering: remove the _OPTIONAL_PresidioPIIMasking isinstance check and proxy import from the SDK handler. Use duck-typed checks (callable check_pii + get_presidio_settings_from_request_data) so any guardrail implementing the interface works, not just Presidio. Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 21 ++++--- litellm/responses/streaming_iterator.py | 61 +++++++++++++++---- 2 files changed, 62 insertions(+), 20 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9f5f97590cd..969909e3434 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5688,27 +5688,31 @@ async def async_responses_websocket( # noqa: PLR0915 _ws_output_guardrail_callbacks: list = [] try: import litellm as _litellm - from litellm.proxy.guardrails.guardrail_hooks.presidio import ( - _OPTIONAL_PresidioPIIMasking, - ) + # 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 isinstance(cb, _OPTIONAL_PresidioPIIMasking) + if callable(getattr(cb, "check_pii", None)) + and callable( + getattr(cb, "get_presidio_settings_from_request_data", None) + ) and getattr(cb, "output_parse_pii", False) ] _ws_output_guardrail_callbacks = [ cb for cb in _litellm.callbacks - if isinstance(cb, _OPTIONAL_PresidioPIIMasking) + 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 ImportError: - pass except Exception as _guardrail_exc: verbose_logger.warning( - "Responses WebSocket: failed to collect Presidio guardrail " + "Responses WebSocket: failed to collect guardrail " "callbacks — PII masking will be skipped. Error: %s", _guardrail_exc, ) @@ -5722,6 +5726,7 @@ async def async_responses_websocket( # noqa: PLR0915 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/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3f5b17f33f8..c18ecd67e0d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1254,6 +1254,7 @@ def __init__( 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 @@ -1265,6 +1266,9 @@ def __init__( 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 @@ -1367,20 +1371,47 @@ async def backend_to_client(self) -> None: finally: await self._log_messages() - async def _mask_response_create(self, message: str) -> str: + def _enforce_authorized_model(self, msg_obj: dict) -> bool: """ - Apply Presidio PII masking to a ``response.create`` message before it - is forwarded to the upstream provider. - - Walks the ``input`` field of the message (string or list of message - items), 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. + 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.guardrail_callbacks: - return message + if not self.authorized_model: + return False + modified = False + if "model" in msg_obj and msg_obj["model"] != self.authorized_model: + msg_obj["model"] = self.authorized_model + modified = True + nested = msg_obj.get("response") + if ( + isinstance(nested, dict) + and "model" in nested + and nested["model"] != self.authorized_model + ): + nested["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`` field, 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): @@ -1389,10 +1420,16 @@ async def _mask_response_create(self, message: str) -> str: 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 = False + modified = model_modified for cb in self.guardrail_callbacks: presidio_config = cb.get_presidio_settings_from_request_data( self.request_data From 0f66b128c84e51f31ebb12afd5c1c9298110f84f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 14:01:16 +0530 Subject: [PATCH 10/26] fix(responses): add _unmask_pii_text to duck-typed contract and mask delta frames - Add callable(_unmask_pii_text) check to the guardrail_callbacks filter so a custom guardrail missing that method cannot cause an AttributeError and silently kill the WebSocket session. - Extend _mask_response_completed to also mask response.output_text.delta (and other delta types) for apply_to_output callbacks, so real-time streaming clients do not receive unredacted model-generated PII in deltas. Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 1 + litellm/responses/streaming_iterator.py | 65 ++++++++++++------- 2 files changed, 43 insertions(+), 23 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 969909e3434..ca61a430abe 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5699,6 +5699,7 @@ async def async_responses_websocket( # noqa: PLR0915 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 = [ diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c18ecd67e0d..53de36654dd 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1559,11 +1559,16 @@ def _unmask_response_event(self, response_str: str) -> str: async def _mask_response_completed(self, response_str: str) -> str: """ - Apply Presidio output masking (apply_to_output=True) to a - ``response.completed`` event before it is forwarded to the client. + Apply Presidio output masking (apply_to_output=True) to backend events + before they are forwarded to the client. - Masks model-generated PII in the output text blocks. Events that are - not ``response.completed`` are returned unchanged. + - ``response.completed``: masks text in every output content block. + - Streaming delta events (``response.output_text.delta``, etc.): masks + the ``delta`` field. Note that PII entities which span multiple delta + fragments may not be fully detected; ``response.completed`` provides + the authoritative masked view of the full output. + + Events of other types are returned unchanged. """ if not self.output_guardrail_callbacks: return response_str @@ -1573,30 +1578,44 @@ async def _mask_response_completed(self, response_str: str) -> str: except (json.JSONDecodeError, TypeError): return response_str - if evt_obj.get("type") != "response.completed": - return response_str - + event_type = evt_obj.get("type") 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", {}) - for output_item in response_obj.get("output", []): - for content_block in output_item.get("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 + + if event_type == "response.completed": + response_obj = evt_obj.get("response", {}) + for output_item in response_obj.get("output", []): + for content_block in output_item.get("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 + + elif event_type in self._DELTA_EVENT_TYPES: + delta = evt_obj.get("delta") + if isinstance(delta, str) and delta: + masked_delta = await cb.check_pii( + text=delta, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=self.request_data, + ) + if masked_delta != delta: + evt_obj["delta"] = masked_delta + modified = True return json.dumps(evt_obj) if modified else response_str From c4a9027fbed3da94c38905e389a4b00856ba1a97 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 9 Jun 2026 10:52:43 +0000 Subject: [PATCH 11/26] Fix Responses WebSocket guardrail edge cases --- .../llms/azure/responses/transformation.py | 15 +++-- .../guardrail_translation/handler.py | 11 ++-- litellm/responses/streaming_iterator.py | 44 ++++++++----- ...test_openai_responses_guardrail_handler.py | 14 +++++ .../test_responses_websocket_all_providers.py | 63 +++++++++++++++++++ 5 files changed, 123 insertions(+), 24 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 1f6ad804107..92ce5b49285 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -201,14 +201,19 @@ def get_websocket_url( if api_base is None: raise ValueError("api_base is required for Azure WebSocket") - base = api_base.rstrip("/") + 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 base.endswith(suffix): - base = base[: -len(suffix)] + if path.endswith(suffix): + path = path[: -len(suffix)] break - ws_url = base.replace("https://", "wss://").replace("http://", "ws://") - return f"{ws_url}/openai/v1/responses" + 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 diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b09d4ef2da7..eaa822dd990 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -502,7 +502,10 @@ async def process_output_streaming_response( # final chunk; iterate output items, apply guardrail, write back. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.completed": - outputs: List[Any] = final_chunk.get("response", {}).get("output", []) + 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] = [] @@ -522,7 +525,7 @@ async def process_output_streaming_response( if request_data is None: request_data = {} if "response" not in request_data: - request_data["response"] = final_chunk.get("response", {}) + 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 @@ -535,7 +538,7 @@ async def process_output_streaming_response( inputs["tool_calls"] = cast( List[ChatCompletionToolCallChunk], tool_calls_to_check ) - response_model = final_chunk.get("response", {}).get("model") + response_model = response_obj.get("model") if response_model: inputs["model"] = response_model @@ -552,7 +555,7 @@ async def process_output_streaming_response( # 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=final_chunk.get("response", {}), + response=response_obj, responses=guardrailed_texts, task_mappings=task_mappings, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 53de36654dd..a0bd1db9c79 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1385,16 +1385,16 @@ def _enforce_authorized_model(self, msg_obj: dict) -> bool: if not self.authorized_model: return False modified = False - if "model" in msg_obj and msg_obj["model"] != self.authorized_model: - msg_obj["model"] = self.authorized_model - modified = True nested = msg_obj.get("response") - if ( - isinstance(nested, dict) - and "model" in nested - and nested["model"] != self.authorized_model - ): - nested["model"] = self.authorized_model + 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 @@ -1534,9 +1534,16 @@ def _unmask_response_event(self, response_str: str) -> str: if event_type == "response.completed": modified = False - response_obj = evt_obj.get("response", {}) - for output_item in response_obj.get("output", []): - for content_block in output_item.get("content", []): + 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") @@ -1587,9 +1594,16 @@ async def _mask_response_completed(self, response_str: str) -> str: ) if event_type == "response.completed": - response_obj = evt_obj.get("response", {}) - for output_item in response_obj.get("output", []): - for content_block in output_item.get("content", []): + 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 + 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") 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 5a4457c7ebd..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 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 966f6688926..d3c5de87d61 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -71,6 +71,14 @@ def test_azure_websocket_url_strips_existing_path(self): ) 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_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() @@ -796,6 +804,61 @@ 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 + + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" From 29cc383e5313b018c3ac16feb15f14d95a0193e1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 17:20:10 +0530 Subject: [PATCH 12/26] fix(responses): log masked output and suppress deltas when apply_to_output active - Move _store_event to after _mask_response_completed so logs receive the redacted form, not raw model output containing PII. - Suppress delta event forwarding when output_guardrail_callbacks are present: per-fragment Presidio cannot catch PII that spans multiple chunks (e.g. "alice@" + "example.com"). Clients receive only the fully-masked response.completed, which Presidio scans on complete text. Co-Authored-By: Claude Sonnet 4.6 --- litellm/responses/streaming_iterator.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index a0bd1db9c79..e16d93e0dff 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1359,9 +1359,25 @@ async def backend_to_client(self) -> None: else: response_str = raw_response - self._store_event(response_str) 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) + + # When apply_to_output masking is active, suppress delta events: + # per-fragment Presidio cannot reliably catch PII that spans + # multiple delta chunks (e.g. "alice@" + "example.com"). + # Clients receive the fully-masked response.completed instead. + if self.output_guardrail_callbacks: + try: + _evt_type = json.loads(output_masked_str).get("type") + except (json.JSONDecodeError, TypeError): + _evt_type = None + if _evt_type in self._DELTA_EVENT_TYPES: + continue + await self.websocket.send_text(output_masked_str) except websockets.exceptions.ConnectionClosed as e: # type: ignore From 1ce3339cf54ea5ef554b4d2489187b1ac7c7434d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 18:04:26 +0530 Subject: [PATCH 13/26] fix(responses): mask and suppress response.output_item.done for apply_to_output response.output_item.done carries completed item text in item.content[*].text before response.completed arrives, allowing unmasked PII to reach the client. - _unmask_response_event: unmask input-PII tokens in item.content[*].text - _mask_response_completed: run check_pii on item.content[*].text for apply_to_output callbacks (same as response.completed handling) - backend_to_client suppression: also skip response.output_item.done when output_guardrail_callbacks are active; client receives only the fully-masked response.completed Co-Authored-By: Claude Sonnet 4.6 --- litellm/responses/streaming_iterator.py | 45 ++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index e16d93e0dff..4314902b65d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1366,16 +1366,20 @@ async def backend_to_client(self) -> None: # guardrails does not appear in success logs. self._store_event(output_masked_str) - # When apply_to_output masking is active, suppress delta events: - # per-fragment Presidio cannot reliably catch PII that spans - # multiple delta chunks (e.g. "alice@" + "example.com"). - # Clients receive the fully-masked response.completed instead. + # When apply_to_output masking is active, suppress delta events + # and response.output_item.done: per-fragment Presidio cannot + # reliably catch PII that spans multiple delta chunks, and + # output_item.done carries the same text that response.completed + # already delivers in fully-masked form. if self.output_guardrail_callbacks: try: _evt_type = json.loads(output_masked_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None - if _evt_type in self._DELTA_EVENT_TYPES: + if ( + _evt_type in self._DELTA_EVENT_TYPES + or _evt_type == "response.output_item.done" + ): continue await self.websocket.send_text(output_masked_str) @@ -1578,6 +1582,20 @@ def _unmask_response_event(self, response_str: str) -> str: evt_obj["delta"] = unmasked return json.dumps(evt_obj) + if event_type == "response.output_item.done": + modified = False + item = evt_obj.get("item") or {} + for content_block in item.get("content") or []: + 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 + return response_str async def _mask_response_completed(self, response_str: str) -> str: @@ -1634,6 +1652,23 @@ async def _mask_response_completed(self, response_str: str) -> str: content_block["text"] = masked modified = True + elif event_type == "response.output_item.done": + item = evt_obj.get("item") or {} + for content_block in item.get("content") or []: + if not isinstance(content_block, dict): + continue + text = content_block.get("text") + if isinstance(text, str) and text: + 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 + elif event_type in self._DELTA_EVENT_TYPES: delta = evt_obj.get("delta") if isinstance(delta, str) and delta: From 562de5e41ab4095a32d66afa3166094f5a76b57e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 18:29:38 +0530 Subject: [PATCH 14/26] fix(types): cast response_obj to ResponsesAPIResponse to satisfy mypy Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/openai/responses/guardrail_translation/handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index eaa822dd990..dc4c540138a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -555,7 +555,7 @@ async def process_output_streaming_response( # 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, + response=cast("ResponsesAPIResponse", response_obj), responses=guardrailed_texts, task_mappings=task_mappings, ) From 2ae01dfd7136dc63d97cadde8935d8ff69be9986 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 22:27:47 +0530 Subject: [PATCH 15/26] Revert "fix(types): cast response_obj to ResponsesAPIResponse to satisfy mypy" This reverts commit d5969557628f9aff58948b9d37cc64d577f95a15. --- litellm/llms/openai/responses/guardrail_translation/handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index dc4c540138a..eaa822dd990 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -555,7 +555,7 @@ async def process_output_streaming_response( # 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=cast("ResponsesAPIResponse", response_obj), + response=response_obj, responses=guardrailed_texts, task_mappings=task_mappings, ) From 19e4c97e0f68c59818531a1658bacb5cb21cd6ea Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 22:27:47 +0530 Subject: [PATCH 16/26] Revert "fix(responses): mask and suppress response.output_item.done for apply_to_output" This reverts commit 219fd54ea3446f4399fde40c07ba0617e2834573. --- litellm/responses/streaming_iterator.py | 45 +++---------------------- 1 file changed, 5 insertions(+), 40 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 4314902b65d..e16d93e0dff 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1366,20 +1366,16 @@ async def backend_to_client(self) -> None: # guardrails does not appear in success logs. self._store_event(output_masked_str) - # When apply_to_output masking is active, suppress delta events - # and response.output_item.done: per-fragment Presidio cannot - # reliably catch PII that spans multiple delta chunks, and - # output_item.done carries the same text that response.completed - # already delivers in fully-masked form. + # When apply_to_output masking is active, suppress delta events: + # per-fragment Presidio cannot reliably catch PII that spans + # multiple delta chunks (e.g. "alice@" + "example.com"). + # Clients receive the fully-masked response.completed instead. if self.output_guardrail_callbacks: try: _evt_type = json.loads(output_masked_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None - if ( - _evt_type in self._DELTA_EVENT_TYPES - or _evt_type == "response.output_item.done" - ): + if _evt_type in self._DELTA_EVENT_TYPES: continue await self.websocket.send_text(output_masked_str) @@ -1582,20 +1578,6 @@ def _unmask_response_event(self, response_str: str) -> str: evt_obj["delta"] = unmasked return json.dumps(evt_obj) - if event_type == "response.output_item.done": - modified = False - item = evt_obj.get("item") or {} - for content_block in item.get("content") or []: - 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 - return response_str async def _mask_response_completed(self, response_str: str) -> str: @@ -1652,23 +1634,6 @@ async def _mask_response_completed(self, response_str: str) -> str: content_block["text"] = masked modified = True - elif event_type == "response.output_item.done": - item = evt_obj.get("item") or {} - for content_block in item.get("content") or []: - if not isinstance(content_block, dict): - continue - text = content_block.get("text") - if isinstance(text, str) and text: - 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 - elif event_type in self._DELTA_EVENT_TYPES: delta = evt_obj.get("delta") if isinstance(delta, str) and delta: From b587c497bfbf4f2443f85eba1d07ff3eff11dde1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 10 Jun 2026 16:35:07 +0530 Subject: [PATCH 17/26] fix(types): accept dict responses in guardrail output write-back Streaming response.completed events pass a dict response object, so widen _apply_guardrail_responses_to_output to match its existing runtime handling. Co-authored-by: Cursor --- litellm/llms/openai/responses/guardrail_translation/handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index eaa822dd990..b5319797cc6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -767,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: From 2b6edfeed6dd6daa31a622eaf54c52b16515e4d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 04:58:30 +0000 Subject: [PATCH 18/26] perf(responses): skip Presidio masking on suppressed WebSocket delta events Delta events are dropped wholesale when apply_to_output masking is active, so masking them first issued a wasted check_pii call per fragment. Move the suppression check ahead of the unmask/mask passes; the event type is invariant across both, so client-visible behavior is unchanged. --- litellm/responses/streaming_iterator.py | 22 +++--- .../test_responses_websocket_all_providers.py | 74 +++++++++++++++++++ 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index e16d93e0dff..b9e124d89f9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1359,25 +1359,27 @@ async def backend_to_client(self) -> None: else: response_str = raw_response - 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) - # When apply_to_output masking is active, suppress delta events: # per-fragment Presidio cannot reliably catch PII that spans - # multiple delta chunks (e.g. "alice@" + "example.com"). - # Clients receive the fully-masked response.completed instead. + # multiple delta chunks (e.g. "alice@" + "example.com"), and the + # client receives the fully-masked response.completed instead. + # Drop them before masking so no Presidio call is wasted on a + # fragment that is never forwarded. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(output_masked_str).get("type") + _evt_type = json.loads(response_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None if _evt_type in self._DELTA_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 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 d3c5de87d61..7cddf2baa9e 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -858,6 +858,80 @@ def _unmask_pii_text(self, text, pii_tokens): 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" + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" From 28b86890af4a0d12ed03d9918f3870161557ae09 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:05:25 +0000 Subject: [PATCH 19/26] test(responses): cover Responses WebSocket PII masking hooks Add regression tests for the native Responses WebSocket guardrail path: input masking and model enforcement in _mask_response_create, token unmasking in _unmask_response_event, apply_to_output masking and delta suppression in _mask_response_completed/backend_to_client, and the get_websocket_url / model_in_websocket_url defaults for the base and Azure configs. Raises diff coverage above the codecov patch target. --- .../test_responses_websocket_all_providers.py | 529 ++++++++++++++++++ 1 file changed, 529 insertions(+) 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 7cddf2baa9e..2780871de5b 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 @@ -79,6 +82,26 @@ def test_azure_websocket_url_strips_query_params(self): ) 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() @@ -933,6 +956,512 @@ async def recv(self, decode=False): assert json.loads(sent_payload)["type"] == "response.completed" +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_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_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_delta(self): + guardrail = _FakeWSGuardrail() + handler = _make_streaming( + request_data={}, output_guardrail_callbacks=[guardrail] + ) + + event = json.dumps( + {"type": "response.output_text.delta", "delta": "alice@example.com"} + ) + masked = json.loads(await handler._mask_response_completed(event)) + assert masked["delta"] == "" + + @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 " + ) + + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" From 1af5562213993a67cc51535f6c1c6f036eb82355 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:27:14 +0000 Subject: [PATCH 20/26] fix(responses): suppress text-bearing done events under output PII masking When apply_to_output masking is active on a native Responses WebSocket, response.output_text.done, response.content_part.done, and response.output_item.done carry the full model output before the masked response.completed arrives, so an authenticated client could read unmasked PII from those events. Suppress them alongside delta events; the client receives only the fully-masked response.completed. --- litellm/responses/streaming_iterator.py | 30 ++++-- .../test_responses_websocket_all_providers.py | 96 +++++++++++++++++++ 2 files changed, 119 insertions(+), 7 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index b9e124d89f9..4bd1f3b74aa 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1359,18 +1359,23 @@ async def backend_to_client(self) -> None: else: response_str = raw_response - # When apply_to_output masking is active, suppress delta events: - # per-fragment Presidio cannot reliably catch PII that spans - # multiple delta chunks (e.g. "alice@" + "example.com"), and the - # client receives the fully-masked response.completed instead. - # Drop them before masking so no Presidio call is wasted on a - # fragment that is never forwarded. + # 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: + 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) @@ -1519,6 +1524,17 @@ async def _mask_response_create(self, message: str) -> str: } ) + # Terminal events that carry the full output text 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", + } + ) + def _unmask_response_event(self, response_str: str) -> str: """ Apply Presidio PII unmasking to backend events before forwarding to 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 2780871de5b..ce088c84af9 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -955,6 +955,102 @@ async def recv(self, decode=False): 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. From 6a60a4c2a9bf63b4e68b75cf85e5f5877428acb1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:30:50 +0000 Subject: [PATCH 21/26] refactor(responses): drop dead delta branch in WebSocket output masking Delta events are suppressed in backend_to_client before _mask_response_completed runs when output masking is active, so the method's delta-handling branch was unreachable. Restrict it to response.completed and cover the Responses API unmask path with a Pydantic ResponseCompletedEvent regression test. --- litellm/responses/streaming_iterator.py | 77 ++++++++----------- .../guardrail_hooks/test_presidio.py | 69 +++++++++++++++++ .../test_responses_websocket_all_providers.py | 5 +- 3 files changed, 101 insertions(+), 50 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 4bd1f3b74aa..9a4375caeab 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1600,16 +1600,13 @@ def _unmask_response_event(self, response_str: str) -> str: async def _mask_response_completed(self, response_str: str) -> str: """ - Apply Presidio output masking (apply_to_output=True) to backend events - before they are forwarded to the client. + Apply Presidio output masking (apply_to_output=True) to the + ``response.completed`` event before it is forwarded to the client. - - ``response.completed``: masks text in every output content block. - - Streaming delta events (``response.output_text.delta``, etc.): masks - the ``delta`` field. Note that PII entities which span multiple delta - fragments may not be fully detected; ``response.completed`` provides - the authoritative masked view of the full output. - - Events of other types are returned unchanged. + Walks ``response.output[*].content[*].text`` and masks every text block. + Delta 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 @@ -1619,51 +1616,37 @@ async def _mask_response_completed(self, response_str: str) -> str: except (json.JSONDecodeError, TypeError): return response_str - event_type = evt_obj.get("type") - modified = False + 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 ) - - if event_type == "response.completed": - response_obj = evt_obj.get("response") or {} - if not isinstance(response_obj, dict): + 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 - 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): + content = output_item.get("content") or [] + if not isinstance(content, list): + continue + for content_block in content: + if not isinstance(content_block, dict): 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 - - elif event_type in self._DELTA_EVENT_TYPES: - delta = evt_obj.get("delta") - if isinstance(delta, str) and delta: - masked_delta = await cb.check_pii( - text=delta, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=self.request_data, - ) - if masked_delta != delta: - evt_obj["delta"] = masked_delta - modified = True + 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 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 9f6084fe9fc..7e74f245098 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2370,6 +2370,75 @@ async def mock_stream(): 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_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 ce088c84af9..f1570c0c50b 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1406,7 +1406,7 @@ async def test_mask_response_completed_event(self): ) @pytest.mark.asyncio - async def test_mask_response_completed_delta(self): + async def test_mask_response_completed_delta_unchanged(self): guardrail = _FakeWSGuardrail() handler = _make_streaming( request_data={}, output_guardrail_callbacks=[guardrail] @@ -1415,8 +1415,7 @@ async def test_mask_response_completed_delta(self): event = json.dumps( {"type": "response.output_text.delta", "delta": "alice@example.com"} ) - masked = json.loads(await handler._mask_response_completed(event)) - assert masked["delta"] == "" + assert await handler._mask_response_completed(event) == event @pytest.mark.asyncio async def test_mask_response_completed_no_guardrails_unchanged(self): From 6d4917790ad9f14172d8c812588bc16e3859d926 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:50:07 +0000 Subject: [PATCH 22/26] fix(presidio): flush buffered chat chunks on mixed unmask stream _stream_pii_unmasking buffered ModelResponseStream chunks but returned early once a /v1/responses event was seen, silently dropping the buffered chat chunks. Flush them in order before switching to passthrough, mirroring _stream_apply_output_masking, and cover it with a regression test. --- .../guardrails/guardrail_hooks/presidio.py | 11 ++++- .../guardrail_hooks/test_presidio.py | 48 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 4b2e8fc59c8..e723c07e3c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1310,7 +1310,10 @@ async def _stream_pii_unmasking( 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] @@ -1319,6 +1322,12 @@ async def _stream_pii_unmasking( 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) 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 7e74f245098..3efc42523f1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2439,6 +2439,54 @@ async def mock_stream(): ) +@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(): """ From 5eb9c805d1060d3c65fa37040faf19415a921cbf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:14:15 +0000 Subject: [PATCH 23/26] fix(responses): mask instructions and tool-call arguments in WebSocket PII path Presidio masking on the native Responses WebSocket path left two gaps. On the request side _mask_response_create only walked the input containers, so PII placed in the instructions field of a response.create frame was forwarded upstream and logged unmasked even with output_parse_pii enabled. Now both the flat and nested instructions strings are masked alongside input. On the response side _mask_response_completed only masked content text blocks, so model-produced PII inside function-call arguments could reach the client when apply_to_output was enabled, both via the standalone response.function_call_arguments.done event and via the function_call output items in response.completed. The done event is now suppressed under output masking and completed function-call arguments are run through check_pii before forwarding or logging. --- litellm/responses/streaming_iterator.py | 66 +++++---- .../test_responses_websocket_all_providers.py | 131 ++++++++++++++++++ 2 files changed, 173 insertions(+), 24 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9a4375caeab..6003315f265 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1429,8 +1429,8 @@ async def _mask_response_create(self, message: str) -> str: - Overwrites any ``model`` field with the connection-authorized model to prevent deployment-substitution attacks (always applied). - - Walks the ``input`` field, calls ``check_pii`` on every text block, - and stores the resulting ``pii_tokens`` map in + - 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. @@ -1457,35 +1457,39 @@ async def _mask_response_create(self, message: str) -> str: presidio_config = cb.get_presidio_settings_from_request_data( self.request_data ) - # response.create supports two shapes: - # flat: {"type": "response.create", "input": [...], ...} - # nested: {"type": "response.create", "response": {"input": [...], ...}} - # Mask both so PII is never forwarded unmasked regardless of shape. + # 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 ) - input_containers: list[tuple[dict, str]] = [] - if "input" in msg_obj: - input_containers.append((msg_obj, "input")) - if nested_response is not None and "input" in nested_response: - input_containers.append((nested_response, "input")) + 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 input_containers: - input_data = container[key] + for container, key in text_containers: + field_value = container[key] - if isinstance(input_data, str): + if isinstance(field_value, str): container[key] = await cb.check_pii( - text=input_data, + text=field_value, output_parse_pii=True, presidio_config=presidio_config, request_data=self.request_data, ) modified = True - elif isinstance(input_data, list): - for item in input_data: + elif isinstance(field_value, list): + for item in field_value: if not isinstance(item, dict): continue content = item.get("content", []) @@ -1524,14 +1528,16 @@ async def _mask_response_create(self, message: str) -> str: } ) - # Terminal events that carry the full output text already delivered by - # ``response.completed``. Suppressed when output masking is active so the - # unmasked copy never reaches the client before the masked completed event. + # 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", } ) @@ -1603,10 +1609,11 @@ 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. - Delta 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. + Walks ``response.output[*].content[*].text`` and masks every text block, + as well as ``response.output[*].arguments`` on function-call 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 @@ -1630,6 +1637,17 @@ async def _mask_response_completed(self, response_str: str) -> str: 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 content = output_item.get("content") or [] if not isinstance(content, list): continue 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 f1570c0c50b..8778fa67417 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1196,6 +1196,50 @@ async def test_mask_response_create_nested_shape(self): 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() @@ -1405,6 +1449,33 @@ async def test_mask_response_completed_event(self): == "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_delta_unchanged(self): guardrail = _FakeWSGuardrail() @@ -1556,6 +1627,66 @@ async def test_backend_to_client_suppresses_deltas_and_masks_completed(self): == "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 + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" From 02f53742622bffeed49bac5052a603f03ae0eba5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:40:18 +0000 Subject: [PATCH 24/26] fix(responses): suppress reasoning_summary_text.done under output PII masking --- litellm/responses/streaming_iterator.py | 1 + .../test_responses_websocket_all_providers.py | 59 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6003315f265..d24c9225e45 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1538,6 +1538,7 @@ async def _mask_response_create(self, message: str) -> str: "response.content_part.done", "response.output_item.done", "response.function_call_arguments.done", + "response.reasoning_summary_text.done", } ) 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 8778fa67417..1b346eec2ac 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1687,6 +1687,65 @@ async def test_backend_to_client_suppresses_function_call_arguments_done(self): ) 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 + class TestWebSocketChunkTypes: """Test handling of different chunk types from streaming responses""" From 846e34fbbdb3ef8c4d395e343793214bfcd9658a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:40:59 +0000 Subject: [PATCH 25/26] fix(responses): mask function_call_output.output in WebSocket PII path response.create input items of type function_call_output carry user-controlled text in output, not content, so the Presidio masking pass forwarded that text upstream unmasked. Mask the output field (string or list of text blocks) alongside content. --- litellm/responses/streaming_iterator.py | 52 ++++++++++-------- .../test_responses_websocket_all_providers.py | 54 +++++++++++++++++++ 2 files changed, 83 insertions(+), 23 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d24c9225e45..5740858813f 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: """ @@ -1492,29 +1496,31 @@ async def _mask_response_create(self, message: str) -> str: for item in field_value: if not isinstance(item, dict): continue - content = item.get("content", []) - if isinstance(content, str): - item["content"] = await cb.check_pii( - text=content, - output_parse_pii=True, - presidio_config=presidio_config, - request_data=self.request_data, - ) - modified = True - elif isinstance(content, list): - for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "input_text" - 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 + 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 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 1b346eec2ac..1b6760746a8 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1174,6 +1174,60 @@ async def test_mask_response_create_input_text_blocks(self): 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() From 495f9bff1ff7a322ead0d3d7d7ff14e0078dc909 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 07:02:31 +0000 Subject: [PATCH 26/26] fix(responses): mask reasoning summary PII in WebSocket output path --- litellm/responses/streaming_iterator.py | 22 ++++- .../test_responses_websocket_all_providers.py | 99 +++++++++++++++++++ 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5740858813f..27ecba83be7 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1545,6 +1545,7 @@ async def _mask_response_create(self, message: str) -> str: "response.output_item.done", "response.function_call_arguments.done", "response.reasoning_summary_text.done", + "response.reasoning_summary_part.done", } ) @@ -1617,8 +1618,9 @@ async def _mask_response_completed(self, response_str: str) -> str: ``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. Delta - and ``*.done`` events are suppressed upstream in ``backend_to_client`` when + 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. """ @@ -1655,6 +1657,22 @@ async def _mask_response_completed(self, response_str: str) -> str: 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 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 1b6760746a8..dc83ebd40d5 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1530,6 +1530,37 @@ async def test_mask_response_completed_masks_function_call_arguments(self): == '{"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() @@ -1800,6 +1831,74 @@ async def test_backend_to_client_suppresses_reasoning_summary_text_done(self): 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"""