diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6b6974becac..7654b56ad9a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/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( diff --git a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py index 5930739d445..8d1b358c8f1 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -12,6 +12,7 @@ import json from collections.abc import Mapping from typing import Any, Final, cast +from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -34,8 +35,13 @@ class FakeAnthropicMessagesStreamIterator: - message_stop """ - def __init__(self, response: AnthropicMessagesResponse): + def __init__( + self, + response: AnthropicMessagesResponse, + message_delta_fields: Mapping[str, object] | None = None, + ): self.response = response + self.message_delta_fields = message_delta_fields self.chunks = self._create_streaming_chunks() self.current_index = 0 @@ -175,14 +181,9 @@ class FakeAnthropicMessagesStreamIterator: cache_read_input_tokens: Final = usage.get("cache_read_input_tokens") if cache_read_input_tokens is not None: delta_usage["cache_read_input_tokens"] = cache_read_input_tokens - safeguard_results: Final = self.response.get("safeguard_results") message_delta: Final = { "type": "message_delta", - "delta": { - "stop_reason": self.response.get("stop_reason"), - "stop_sequence": self.response.get("stop_sequence"), - **({"safeguard_results": safeguard_results} if safeguard_results is not None else {}), - }, + "delta": anthropic_message_delta(response_dict, self.message_delta_fields), "usage": delta_usage, } chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode()) diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index e65fe596da7..44db193ea83 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -18,7 +18,11 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP -from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE +from litellm.llms.anthropic.pass_through.messages.utils import ( + INCOMPLETE_STREAM_ERROR_MESSAGE, + MESSAGE_DELTA_OPTIONAL_KEYS, + anthropic_message_delta, +) from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -280,30 +284,32 @@ def anthropic_messages_response_as_sse_events(response: AnthropicMessagesRespons content_events: Final = ( event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block) ) - # A real message_start always carries a null stop_reason/stop_sequence and - # a zero output_tokens - those are only known once generation finishes, so - # copying the completed response's final values here would let a client - # treat the message as already finished, or double-count output tokens. + # A real message_start always carries a null stop_reason/stop_sequence/stop_details, + # a zero output_tokens and no safeguard_results - those are only known once + # generation finishes, so copying the completed response's final values here + # would let a client treat the message as already finished, or double-count + # output tokens. message_start_usage: Final = { **(response.get("usage") or {}), "output_tokens": 0, } + message_start_fields: Final = { + key: value for key, value in response.items() if key not in MESSAGE_DELTA_OPTIONAL_KEYS + } message_start_payload: Final = { "type": "message_start", "message": { - **response, + **message_start_fields, "content": [], "stop_reason": None, "stop_sequence": None, + "stop_details": None, "usage": message_start_usage, }, } message_delta_payload: Final = { "type": "message_delta", - "delta": { - "stop_reason": response.get("stop_reason"), - "stop_sequence": response.get("stop_sequence"), - }, + "delta": anthropic_message_delta(response), "usage": response.get("usage") or {}, } return ( diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index d8b96c1b7a0..d91a5e2fb5c 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/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 @@ -49,6 +49,95 @@ def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: return AnthropicStopDetails(type="refusal", category=None, explanation=explanation) +MESSAGE_DELTA_OPTIONAL_KEYS: Final = ("stop_details", "safeguard_results") + + +def anthropic_message_delta( + response: Mapping[str, object], optional_fields: Mapping[str, object] | None = None +) -> Mapping[str, object]: + """ + The ``delta`` of the ``message_delta`` event synthesized from a complete + Anthropic Messages response. A real stream carries ``stop_details`` and + ``safeguard_results`` there, so a client that asked to stream reads them + from this event and nowhere else. ``optional_fields`` replaces the + response's own when it was rebuilt from a stream that carried them. + """ + optional_source: Final = response if optional_fields is None else optional_fields + present_optional: Final = { + key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := optional_source.get(key)) is not None + } + return { + "stop_reason": response.get("stop_reason"), + "stop_sequence": response.get("stop_sequence"), + **present_optional, + } + + +_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 if keep_explanations else None + 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): diff --git a/litellm/llms/anthropic/pass_through/stream_assembly.py b/litellm/llms/anthropic/pass_through/stream_assembly.py index 89961bbe255..90cb9f4c335 100644 --- a/litellm/llms/anthropic/pass_through/stream_assembly.py +++ b/litellm/llms/anthropic/pass_through/stream_assembly.py @@ -1,16 +1,18 @@ import codecs import json from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.pass_through.messages.utils import MESSAGE_DELTA_OPTIONAL_KEYS from litellm.types.utils import ( Choices, Message, @@ -447,6 +449,27 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: ) +def anthropic_message_delta_fields(all_chunks: Sequence[object]) -> Mapping[str, object]: + """The ``stop_details`` and ``safeguard_results`` of the stream's ``message_delta``. + + The assembled OpenAI shape has no slot for them, and a client that asked to stream reads them + from that event and nowhere else (a Claude Code auto-mode session takes its classifier verdict + from ``safeguard_results``), so a caller re-emitting a rewritten response merges them back in. + """ + sse_stream: Final = _joined_sse_stream(all_chunks) + if sse_stream is None: + return MappingProxyType({}) + message_delta: Final = next( + (event for event in reversed(_parsed_sse_events(sse_stream)) if event.get("type") == "message_delta"), None + ) + delta: Final = as_str_mapping(message_delta.get("delta")) if message_delta is not None else None + if delta is None: + return MappingProxyType({}) + return MappingProxyType( + {key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := delta.get(key)) is not None} + ) + + def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: """Whether raw SSE frames are Anthropic Messages events. diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index c3a6f259f95..2baeabdce7b 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -8,8 +8,13 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response. from __future__ import annotations import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final +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 ) @@ -50,7 +55,62 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 ) -def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]: +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) + return MappingProxyType( + { + tool_use_id: block.get("input") + for block in blocks + if block.get("type") == "tool_use" and isinstance(tool_use_id := block.get("id"), str) + } + ) + + +def _unchanged_tool_use_ids(source_content: Sequence[object], delivered_content: Sequence[object]) -> frozenset[str]: + source_inputs: Final = _tool_use_inputs(source_content) + return frozenset( + tool_use_id + for tool_use_id, tool_input in _tool_use_inputs(delivered_content).items() + if tool_use_id in source_inputs and source_inputs[tool_use_id] == tool_input + ) + + +def _rewritten_message_delta_fields( + source_chunks: Sequence[object], unchanged_tool_use_ids: frozenset[str], keep_explanations: bool +) -> Mapping[str, object]: + source_fields: Final = anthropic_message_delta_fields(source_chunks) + stop_details: Final = source_fields.get("stop_details") + safeguard_results: Final = source_fields.get("safeguard_results") + return MappingProxyType( + { + **( + {"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 {} + ), + } + ) + + +def anthropic_sse_chunks_from_response( + 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. + + ``message_delta`` carries the source stream's ``stop_details`` and ``safeguard_results`` and + nothing the OpenAI round trip made up, with a verdict kept only for a tool call delivered with + the input the classifier judged. Their free-text ``explanation`` was written about the text the + guardrail rewrote, so it is dropped unless the caller rewrites no text (``keep_explanations``). + """ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, ) @@ -58,7 +118,22 @@ def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, FakeAnthropicMessagesStreamIterator, ) - anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( - response=assembled + adapter: Final = LiteLLMAnthropicMessagesAdapter() + anthropic_response: Final = adapter.translate_openai_response_to_anthropic(response=assembled) + source_response: Final = assemble_anthropic_sse_stream(source_chunks) + source_content: Final = ( + adapter.translate_openai_response_to_anthropic(response=source_response).get("content") + if source_response is not None + else None + ) + unchanged_tool_use_ids: Final = _unchanged_tool_use_ids( + source_content or (), anthropic_response.get("content") or () + ) + return tuple( + FakeAnthropicMessagesStreamIterator( + response=anthropic_response, + message_delta_fields=_rewritten_message_delta_fields( + source_chunks, unchanged_tool_use_ids, keep_explanations + ), + ).chunks ) - return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b92020c89be..0672ccf6a66 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -3058,7 +3058,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################################### if raw_sse: for sse_chunk in ( - anthropic_sse_chunks_from_response(assembled_model_response) + anthropic_sse_chunks_from_response(assembled_model_response, all_chunks) if model_response_text(assembled_model_response) != pre_guardrail_text else all_chunks ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index e590e5c2e43..b9871b1e276 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1263,7 +1263,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Return sanitized stream if surface is _StreamSurface.ANTHROPIC_MESSAGES: - for sse_chunk in anthropic_sse_chunks_from_response(assembled_response): + for sse_chunk in anthropic_sse_chunks_from_response(assembled_response, all_chunks): yield sse_chunk return mock_response: Final = MockResponseIterator(model_response=assembled_response) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 01a9259203a..f21aabc7f33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1532,7 +1532,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask") if model_response_text(assembled) == original_text: return chunks - return anthropic_sse_chunks_from_response(assembled) + return anthropic_sse_chunks_from_response(assembled, chunks) @staticmethod def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index ebd50c8cd39..fa92a0c90e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -971,7 +971,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return self._modify_response_with_permission_errors(anthropic_response, anthropic_denials) - for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response): + for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response, all_chunks, keep_explanations=True): yield sse_chunk def _check_assembled_stream( diff --git a/litellm/router.py b/litellm/router.py index 99adc6c5fa6..1b3857dfb71 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6055,6 +6055,9 @@ class Router: including pings and provider error frames, is forwarded live and verbatim instead. """ + from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( aclose_if_supported, aclose_superseded_attempt, @@ -6066,6 +6069,7 @@ class Router: ) source_iterator: Final = response + replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) async def close_source() -> None: with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): @@ -6126,7 +6130,11 @@ class Router: if not has_generated_content and error_event is None else None ) - if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs): + if ( + refusal_stop_details is not None + and not replays_a_finished_agentic_turn + and self._refusal_fallback_available(model, initial_kwargs) + ): refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) raise MidStreamFallbackError( message=refusal_error.message, diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py index c9da24484e2..57ab6e0e745 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py @@ -4,28 +4,35 @@ with no verdicts.""" from __future__ import annotations +import json from collections.abc import Callable, Iterator, Mapping +from itertools import chain from typing import Final import pytest -from integration._support.client import object_value +from integration._support.anthropic_refusal import DOCS_STOP_DETAILS +from integration._support.client import list_value, object_value from integration._support.mcp import McpPeer, mcp_peer, tool_calls from integration._support.process import graceful_stop_seconds from integration._support.safeguards_rig import ( BASH_TOOL, CHAT_MODEL, + CLAUDE_MODEL, Responder, Rig, auto_mode_body, available, chat_backend, json_body, + native_message, nonce, request_json, safeguards_rig, + sse_events, stream_results, stream_tool_use_ids, tool_id, + tool_use_block, tool_use_ids, ) from integration._support.wire import Reply, Request @@ -115,3 +122,45 @@ def test_auto_executed_mcp_tool_costs_no_classifier_call(rig: Rig, peer: McpPeer received: Final = rig.received(request_nonce) assert rig.classifier_calls(received) == (), [request.target for request in received] assert len(rig.backend_calls(received)) == 2, [request.target for request in received] + + +def carries_a_tool_result_block(request: Request) -> bool: + messages: Final = list_value(request_json(request).get("messages", [])) + contents: Final = tuple(object_value(message).get("content") for message in messages) + blocks: Final = chain.from_iterable(content for content in contents if isinstance(content, list)) + return any(object_value(block).get("type") == "tool_result" for block in blocks) + + +def claude_mcp_tool_then_refusal(request_nonce: str) -> Responder: + identity: Final = tool_id(request_nonce) + tool_turn: Final = native_message( + identity, + content=(tool_use_block(f"{identity}m", name=f"{MCP_ALIAS}-add", arguments={"a": 2, "b": 3}),), + model="claude-sonnet-4-6", + ) + refusal_turn: Final = { + **native_message(f"{identity}r", content=(), stop_reason="refusal", model="claude-sonnet-4-6"), + "stop_details": DOCS_STOP_DETAILS, + } + + def respond(request: Request) -> Reply: + turn: Final = refusal_turn if carries_a_tool_result_block(request) else tool_turn + return Reply(body=json.dumps(turn).encode()) + + return respond + + +def test_auto_executed_mcp_loop_streams_the_final_turn_refusal_stop_details(rig: Rig, peer: McpPeer) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, claude_mcp_tool_then_refusal(request_nonce)) + peer.drain() + status, raw = rig.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, tools=[dict(AUTO_EXECUTED), dict(BASH_TOOL)])) + assert status == 200, raw + assert executed_tools(peer) == ("add",) + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + assert deltas[0]["stop_reason"] == "refusal", raw + assert deltas[0].get("stop_details") == DOCS_STOP_DETAILS, raw + assert len(rig.backend_calls(rig.received(request_nonce))) == 2 diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py new file mode 100644 index 00000000000..99a6dd18d48 --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py @@ -0,0 +1,470 @@ +"""A /v1/messages stream LiteLLM builds from a complete reply (web search interception, a fallback's answer, a guardrail +rewrite) carries the ``stop_details`` and ``safeguard_results`` a real Anthropic stream puts on its ``message_delta``.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Iterator, Mapping, Sequence +from typing import Final +from urllib.parse import urlsplit + +import pytest +from integration._support.client import eventually, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + BASH_TOOL, + CLAUDE_MODEL, + CLAUDE_TARGET, + TAVILY_TARGET, + Responder, + Rig, + auto_mode_body, + available, + deployment, + native_message, + native_stream_frames, + nonce, + request_json, + safeguards_rig, + sse_events, + stream_results, + stream_tool_use_ids, + text_block, + tool_id, + tool_use_block, + unavailable, +) +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +CLAUDE: Final = "anthropic/claude-sonnet-4-6" +PRIMARY_MODEL: Final = "synthesized-primary" +FALLBACK_MODEL: Final = "synthesized-fallback" +REFUSED_MODEL: Final = "synthesized-refused" +MOCK_MODEL: Final = "synthesized-mock" +PIPELINE_MODEL: Final = "synthesized-pipeline" +AZURE_REFUSED_MODEL: Final = "synthesized-azure-refused" +CACHED_FALLBACK_MODEL: Final = "synthesized-cached-fallback" +PRIMARY_TARGET: Final = "/primary/v1/messages" +FALLBACK_TARGET: Final = "/fallback/v1/messages" +REFUSED_TARGET: Final = "/refused/v1/messages" +PIPELINE_TARGET: Final = "/pipeline/v1/messages" +AZURE_REFUSED_TARGET: Final = "/azure/anthropic/v1/messages" +CACHED_FALLBACK_TARGET: Final = "/cached-fallback/v1/messages" +MOCK_ANSWER: Final = "mocked fallback answer" +TOOL_PERMISSION: Final = "synthesized-tool-permission" +PIPELINE_REWRITE: Final = "synthesized-pipeline-rewrite" +SECRET: Final = "SECRET" +MASKED: Final = "[MASKED]" +REWRITE_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " if input_type != 'response':\n" + " return allow()\n" + f" texts = [text.replace({SECRET!r}, {MASKED!r}) for text in inputs.get('texts') or []]\n" + " calls = [\n" + " {'id': call['id'], 'type': 'function', 'function': {'name': call['function']['name'],\n" + f" 'arguments': call['function']['arguments'].replace({SECRET!r}, {MASKED!r})}}}}\n" + " for call in inputs.get('tool_calls') or []\n" + " ]\n" + " return modify(texts=texts, tool_calls=calls)\n" +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +READ_TOOL: Final[Mapping[str, JsonValue]] = { + "name": "Read", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, +} +NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +REFUSAL_STOP_DETAILS: Final[Mapping[str, JsonValue]] = { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", +} +_FRAME_BOUNDARY: Final = re.compile(rb"(?<=\n\n)") + +LITELLM_SETTINGS: Final[Mapping[str, JsonValue]] = { + "callbacks": ["websearch_interception"], + "websearch_interception_params": {"enabled_providers": ["anthropic"], "search_tool_name": "tavily-search"}, +} + + +def extra_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + deployment(PRIMARY_MODEL, CLAUDE, f"{wire_url}/primary"), + deployment(FALLBACK_MODEL, CLAUDE, f"{wire_url}/fallback"), + deployment(REFUSED_MODEL, CLAUDE, f"{wire_url}/refused"), + deployment(MOCK_MODEL, CLAUDE, f"{wire_url}/mock", mock_response=MOCK_ANSWER), + deployment(PIPELINE_MODEL, CLAUDE, f"{wire_url}/pipeline"), + deployment(AZURE_REFUSED_MODEL, "azure_ai/claude-sonnet-4-6", f"{wire_url}/azure/anthropic"), + deployment(CACHED_FALLBACK_MODEL, CLAUDE, f"{wire_url}/cached-fallback"), + ) + + +def top_level(wire_url: str) -> Mapping[str, JsonValue]: + tavily: Final[Mapping[str, JsonValue]] = { + "search_tool_name": "tavily-search", + "litellm_params": { + "search_provider": "tavily", + "api_key": "synthetic-tavily-key", + "api_base": f"{wire_url}/tavily", + }, + } + tool_permission: Final[Mapping[str, JsonValue]] = { + "guardrail_name": TOOL_PERMISSION, + "litellm_params": { + "guardrail": "tool_permission", + "mode": "post_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": "rewrite", + "rules": [{"id": "deny_bash", "tool_name": "^Bash$", "decision": "deny"}], + }, + } + pipeline_rewrite: Final[Mapping[str, JsonValue]] = { + "guardrail_name": PIPELINE_REWRITE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": REWRITE_CODE, + }, + } + masking_policy: Final[Mapping[str, JsonValue]] = { + "guardrails": {"add": [PIPELINE_REWRITE]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": PIPELINE_REWRITE, "on_pass": "allow", "on_fail": "block"}], + }, + } + return { + "search_tools": [dict(tavily)], + "guardrails": [dict(tool_permission), dict(pipeline_rewrite)], + "policies": {"synthesized-masking": dict(masking_policy)}, + "policy_attachments": [{"policy": "synthesized-masking", "models": [PIPELINE_MODEL]}], + "router_settings": { + "disable_cooldowns": True, + "content_policy_fallbacks": [ + {PRIMARY_MODEL: [FALLBACK_MODEL]}, + {REFUSED_MODEL: [MOCK_MODEL]}, + {AZURE_REFUSED_MODEL: [CACHED_FALLBACK_MODEL]}, + ], + }, + } + + +@pytest.fixture(scope="module") +def synthesizing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("synthesized-streams") + with safeguards_rig( + directory, litellm_settings=LITELLM_SETTINGS, top_level=top_level, extra_models=extra_models + ) as rig: + yield rig + + +def claude_backend(message: Mapping[str, JsonValue], delta_fields: Mapping[str, JsonValue]) -> Responder: + def respond(request: Request) -> Reply: + if request_json(request).get("stream") is True: + return Reply( + content_type="text/event-stream", chunks=native_stream_frames(message, delta_extra=delta_fields) + ) + return Reply(body=json.dumps({**message, **delta_fields}).encode()) + + return respond + + +def by_target(backends: Mapping[str, Responder]) -> Responder: + def respond(request: Request) -> Reply: + return backends[urlsplit(request.target).path](request) + + return respond + + +def refusal(identity: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(), stop_reason="refusal", model="claude-sonnet-4-6") + + +def answer(identity: str, text: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(text_block(text),), stop_reason="end_turn", model="claude-sonnet-4-6") + + +def refusal_fields() -> Mapping[str, JsonValue]: + return {"stop_details": dict(REFUSAL_STOP_DETAILS), "safeguard_results": list(available({}))} + + +def search_body(model: str, request_nonce: str) -> Mapping[str, JsonValue]: + return auto_mode_body(model, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for", **NO_CACHE) + + +def message_delta(raw: bytes) -> Mapping[str, JsonValue]: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + return deltas[0] + + +def message_start(raw: bytes) -> Mapping[str, JsonValue]: + starts: Final = tuple( + object_value(event["message"]) for event in sse_events(raw) if event.get("type") == "message_start" + ) + assert len(starts) == 1, raw + return starts[0] + + +def streamed_text(raw: bytes) -> str: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "content_block_delta" + ) + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +def delta_frame(raw: bytes) -> bytes: + frames: Final = tuple(frame for frame in _FRAME_BOUNDARY.split(raw) if b'"message_delta"' in frame) + assert len(frames) == 1, raw + return frames[0] + + +def targets(rig: Rig, request_nonce: str) -> tuple[str, ...]: + return tuple(urlsplit(request.target).path for request in rig.received(request_nonce)) + + +def test_intercepted_refusal_stream_delivers_the_upstream_stop_details(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_refusal_stream_with_a_content_policy_fallback_delivers_the_refusal_without_falling_back( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + synthesizing.register( + request_nonce, + by_target( + { + PRIMARY_TARGET: claude_backend(refusal(identity), refusal_fields()), + FALLBACK_TARGET: claude_backend(answer(f"{identity}f", f"fallback answer {request_nonce}"), {}), + } + ), + ) + status, raw = synthesizing.stream(search_body(PRIMARY_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (PRIMARY_TARGET,) + + +def test_tool_permission_rewrite_keeps_only_the_delivered_tool_verdict_with_its_explanation(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + denied: Final = tool_id(request_nonce) + allowed: Final = f"{denied}r" + allowed_verdict: Final[Mapping[str, JsonValue]] = { + "type": "evaluated", + "outcome": "flagged", + "explanation": f"reads a secret {allowed}", + } + verdicts: Final[Mapping[str, JsonValue]] = { + denied: {"type": "evaluated", "outcome": "flagged", "explanation": f"destructive {denied}"}, + allowed: dict(allowed_verdict), + } + message: Final = native_message( + denied, + content=( + text_block("Cleaning up."), + tool_use_block(denied), + tool_use_block(allowed, name="Read", arguments={"path": "/work/repo/.env"}), + ), + model="claude-sonnet-4-6", + ) + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body( + CLAUDE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], guardrails=[TOOL_PERMISSION], **NO_CACHE + ) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (allowed,), raw + assert stream_results(raw) == available({allowed: dict(allowed_verdict)}), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_answer_stream_keeps_its_safeguard_results(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + text: Final = f"found it {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + claude_backend(answer(tool_id(request_nonce), text), {"stop_details": None, "safeguard_results": results}), + ) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert streamed_text(raw) == text, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + received: Final = targets(synthesizing, request_nonce) + assert received == (CLAUDE_TARGET,) and TAVILY_TARGET not in received, received + + +def test_native_refusal_stream_without_a_guardrail_keeps_the_upstream_message_delta_bytes(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + message: Final = native_message( + identity, content=(text_block("Partial."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + fields: Final = {**refusal_fields(), "safeguard_results": list(unavailable((identity,), "error"))} + sent: Final = native_stream_frames(message, delta_extra=fields) + synthesizing.register(request_nonce, claude_backend(message, fields)) + status, raw = synthesizing.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + assert delta_frame(raw) == delta_frame(b"".join(sent)), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_native_refusal_falling_back_to_a_mock_reply_renders_a_fresh_message_lifecycle(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(auto_mode_body(REFUSED_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + started: Final = message_start(raw) + assert started["stop_reason"] is None, raw + assert started.get("stop_details") is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None}, raw + assert streamed_text(raw) == MOCK_ANSWER, raw + assert b'"refusal"' not in raw, raw + assert targets(synthesizing, request_nonce) == (REFUSED_TARGET,) + + +def test_native_refusal_falling_back_to_a_cached_intercepted_reply_streams_its_fields_in_message_delta( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + text: Final = f"cached answer {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + by_target( + { + AZURE_REFUSED_TARGET: claude_backend(refusal(identity), refusal_fields()), + CACHED_FALLBACK_TARGET: claude_backend( + answer(f"{identity}c", text), {"stop_details": None, "safeguard_results": results} + ), + } + ), + ) + cacheable: Final = auto_mode_body( + CACHED_FALLBACK_MODEL, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for" + ) + + def prime() -> tuple[str, ...]: + primed: Final = synthesizing.post({**cacheable, "stream": False}) + assert primed.status_code == 200, primed.text + return targets(synthesizing, request_nonce) + + first_prime: Final = prime() + assert first_prime == (CACHED_FALLBACK_TARGET,), first_prime + eventually(prime, lambda received: received == ()) + status, raw = synthesizing.stream({**cacheable, "model": AZURE_REFUSED_MODEL}) + assert status == 200, raw + assert streamed_text(raw) == text, raw + started: Final = message_start(raw) + assert started["stop_details"] is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + assert targets(synthesizing, request_nonce) == (AZURE_REFUSED_TARGET,) + + +def flagged_verdict(explanation: str) -> Mapping[str, JsonValue]: + return {"type": "evaluated", "outcome": "flagged", "explanation": explanation} + + +FLAGGED_WITHOUT_EXPLANATION: Final[Mapping[str, JsonValue]] = {"type": "evaluated", "outcome": "flagged"} + + +def masked_refusal(identity: str) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, content=(text_block(f"Partial {SECRET}."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {**REFUSAL_STOP_DETAILS, "explanation": None}, + "safeguard_results": list(available({})), + } + return message, refusal_fields(), expected + + +def masked_tool_turn( + identity: str, +) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, + content=(text_block(f"Reading {SECRET}."), tool_use_block(identity, name="Read", arguments={"path": "/work"})), + model="claude-sonnet-4-6", + ) + sent: Final[Mapping[str, JsonValue]] = { + "safeguard_results": list(available({identity: dict(flagged_verdict(f"reads {identity}"))})) + } + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": list(available({identity: dict(FLAGGED_WITHOUT_EXPLANATION)})), + } + return message, sent, expected + + +MASKED_STREAMS: Final = {"refusal_stop_details": masked_refusal, "verdict_explanation": masked_tool_turn} + + +@pytest.mark.parametrize("shape", tuple(MASKED_STREAMS)) +def test_pipeline_text_rewrite_drops_explanations_from_the_upstream_message_delta( + synthesizing: Rig, shape: str +) -> None: + request_nonce: Final = nonce() + message, sent, expected = MASKED_STREAMS[shape](tool_id(request_nonce)) + synthesizing.register(request_nonce, claude_backend(message, sent)) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert MASKED in streamed_text(raw) and SECRET not in streamed_text(raw), raw + assert message_delta(raw) == expected, raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) + + +def test_pipeline_tool_call_rewrite_keeps_only_the_untouched_tool_verdict(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + rewritten: Final = tool_id(request_nonce) + untouched: Final = f"{rewritten}r" + message: Final = native_message( + rewritten, + content=( + tool_use_block(rewritten, arguments={"command": f"cat {SECRET}"}), + tool_use_block(untouched, name="Read", arguments={"path": "/work"}), + ), + model="claude-sonnet-4-6", + ) + verdicts: Final[Mapping[str, JsonValue]] = { + rewritten: dict(flagged_verdict(f"reads a file {rewritten}")), + untouched: dict(flagged_verdict(f"reads a directory {untouched}")), + } + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (rewritten, untouched), raw + assert SECRET.encode() not in raw and MASKED.encode() in raw, raw + assert stream_results(raw) == available({untouched: dict(FLAGGED_WITHOUT_EXPLANATION)}), raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index d784fecd146..d69866ec43d 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -2,7 +2,7 @@ import json import re import threading import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager from dataclasses import dataclass @@ -10,6 +10,7 @@ from pathlib import Path from typing import Final import httpx +import pytest import yaml from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy_process @@ -457,6 +458,96 @@ def test_anthropic_messages_per_request_guardrails_selects_masking(gateway: Gate assert len(rig.analyzer.drain()) == 1 +READ_TOOL_USE_ID: Final = "toolu_presidio_read" +READ_INPUT: Final = {"path": "/work/report.txt"} + + +def anthropic_tool_turn_stream(identity: str, text: str, delta: Mapping[str, object]) -> tuple[bytes, ...]: + opening, text_start, text_delta, text_stop, _, closing = anthropic_stream(identity, text) + tool_use: Final = {"type": "tool_use", "id": READ_TOOL_USE_ID, "name": "Read", "input": {}} + return ( + opening, + text_start, + text_delta, + text_stop, + anthropic_frame("content_block_start", {"type": "content_block_start", "index": 1, "content_block": tool_use}), + anthropic_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": json.dumps(READ_INPUT)}, + }, + ), + anthropic_frame("content_block_stop", {"type": "content_block_stop", "index": 1}), + anthropic_frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 4}}), + closing, + ) + + +def read_verdicts(verdict: Mapping[str, object]) -> Sequence[Mapping[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {READ_TOOL_USE_ID: verdict}}}] + + +def anthropic_deltas(received: Received, event_type: str) -> tuple[Mapping[str, object], ...]: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ") + ) + return tuple(event["delta"] for event in events if event.get("type") == event_type) + + +def anthropic_message_delta(received: Received) -> Mapping[str, object]: + deltas: Final = anthropic_deltas(received, "message_delta") + assert len(deltas) == 1, received.text + return deltas[0] + + +def anthropic_text_deltas(received: Received) -> str: + deltas: Final = anthropic_deltas(received, "content_block_delta") + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +@pytest.mark.parametrize( + ("delta", "field", "kept"), + [ + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": read_verdicts( + {"type": "evaluated", "outcome": "flagged", "explanation": f"{PERSON} asked for the report"} + ), + }, + "safeguard_results", + read_verdicts({"type": "evaluated", "outcome": "flagged"}), + ), + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"declined for {PERSON}"}, + }, + "stop_details", + {"type": "refusal", "category": "cyber", "explanation": None}, + ), + ], + ids=["verdict_explanation", "stop_details_explanation"], +) +def test_anthropic_messages_masked_stream_keeps_upstream_safeguard_fields_without_explanations( + gateway: Gateway, tmp_path: Path, delta: Mapping[str, object], field: str, kept: object +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_tool_turn_stream(identity, f"{PERSON} designed it.", delta)) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text_deltas(received) == f"{MASK} designed it.", received.text + assert anthropic_message_id(received) == identity, received.text + assert anthropic_message_delta(received).get(field) == kept, received.text + assert PERSON not in received.text, received.text + assert len(rig.upstream.drain()) == 1 + + def openai_provider(identity: str) -> Callable[[Request], Reply]: def provider(request: Request) -> Reply: assert request.target == "/v1/chat/completions", request.target diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 892736b5043..0aca60351bf 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/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() diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py new file mode 100644 index 00000000000..476cea2833c --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py @@ -0,0 +1,93 @@ +import json +from collections.abc import Iterable, Mapping +from typing import Final + +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, +) + +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _decoded_event(chunk: object) -> tuple[str, dict[str, object]]: + assert isinstance(chunk, bytes) + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decoded_events(chunks: Iterable[object]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response)) + assert tuple(event_type for event_type, _ in decoded)[-2:] == ("message_delta", "message_stop") + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_message_delta_carries_safeguard_results(): + """Websearch interception and the MCP loop render a complete response as + this stream for a client that asked to stream (the guardrail path does too, + after merging the upstream fields back in, which + tests/unit/proxy/guardrails/test_anthropic_sse.py covers); a Claude Code + auto-mode client reads the classifier verdict from message_delta.delta, + where a real Anthropic stream puts it.""" + response: Final = { + "id": "msg_1", + "model": "claude-sonnet-4-6", + "role": "assistant", + "type": "message", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + assert _message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_message_delta_carries_refusal_stop_details(): + stop_details: Final = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response: Final = { + "id": "msg_2", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "usage": {"input_tokens": 10, "output_tokens": 0}, + } + assert _message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_message_delta_omits_absent_optional_fields(): + response: Final = { + "id": "msg_3", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": {"input_tokens": 10, "output_tokens": 2}, + } + assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py index 7e983be835b..b995c6dba9b 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py @@ -1306,6 +1306,33 @@ def test_anthropic_messages_streaming_response_reports_withheld_output_of_its_st ) +def test_anthropic_messages_response_as_sse_events_keeps_message_delta_fields_out_of_message_start(): + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}] + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_1", + "model": "claude-haiku", + "role": "assistant", + "type": "message", + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + "content": [{"type": "text", "text": "no"}], + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + message_start = decoded[0][1]["message"] + assert message_start["stop_details"] is None + assert "safeguard_results" not in message_start + assert decoded[-2][1]["delta"] == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + } + + def test_anthropic_messages_response_as_sse_events_no_content_blocks(): response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) @@ -1530,6 +1557,72 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co assert deferred_fired.is_set() +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _sse_events_message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_safeguard_results(): + """A Claude Code auto-mode request's classifier verdict rides in + message_delta.delta on a real Anthropic stream; a synthesized stream that + drops it leaves the client without the verdict.""" + response = { + "id": "msg_7", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_refusal_stop_details(): + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_8", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_omits_absent_optional_fields(): + response = { + "id": "msg_9", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + } + assert _sse_events_message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} + + class _OneVerdictClassifier: def __init__(self, reply: str) -> None: self._reply = reply diff --git a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py index 33f95ccc9ff..dbd49064e4b 100644 --- a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py @@ -4,6 +4,7 @@ from typing import Final import pytest from litellm.llms.anthropic.pass_through.stream_assembly import ( + anthropic_message_delta_fields, assemble_anthropic_sse_stream, is_anthropic_sse_stream, is_raw_sse_stream, @@ -55,6 +56,35 @@ ANTHROPIC_STREAM: Final = ( _frame("message_stop", {"type": "message_stop"}), ) +SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] +SAFEGUARDED_STREAM: Final = ( + *ANTHROPIC_STREAM[:5], + _frame( + "message_delta", + { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": SAFEGUARD_RESULTS, + }, + "usage": {"output_tokens": 2}, + }, + ), + _frame("message_stop", {"type": "message_stop"}), +) + GOOGLE_STREAM: Final = (b'data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}\n\n',) ANTHROPIC_ERROR: Final = _frame("error", {"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}) CHAT_ERROR: Final = b'data: {"error": {"message": "blocked by guardrail"}}\n\n' @@ -108,6 +138,12 @@ def test_assemble_anthropic_sse_stream_restores_upstream_identity_only_when_aske assert (restored.id, restored.model) == (MESSAGE_ID, MODEL) +def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops(): + assert dict(anthropic_message_delta_fields(ANTHROPIC_STREAM)) == {} + assert dict(anthropic_message_delta_fields(SAFEGUARDED_STREAM)) == {"safeguard_results": SAFEGUARD_RESULTS} + assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 4e19c7ed328..2ee4c8c72c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/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 @@ -1379,6 +1380,54 @@ class TestToolPermissionGuardrailAnthropicMessages: assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served" + @pytest.mark.asyncio + 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.""" + 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": { + "type": "available", + "tool_uses": {"tu_ok": bash_verdict, "tu_bad": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + 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: Final = [ + *bash_chunks[:7], + *read_tool_use, + f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(), + bash_chunks[-1], + ] + + 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" + ) + assert "Permission denied" in body + assert delta["safeguard_results"] == [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {"tu_ok": bash_verdict}}} + ] + @pytest.mark.asyncio async def test_message_start_without_a_dict_message_fails_closed(self): """Malformed SSE must not be forwarded unscanned. diff --git a/tests/unit/proxy/guardrails/test_anthropic_sse.py b/tests/unit/proxy/guardrails/test_anthropic_sse.py new file mode 100644 index 00000000000..95f596aab11 --- /dev/null +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -0,0 +1,316 @@ +import json +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm.proxy.guardrails.anthropic_sse import ( + anthropic_sse_chunks_from_response, + assemble_anthropic_sse_stream, +) +from litellm.types.utils import Choices, ModelResponse + + +def _frame(event: str, data: dict[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +def _decoded_event(chunk: bytes) -> tuple[str, dict[str, object]]: + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decode(chunks: tuple[bytes, ...]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _rewritten_delta( + upstream: tuple[bytes, ...], rewrite: Callable[[ModelResponse], None], keep_explanations: bool = False +) -> object: + assembled: Final = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + rewrite(assembled) + rewritten: Final = anthropic_sse_chunks_from_response(assembled, upstream, keep_explanations=keep_explanations) + return next(payload["delta"] for event, payload in _decode(rewritten) if event == "message_delta") + + +def _set_text(text: str) -> Callable[[ModelResponse], None]: + def rewrite(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = text + + return rewrite + + +def _tool_use_frames(index: int, tool_use_id: str, name: str) -> tuple[bytes, ...]: + return ( + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": index, + "content_block": {"type": "tool_use", "id": tool_use_id, "name": name, "input": {}}, + }, + ), + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": '{"command": "ls"}'}, + }, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": index}), + ) + + +def _upstream_stream( + delta: dict[str, object], text: str = "call 555-0100", tool_uses: tuple[tuple[str, str], ...] = () +) -> tuple[bytes, ...]: + tool_use_frames = tuple( + chain.from_iterable( + _tool_use_frames(index, tool_use_id, name) for index, (tool_use_id, name) in enumerate(tool_uses, start=1) + ) + ) + return ( + _frame( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_upstream_1", + "type": "message", + "role": "assistant", + "model": "claude-upstream", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + }, + ), + _frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + *tool_use_frames, + _frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 2}}), + _frame("message_stop", {"type": "message_stop"}), + ) + + +def _safeguard_results(tool_uses: dict[str, dict[str, str]]) -> list[dict[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_uses}}] + + +SAFEGUARD_RESULTS: Final = _safeguard_results({}) +REFUSAL_STOP_DETAILS: Final = {"type": "refusal", "category": "cyber", "explanation": None} + + +@pytest.mark.parametrize( + ("stop_reason", "delta_fields"), + [ + ("end_turn", {"safeguard_results": SAFEGUARD_RESULTS}), + ("refusal", {"stop_details": REFUSAL_STOP_DETAILS}), + ], + ids=["safeguard_results", "refusal_stop_details"], +) +def test_rewritten_stream_keeps_the_upstream_message_delta_fields(stop_reason: str, delta_fields: dict[str, object]): + """A guardrail that rewrites a streamed response assembles it into the OpenAI shape, which has no + slot for ``safeguard_results`` or ``stop_details``, so the re-emitted ``message_delta`` used to + drop the classifier verdict a Claude Code auto-mode session reads from it.""" + upstream_delta: dict[str, object] = {"stop_reason": stop_reason, "stop_sequence": None, **delta_fields} + upstream = _upstream_stream(upstream_delta) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "call " + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + payloads = {event: payload for event, payload in decoded} + + assert payloads["content_block_delta"]["delta"] == {"type": "text_delta", "text": "call "} + assert payloads["message_delta"]["delta"] == upstream_delta + assert not set(delta_fields) & set(payloads["message_start"]["message"]) + + +SSN: Final = "123-45-6789" + + +@pytest.mark.parametrize( + ("upstream_delta", "tool_uses", "expected_fields"), + [ + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"Request involves SSN {SSN}"}, + }, + (), + {"stop_details": {"type": "refusal", "category": "cyber", "explanation": None}}, + ), + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": f"leaks SSN {SSN}"}} + ), + }, + (("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", "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] +): + upstream = _upstream_stream(upstream_delta, text=f"my SSN is {SSN}", tool_uses=tool_uses) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "my SSN is " + + rewritten = anthropic_sse_chunks_from_response(assembled, upstream) + message_delta = next(payload for event, payload in _decode(rewritten) if event == "message_delta") + + assert not [chunk for chunk in rewritten if SSN.encode() in chunk] + assert {key: message_delta["delta"][key] for key in expected_fields} == expected_fields + + +def test_masked_stream_drops_safeguard_results_it_cannot_read(): + upstream_delta: Final = { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": {"type": "dangerous_tool_use", "explanation": f"saw SSN {SSN}"}, + } + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("my SSN is ")) == { + "stop_reason": "end_turn", + "stop_sequence": None, + } + + +def test_rewritten_stream_keeps_only_the_verdicts_of_tool_uses_it_still_delivers(): + kept_verdict = {"type": "evaluated", "outcome": "not_flagged"} + upstream = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_glob": {"type": "evaluated", "outcome": "flagged"}, "toolu_bash": kept_verdict} + ), + }, + tool_uses=(("toolu_glob", "Glob"), ("toolu_bash", "Bash")), + ) + assembled = assemble_anthropic_sse_stream(upstream) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.tool_calls = [call for call in choice.message.tool_calls or [] if call.id == "toolu_bash"] + choice.message.content = "Permission denied: Tool 'Glob' denied" + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + message_delta = next(payload for event, payload in decoded if event == "message_delta") + delivered_tool_use_ids = [ + payload["content_block"]["id"] + for event, payload in decoded + if event == "content_block_start" and payload["content_block"]["type"] == "tool_use" + ] + + assert delivered_tool_use_ids == ["toolu_bash"] + assert message_delta["delta"]["safeguard_results"] == _safeguard_results({"toolu_bash": kept_verdict}) + + +def test_rewritten_refusal_gets_no_stop_details_the_upstream_did_not_send(): + """The OpenAI round trip sets ``stop_details`` on every refusal, and the repo reads a refusal + carrying one as a safeguard refusal, so a plain refusal must not come back as one.""" + upstream_delta: Final = {"stop_reason": "refusal", "stop_sequence": None} + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("call ")) == upstream_delta + + +def test_rewritten_stream_keeps_safeguard_shapes_it_does_not_know(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": {"type": "pause_turn", "category": None, "explanation": None, "policy_url": "https://x.test"}, + "safeguard_results": [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_bash": {"type": "unavailable", "reason": "overloaded"}}, + }, + "classifier_version": "2026-10", + }, + {"type": "prompt_injection", "status": {"type": "skipped"}}, + ], + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("call ")) == upstream_delta + + +def test_rewrite_that_changes_no_text_keeps_the_upstream_explanations(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}} + ), + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("Permission denied"), keep_explanations=True) == upstream_delta + + +def test_rewritten_stream_drops_the_verdict_of_a_tool_use_whose_input_the_guardrail_changed(): + """A verdict scored on the upstream input must not vouch for the input the client runs.""" + verdict: Final = {"type": "evaluated", "outcome": "not_flagged"} + upstream: Final = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results({"toolu_masked": verdict, "toolu_untouched": verdict}), + }, + tool_uses=(("toolu_masked", "Bash"), ("toolu_untouched", "Bash")), + ) + + def mask_first_tool_input(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + tool_calls: Final = choice.message.tool_calls or [] + assert [call.id for call in tool_calls] == ["toolu_masked", "toolu_untouched"] + tool_calls[0].function.arguments = '{"command": ""}' + + rewritten_delta: Final = _rewritten_delta(upstream, mask_first_tool_input) + + assert isinstance(rewritten_delta, dict) + assert rewritten_delta["safeguard_results"] == _safeguard_results({"toolu_untouched": verdict}) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 56184fa299a..637fdfa6873 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -14,6 +14,7 @@ import respx import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import FakeAnthropicMessagesStreamIterator from litellm.litellm_core_utils import get_llm_provider_logic from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import ( @@ -32,6 +33,7 @@ from litellm.router_utils.fallback_event_handlers import ( is_bad_request_without_fallback, log_failure_fallback_event, log_success_fallback_event, + mid_stream_fallback_hop_kwargs, mid_stream_fallback_snapshot_kwargs, mid_stream_retry_kwargs, record_pre_routing_selection, @@ -1865,6 +1867,41 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through(): assert len(fake.calls) == 1 +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.asyncio +async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback(): + """Websearch, MCP and advisor loops run their tools before streaming the final reply, so a + fallback would run every tool call again for a reply the client then gets from another model.""" + fake: Final = FakeAnthropicUpstream() + router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream: Final = await router._aanthropic_messages_streaming_iterator( + response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE), + initial_kwargs=mid_stream_fallback_hop_kwargs( + model="fable-tier", + original_generic_function=litellm.anthropic_messages, + controls=None, + kwargs={ + "model": "fable-tier", + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ), + ) + events: Final = [ + json.loads(line[len(b"data: ") :]) + for line in (await _collect(stream)).splitlines() + if line.startswith(b"data: ") + ] + + assert fake.calls == [] + message_delta: Final = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():