Merge pull request #41558 from BerriAI/litellm_lit_6568_streaming_redaction

fix(guardrails): stream Prompt Security post_call redactions in incremental_diff mode

litellm/integrations/custom_guardrail.py

@@ -1379,8 +1379,9 @@ class CustomGuardrail(CustomLogger):
         raise e
 
     def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool:
-        """True when any key of either mapping differs between them (mask), False otherwise (allow)."""
-        return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys())
+        """True when any content key of either mapping differs between them (mask), False otherwise (allow)."""
+        compared_keys: Final = (original_inputs.keys() | response.keys()) - _STREAM_CONTROL_KEYS
+        return any(original_inputs.get(key) != response.get(key) for key in compared_keys)
 
     def mask_content_in_string(
         self,
@@ -1490,6 +1491,7 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object)
 _PRE_CALL_CONTENT_KEYS: Final = frozenset(
     {"messages", "input", "prompt", "system", "instructions", "tools", "functions", "function_call", "tool_choice"}
 )
+_STREAM_CONTROL_KEYS: Final = frozenset({"stream_holdback_chars"})
 
 
 def _original_inputs_for(

litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py

@@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         guardrail_name=guardrail.get("guardrail_name", ""),
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
+        streaming_transform_mode=getattr(litellm_params, "streaming_transform_mode", None),
         file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
         block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None),
     )

litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py

@@ -38,6 +38,11 @@ class PromptSecurityGuardrailMissingSecrets(Exception):
     pass
 
 
+def _modified_or_original(text: str, verdict: "_ProtectVerdict") -> str:
+    modified_text: Final = verdict.get("modified_text") if verdict.get("action") == "modify" else None
+    return text if modified_text is None else modified_text
+
+
 def _inputs_with_structured_messages(
     inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
 ) -> GenericGuardrailAPIInputs:
@@ -119,6 +124,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         user: str | None = None,
         system_prompt: str | None = None,
         check_tool_results: bool | None = None,
+        streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
         file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
         file_sanitization_fail_open: bool | None = None,
         block_on_file_modify: bool | None = None,
@@ -148,6 +154,10 @@ class PromptSecurityGuardrail(CustomGuardrail):
             )
             raise PromptSecurityGuardrailMissingSecrets(msg)
 
+        self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = (
+            "block_only" if streaming_transform_mode is None else streaming_transform_mode
+        )
+
         # Configuration for file sanitization
         self.max_poll_attempts = 30  # Maximum number of polling attempts
         self.poll_interval = 2  # Seconds between polling attempts
@@ -342,16 +352,46 @@ class PromptSecurityGuardrail(CustomGuardrail):
         texts: list[str],
         user_api_key_alias: str | None,
     ) -> GenericGuardrailAPIInputs:
-        """Handle response-side guardrail checks."""
+        """Handle response-side guardrail checks, one protect verdict per text.
+
+        Prompt Security rewrites a single string, so texts from several choices must be scanned separately
+        or one ``modified_text`` cannot be mapped back onto the choice it came from. It also returns no span
+        offsets, so on a stream every text is held back in full until the final verdict: a value the vendor
+        redacts later may start anywhere in text that looked clean so far, and streamed bytes cannot be recalled.
+        """
         if not texts:
             return inputs
 
-        # Combine all texts for response checking
-        combined_text: Final = "\n".join(texts)
+        verdicts: Final = await asyncio.gather(
+            *(self._protect_response_text(text, user_api_key_alias) for text in texts)
+        )
+        violations: Final = tuple(
+            violation
+            for verdict in verdicts
+            if verdict.get("action") == "block"
+            for violation in verdict.get("violations", ())
+        )
+        if any(verdict.get("action") == "block" for verdict in verdicts):
+            raise HTTPException(
+                status_code=400,
+                detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
+            )
+        returned_texts: Final = [  # mutable-ok: GenericGuardrailAPIInputs.texts is list[str]
+            _modified_or_original(text, verdict) for text, verdict in zip(texts, verdicts, strict=True)
+        ]
+        patched: Final[GenericGuardrailAPIInputs] = {
+            **inputs,
+            "texts": returned_texts,
+            "stream_holdback_chars": [  # mutable-ok: GenericGuardrailAPIInputs.stream_holdback_chars is list[int]
+                len(text) for text in returned_texts
+            ],
+        }
+        return patched
 
