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 a61829d79cb..8d1b358c8f1 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -10,7 +10,6 @@ the LLM doesn't make a tool call, and we need to return a stream to the user. import json from collections.abc import Mapping -from types import MappingProxyType from typing import Any, Final, cast from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta @@ -39,7 +38,7 @@ class FakeAnthropicMessagesStreamIterator: def __init__( self, response: AnthropicMessagesResponse, - message_delta_fields: Mapping[str, object] = MappingProxyType({}), + message_delta_fields: Mapping[str, object] | None = None, ): self.response = response self.message_delta_fields = message_delta_fields @@ -184,7 +183,7 @@ class FakeAnthropicMessagesStreamIterator: delta_usage["cache_read_input_tokens"] = cache_read_input_tokens message_delta: Final = { "type": "message_delta", - "delta": {**anthropic_message_delta(response_dict), **self.message_delta_fields}, + "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/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index 20242de9d2f..67d3affd7c0 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -52,15 +52,19 @@ def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: MESSAGE_DELTA_OPTIONAL_KEYS: Final = ("stop_details", "safeguard_results") -def anthropic_message_delta(response: Mapping[str, object]) -> Mapping[str, object]: +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. + 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: response[key] for key in MESSAGE_DELTA_OPTIONAL_KEYS if response.get(key) is not None + 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"), diff --git a/litellm/llms/anthropic/pass_through/stream_assembly.py b/litellm/llms/anthropic/pass_through/stream_assembly.py index 654dddd00d5..90cb9f4c335 100644 --- a/litellm/llms/anthropic/pass_through/stream_assembly.py +++ b/litellm/llms/anthropic/pass_through/stream_assembly.py @@ -4,11 +4,9 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter, ValidationError - 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, @@ -451,16 +449,6 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: ) -_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) - - -def _mapping_or_none(value: object) -> Mapping[str, object] | None: - try: - return _MAPPING_ADAPTER.validate_python(value) - except ValidationError: - return 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``. @@ -474,7 +462,7 @@ def anthropic_message_delta_fields(all_chunks: Sequence[object]) -> Mapping[str, message_delta: Final = next( (event for event in reversed(_parsed_sse_events(sse_stream)) if event.get("type") == "message_delta"), None ) - delta: Final = _mapping_or_none(message_delta.get("delta")) if message_delta is not None else 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( diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index cea4f88e4e1..8e363ff92f0 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -14,6 +14,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.pass_through.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 @@ -24,14 +25,6 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( from litellm.llms.anthropic.pass_through.stream_assembly import ( is_raw_sse_stream as is_raw_sse_stream, # noqa: PLC0414 # public re-export ) -from litellm.types.llms.anthropic import ( - AnthropicStopDetails, - SafeguardResult, - SafeguardStatusAvailable, - SafeguardStatusUnsupported, - SafeguardToolUseEvaluated, - SafeguardToolUseVerdict, -) from litellm.types.utils import Choices, ModelResponse @@ -62,87 +55,101 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 is_sse_error_stream as is_sse_error_stream, # noqa: PLC0414 # public re-export ) -_STOP_DETAILS: Final = TypeAdapter(AnthropicStopDetails) -_SAFEGUARD_RESULTS: Final = TypeAdapter(tuple[SafeguardResult, ...]) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) -def _stop_details_without_explanation(raw: object) -> AnthropicStopDetails | None: - try: - stop_details: Final = _STOP_DETAILS.validate_python(raw) - except ValidationError: - return None +def _stop_details(raw: object, keep_explanations: bool) -> object: + stop_details: Final = as_str_mapping(raw) + if keep_explanations or stop_details is None or "explanation" not in stop_details: + return raw return {**stop_details, "explanation": None} -def _verdict_without_explanation(verdict: SafeguardToolUseVerdict) -> SafeguardToolUseVerdict: - match verdict: - case {"type": "evaluated"}: - return SafeguardToolUseEvaluated(type="evaluated", outcome=verdict["outcome"]) - case _: - return verdict +def _verdict(raw: object, keep_explanations: bool) -> object: + verdict: Final = as_str_mapping(raw) + if keep_explanations or verdict is None: + return raw + return {key: value for key, value in verdict.items() if key != "explanation"} -def _status_for_delivered_tool_uses( - status: SafeguardStatusAvailable | SafeguardStatusUnsupported, delivered_tool_use_ids: frozenset[str] -) -> SafeguardStatusAvailable | SafeguardStatusUnsupported: - match status: - case {"type": "available"}: - return SafeguardStatusAvailable( - type="available", - tool_uses={ - tool_use_id: _verdict_without_explanation(verdict) - for tool_use_id, verdict in status["tool_uses"].items() - if tool_use_id in delivered_tool_use_ids - }, - ) - case _: - return status +def _status(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: + status: Final = as_str_mapping(raw) + tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None + if status is None or tool_uses is None: + return raw + return { + **status, + "tool_uses": { + tool_use_id: _verdict(verdict, keep_explanations) + for tool_use_id, verdict in tool_uses.items() + if tool_use_id in unchanged_tool_use_ids + }, + } -def _safeguard_results_for_delivered_tool_uses( - raw: object, delivered_tool_use_ids: frozenset[str] -) -> tuple[SafeguardResult, ...] | None: +def _safeguard_result(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: + result: Final = as_str_mapping(raw) + if result is None or "status" not in result: + return raw + return {**result, "status": _status(result["status"], unchanged_tool_use_ids, keep_explanations)} + + +def _safeguard_results(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: try: - results: Final = _SAFEGUARD_RESULTS.validate_python(raw) + results: Final = _ITEMS.validate_python(raw) except ValidationError: - return None - return tuple( - SafeguardResult( - type=result["type"], status=_status_for_delivered_tool_uses(result["status"], delivered_tool_use_ids) - ) - for result in results - ) + return raw + return tuple(_safeguard_result(result, unchanged_tool_use_ids, keep_explanations) for result in results) -def _tool_use_id(block: object) -> str | None: - match block: - case {"type": "tool_use", "id": str() as tool_use_id}: - return tool_use_id - case _: - return None - - -def _rewritten_message_delta_fields( - source_chunks: Sequence[object], delivered_content: Sequence[object] -) -> Mapping[str, object]: - source_fields: Final = anthropic_message_delta_fields(source_chunks) - delivered_tool_use_ids: Final = frozenset( - tool_use_id for block in delivered_content if (tool_use_id := _tool_use_id(block)) is not None - ) - stop_details: Final = _stop_details_without_explanation(source_fields.get("stop_details")) - safeguard_results: Final = _safeguard_results_for_delivered_tool_uses( - source_fields.get("safeguard_results"), delivered_tool_use_ids - ) +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( { - **({"stop_details": stop_details} if stop_details is not None else {}), - **({"safeguard_results": safeguard_results} if safeguard_results is not None else {}), + 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 anthropic_sse_chunks_from_response(assembled: ModelResponse, source_chunks: Sequence[object]) -> tuple[bytes, ...]: - """Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces.""" +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": _stop_details(stop_details, keep_explanations)} if stop_details is not None else {}), + **( + {"safeguard_results": _safeguard_results(safeguard_results, unchanged_tool_use_ids, keep_explanations)} + 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, ) @@ -150,14 +157,22 @@ def anthropic_sse_chunks_from_response(assembled: ModelResponse, source_chunks: 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, anthropic_response.get("content") or () + source_chunks, unchanged_tool_use_ids, keep_explanations ), ).chunks ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index bd5744a934e..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, all_chunks): + 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/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 index 6ec94c1d994..476cea2833c 100644 --- 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 @@ -1,10 +1,12 @@ 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 = [ +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ { "type": "dangerous_tool_use", "status": { @@ -17,20 +19,21 @@ _CLAUDE_CODE_SAFEGUARD_RESULTS = [ ] -def _decode_sse_chunks(chunks: list[bytes]) -> list[tuple[str, dict]]: - decoded = [] - for chunk in chunks: - assert isinstance(chunk, bytes) - lines = chunk.decode().split("\n") - assert lines[0].startswith("event: ") - decoded.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: ")))) - return decoded +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 _message_delta(response: dict) -> dict: - decoded = _decode_sse_chunks(list(FakeAnthropicMessagesStreamIterator(response=response))) - assert [event_type for event_type, _ in decoded][-2:] == ["message_delta", "message_stop"] - deltas = [payload["delta"] for event_type, payload in decoded if event_type == "message_delta"] +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] @@ -42,7 +45,7 @@ def test_message_delta_carries_safeguard_results(): 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 = { + response: Final = { "id": "msg_1", "model": "claude-sonnet-4-6", "role": "assistant", @@ -62,8 +65,8 @@ def test_message_delta_carries_safeguard_results(): def test_message_delta_carries_refusal_stop_details(): - stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} - response = { + stop_details: Final = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response: Final = { "id": "msg_2", "content": [], "stop_reason": "refusal", @@ -79,7 +82,7 @@ def test_message_delta_carries_refusal_stop_details(): def test_message_delta_omits_absent_optional_fields(): - response = { + response: Final = { "id": "msg_3", "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", 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 039835518e6..fab79ea2fce 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 @@ -1557,7 +1557,7 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co assert deferred_fired.is_set() -_CLAUDE_CODE_SAFEGUARD_RESULTS = [ +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ { "type": "dangerous_tool_use", "status": { @@ -1570,9 +1570,9 @@ _CLAUDE_CODE_SAFEGUARD_RESULTS = [ ] -def _sse_events_message_delta(response: dict) -> dict: - decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) - deltas = [payload["delta"] for event_type, payload in decoded if event_type == "message_delta"] +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] 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..f439bd40059 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -1379,6 +1379,51 @@ 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.""" + bash_verdict = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"} + safeguard_results = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"tu_ok": bash_verdict, "tu_bad": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + bash_chunks = self._sse_chunks("Bash", "tu_ok") + read_chunks = self._sse_chunks("Read", "tu_bad") + read_tool_use = [ + chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7] + ] + message_delta = { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 5}, + } + chunks = [ + *bash_chunks[:7], + *read_tool_use, + f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(), + bash_chunks[-1], + ] + + with patch.object(self.rewriting, "should_run_guardrail", return_value=True): + out = await self._drain(self.rewriting, chunks) + + body = b"".join(out).decode() + delta = next( + 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 index e998d36d375..000ecaadcb9 100644 --- a/tests/unit/proxy/guardrails/test_anthropic_sse.py +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -1,4 +1,5 @@ import json +from collections.abc import Callable from itertools import chain from typing import Final @@ -8,20 +9,40 @@ from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, ) -from litellm.types.utils import Choices +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 _decode(chunks: tuple[bytes, ...]) -> list[tuple[str, dict[str, object]]]: - decoded = [] - for chunk in chunks: - lines = chunk.decode().split("\n") - assert lines[0].startswith("event: ") - decoded.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: ")))) - return decoded +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, ...]: @@ -197,3 +218,71 @@ def test_rewritten_stream_keeps_only_the_verdicts_of_tool_uses_it_still_delivers 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})