fix(guardrails): apply the safeguard_results and stop_details rules to policy pipeline stream rewrites

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -30,6 +30,7 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import (
     LiteLLMAnthropicMessagesAdapter,
     is_provider_native_tool_dict,
 )
+from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details
 from litellm.llms.anthropic.pass_through.stream_assembly import (
     build_complete_streaming_response,
     build_usage_only_response_from_chunks,
@@ -1456,6 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation):
             return _SSEFieldRewrite("delta", "text", next(replacements))
 
         cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
+        cls._redact_ended_stream_message_delta(
+            responses_so_far, kept_tool_use_ids=frozenset(cls._streamed_tool_use_ids(responses_so_far).values())
+        )
 
     @classmethod
     def _write_ended_stream_tool_call_rewrites(
@@ -1523,6 +1527,64 @@ class AnthropicMessagesHandler(BaseTranslation):
                     return None
 
         cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
+        streamed_tool_use_ids: Final = cls._streamed_tool_use_ids(responses_so_far)
+        cls._redact_ended_stream_message_delta(
+            responses_so_far,
+            kept_tool_use_ids=frozenset(
+                tool_use_id for index, tool_use_id in streamed_tool_use_ids.items() if index not in rewrites_by_block
+            ),
+        )
+
+    @classmethod
+    def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]:
+        """The id of each buffered ``tool_use`` block, by its content block index."""
+        events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
+        tool_use_blocks: Final = (
+            (event.get("index"), block)
+            for event in events
+            if event.get("type") == "content_block_start"
+            and (block := as_str_mapping(event.get("content_block"))) is not None
+            and block.get("type") == "tool_use"
+        )
+        return MappingProxyType(
+            {
+                index: tool_use_id
+                for index, block in tool_use_blocks
+                if isinstance(index, int) and isinstance(tool_use_id := block.get("id"), str)
+            }
+        )
+
+    @classmethod
+    def _redact_ended_stream_message_delta(
+        cls,
+        responses_so_far: MutableSequence[object],  # mutable-ok: rewrites the caller's buffered chunks in place
+        *,
+        kept_tool_use_ids: frozenset[str],
+    ) -> None:
+        """Hold the buffered ``message_delta`` to the rules a re-emitted rewrite follows: a verdict survives
+        only for a tool call delivered unchanged, and no free-text ``explanation`` describes the text the
+        guardrail rewrote."""
+
+        def rewrite_stop_details(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
+            delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None
+            if delta is None or "stop_details" not in delta:
+                return None
+            return _SSEFieldRewrite(
+                "delta", "stop_details", redacted_stop_details(delta["stop_details"], keep_explanations=False)
+            )
+
+        def rewrite_safeguard_results(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
+            delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None
+            if delta is None or "safeguard_results" not in delta:
+                return None
+            return _SSEFieldRewrite(
+                "delta",
+                "safeguard_results",
+                filtered_safeguard_results(delta["safeguard_results"], kept_tool_use_ids, keep_explanations=False),
+            )
+
+        cls._rewrite_ended_stream_events(responses_so_far, rewrite_stop_details)
+        cls._rewrite_ended_stream_events(responses_so_far, rewrite_safeguard_results)
 
     @staticmethod
     def _rewrite_ended_stream_events(

litellm/llms/anthropic/pass_through/messages/utils.py

@@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping, Sequence
 from functools import lru_cache
 from typing import TYPE_CHECKING, Final, cast, get_type_hints
 
-from pydantic import JsonValue
+from pydantic import JsonValue, TypeAdapter, ValidationError
 
 from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
@@ -73,6 +73,71 @@ def anthropic_message_delta(
     }
 
 
+_SAFEGUARD_RESULT_ITEMS: Final = TypeAdapter(tuple[object, ...])
+
+
+def redacted_stop_details(raw: object, keep_explanations: bool) -> object:
+    """``stop_details`` as a rewritten stream may deliver it: its free-text ``explanation`` can quote
+    the text a guardrail rewrote, so it is nulled unless the rewrite changed no text."""
+    stop_details: Final = as_str_mapping(raw)
+    if keep_explanations or stop_details is None or "explanation" not in stop_details:
+        return raw
+    return {**stop_details, "explanation": None}
+
+
+def _without_explanation(raw: object) -> object:
+    mapping: Final = as_str_mapping(raw)
+    if mapping is None:
+        return raw
+    return {key: value for key, value in mapping.items() if key != "explanation"}
+
+
+def _status_without_explanations(raw: object) -> object:
+    status: Final = as_str_mapping(_without_explanation(raw))
+    if status is None:
+        return raw
+    tool_uses: Final = as_str_mapping(status.get("tool_uses"))
+    if tool_uses is None:
+        return status
+    return {
+        **status,
+        "tool_uses": {tool_use_id: _without_explanation(verdict) for tool_use_id, verdict in tool_uses.items()},
+    }
+
+
+def _result_without_explanations(raw: object) -> object:
+    result: Final = as_str_mapping(_without_explanation(raw))
+    if result is None:
+        return raw
+    if "status" not in result:
+        return result
+    return {**result, "status": _status_without_explanations(result["status"])}
+
+
+def _kept_tool_use_verdicts(raw: object, kept_tool_use_ids: frozenset[str]) -> object:
+    result: Final = as_str_mapping(raw)
+    status: Final = as_str_mapping(result.get("status")) if result is not None else None
+    tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None
+    if result is None or status is None or tool_uses is None:
+        return raw
+    kept_tool_uses: Final = {
+        tool_use_id: verdict for tool_use_id, verdict in tool_uses.items() if tool_use_id in kept_tool_use_ids
+    }
+    return {**result, "status": {**status, "tool_uses": kept_tool_uses}}
+
+
+def filtered_safeguard_results(raw: object, kept_tool_use_ids: frozenset[str], keep_explanations: bool) -> object:
+    """``safeguard_results`` as a rewritten stream may deliver them: a ``tool_uses`` verdict survives only
+    for a tool call the client gets with the input the classifier judged, and the free-text ``explanation`` of a
+    result, its status and each verdict is dropped unless the rewrite changed no text."""
+    try:
+        results: Final = _SAFEGUARD_RESULT_ITEMS.validate_python(raw)
+    except ValidationError:
+        return raw
+    kept: Final = tuple(_kept_tool_use_verdicts(result, kept_tool_use_ids) for result in results)
+    return kept if keep_explanations else tuple(_result_without_explanations(result) for result in kept)
+
+
 def _mapping_field(container: object, key: str) -> object | None:
     """One key of a raw provider payload, or None when the payload is not a mapping."""
     if not isinstance(container, Mapping):

litellm/proxy/guardrails/anthropic_sse.py

@@ -12,9 +12,8 @@ from collections.abc import Mapping, Sequence
 from types import MappingProxyType
 from typing import Final
 
-from pydantic import TypeAdapter, ValidationError
-
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
+from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details
 from litellm.llms.anthropic.pass_through.stream_assembly import anthropic_message_delta_fields
 from litellm.llms.anthropic.pass_through.stream_assembly import (
     assemble_anthropic_sse_stream as assemble_anthropic_sse_stream,  # noqa: PLC0414  # public re-export
@@ -55,52 +54,6 @@ from litellm.llms.anthropic.pass_through.stream_assembly import (  # noqa: E402
     is_sse_error_stream as is_sse_error_stream,  # noqa: PLC0414  # public re-export
 )
 
-_ITEMS: Final = TypeAdapter(tuple[object, ...])
-
-
-def _stop_details(raw: object, keep_explanations: bool) -> object:
-    stop_details: Final = as_str_mapping(raw)
-    if keep_explanations or stop_details is None or "explanation" not in stop_details:
-        return raw
-    return {**stop_details, "explanation": None}
-
-
-def _verdict(raw: object, keep_explanations: bool) -> object:
-    verdict: Final = as_str_mapping(raw)
-    if keep_explanations or verdict is None:
-        return raw
-    return {key: value for key, value in verdict.items() if key != "explanation"}
-
-
-def _status(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object:
-    status: Final = as_str_mapping(raw)
-    tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None
-    if status is None or tool_uses is None:
-        return raw
-    return {
-        **status,
-        "tool_uses": {
-            tool_use_id: _verdict(verdict, keep_explanations)
-            for tool_use_id, verdict in tool_uses.items()
-            if tool_use_id in unchanged_tool_use_ids
-        },
-    }
-
-
-def _safeguard_result(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object:
-    result: Final = as_str_mapping(raw)
-    if result is None or "status" not in result:
-        return raw
-    return {**result, "status": _status(result["status"], unchanged_tool_use_ids, keep_explanations)}
-
-
-def _safeguard_results(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object:
-    try:
-        results: Final = _ITEMS.validate_python(raw)
-    except ValidationError:
-        return raw
-    return tuple(_safeguard_result(result, unchanged_tool_use_ids, keep_explanations) for result in results)
-
 
 def _tool_use_inputs(content: Sequence[object]) -> Mapping[str, object]:
     blocks: Final = (block for raw in content if (block := as_str_mapping(raw)) is not None)
@@ -130,9 +83,17 @@ def _rewritten_message_delta_fields(
     safeguard_results: Final = source_fields.get("safeguard_results")
     return MappingProxyType(
         {
-            **({"stop_details": _stop_details(stop_details, keep_explanations)} if stop_details is not None else {}),
             **(
-                {"safeguard_results": _safeguard_results(safeguard_results, unchanged_tool_use_ids, keep_explanations)}
+                {"stop_details": redacted_stop_details(stop_details, keep_explanations)}
+                if stop_details is not None
+                else {}
+            ),
+            **(
+                {
+                    "safeguard_results": filtered_safeguard_results(
+                        safeguard_results, unchanged_tool_use_ids, keep_explanations
+                    )
+                }
                 if safeguard_results is not None
                 else {}
             ),
@@ -141,7 +102,7 @@ def _rewritten_message_delta_fields(
 
 
 def anthropic_sse_chunks_from_response(
-    assembled: ModelResponse, source_chunks: Sequence[object], *, keep_explanations: bool = False
+    assembled: ModelResponse, source_chunks: Sequence[object] = (), *, keep_explanations: bool = False
 ) -> tuple[bytes, ...]:
     """Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces.
 

tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -6,7 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic
 """
 
 import json
-from typing import Any, Literal, Optional
+from typing import Any, Literal, Optional, Final
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -18,7 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami
 from litellm.llms.anthropic.chat.guardrail_translation.handler import (
     AnthropicMessagesHandler,
 )
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
 
 
 class MockPassThroughGuardrail(CustomGuardrail):
@@ -383,6 +383,97 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
         assert '"name": "lookup_fruit"' not in raw
         assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"}
 
+    @staticmethod
+    def _message_delta(chunks: list[bytes]) -> dict[str, object]:
+        lines: Final = b"".join(chunks).decode().split("\n")
+        events: Final = (json.loads(line[len("data:") :]) for line in lines if line.startswith("data:"))
+        return next(event["delta"] for event in events if event.get("type") == "message_delta")
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_text_rewrite_drops_the_upstream_explanations(self):
+        delta_fields: Final = {
+            "stop_details": {"type": "refusal", "category": "cyber", "explanation": "the reply says world"},
+            "safeguard_results": [
+                {"type": "prompt_injection", "status": {"type": "available", "explanation": "quotes world"}}
+            ],
+        }
+        source: Final[list[bytes]] = self._ended_sse_chunks()
+        chunks: Final = [
+            chunk.replace(
+                b'"stop_sequence": null}', b'"stop_sequence": null, ' + json.dumps(delta_fields)[1:-1].encode() + b"}"
+            )
+            for chunk in source
+        ]
+
+        await AnthropicMessagesHandler().process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=self._masking_guardrail(),
+            litellm_logging_obj=MagicMock(),
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert self._delta_texts(chunks) == ["hello [MASKED]", ""]
+        assert self._message_delta(chunks) == {
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "stop_details": {"type": "refusal", "category": "cyber", "explanation": None},
+            "safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}],
+        }
+
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_tool_use_rewrite_drops_the_verdict_of_the_rewritten_call_only(self):
+        class MaskPersimmon(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict[str, object],
+                input_type: Literal["request", "response"],
+                logging_obj: Optional[LiteLLMLoggingObj] = None,
+            ) -> GenericGuardrailAPIInputs:
+                for tool_call in inputs.get("tool_calls", []):
+                    if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments:
+                        tool_call.function.arguments = '{"fruit": "[MASKED]"}'
+                return inputs
+
+        first_call: Final[list[bytes]] = self._ended_tool_use_sse_chunks()
+        second_call: Final = [
+            chunk.replace(b'"index": 0', b'"index": 1')
+            .replace(b"toolu_1", b"toolu_2")
+            .replace(b"persim", b"pear")
+            .replace(b'mon\\"}', b'\\"}')
+            for chunk in first_call[1:6]
+        ]
+        verdicts: Final = {
+            "toolu_1": {"type": "evaluated", "outcome": "flagged", "explanation": "looks up persimmon"},
+            "toolu_2": {"type": "evaluated", "outcome": "not_flagged", "explanation": "looks up pear"},
+        }
+        safeguard_results: Final = [
+            {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": verdicts}}
+        ]
+        message_delta: Final = first_call[6].replace(
+            b'"stop_sequence": null}',
+            b'"stop_sequence": null, "safeguard_results": ' + json.dumps(safeguard_results).encode() + b"}",
+        )
+        chunks: Final = [*first_call[:6], *second_call, message_delta, first_call[7]]
+
+        await AnthropicMessagesHandler().process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=MaskPersimmon(guardrail_name="test"),
+            litellm_logging_obj=MagicMock(),
+            deliver_ended_stream_rewrites=True,
+        )
+
+        assert self._partial_jsons(chunks)[:3] == ['{"fruit": "[MASKED]"}', "", ""]
+        assert self._message_delta(chunks)["safeguard_results"] == [
+            {
+                "type": "dangerous_tool_use",
+                "status": {
+                    "type": "available",
+                    "tool_uses": {"toolu_2": {"type": "evaluated", "outcome": "not_flagged"}},
+                },
+            }
+        ]
+
     @pytest.mark.asyncio
     async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self):
         handler = AnthropicMessagesHandler()

tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py

@@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
 import json
 import logging
 import re
+from typing import Final
 from unittest.mock import patch
 
 import pytest
@@ -1383,8 +1384,16 @@ class TestToolPermissionGuardrailAnthropicMessages:
     async def test_rewrite_mode_keeps_the_allowed_tool_use_verdict_with_its_explanation(self):
         """Claude Code auto mode reads the classifier verdict from message_delta. The rewrite changes
         no text the explanation could repeat, so the allowed call keeps its verdict whole."""
-        bash_verdict = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}
-        safeguard_results = [
+        default_on_rewriting: Final = ToolPermissionGuardrail(
+            guardrail_name="anthropic-rewrite-default-on",
+            rules=self.rules,
+            default_action="deny",
+            on_disallowed_action="rewrite",
+            event_hook=GuardrailEventHooks.post_call,
+            default_on=True,
+        )
+        bash_verdict: Final = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}
+        safeguard_results: Final = [
             {
                 "type": "dangerous_tool_use",
                 "status": {
@@ -1393,28 +1402,23 @@ class TestToolPermissionGuardrailAnthropicMessages:
                 },
             }
         ]
-        bash_chunks = self._sse_chunks("Bash", "tu_ok")
-        read_chunks = self._sse_chunks("Read", "tu_bad")
-        read_tool_use = [
-            chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]
-        ]
-        message_delta = {
+        bash_chunks: Final = self._sse_chunks("Bash", "tu_ok")
+        read_chunks: Final = self._sse_chunks("Read", "tu_bad")
+        read_tool_use: Final = [chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]]
+        message_delta: Final = {
             "type": "message_delta",
             "delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results},
             "usage": {"output_tokens": 5},
         }
-        chunks = [
+        chunks: Final = [
             *bash_chunks[:7],
             *read_tool_use,
             f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(),
             bash_chunks[-1],
         ]
 
-        with patch.object(self.rewriting, "should_run_guardrail", return_value=True):
-            out = await self._drain(self.rewriting, chunks)
-
-        body = b"".join(out).decode()
-        delta = next(
+        body: Final = b"".join(await self._drain(default_on_rewriting, chunks)).decode()
+        delta: Final = next(
             json.loads(line[6:])["delta"]
             for line in body.splitlines()
             if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_delta"

tests/unit/proxy/guardrails/test_anthropic_sse.py

@@ -169,8 +169,23 @@ SSN: Final = "123-45-6789"
             (("toolu_bash", "Bash"),),
             {"safeguard_results": _safeguard_results({"toolu_bash": {"type": "evaluated", "outcome": "flagged"}})},
         ),
+        (
+            {
+                "stop_reason": "end_turn",
+                "stop_sequence": None,
+                "safeguard_results": [
+                    {
+                        "type": "prompt_injection",
+                        "status": {"type": "available", "explanation": f"quotes SSN {SSN}"},
+                        "explanation": f"saw SSN {SSN}",
+                    }
+                ],
+            },
+            (),
+            {"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}]},
+        ),
     ],
-    ids=["stop_details_explanation", "verdict_explanation"],
+    ids=["stop_details_explanation", "verdict_explanation", "unknown_shape_explanation"],
 )
 def test_masked_stream_drops_the_unmasked_upstream_explanations(
     upstream_delta: dict[str, object], tool_uses: tuple[tuple[str, str], ...], expected_fields: dict[str, object]