+    async def _protect_response_text(self, text: str, user_api_key_alias: str | None) -> _ProtectVerdict:
         headers: Final = self._build_headers(user_api_key_alias)
         payload: Final = {
-            "response": combined_text,
+            "response": text,
             "user": user_api_key_alias or self.user,
             "system_prompt": self.system_prompt,
         }
@@ -360,7 +400,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
             method="POST",
             url=f"{self.api_base}/api/protect",
             headers=headers,
-            payload={"response_length": len(combined_text)},
+            payload={"response_length": len(text)},
         )
 
         response: Final = await self.async_handler.post(
@@ -377,26 +417,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
             payload={"result": res.get("result")},
         )
 
-        result: Final = res.get("result", {}).get("response", {})
-        if result is None:
-            return inputs
-
-        action: Final = result.get("action")
-        violations: Final = result.get("violations", [])
-
-        if action == "block":
-            raise HTTPException(
-                status_code=400,
-                detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
-            )
-        elif action == "modify":
-            modified_text: Final = result.get("modified_text")
-            if modified_text is not None:
-                # If we combined multiple texts, return the modified version as single text
-                # The framework will handle distributing it back
-                inputs["texts"] = [modified_text]
-
-        return inputs
+        verdict: Final = res.get("result", {}).get("response", {})
+        return {} if verdict is None else verdict
 
     def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
         return [text for message in messages for text in message_slot_texts(message)]

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -104,6 +104,10 @@ def _chunk_choices(item: object) -> Sequence[object]:
     return choices
 
 
+def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
+    return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
+
+
 def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
     if scan_key is None:
         return False
@@ -472,6 +476,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         emitted_text_per_choice: dict[int, str],
         holdback_per_choice: dict[int, int],
         finish_reason_per_choice: dict[int, str | None],
+        held_chars_per_choice: dict[int, int],
         is_final: bool,
     ) -> ModelResponseStream | None:
         """Build the synthetic chunk carrying the newly-guardrailed deltas.
@@ -479,7 +484,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         For each choice, the new delta is the mutated accumulated text past what
         has already been emitted, minus a trailing holdback (forced to 0 on the
         final flush). ``emitted_text_per_choice`` holds the exact bytes already
-        sent per choice and is extended in place. Returns None when there is no
+        sent per choice and is extended in place; ``held_chars_per_choice`` is
+        updated in place with how many mutated chars per choice are still withheld
+        after this round. Returns None when there is no
         text to emit (e.g. a tool-call-only turn) or nothing new and this is not
         the final chunk.
 
@@ -536,6 +543,7 @@ class UnifiedLLMGuardrails(CustomLogger):
             holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
             end = max(len(already), len(text) - holdback)
             deltas[choice_idx] = text[len(already) : end]
+            held_chars_per_choice[choice_idx] = len(text) - end
 
         # Iterate the mutated choices (not just those in reference_chunk) so a
         # choice with pending text is never dropped for n > 1. finish_reason is
@@ -590,6 +598,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         responses_yielded: list[object],
         emitted_text_per_choice: dict[int, str],
         finish_reason_per_choice: dict[int, str | None],
+        held_chars_per_choice: dict[int, int],
         is_final: bool,
     ) -> AsyncGenerator[object, None]:
         """Run one guardrail processing round and emit the resulting diff chunk.
@@ -618,6 +627,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                 emitted_text_per_choice=emitted_text_per_choice,
                 holdback_per_choice=sink.holdback_per_choice,
                 finish_reason_per_choice=finish_reason_per_choice,
+                held_chars_per_choice=held_chars_per_choice,
                 is_final=is_final,
             )
         except ModifyResponseException as e:
@@ -673,6 +683,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         responses_yielded: Final[list[object]] = []
         emitted_text_per_choice: Final[dict[int, str]] = {}
         finish_reason_per_choice: Final[dict[int, str | None]] = {}
+        held_chars_per_choice: Final[dict[int, int]] = {}
         chunk_counter = 0
         last_chunk: object | None = None
 
@@ -688,6 +699,7 @@ class UnifiedLLMGuardrails(CustomLogger):
                 responses_yielded=responses_yielded,
                 emitted_text_per_choice=emitted_text_per_choice,
                 finish_reason_per_choice=finish_reason_per_choice,
