From d9c793bdcdf717fa102f171ee533033769654c91 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 30 Aug 2026 04:19:43 +0000 Subject: [PATCH] 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. --- .../proxy/policy_engine/pipeline_executor.py | 36 ++++++---- .../policy_engine/test_pipeline_executor.py | 70 +++++++++++++++++++ 2 files changed, 94 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 21f3aca3f58..64895ba5b18 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/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 diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 52fd8777a19..37c77761bb2 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/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"