fix(guardrails): hold tool-call windows until the final scan and expose Bedrock streaming flags to the UI

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation):
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
+        tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_use_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
         )
 
     @classmethod

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -40,11 +40,15 @@ class StreamingScanKey:
     """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
     compare equal when the round would scan the same content again; ``stream_ended``
     stays out of the comparison and only says whether the handler is on its
-    end-of-stream path, where an empty payload is still scanned today."""
+    end-of-stream path, where an empty payload is still scanned today.
+    ``tool_calls_in_flight`` also stays out of the comparison: it flags that tool
+    calls have streamed which this round cannot scan yet, so a buffered window
+    holding them must stay withheld until the end-of-stream scan covers them."""
 
     texts: tuple[str, ...]
     tool_calls: tuple[str, ...] = ()
     stream_ended: bool = field(default=False, compare=False)
+    tool_calls_in_flight: bool = field(default=False, compare=False)
 
     @property
     def has_nothing_to_scan(self) -> bool:

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

@@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
         stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=tuple(self._combine_streaming_texts(chunks).values()),
-            tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_call_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended,
         )
 
     @staticmethod

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -1177,9 +1177,22 @@ class OpenAIResponsesHandler(BaseTranslation):
             return None
         if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
             return self._completed_response_scan_key(stream_item_field(last_event, "response"))
+        stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            stream_ended=self._check_streaming_has_ended(responses_so_far),
+            stream_ended=stream_ended,
+            tool_calls_in_flight=not stream_ended and self._has_streamed_tool_call_events(responses_so_far),
+        )
+
+    @staticmethod
+    def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool:
+        return any(
+            stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
+            or (
+                stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
+                and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES
+            )
+            for event in responses_so_far
         )
 
     @staticmethod

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -37,6 +37,7 @@ from litellm.types.guardrails import (
     ApplyGuardrailResponse,
     BaseLitellmParams,
     BedrockGuardrailConfigModel,
+    BedrockGuardrailStreamingParams,
     Guardrail,
     GuardrailEventHooks,
     GuardrailInfoResponse,
@@ -1959,7 +1960,10 @@ async def get_provider_specific_params():
     ```
     """
     # Get fields from the models
-    bedrock_fields: Final = _get_fields_from_model(BedrockGuardrailConfigModel)
+    bedrock_fields: Final = {
+        **_get_fields_from_model(BedrockGuardrailConfigModel),
+        **_get_fields_from_model(BedrockGuardrailStreamingParams),
+    }
     presidio_fields: Final = _get_fields_from_model(PresidioPresidioConfigModelUserInterface)
     lakera_v2_fields: Final = _get_fields_from_model(LakeraV2GuardrailConfigModel)
     tool_permission_fields: Final = _get_fields_from_model(ToolPermissionGuardrailConfigModel)

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -330,7 +330,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
 
     def _streams_incrementally(self) -> bool:
-        return not self.streaming_buffer_until_moderated and not self.mask_response_content
+        if self.mask_response_content:
+            return False
+        if not self.streaming_buffer_until_moderated:
+            return True
+        return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only
 
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:

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

@@ -1032,6 +1032,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         # message (True) vs emit a standalone block message (False, buffered).
         chunks_yielded = False
         last_scan_key: StreamingScanKey | None = None  # rebind-ok: replaced after every scan round
+        tool_calls_in_flight = False  # rebind-ok: tracks the latest scan key's unscanned tool calls
 
         async for item in response:
             chunk_counter += 1
@@ -1079,6 +1080,9 @@ class UnifiedLLMGuardrails(CustomLogger):
             if chunk_counter % sampling_rate == 0:
                 endpoint_translation = mappings[CallTypes(call_type)]()
                 scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
+                if scan_key is not None:
+                    tool_calls_in_flight = scan_key.tool_calls_in_flight
+                hold_window = buffer_until_moderated and tool_calls_in_flight
                 if _is_redundant_scan(scan_key, last_scan_key):
                     verbose_proxy_logger.debug(
                         "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
@@ -1086,6 +1090,8 @@ class UnifiedLLMGuardrails(CustomLogger):
                         guardrail_to_apply.guardrail_name,
                     )
                     if buffer_until_moderated:
+                        if hold_window:
+                            continue
                         for withheld_item in withheld_items:
                             chunks_yielded = True
                             responses_yielded.append(withheld_item)
