fix(policy_engine): snapshot stream observer inputs before dispatch

Guardrails that rewrite in place (`inputs["texts"] = masked; return
inputs` or mutating the same list/tool-call objects) left the observer's
pre and post views pointing at the already-rewritten values, so the
rewrote flag stayed false and `_pipeline_gated_stream` released the
original buffered chunks unredacted. Snapshot texts and tool-call shapes
before delegating to the inner guardrail so the diff sees the true
before-and-after.

litellm/proxy/policy_engine/pipeline_executor.py

@@ -57,16 +57,14 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]:
     return (function.get("name"), function.get("arguments"))
 
 
-def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool:
-    return sent is not None and returned is not None and tuple(returned) != tuple(sent)
+def _texts_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
+    return None if texts is None else tuple(texts)
 
 
-def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool:
-    if sent is None or returned is None:
-        return False
-    return tuple(_tool_call_shape(tool_call) for tool_call in returned) != tuple(
-        _tool_call_shape(tool_call) for tool_call in sent
-    )
+def _tool_call_shapes_snapshot(
+    tool_calls: Sequence[object] | None,
+) -> tuple[tuple[object, object], ...] | None:
+    return None if tool_calls is None else tuple(_tool_call_shape(tc) for tc in tool_calls)
 
 
 class _StreamRewriteObserver(CustomGuardrail):
@@ -90,14 +88,28 @@ class _StreamRewriteObserver(CustomGuardrail):
         input_type: Literal["request", "response"],
         logging_obj: "LiteLLMLoggingObj | None" = None,
     ) -> GenericGuardrailAPIInputs:
+        # Snapshot inputs *before* dispatching. Guardrails that rewrite in
+        # place (``inputs["texts"] = masked; return inputs`` or mutating the
+        # same list/tool-call objects) leave the pre and post views pointing
+        # at the already-rewritten values, which would hide the rewrite from
+        # a post-call diff and let the original buffered chunks reach the
+        # client unredacted.
+        sent_texts: Final = _texts_snapshot(inputs.get("texts"))
+        sent_tool_call_shapes: Final = _tool_call_shapes_snapshot(inputs.get("tool_calls"))
         outputs: Final = await self.inner.apply_guardrail(
             inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
         )
-        self.rewrote = (
-            self.rewrote
-            or _rewrote_texts(inputs.get("texts"), outputs.get("texts"))
-            or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls"))
+        returned_texts: Final = _texts_snapshot(outputs.get("texts"))
+        returned_tool_call_shapes: Final = _tool_call_shapes_snapshot(outputs.get("tool_calls"))
+        texts_rewrote: Final = (
+            sent_texts is not None and returned_texts is not None and returned_texts != sent_texts
         )
+        tool_calls_rewrote: Final = (
+            sent_tool_call_shapes is not None
+            and returned_tool_call_shapes is not None
+            and returned_tool_call_shapes != sent_tool_call_shapes
+        )
+        self.rewrote = self.rewrote or texts_rewrote or tool_calls_rewrote
         return outputs
 
 

tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py

@@ -872,3 +872,73 @@ async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeyp
 
     assert result.terminal_action == "allow"
     assert [step.outcome for step in result.step_results] == ["pass"]
+
+
+class _InPlaceTextRewritingGuardrail(CustomGuardrail):
+    """Rewrites ``inputs["texts"]`` in place and returns the same dict, the common
+    pattern the stream-rewrite observer must catch."""
+
+    def __init__(self, replacement_texts):
+        super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
+        self.replacement_texts = replacement_texts
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        inputs["texts"] = self.replacement_texts
+        return inputs
+
+
+class _InPlaceListMutatingTextGuardrail(CustomGuardrail):
+    """Mutates the same texts list in place (``texts[i] = masked``) and returns
+    the original inputs dict, exercising the observer's snapshot logic against
+    guardrails that never rebind the ``texts`` key at all."""
+
+    def __init__(self, replacement_texts):
+        super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
+        self.replacement_texts = list(replacement_texts)
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        texts = inputs["texts"]
+        for i, replacement in enumerate(self.replacement_texts):
+            texts[i] = replacement
+        return inputs
+
+
+class _MutableInputsTextTranslation:
+    """Sends a mutable ``texts`` list to the guardrail so an in-place rewrite is
+    visible on the same object the observer holds a reference to."""
+
+    async def process_output_streaming_response(
+        self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
+    ):
+        await guardrail_to_apply.apply_guardrail(
+            inputs={"texts": ["hello world"]},
+            request_data=request_data or {},
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+        return responses_so_far
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_detects_in_place_texts_assignment(monkeypatch):
+    """Regression: a guardrail that assigns ``inputs["texts"] = masked; return inputs``
+    must still trip the observer, since after the call ``inputs.get("texts")`` and
+    ``outputs.get("texts")`` point at the same rewritten value."""
+    monkeypatch.setattr(litellm, "callbacks", [_InPlaceTextRewritingGuardrail(["hello [MASKED]"])])
+
+    with pytest.raises(UndeliverableStreamRewrite) as info:
+        await _run_streaming_step(["hello [MASKED]"], _MutableInputsTextTranslation())
+
+    assert info.value.guardrail_name == "masker"
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_detects_in_place_texts_list_mutation(monkeypatch):
+    """Regression: a guardrail that mutates the same ``texts`` list in place must
+    still trip the observer, since the ``texts`` key is never rebound."""
+    monkeypatch.setattr(litellm, "callbacks", [_InPlaceListMutatingTextGuardrail(["hello [MASKED]"])])
+
+    with pytest.raises(UndeliverableStreamRewrite) as info:
+        await _run_streaming_step(["hello [MASKED]"], _MutableInputsTextTranslation())
+
+    assert info.value.guardrail_name == "masker"