+                held_chars_per_choice=held_chars_per_choice,
                 is_final=is_final,
             )
 
@@ -724,12 +736,18 @@ class UnifiedLLMGuardrails(CustomLogger):
                     # finish_reason to the final text terminator (see the
                     # _tool_call_passthrough_chunk docstring).
                     tool_only = self._tool_call_passthrough_chunk(
-                        item, finish_reason_per_choice=finish_reason_per_choice
+                        item,
+                        finish_reason_per_choice=finish_reason_per_choice,
+                        held_choices=_held_choices(held_chars_per_choice),
                     )
                     responses_yielded.append(tool_only)
                     yield tool_only
                     continue
 
+                if self._is_trailing_metadata_chunk(item):
+                    responses_so_far.append(item)
+                    continue
+
                 chunk_counter += 1
                 responses_so_far.append(item)
                 last_chunk = item
@@ -773,12 +791,33 @@ class UnifiedLLMGuardrails(CustomLogger):
                 ):
                     yield out
 
-            if last_chunk is not None:
-                async for out in _round(last_chunk, is_final=True):
-                    yield out
+            async for out in self._emit_stream_tail(
+                last_chunk=last_chunk,
+                final_round=_round,
+                responses_so_far=responses_so_far,
+                responses_yielded=responses_yielded,
+            ):
+                yield out
         except _StreamTerminated:
             return
 