@@ -1151,6 +1157,14 @@ class UnifiedLLMGuardrails(CustomLogger):
                     return
                 if scan_key is not None:
                     last_scan_key = scan_key
+                if hold_window:
+                    verbose_proxy_logger.debug(
+                        "Holding %s buffered chunks for guardrail %s: streamed tool calls await the end-of-stream scan",
+                        len(withheld_items),
+                        guardrail_to_apply.guardrail_name,
+                    )
+                    withheld_items[:] = original_items
+                    continue
                 for original_item in original_items:
                     chunks_yielded = True
                     responses_yielded.append(original_item)

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

@@ -2482,7 +2482,10 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
         open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
         ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
         assert open_key == StreamingScanKey(texts=("hi",))
+        assert open_key.tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._text_delta("hi")]).tool_calls_in_flight is False
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key.tool_calls_in_flight is False
         assert ended_key != open_key
 
 

tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -2206,8 +2206,11 @@ class TestStreamingScanKey:
             [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
         )
         assert open_key == StreamingScanKey(texts=("hi",))
+        assert open_key.tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._chunk("hi")]).tool_calls_in_flight is False
         assert ended_key.texts == ("hi",)
         assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key.tool_calls_in_flight is False
         assert ended_key != open_key
 
     def test_text_after_the_first_choice_finishes_still_changes_the_key(self):

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -3211,3 +3211,24 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
     def test_output_item_done_round_is_never_deduped(self):
         done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
         assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None
+
+    def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self):
+        handler = OpenAIResponsesHandler()
+        added = {
+            "type": "response.output_item.added",
+            "sequence_number": 1,
+            "item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather"},
+        }
+        arguments_delta = {
+            "type": "response.function_call_arguments.delta",
+            "sequence_number": 2,
+            "item_id": "fc_1",
+            "delta": '{"city":',
+        }
+        function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
+        assert handler.get_streaming_scan_key([self._delta(0, "hi")]).tool_calls_in_flight is False
+        assert handler.get_streaming_scan_key([self._delta(0, "hi"), added]).tool_calls_in_flight is True
+        assert handler.get_streaming_scan_key([self._delta(0, "hi"), arguments_delta]).tool_calls_in_flight is True
+        ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), added, self._completed(3, [function_call])])
+        assert ended_key.tool_calls_in_flight is False
+        assert len(ended_key.tool_calls) == 1

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -5724,6 +5724,44 @@ async def test_buffered_default_hook_scans_before_any_chunk():
     assert len([e for e in events if e != "scan"]) >= 1
 
 
+@pytest.mark.asyncio
+async def test_buffered_release_on_scan_hook_releases_each_window_after_its_scan():
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-release-on-scan",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+        streaming_buffer_release_on_scan=True,
+        streaming_sampling_rate=1,
+    )
+
+    assert guardrail._streams_incrementally() is True
+    events = await _run_streaming_hook_recording_order(guardrail)
+
+    assert events == ["scan", ("chunk", "Hello"), "scan", ("chunk", " world"), ("chunk", "")]
+
+
+@pytest.mark.asyncio
+async def test_buffered_release_on_scan_defers_to_end_of_stream_only():
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-release-on-scan-end-only",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+        streaming_buffer_release_on_scan=True,
+        streaming_end_of_stream_only=True,
+        streaming_sampling_rate=1,
+    )
+
+    assert guardrail._streams_incrementally() is False
+    events = await _run_streaming_hook_recording_order(guardrail)
+
+    assert events.count("scan") == 1
+    assert events[0] == "scan"
+
+
 @pytest.mark.asyncio
 async def test_masking_keeps_buffered_path_even_when_unbuffered_configured():
     guardrail = BedrockGuardrail(

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

@@ -23,10 +23,18 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
 )
-from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices
+from litellm.types.utils import (
+    ChatCompletionDeltaToolCall,
+    Delta,
+    Function,
+    GenericGuardrailAPIInputs,
+    ModelResponseStream,
+    StreamingChoices,
+)
 
 BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
 ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER"
+TOOL_ARGUMENTS_MARKER = "TOOL-ARGS-SECRET"
 
 
 class _BlockingGuardrail(CustomGuardrail):
@@ -76,6 +84,24 @@ class _CountingPassingGuardrail(_PassingGuardrail):
         return inputs
 
 
