From 194c3de14295858c9b02083b769dd52ced8ab172 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 3 Nov 2026 18:07:17 +0000 Subject: [PATCH] fix(guardrails): deliver post_call rewrites into buffered streams on the guardrails config path A post_call guardrail that rewrites text or tool-call arguments never reached the client on a streamed response under the plain guardrails: list, even with streaming_buffer_until_moderated: true, because the end-of-stream scan never asked the translation to write the rewrite back and the buffered originals were replayed. The buffered end scan now passes deliver_ended_stream_rewrites, releases the rewritten chunks, and falls back to the originals when the rewrite cannot be placed. The guard that disabled buffering for mask_response_content guardrails is removed. --- .../chat/guardrail_translation/handler.py | 30 +- .../chat/guardrail_translation/handler.py | 16 + .../guardrail_translation/handler.py | 15 + .../unified_guardrail/unified_guardrail.py | 108 +++-- .../test_streaming_buffer_until_moderated.py | 379 +++++++++++++++++- 5 files changed, 501 insertions(+), 47 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..ac6949b4dbb 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -519,6 +519,26 @@ def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewrit return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} +def _ended_stream_tool_calls( + returned_tool_calls: Sequence[object] | None, + *, + streamed_tool_calls: Sequence[object], + expected_count: int, + guardrail_name: str, +) -> Sequence[object]: + """A guardrail that adds or drops a tool call on an ended stream cannot be written back, block for block.""" + if returned_tool_calls is None: + return streamed_tool_calls + if len(returned_tool_calls) != expected_count: + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried {expected_count}", + ) + return returned_tool_calls + + def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: """The guardrail-visible shape of each tool call, whether the guardrail handed back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" @@ -1611,10 +1631,12 @@ class AnthropicMessagesHandler(BaseTranslation): pre_guardrail_tool_calls=pre_guardrail_tool_calls, post_guardrail_tool_calls=_with_query_rewrites( _tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) - and len(returned_tool_calls) == len(pre_guardrail_tool_calls) - else tool_calls_list or () + _ended_stream_tool_calls( + returned_tool_calls, + streamed_tool_calls=tool_calls_list or (), + expected_count=len(pre_guardrail_tool_calls), + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) ), tool_call_blocks, _query_rewrites( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..fab41584090 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -393,6 +393,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + *, + deliver_ended_stream_rewrites: bool = False, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -458,6 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if hasattr(response, "model") and response.model: inputs["model"] = response.model + pre_guardrail_tool_call_count: Final = len(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -467,6 +470,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") + if ( + deliver_ended_stream_rewrites + and isinstance(returned_tool_calls, list) + and len(returned_tool_calls) != pre_guardrail_tool_call_count + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried " + f"{pre_guardrail_tool_call_count}", + ) guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) @@ -667,6 +682,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) if not deliver_ended_stream_rewrites: return diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a18cbe13925..3720027fa37 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -177,9 +177,23 @@ def _post_guardrail_tool_call_shapes( returned_tool_calls: Sequence[object] | None, pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], guardrail_name: str | None, + *, + deliver_ended_stream_rewrites: bool = False, ) -> tuple[_ToolCallShape, ...]: if not pre_guardrail_tool_calls: return pre_guardrail_tool_calls + if ( + deliver_ended_stream_rewrites + and returned_tool_calls is not None + and len(returned_tool_calls) != len(pre_guardrail_tool_calls) + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name or "unknown", + f"the guardrail returned {len(returned_tool_calls)} tool calls for a stream that carried " + f"{len(pre_guardrail_tool_calls)}", + ) if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls): verbose_proxy_logger.warning( "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, " @@ -1284,6 +1298,7 @@ class OpenAIResponsesHandler(BaseTranslation): returned_tool_calls=guardrailed_inputs.get("tool_calls"), pre_guardrail_tool_calls=pre_guardrail_tool_calls, guardrail_name=guardrail_to_apply.guardrail_name, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) # Write guardrailed texts back into the output items in-place. diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 5ce2dfb3b87..8df475d0d48 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -9,6 +9,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import asyncio import contextlib import copy +import functools import json from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable @@ -234,6 +235,57 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N data["litellm_metadata"] = user_metadata +async def _scan_ended_stream( + endpoint_translation: "BaseTranslation", + *, + guardrail_to_apply: CustomGuardrail, + responses_so_far: list[object], # mutable-ok: the translation writes the rewrites into the caller's buffer + user_api_key_dict: UserAPIKeyAuth, + request_data: dict[str, object], # mutable-ok: handed through to the translation's dict parameter + buffered_items: tuple[object, ...] | None, + deliver_rewrites: bool, + shield: bool, +) -> tuple[object, ...] | None: + """Run the end-of-stream scan and return the chunks to release to the client. + + A fully withheld stream holds every chunk in ``responses_so_far``, so the guardrail's text and tool-call rewrites + are released from there, whether the translation rewrote the chunks in place or replaced them (SSE bytes). + A rewrite the translation cannot place (a changed tool-call count) fails closed: nothing is released and the + client gets a 500, since the originals carry the very content the guardrail rewrote. + """ + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + logging_obj: Final = request_data.get("litellm_logging_obj") + scan: Final = functools.partial( + endpoint_translation.process_output_streaming_response, + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + try: + with anyio.CancelScope(shield=shield): + await (scan(deliver_ended_stream_rewrites=True) if deliver_rewrites else scan()) + except UndeliverableStreamRewrite as undeliverable: + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails: guardrail %s rewrote the buffered stream but the rewrite cannot be written back " + "(%s). The response was withheld", + undeliverable.guardrail_name, + undeliverable.reason, + ) + raise HTTPException( + status_code=500, + detail={ + "error": f"Guardrail {undeliverable.guardrail_name} rewrote the response but the rewrite cannot be " + f"delivered on this stream ({undeliverable.reason}). The response was withheld", + "guardrail": undeliverable.guardrail_name, + }, + ) from undeliverable + return tuple(responses_so_far) if deliver_rewrites else buffered_items + + class UnifiedLLMGuardrails(CustomLogger): def __init__( self, @@ -1158,31 +1210,13 @@ class UnifiedLLMGuardrails(CustomLogger): # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). streaming_transform_mode: Final = _streaming_setting("streaming_transform_mode", "block_only") - # Withhold every chunk until end-of-stream moderation passes, then - # release the original chunks (clean) or only the block message - # (blocked) -- moderating the whole response *before* any content - # reaches the client. Only safe for allow/block guardrails: on - # release the original chunks are replayed as-is, so a - # content-rewriting guardrail (e.g. PII masking) would leak - # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated: object = _streaming_setting( + # Withhold every chunk until end-of-stream moderation passes, then release the chunks + # (with the guardrail's rewrite written into them) or only the block message. + buffer_until_moderated: Final = _streaming_setting( "streaming_buffer_until_moderated", buffer_until_moderated_default ) release_on_scan: Final = _streaming_setting("streaming_buffer_release_on_scan", False) - if ( - buffer_until_moderated - and guardrail_to_apply is not None - and getattr(guardrail_to_apply, "mask_response_content", False) - ): - verbose_proxy_logger.warning( - "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s " - "because mask_response_content=True -- buffered replay would release " - "unredacted original chunks instead of the moderated output.", - guardrail_to_apply.guardrail_name, - ) - buffer_until_moderated = False - if buffer_until_moderated and not release_on_scan: end_of_stream_only = True @@ -1476,19 +1510,25 @@ class UnifiedLLMGuardrails(CustomLogger): return try: - with anyio.CancelScope(shield=chunks_yielded): - if pending_scan is not None: - pending_scan.result() # a failed background scan raises into the handlers below - await endpoint_translation.process_output_streaming_response( - responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=request_data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - request_data=request_data, - ) - # Moderation passed: release the withheld original chunks. - if buffered_items is not None: - for buffered_item in buffered_items: + if pending_scan is not None: + pending_scan.result() # a failed background scan raises into the handlers below + released_items: Final = await _scan_ended_stream( + endpoint_translation, + guardrail_to_apply=guardrail_to_apply, + responses_so_far=responses_so_far, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + buffered_items=buffered_items, + deliver_rewrites=bool( + buffer_until_moderated + and end_of_stream_only + and not responses_yielded + and endpoint_translation.delivers_ended_stream_rewrites + ), + shield=chunks_yielded, + ) + if released_items is not None: + for buffered_item in released_items: yield buffered_item for pending_item in pending_end_of_stream_items: responses_yielded.append(pending_item) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 2436c85fedd..c2715edec34 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -14,6 +14,7 @@ import json from typing import Any, AsyncGenerator, List, Literal, Optional import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import ( @@ -30,6 +31,8 @@ from litellm.proxy.utils import _pipelines_stream_live from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import ( ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, + ChatCompletionToolCallChunk, Delta, Function, FunctionCall, @@ -562,20 +565,378 @@ async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buff assert guardrail.scan_count == 1 +MASKED_MARKER = "[MASKED]" + + +class _MaskingGuardrail(_CountingPassingGuardrail): + """Rewrites the markers out of the response text and tool-call arguments, as a PII masker does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + masked: GenericGuardrailAPIInputs = { + "texts": [text.replace(ORIGINAL_MARKER, MASKED_MARKER) for text in inputs.get("texts") or []] + } + tool_calls = inputs.get("tool_calls") + if tool_calls: + masked["tool_calls"] = [_mask_tool_call(tool_call) for tool_call in tool_calls] + return masked + + +def _mask_tool_call( + tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall, +) -> ChatCompletionToolCallChunk | ChatCompletionMessageToolCall: + """Handlers hand chat tool calls over as dicts and /v1/messages tool calls as pydantic objects.""" + if isinstance(tool_call, dict): + function = tool_call["function"] + arguments = function["arguments"].replace(TOOL_ARGUMENTS_MARKER, MASKED_MARKER) + return {**tool_call, "function": {**function, "arguments": arguments}} + tool_call.function.arguments = tool_call.function.arguments.replace(TOOL_ARGUMENTS_MARKER, MASKED_MARKER) + return tool_call + + +class _ToolDroppingGuardrail(_MaskingGuardrail): + """Masks like _MaskingGuardrail but returns no tool calls at all, as a guardrail that strips a disallowed tool does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + masked = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + return {**masked, "tool_calls": []} + + +class _ToolClearingGuardrail(_MaskingGuardrail): + """Empties the handler's own tool-call list in place and hands that same list back.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + masked = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + tool_calls = inputs.get("tool_calls") + if tool_calls is not None: + tool_calls.clear() + return {**masked, "tool_calls": tool_calls or []} + + +async def _anthropic_tool_use_stream(): + """A complete /v1/messages SSE stream with a text block and a tool_use block whose input carries the marker.""" + yield _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_tool", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ) + yield _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + yield _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": f"Sending {ORIGINAL_MARKER}"}, + }, + ) + yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + yield _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "send_email", "input": {}}, + }, + ) + for partial in ['{"to": "', TOOL_ARGUMENTS_MARKER, '"}']: + yield _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": partial}}, + ) + yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 1}) + yield _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + yield _sse_event("message_stop", {"type": "message_stop"}) + + +async def _run_messages_stream(guardrail: CustomGuardrail, stream: AsyncGenerator[bytes, None]) -> str: + guardrail.streaming_end_of_stream_only = True + guardrail.streaming_buffer_until_moderated = True + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=stream, + request_data=request_data, + ): + collected.append(chunk) + return _decode(collected) + + +async def _run_buffered_chat( + guardrail: CustomGuardrail, content_chunks: list[str], tool_argument_chunks: list[str] | None = None +) -> list[Any]: + """The plain ``streaming_buffer_until_moderated: true`` config: hold the whole stream, scan once at the end.""" + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = False + guardrail.streaming_end_of_stream_only = False + guardrail.streaming_sampling_rate = 2 + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + 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, tool_argument_chunks, False), + request_data=request_data, + ): + collected.append(chunk) + assert yielded_count == [0] * len(yielded_count), "buffered stream released a chunk before the end scan" + return collected + + @pytest.mark.asyncio -async def test_buffered_mode_disabled_for_content_rewriting_guardrail(): - """Buffered replay yields the withheld *original* chunks verbatim, which - is unsafe for a guardrail that rewrites response text (e.g. PII masking): - the client would get the unredacted original instead of the moderated - output. mask_response_content=True must force buffering off so the - request falls back to the (correctly moderated) non-buffered path.""" - guardrail = _PassingGuardrail(guardrail_name="masker", event_hook="post_call", mask_response_content=True) +async def test_buffered_messages_stream_delivers_masked_text_to_the_client(): + """A masking guardrail on a fully buffered /v1/messages stream, the client gets the masked text, never the original.""" + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call", mask_response_content=True) raw = await _run(guardrail) - assert guardrail.streaming_buffer_until_moderated is True # request asked for buffering - assert ORIGINAL_MARKER in raw + assert guardrail.scan_count == 1 + assert MASKED_MARKER in raw, raw + assert ORIGINAL_MARKER not in raw, raw assert BLOCK_MESSAGE not in raw +@pytest.mark.asyncio +async def test_buffered_messages_stream_delivers_masked_tool_call_input(): + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call") + raw = await _run_messages_stream(guardrail, _anthropic_tool_use_stream()) + assert guardrail.scan_count == 1 + assert TOOL_ARGUMENTS_MARKER not in raw, raw + assert ORIGINAL_MARKER not in raw, raw + assert raw.count(MASKED_MARKER) == 2, raw + assert "message_stop" in raw + + +@pytest.mark.asyncio +async def test_buffered_chat_stream_delivers_masked_text_and_tool_arguments(): + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call") + collected = await _run_buffered_chat( + guardrail, [f"call {ORIGINAL_MARKER} now"], ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'] + ) + assert _chat_text(collected) == f"call {MASKED_MARKER} now" + assert _tool_argument_text(collected) == f'{{"cmd": "{MASKED_MARKER}"}}' + assert collected[-1].choices[0].finish_reason == "tool_calls" + + +def _two_choice_tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: + """An ``n=2`` chunk: tool-call rewrites cannot be written back onto a multi-choice stream.""" + return ModelResponseStream( + id="chatcmpl-two-choices", + created=1724900000, + model="gpt-4", + choices=[ + StreamingChoices( + index=index, + finish_reason=finish_reason, + delta=Delta( + role="assistant", + content=f"call {ORIGINAL_MARKER} now", + tool_calls=[ + ChatCompletionDeltaToolCall( + id=f"call_{index}", + type="function", + index=0, + function=Function(name="run_shell", arguments=f'{{"cmd": "{TOOL_ARGUMENTS_MARKER}"}}'), + ) + ], + ), + ) + for index in (0, 1) + ], + ) + + +async def _two_choice_tool_call_stream() -> AsyncGenerator[ModelResponseStream, None]: + yield _two_choice_tool_call_chunk() + yield _two_choice_tool_call_chunk(finish_reason="tool_calls") + + +@pytest.mark.asyncio +async def test_buffered_messages_stream_withholds_everything_when_the_guardrail_drops_a_tool_call(): + """A tool call the guardrail removed must not reach the client as the original, the request fails with a 500.""" + guardrail = _ToolDroppingGuardrail(guardrail_name="dropper", event_hook="post_call") + with pytest.raises(HTTPException) as raised: + await _run_messages_stream(guardrail, _anthropic_tool_use_stream()) + assert raised.value.status_code == 500 + assert "dropper" in str(raised.value.detail) + assert guardrail.scan_count == 1 + + +@pytest.mark.asyncio +async def test_buffered_chat_stream_withholds_everything_when_the_tool_call_rewrite_cannot_be_placed(): + """Neither the originals nor a half-applied rewrite reach the client, the request fails with a 500 instead.""" + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call") + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = False + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + + async def drain() -> None: + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, response=_two_choice_tool_call_stream(), request_data=request_data + ): + collected.append(chunk) + + with pytest.raises(HTTPException) as raised: + await drain() + assert raised.value.status_code == 500 + assert "masker" in str(raised.value.detail) + assert guardrail.scan_count == 1 + assert collected == [] + + +@pytest.mark.asyncio +async def test_buffered_responses_stream_delivers_masked_text_to_the_client(): + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call") + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = False + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_replay(_responses_message_stream_events(["Here is ", ORIGINAL_MARKER])), + request_data=request_data, + ): + collected.append(chunk) + raw = json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected]) + assert guardrail.scan_count == 1 + assert MASKED_MARKER in raw, raw + assert ORIGINAL_MARKER not in raw, raw + assert "response.completed" in raw + + +@pytest.mark.asyncio +async def test_buffered_responses_stream_withholds_everything_when_the_guardrail_drops_a_tool_call(): + """A function call the guardrail removed must not reach the client as the original, the request fails with a 500.""" + guardrail = _ToolDroppingGuardrail(guardrail_name="dropper", event_hook="post_call") + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = False + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + + async def drain() -> None: + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_replay( + _responses_truncated_function_call_events("Running", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']) + ), + request_data=request_data, + ): + collected.append(chunk) + + with pytest.raises(HTTPException) as raised: + await drain() + + assert raised.value.status_code == 500 + assert "dropper" in str(raised.value.detail) + assert guardrail.scan_count == 1 + assert collected == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "guardrail_cls", [_ToolDroppingGuardrail, _ToolClearingGuardrail], ids=["returns fewer", "clears in place"] +) +async def test_buffered_chat_stream_withholds_everything_when_the_guardrail_drops_a_tool_call(guardrail_cls: type): + guardrail = guardrail_cls(guardrail_name="dropper", event_hook="post_call") + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = False + with pytest.raises(HTTPException) as raised: + await _run_buffered_chat(guardrail, ["Running"], ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']) + assert raised.value.status_code == 500 + assert "dropper" in str(raised.value.detail) + assert guardrail.scan_count == 1 + + +@pytest.mark.asyncio +async def test_live_messages_stream_keeps_block_only_for_a_masking_guardrail(): + """Without buffering the chunks are already out, so a rewrite cannot reach the client and nothing is replaced.""" + guardrail = _MaskingGuardrail(guardrail_name="masker", event_hook="post_call") + guardrail.streaming_buffer_until_moderated = False + guardrail.streaming_end_of_stream_only = True + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: list[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, response=_anthropic_stream(), request_data=request_data + ): + collected.append(chunk) + raw = _decode(collected) + assert ORIGINAL_MARKER in raw + assert MASKED_MARKER not in raw + + def _detect_only_pipeline(guardrail: CustomGuardrail) -> tuple[tuple[str, GuardrailPipeline], ...]: step = PipelineStep(guardrail=guardrail.guardrail_name, on_pass="allow", on_fail="next") return (("live-policy", GuardrailPipeline(mode="post_call", steps=[step])),)