+    async def _emit_stream_tail(
+        self,
+        *,
+        last_chunk: object | None,
+        final_round: Callable[[object, bool], AsyncGenerator[object, None]],
+        responses_so_far: Sequence[object],
+        responses_yielded: list[object],
+    ) -> AsyncGenerator[object, None]:
+        """Flush the held text with holdback 0, then replay metadata-only chunks
+        (usage) so they land after the text and its finish_reason, as upstream sent them."""
+        if last_chunk is not None:
+            async for out in final_round(last_chunk, True):
+                yield out
+        for trailing in self._trailing_metadata_chunks(responses_so_far):
+            responses_yielded.append(trailing)
+            yield trailing
+
     async def _inspect_full_response_for_block(
         self,
         *,
@@ -829,6 +868,23 @@ class UnifiedLLMGuardrails(CustomLogger):
                 return True
         return False
 
+    @classmethod
+    def _is_trailing_metadata_chunk(cls, item: object) -> bool:
+        """True for a chunk that carries only stream metadata (no choices, or a
+        ``usage`` chunk whose deltas are empty); such chunks are replayed after
+        the final text flush instead of being folded into the transform."""
+        if not _chunk_choices(item):
+            return True
+        return (
+            getattr(item, "usage", None) is not None
+            and not cls._chunk_carries_text(item)
+            and not cls._chunk_has_finish_reason(item)
+        )
+
+    @classmethod
+    def _trailing_metadata_chunks(cls, items: Sequence[object]) -> tuple[object, ...]:
+        return tuple(item for item in items if cls._is_trailing_metadata_chunk(item))
+
     @staticmethod
     def _chunk_carries_text(item: object) -> bool:
         """True if any choice in this chunk has non-empty string ``delta.content``."""
@@ -843,6 +899,7 @@ class UnifiedLLMGuardrails(CustomLogger):
     def _tool_call_passthrough_chunk(
         item: object,
         finish_reason_per_choice: "dict[int, str | None] | None" = None,
+        held_choices: frozenset[int] = frozenset(),
     ) -> ModelResponseStream:
         """Copy of a chunk carrying tool calls with all text content stripped.
 
@@ -851,8 +908,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         transform instead). Applies per choice so an n>1 chunk mixing a text
         choice and a tool-call choice does not leak the text choice.
 
-        For a choice that carries BOTH text content AND tool_calls, ``finish_reason``
-        is suppressed on the passthrough and recorded on
+        For a choice that carries BOTH text content AND tool_calls, or whose earlier
+        text is still withheld (``held_choices``), ``finish_reason`` is suppressed on
+        the passthrough and recorded on
         ``finish_reason_per_choice`` (when provided) so the final synthetic text
         chunk delivers it. Emitting the passthrough's ``finish_reason`` before the
         text flush would let a spec-compliant SSE client stop reading at
@@ -865,7 +923,8 @@ class UnifiedLLMGuardrails(CustomLogger):
             idx = getattr(choice, "index", 0) or 0
             original_finish = getattr(choice, "finish_reason", None)
             has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != ""
-            if has_text and original_finish is not None and finish_reason_per_choice is not None:
+            text_pending = has_text or idx in held_choices
+            if text_pending and original_finish is not None and finish_reason_per_choice is not None:
                 finish_reason_per_choice[idx] = original_finish
                 passthrough_finish: str | None = None
             else:

litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import Field
 
 from .base import GuardrailConfigModel
@@ -20,6 +22,16 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel):
         default=True,
         description="Whether a file sanitization `modify` verdict blocks the request instead of replacing the file content.",
     )
+    streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = Field(
+        default=None,
+        description=(
+            "How post_call `modify` verdicts reach a streaming client. `block_only` (default) streams the raw upstream "
+            "chunks and only a `block` verdict ends the stream, so `modified_text` is dropped. `incremental_diff` "
+            "buffers the whole response and sends the redacted text once the final verdict is in, so the first token "
+            "arrives with the last, while a `block` verdict still ends the stream early. "
+            "OpenAI chat completions streaming only."
+        ),
+    )
 
     @staticmethod
     def ui_friendly_name() -> str:

tests/test_litellm/integrations/test_custom_guardrail.py

@@ -3130,3 +3130,22 @@ class TestPreCallHookResponseIsNotLoggedVerbatim:
         )
 
         assert self._logged_response(data) == "mask"
+
+    @pytest.mark.asyncio
+    async def test_apply_guardrail_adding_only_stream_holdback_logs_allow(self):
+        class HoldbackOnlyGuardrail(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict[str, object],
+                input_type: Literal["request", "response"],
+                logging_obj: Optional["LiteLLMLoggingObj"] = None,
+            ) -> GenericGuardrailAPIInputs:
+                return {**inputs, "stream_holdback_chars": [6]}
+
+        data = self._request()
+        await HoldbackOnlyGuardrail(guardrail_name="g").apply_guardrail(
+            inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="response"
+        )
+
+        assert self._logged_response(data) == "allow"

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

@@ -1119,6 +1119,7 @@ class TestStreamingTransform:
             emitted_text_per_choice={},
             holdback_per_choice={},
             finish_reason_per_choice={0: "stop", 1: "length"},
+            held_chars_per_choice={},
             is_final=True,
         )
 
@@ -1157,6 +1158,7 @@ class TestStreamingTransform:
             emitted_text_per_choice={},
             holdback_per_choice={},
             finish_reason_per_choice={},
+            held_chars_per_choice={},
             is_final=False,
         )
 
@@ -1179,6 +1181,7 @@ class TestStreamingTransform:
                 emitted_text_per_choice={0: "My SSN is 123"},
                 holdback_per_choice={},
                 finish_reason_per_choice={},
+                held_chars_per_choice={},
                 is_final=False,
             )
 
@@ -1312,6 +1315,65 @@ class TestStreamingTransform:
         assert out[1].choices[0].delta.tool_calls
         assert out[1].choices[0].finish_reason == "tool_calls"
 
+    @pytest.mark.asyncio
+    async def test_held_text_flushes_before_tool_call_finish_reason(self):
+        """Text still held back when a separate terminal tool-call chunk arrives is
+        delivered before the stream's finish_reason, not after it."""
+        guardrail = _StreamingTextGuardrail(holdback_schedule=[100, 100, 100])
+
+        tool_chunk = ModelResponseStream(
+            choices=[
+                StreamingChoices(
+                    index=0,
+                    delta=Delta(
+                        content=None,
+                        tool_calls=[
+                            {
+                                "index": 0,
+                                "id": "call_1",
+                                "type": "function",
+                                "function": {"name": "get_weather", "arguments": "{}"},
+                            }
+                        ],
+                    ),
+                    finish_reason="tool_calls",
+                )
+            ],
+        )
+        chunks = [_stream_chunk("let me check "), tool_chunk]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        finished_at = [i for i, item in enumerate(out) if item.choices[0].finish_reason is not None]
+        assert finished_at == [len(out) - 1]
+        assert out[-1].choices[0].finish_reason == "tool_calls"
+        assert "".join(_delta_text(i) for i in out) == "LET ME CHECK "
+        assert any(item.choices[0].delta.tool_calls for item in out)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "usage_choices",
+        [[], [StreamingChoices(index=0, delta=Delta(), finish_reason=None)]],
+        ids=["choiceless", "empty-delta"],
+    )
+    async def test_usage_chunk_is_forwarded_after_final_text(self, usage_choices):
+        """A trailing usage chunk (stream_options.include_usage) is delivered after
+        the transformed text instead of being swallowed, whether it arrives with
+        no choices or, as CustomStreamWrapper emits it, with one empty delta."""
+        guardrail = _StreamingTextGuardrail(holdback_schedule=[100, 100, 100])
+        usage_chunk = ModelResponseStream(
+            choices=usage_choices,
+            usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+        )
+        chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop"), usage_chunk]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert "".join(_delta_text(i) for i in out) == "HELLO WORLD"
+        assert out[-1].usage.total_tokens == 5
+        assert not _delta_text(out[-1])
+        assert out[-2].choices[0].finish_reason == "stop"
+
     @pytest.mark.asyncio
     async def test_tool_call_blocking_guardrail_is_enforced(self):
         """A guardrail that blocks on tool calls must terminate the incremental_diff

tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py

@@ -8,12 +8,15 @@ from fastapi.exceptions import HTTPException
 from httpx import ReadTimeout, Request, Response
 
 import litellm
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import (
     PromptSecurityGuardrail,
     PromptSecurityGuardrailMissingSecrets,
 )
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
 from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
 
 
 def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
@@ -415,6 +418,199 @@ async def test_apply_guardrail_modify_response(monkeypatch: pytest.MonkeyPatch):
     assert result["texts"] == ["Your SSN is [REDACTED]"]
 
 
+@pytest.mark.asyncio
+async def test_apply_guardrail_modify_response_keeps_multi_choice_texts_aligned():
+    """With n>1 each choice text gets its own verdict, so a rewrite lands on the choice it came from."""
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="test-guard",
+        event_hook="post_call",
+        default_on=True,
+        api_key="test-key",
+        api_base="https://test.prompt.security",
+    )
+
+    async def mock_post(*args, **kwargs):
+        text = kwargs["json"]["response"]
+        redacted = text.replace("123-45-6789", "[REDACTED]")
+        mock_response = Response(
+            json={
+                "result": {
+                    "response": {
+                        "action": "modify" if redacted != text else "log",
+                        "violations": [],
+                        "modified_text": redacted,
+                    }
+                }
+            },
+            status_code=200,
+            request=Request(method="POST", url="https://test.prompt.security/api/protect"),
+        )
+        mock_response.raise_for_status = lambda: None
+        return mock_response
+
+    with patch.object(guardrail.async_handler, "post", side_effect=mock_post):
+        result = await guardrail.apply_guardrail(
+            inputs={"texts": ["all clear", "SSN 123-45-6789 on file"]},
+            request_data={},
+            input_type="response",
+        )
+
+    assert result["texts"] == ["all clear", "SSN [REDACTED] on file"]
+    assert result["stream_holdback_chars"] == [len("all clear"), len("SSN [REDACTED] on file")]
+
+
+def test_prompt_security_streaming_transform_mode_from_config(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
+    monkeypatch.setattr(litellm, "callbacks", [])
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    init_guardrails_v2(
+        all_guardrails=[
+            {
+                "guardrail_name": "prompt_security_streaming",
+                "litellm_params": {
+                    "guardrail": "prompt_security",
+                    "mode": "post_call",
+                    "default_on": True,
+                    "streaming_transform_mode": "incremental_diff",
+                },
+            }
+        ],
+        config_file_path="",
+    )
+
+    registered = [c for c in litellm.callbacks if isinstance(c, PromptSecurityGuardrail)]
+    assert len(registered) == 1
+    assert registered[0].streaming_transform_mode == "incremental_diff"
+    assert PromptSecurityGuardrail(api_key="k", api_base="https://b").streaming_transform_mode == "block_only"
+
+
+def _stream_chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        choices=[StreamingChoices(index=0, delta=Delta(content=content, role="assistant"), finish_reason=finish_reason)]
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("chunks", "secret", "redacted_output"),
+    [
+        pytest.param(
+            (
+                "Sure. I checked the billing record for this account and confirmed the details below. Card 4111 1111 ",
+                "1111 1111 is on file.",
+            ),
+            "4111 1111 1111 1111",
+            "Sure. I checked the billing record for this account and confirmed the details below. "
+            "Card [REDACTED] is on file.",
+            id="spaced_value_after_full_sentence",
+        ),
+        pytest.param(
+            ("Ship to 12 Main St. ", "Springfield 62704 today."),
+            "12 Main St. Springfield 62704",
+            "Ship to [REDACTED] today.",
+            id="value_spanning_abbreviation_period",
+        ),
+        pytest.param(
+            (
+                "Customer record follows.\nName: John Smith\n"
+                "Address: 12 Main St, Springfield IL 62704, United States\n",
+                "SSN: 123-45-6789\nThat is all.",
+            ),
+            "Name: John Smith\nAddress: 12 Main St, Springfield IL 62704, United States\nSSN: 123-45-6789",
+            "Customer record follows.\n[REDACTED]\nThat is all.",
+            id="multi_line_record_redacted_as_one_span",
+        ),
+    ],
+)
+async def test_prompt_security_incremental_diff_redacts_value_split_across_chunks(
+    chunks: tuple[str, ...],
+    secret: str,
+    redacted_output: str,
+):
+    """A modify verdict reaches the client redacted even when the value straddles a sampled scan."""
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="prompt_security_streaming",
+        event_hook="post_call",
+        default_on=True,
+        api_key="test-key",
+        api_base="https://test.prompt.security",
+        streaming_transform_mode="incremental_diff",
+    )
+    guardrail.streaming_sampling_rate = 1
+
+    async def mock_post(*args, **kwargs):
+        text = kwargs["json"]["response"]
+        redacted = text.replace(secret, "[REDACTED]")
+        mock_response = Response(
+            json={
+                "result": {
+                    "response": {
+                        "action": "modify" if redacted != text else "log",
+                        "violations": ["pii"] if redacted != text else [],
+                        "modified_text": redacted,
+                    }
+                }
+            },
+            status_code=200,
+            request=Request(method="POST", url="https://test.prompt.security/api/protect"),
+        )
+        mock_response.raise_for_status = lambda: None
+        return mock_response
+
+    async def _upstream():
+        for chunk in chunks:
+            yield _stream_chunk(chunk)
+        yield _stream_chunk("", finish_reason="stop")
+
+    with patch.object(guardrail.async_handler, "post", side_effect=mock_post):
+        out = [
+            item
+            async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+                user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
+                response=_upstream(),
+                request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"},
+            )
+        ]
+
+    assert all(isinstance(item, ModelResponseStream) for item in out)
+    deltas = [item.choices[0].delta.content for item in out if item.choices and item.choices[0].delta.content]
+    assert deltas == [redacted_output]
+    assert all(secret[:6] not in delta for delta in deltas)
+
+
+@pytest.mark.asyncio
+async def test_prompt_security_clean_non_streaming_response_logs_allow():
+    """A log verdict keeps the text (even if modified_text is present) and is logged as allow."""
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="prompt_security_streaming",
+        event_hook="post_call",
+        default_on=True,
+        api_key="test-key",
+        api_base="https://test.prompt.security",
+        streaming_transform_mode="incremental_diff",
+    )
+    mock_response = Response(
+        json={"result": {"response": {"action": "log", "violations": [], "modified_text": "order noted"}}},
+        status_code=200,
+        request=Request(method="POST", url="https://test.prompt.security/api/protect"),
+    )
+    mock_response.raise_for_status = lambda: None
+    request_data = {"metadata": {}}
+
+    with patch.object(guardrail.async_handler, "post", return_value=mock_response):
+        result = await guardrail.apply_guardrail(
+            inputs={"texts": ["order confirmed"]},
+            request_data=request_data,
+            input_type="response",
+        )
+
+    assert result["texts"] == ["order confirmed"]
+    info = request_data["metadata"]["standard_logging_guardrail_information"]
+    assert [entry["guardrail_response"] for entry in info] == ["allow"]
+
+
 @pytest.mark.asyncio
 async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
     """Test file sanitization for images"""