+class _ToolCallRecordingGuardrail(_CountingPassingGuardrail):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tool_call_scan_indexes: List[int] = []
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if inputs.get("tool_calls"):
+            self.tool_call_scan_indexes.append(self.scan_count)
+        return inputs
+
+
 class _SecondScanBlockingGuardrail(_CountingPassingGuardrail):
     async def apply_guardrail(
         self,
@@ -165,20 +191,63 @@ def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelRes
     )
 
 
+def _tool_call_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-windowed",
+        created=1724900000,
+        model="gpt-4",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(
+                    role="assistant",
+                    content=None,
+                    tool_calls=[
+                        ChatCompletionDeltaToolCall(
+                            id="call_1",
+                            type="function",
+                            index=0,
+                            function=Function(name="run_shell", arguments=arguments),
+                        )
+                    ],
+                ),
+                finish_reason=finish_reason,
+            )
+        ],
+    )
+
+
 async def _windowed_chat_stream(
-    yielded_count: List[int], collected: List[Any], content_chunks: List[str]
+    yielded_count: List[int],
+    collected: List[Any],
+    content_chunks: List[str],
+    tool_argument_chunks: List[str] | None = None,
 ) -> AsyncGenerator[ModelResponseStream, None]:
     for content in content_chunks:
         yielded_count.append(len(collected))
         yield _chat_chunk(content)
+    for arguments in tool_argument_chunks or []:
+        yielded_count.append(len(collected))
+        yield _tool_call_chunk(arguments)
     yielded_count.append(len(collected))
-    yield _chat_chunk(finish_reason="stop")
+    yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop")
+
+
+def _tool_argument_text(chunks: List[Any]) -> str:
+    return "".join(
+        tool_call.function.arguments or ""
+        for chunk in chunks
+        if isinstance(chunk, ModelResponseStream)
+        for choice in chunk.choices
+        for tool_call in choice.delta.tool_calls or []
+    )
 
 
 async def _run_windowed(
     guardrail: CustomGuardrail,
     content_chunks: List[str],
     end_of_stream_only: bool = False,
+    tool_argument_chunks: List[str] | None = None,
 ) -> tuple[List[Any], List[int]]:
     guardrail.streaming_buffer_until_moderated = True
     guardrail.streaming_buffer_release_on_scan = True
@@ -195,7 +264,7 @@ async def _run_windowed(
     yielded_count: List[int] = []
     async for chunk in unified.async_post_call_streaming_iterator_hook(
         user_api_key_dict=user_api_key_dict,
-        response=_windowed_chat_stream(yielded_count, collected, content_chunks),
+        response=_windowed_chat_stream(yielded_count, collected, content_chunks, tool_argument_chunks),
         request_data=request_data,
     ):
         collected.append(chunk)
@@ -282,6 +351,20 @@ async def test_windowed_buffer_drops_blocked_window():
     assert '"error"' not in raw
 
 
+@pytest.mark.asyncio
+async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan():
+    guardrail = _ToolCallRecordingGuardrail(guardrail_name="windowed-tools", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three "]
+    tool_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks, tool_argument_chunks=tool_argument_chunks)
+
+    assert yielded_count == [0, 0, 2, 2, 2, 2, 2]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert _tool_argument_text(collected) == "".join(tool_argument_chunks)
+    assert guardrail.tool_call_scan_indexes == [guardrail.scan_count]
+
+
 @pytest.mark.asyncio
 async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered():
     guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call")

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -682,6 +682,22 @@ async def test_provider_specific_params_includes_embedding_toggle():
         assert field["default_value"] is False
 
 
+@pytest.mark.asyncio
+async def test_provider_specific_params_exposes_bedrock_streaming_flags():
+    from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params
+
+    provider_params = await get_provider_specific_params()
+
+    bedrock = provider_params["bedrock"]
+    assert "guardrailIdentifier" in bedrock
+    assert "guardrailVersion" in bedrock
+    assert bedrock["streaming_buffer_release_on_scan"]["type"] == "boolean"
+    assert bedrock["streaming_buffer_release_on_scan"]["default_value"] is False
+    assert bedrock["streaming_buffer_until_moderated"]["default_value"] is True
+    assert bedrock["streaming_end_of_stream_only"]["type"] == "boolean"
+    assert bedrock["streaming_sampling_rate"]["type"] == "number"
+
+
 @pytest.mark.asyncio
 async def test_provider_specific_params_includes_hide_secrets():
     """hide-secrets lives in the enterprise package so it is not in