From d725c576a02abb679cd80898bdff59b08f48bf31 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:43:31 -0700 Subject: [PATCH] fix(bedrock): keep Converse tool history as tool messages until send and match chat on emptied turns When a guardrail removes every tool, the tool calls and results stay tool messages for the guardrails after it, so a scan_only_tool_results guardrail still masks them. The toolConfig with no tool is dropped and the history turned into text only when the request is signed, as chat completions does at request build. A caller's toolConfig with no tool gets the same treatment, where Bedrock returned 400 Request-field strings join the latest user message that has text when a tool-result-only turn follows it, so a guardrail judging only that message reads the caller's text with them A turn the vendor empties entirely is sent as written, as on main and chat completions, instead of a 500 Messages go to the guardrail through a JSON round trip instead of a deep copy, and a reply that changes nothing returns early, so a long Converse conversation guards about as fast as on /v1/messages --- .../llms/bedrock/chat/converse_to_openai.py | 23 ++- .../guardrail_translation/handler.py | 105 +++---------- .../bedrock/passthrough/transformation.py | 78 +++++++++- ...est_bedrock_converse_guardrail_rewrites.py | 59 ++++++-- .../guardrail_translation/test_handler.py | 141 ++++++++++++++++-- 5 files changed, 292 insertions(+), 114 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_to_openai.py b/litellm/llms/bedrock/chat/converse_to_openai.py index 02b94453e54..ce80a2dc8bb 100644 --- a/litellm/llms/bedrock/chat/converse_to_openai.py +++ b/litellm/llms/bedrock/chat/converse_to_openai.py @@ -295,26 +295,37 @@ def _user_row(slots: Sequence[Slot]) -> ConverseRow: return ConverseRow(ChatCompletionUserMessage(role="user", content=parts), tuple(slots)) -def _with_request_fields(rows: tuple[ConverseRow, ...], value: object) -> tuple[ConverseRow, ...]: +def _with_request_fields(rows: tuple[ConverseRow, ...], value: object, user_turn: int) -> tuple[ConverseRow, ...]: """ ``rows`` with the ``additionalModelRequestFields`` strings after the text of the latest user row, so a guardrail - that judges only the latest user message reads both. They get a user row of their own when there is none. + that judges only the latest user message reads both. When no user row has text they get a user row of their own + at ``user_turn``, where the rows of the latest user turn that is not only tool results start. """ slots: Final = _string_leaves(("additionalModelRequestFields",), value) if not slots: return rows latest: Final = next((index for index in reversed(range(len(rows))) if rows[index].message["role"] == "user"), None) if latest is None: - return (*rows, _user_row(slots)) + return (*rows[:user_turn], _user_row(slots), *rows[user_turn:]) return (*rows[:latest], _user_row((*rows[latest].content, *slots)), *rows[latest + 1 :]) +def _is_user_turn(message: _Message) -> bool: + return message.role != "assistant" and any(block.tool_result is None for block in message.content) + + def converse_rows(request: ConverseRequest, include_tool_calls: bool) -> tuple[ConverseRow, ...]: - messages: Final = ( + system: Final = _system_rows(request.system) + messages: Final = tuple( _message_rows(index, message, include_tool_calls) for index, message in enumerate(request.messages) ) - rows: Final = (*_system_rows(request.system), *chain.from_iterable(messages)) - return _with_request_fields(rows, request.additional_model_request_fields) + starts: Final = tuple(accumulate(map(len, messages), initial=len(system))) + rows: Final = (*system, *chain.from_iterable(messages)) + user_turn: Final = next( + (starts[index] for index in reversed(range(len(messages))) if _is_user_turn(request.messages[index])), + len(rows), + ) + return _with_request_fields(rows, request.additional_model_request_fields, user_turn) @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 340de027bfa..a90b4829634 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -105,12 +105,7 @@ class _Remove: holder: _StringHolder -@dataclass(frozen=True, slots=True) -class _DropToolConfig: - pass - - -_Write: TypeAlias = _Assign | _DropToolConfig +_Write: TypeAlias = _Assign _ContentHolder: TypeAlias = tuple[object, str] _BlankDrop: TypeAlias = tuple[_ContentHolder, frozenset[int]] _BODY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[dict[str, object]]) @@ -332,8 +327,8 @@ def _replaced_tools_write(body: Mapping[str, object], edit: ReplacedTools) -> _W """ Rebuild ``toolConfig.tools`` in the order the guardrail returned. Entries it was never sent (cache points, system tools) stay right after the tool they followed, and a cache point goes away with its tool. Refuses when - ``toolChoice`` names a tool that is no longer there, drops ``toolConfig`` when nothing is left, and creates it - when the request had none. + ``toolChoice`` names a tool that is no longer there, and creates ``toolConfig`` when the request had none. A + ``toolConfig`` left with no tool stays until the request is sent, so later guardrails still see the tool history. """ tool_config: Final = _child(body, "toolConfig") tools: Final = () if tool_config is None else as_json_array(_child(tool_config, "tools")) @@ -357,7 +352,7 @@ def _replaced_tools_write(body: Mapping[str, object], edit: ReplacedTools) -> _W return None if tool_config is None: return None if all(map(_is_cache_point, rebuilt)) else _Assign((body, "toolConfig"), {"tools": rebuilt}) - return _DropToolConfig() if all(map(_is_cache_point, rebuilt)) else _Assign((tool_config, "tools"), rebuilt) + return _Assign((tool_config, "tools"), rebuilt) def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None: @@ -372,11 +367,7 @@ def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None: def _kept_write(write: _Write, flat_changes: frozenset[_LeafIdentity]) -> bool: - match write: - case _Assign(): - return _leaf_identity(write.holder) not in flat_changes - case _DropToolConfig(): - return True + return _leaf_identity(write.holder) not in flat_changes def _holder_edits( @@ -397,52 +388,6 @@ def _apply(writes: Sequence[_Assign | _Remove]) -> None: del container[key] -def _tool_result_part(item: object) -> str | None: - text: Final = _child(item, "text") - if isinstance(text, str): - return text - fields: Final = as_json_object(item) - return json.dumps(fields["json"]) if fields is not None and "json" in fields else None - - -def _tool_result_text(tool_result: object) -> str: - parts: Final = as_json_array(_child(tool_result, "content")) or () - rendered: Final = "\n".join(part for item in parts if (part := _tool_result_part(item)) is not None).strip() - return rendered or "" - - -def _neutralized(block: object) -> object: - tool_use: Final = _child(block, "toolUse") - if as_json_object(tool_use) is not None: - arguments: Final = json.dumps(_child(tool_use, "input")) - return {"text": f"[tool call {_child(tool_use, 'toolUseId')}: {_child(tool_use, 'name')}({arguments})]"} - tool_result: Final = _child(block, "toolResult") - if as_json_object(tool_result) is None: - return block - return {"text": f"[tool result for {_child(tool_result, 'toolUseId')}: {_tool_result_text(tool_result)}]"} - - -def _has_tool_block(content: Sequence[object]) -> bool: - return any(_child(block, "toolUse") is not None or _child(block, "toolResult") is not None for block in content) - - -def _tool_config_drop(body: Mapping[str, object]) -> tuple[_Assign | _Remove, ...]: - """Remove ``toolConfig`` and turn every ``toolUse`` and ``toolResult`` block into text.""" - turns: Final = as_json_array(body.get("messages")) or () - neutralized: Final = tuple( - _Assign((message, "content"), [_neutralized(block) for block in content]) - for message in turns - if (content := as_json_array(_child(message, "content"))) is not None and _has_tool_block(content) - ) - return (*neutralized, _Remove((body, "toolConfig"))) - - -def _write_back_structured_edits(body: Mapping[str, object], writes: Sequence[_Write]) -> None: - _apply(tuple(write for write in writes if isinstance(write, _Assign))) - if any(isinstance(write, _DropToolConfig) for write in writes): - _apply(_tool_config_drop(body)) - - def _block_text_holder(block: object) -> _ContentHolder | None: if isinstance(_child(block, "text"), str): return block, "text" @@ -454,7 +399,7 @@ def _final_texts( holders: Sequence[_StringHolder], applied_texts: Sequence[str], writes: Sequence[_Write] ) -> Mapping[_LeafIdentity, object]: flat: Final = {_leaf_identity(holder): text for holder, text in zip(holders, applied_texts)} - structured: Final = {_leaf_identity(write.holder): write.value for write in writes if isinstance(write, _Assign)} + structured: Final = {_leaf_identity(write.holder): write.value for write in writes} return MappingProxyType({**flat, **structured}) @@ -478,16 +423,15 @@ def _empties_turn(drop: _BlankDrop) -> bool: return bool(dropped) and all(id(block) in dropped or _is_cache_point(block) for block in blocks) -def _blank_drops(body: Mapping[str, object], final: Mapping[_LeafIdentity, object]) -> tuple[_BlankDrop, ...] | None: - """The system and message text blocks the guardrail left blank, or None when dropping them empties a turn.""" - turns: Final = tuple( - ((message, "content"), _blank_blocks(final, (message, "content"))) - for message in as_json_array(body.get("messages")) or () - ) - if any(_empties_turn(turn) for turn in turns): - return None +def _turn_drop(final: Mapping[_LeafIdentity, object], message: object) -> _BlankDrop: + return (message, "content"), _blank_blocks(final, (message, "content")) + + +def _blank_drops(body: Mapping[str, object], final: Mapping[_LeafIdentity, object]) -> tuple[_BlankDrop, ...]: + """The system and message text blocks the guardrail left blank, except in a turn they would leave empty.""" + turns: Final = (_turn_drop(final, message) for message in as_json_array(body.get("messages")) or ()) system: Final = ((body, "system"), _blank_blocks(final, (body, "system"))) - return tuple(drop for drop in (system, *turns) if drop[1]) + return tuple(drop for drop in (system, *(turn for turn in turns if not _empties_turn(turn))) if drop[1]) def _blank_drop_write(drop: _BlankDrop) -> _Assign | _Remove: @@ -519,7 +463,7 @@ def _edits_drop_a_text_change( writes: Final = _holder_edits(rewritten, edits, frozenset()) if writes is None: return True - _write_back_structured_edits(rewritten, writes) + _apply(writes) edited: Final = [container[key] for container, key in holders] return any( returned != before and after == before for before, returned, after in zip(texts, guardrailed_texts, edited) @@ -1375,7 +1319,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): sent_messages: Final = [rows[index].message for index in scoped_indices] inputs: Final = GenericGuardrailAPIInputs(texts=list(texts)) if sent_messages: - inputs["structured_messages"] = copy.deepcopy(sent_messages) + inputs["structured_messages"] = json.loads(json.dumps(sent_messages)) tools: Final = converse_tools(request) if request and not scope.scan_only_tool_results else () if tools: inputs["tools"] = copy.deepcopy([tool.param for tool in tools]) @@ -1437,17 +1381,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): request_body, (*row_edits, *tool_changes), _changed_leaves(holders, texts, applied_texts) ) ) - blank_drops: Final = ( - None - if holder_edits is None - else _blank_drops(request_body, _final_texts(holders, applied_texts, holder_edits)) - ) - if holder_edits is None or blank_drops is None: + if holder_edits is None: raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) - if applied_texts: - _write_back_texts(applied_texts, holders) - _write_back_structured_edits(request_body, holder_edits) - _apply(tuple(_blank_drop_write(drop) for drop in blank_drops)) + if holder_edits or (applied_texts and list(applied_texts) != texts): + blank_drops: Final = _blank_drops(request_body, _final_texts(holders, applied_texts, holder_edits)) + if applied_texts: + _write_back_texts(applied_texts, holders) + _apply(holder_edits) + _apply(tuple(_blank_drop_write(drop) for drop in blank_drops)) return data diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index fb0ecd5b825..d7d6341fe7a 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector +from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -34,6 +35,78 @@ def is_bedrock_streaming_endpoint(endpoint: str) -> bool: return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS +_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) + + +def _is_converse_url(url: str) -> bool: + return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS + + +def _is_cache_point(entry: object) -> bool: + fields: Final = as_json_object(entry) + return fields is not None and "cachePoint" in fields + + +def _tool_result_part(item: object) -> str | None: + fields: Final = as_json_object(item) + if fields is None: + return None + text: Final = fields.get("text") + if isinstance(text, str): + return text + return json.dumps(fields["json"]) if "json" in fields else None + + +def _tool_result_text(tool_result: Mapping[str, object]) -> str: + parts: Final = as_json_array(tool_result.get("content")) or () + rendered: Final = "\n".join(part for item in parts if (part := _tool_result_part(item)) is not None).strip() + return rendered or "" + + +def _block_field(block: object, key: str) -> Mapping[str, object] | None: + fields: Final = as_json_object(block) + return None if fields is None else as_json_object(fields.get(key)) + + +def _as_text_block(block: object) -> object: + tool_use: Final = _block_field(block, "toolUse") + if tool_use is not None: + arguments: Final = json.dumps(tool_use.get("input")) + return {"text": f"[tool call {tool_use.get('toolUseId')}: {tool_use.get('name')}({arguments})]"} + tool_result: Final = _block_field(block, "toolResult") + if tool_result is None: + return block + return {"text": f"[tool result for {tool_result.get('toolUseId')}: {_tool_result_text(tool_result)}]"} + + +def _has_tool_block(content: Sequence[object]) -> bool: + return any( + _block_field(block, "toolUse") is not None or _block_field(block, "toolResult") is not None for block in content + ) + + +def _as_text_message(message: object) -> object: + fields: Final = as_json_object(message) + content: Final = None if fields is None else as_json_array(fields.get("content")) + if fields is None or content is None or not _has_tool_block(content): + return message + return {**fields, "content": [_as_text_block(block) for block in content]} + + +def _converse_body_to_send(body: Mapping[str, object]) -> Mapping[str, object]: + """ + ``body`` without a Converse ``toolConfig`` that holds no tool, with every ``toolUse`` and ``toolResult`` turned into + text, since Bedrock rejects tool blocks without a tool. Any other ``body`` comes back as it is. + """ + tool_config: Final = as_json_object(body.get("toolConfig")) + tools: Final = None if tool_config is None else as_json_array(tool_config.get("tools")) + if tools is None or not all(map(_is_cache_point, tools)): + return body + messages: Final = as_json_array(body.get("messages")) + rest: Final = {key: value for key, value in body.items() if key != "toolConfig"} + return rest if messages is None else {**rest, "messages": [_as_text_message(message) for message in messages]} + + _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) @@ -286,16 +359,17 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD self, headers: dict, litellm_params: dict, - request_data: dict | None, + request_data: dict[str, object] | None, api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: + body: Final[Mapping[str, object]] = request_data or {} optional_params: Final = litellm_params.copy() return self._sign_request( service_name="bedrock", headers=headers, optional_params=optional_params, - request_data=request_data or {}, + request_data=dict(_converse_body_to_send(body) if _is_converse_url(api_base) else body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py index 23a2951d1fb..08303236c09 100644 --- a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py +++ b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py @@ -261,6 +261,7 @@ def _guardrails(vendor_url: str) -> tuple[dict[str, JsonValue], ...]: _guardrail(_generic(vendor), {**generic, "api_base": f"{vendor_url}/gga/{vendor.value}"}) for vendor in _Vendor ), + _guardrail("mcp-end-user", {"guardrail": "mcp_end_user_permission"}), _guardrail( "gga-block-skip-tool", {**generic, "api_base": f"{vendor_url}/gga/{_Vendor.BLOCKS.value}", "skip_tool_message_in_guardrail": True}, @@ -274,7 +275,6 @@ def _guardrails(vendor_url: str) -> tuple[dict[str, JsonValue], ...]: }, ), _guardrail("tool-policy", {"guardrail": "tool_policy"}), - _guardrail("mcp-end-user", {"guardrail": "mcp_end_user_permission"}), ) @@ -329,7 +329,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: owned_proxy(gateway, directory, {}, config=_config(directory, guardrails, bedrock), workers=2) as owned, owned.scenario() as scenario, ): - yield _Rig(owned, vendor, bedrock, {name: scenario.key(guardrails=[name]) for name in names}) + keys: Final = {name: scenario.key(guardrails=[name]) for name in names} + tool_remover_then_masker: Final = ["mcp-end-user", "gga-mask-tool-only"] + yield _Rig( + owned, + vendor, + bedrock, + {**keys, "+".join(tool_remover_then_masker): scenario.key(guardrails=tool_remover_then_masker)}, + ) @pytest.fixture(autouse=True) @@ -579,9 +586,10 @@ def test_text_block_the_guardrail_cleared_is_dropped_before_bedrock(rig: _Rig) - assert sent["messages"] == [_user({"text": "Order details below."})] -def test_turn_the_guardrail_cleared_entirely_is_refused_before_bedrock(rig: _Rig) -> None: +def test_turn_the_guardrail_cleared_entirely_is_sent_as_written(rig: _Rig) -> None: body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} - assert _UNAPPLIABLE in _refused(rig, _converse(rig, _generic(_Vendor.CLEARS_CARD_TEXTS), body)) + sent: Final = _forwarded(rig, _converse(rig, _generic(_Vendor.CLEARS_CARD_TEXTS), body)) + assert sent["messages"] == [_user({"text": ""})] def test_deeply_nested_tool_result_json_reaches_bedrock_with_the_message_mask(rig: _Rig) -> None: @@ -644,18 +652,12 @@ def test_mcp_end_user_permission_drops_tools_of_servers_the_end_user_cannot_use( @pytest.mark.parametrize( - ("guardrail_config", "result_block"), - [ - ({}, {"text": "[tool result for tooluse_post: posted to #ops]"}), - ( - {"guardrailConfig": {"guardrailIdentifier": "gr-scripted", "guardrailVersion": "1"}}, - {"guardContent": {"text": {"text": "[tool result for tooluse_post: posted to #ops]"}}}, - ), - ], + "guardrail_config", + [{}, {"guardrailConfig": {"guardrailIdentifier": "gr-scripted", "guardrailVersion": "1"}}], ids=["plain", "with-guardrail-config"], ) def test_mcp_end_user_permission_removing_every_tool_drops_the_tool_config_and_turns_tool_history_into_text( - rig: _Rig, guardrail_config: Mapping[str, JsonValue], result_block: Mapping[str, JsonValue] + rig: _Rig, guardrail_config: Mapping[str, JsonValue] ) -> None: tool_use: Final = {"toolUseId": "tooluse_post", "name": "slack-post_message", "input": {"channel": "#ops"}} body: Final = { @@ -682,6 +684,35 @@ def test_mcp_end_user_permission_removing_every_tool_drops_the_tool_config_and_t "messages": [ _user({"text": "Post the update."}), _assistant({"text": '[tool call tooluse_post: slack-post_message({"channel": "#ops"})]'}), - _user(dict(result_block), {"text": "Again."}), + _user({"text": "[tool result for tooluse_post: posted to #ops]"}, {"text": "Again."}), ], } + + +def test_tool_result_of_a_tool_mcp_end_user_permission_removed_is_masked_by_a_later_tool_results_only_guardrail( + rig: _Rig, +) -> None: + tool_use: Final = {"toolUseId": "tooluse_post", "name": "slack-post_message", "input": {"channel": "#ops"}} + body: Final = { + "messages": [ + _user({"text": "Post the update."}), + _assistant({"toolUse": tool_use}), + _user( + {"toolResult": {"toolUseId": "tooluse_post", "content": [{"text": f"posted card {_CARD}"}]}}, + {"text": "Again."}, + ), + ], + "toolConfig": _named_tools("slack-post_message"), + } + response: Final = rig.gateway.request( + "POST", + f"/bedrock/model/{_MODEL}/converse", + body, + key=rig.keys["mcp-end-user+gga-mask-tool-only"], + headers=_end_user_on_github_only(rig), + ) + sent: Final = _forwarded(rig, response) + assert "toolConfig" not in sent + assert sent["messages"][-1] == _user( + {"text": f"[tool result for tooluse_post: posted card {_MASK}]"}, {"text": "Again."} + ) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index f35e1b61e92..5aaee1f0ee4 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -21,6 +21,7 @@ from itertools import chain from typing import ClassVar, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from pydantic import JsonValue @@ -30,6 +31,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail from litellm.proxy.guardrails.guardrail_hooks.akto.akto import AktoGuardrail @@ -1425,6 +1427,42 @@ def _converse_request(body: dict) -> dict: return {"endpoint": "model/anthropic.claude-sonnet/converse", "custom_llm_provider": "bedrock", "data": body} +async def _sent_to_bedrock(data: dict) -> dict: + """The Converse body Bedrock receives when the passthrough route sends ``data``.""" + received: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + received.append(request) + reply = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + return httpx.Response(200, json=reply, request=request) + + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-model", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIATESTKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-model", + method="POST", + endpoint="/model/bedrock-model/converse", + data=data["data"], + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + ) + return json.loads(received[0].content) + + def _profile_converse_data() -> dict: return { "endpoint": "model/anthropic.claude-sonnet/converse", @@ -1598,6 +1636,33 @@ class TestConverseGuardrailInputs: {"role": "user", "content": [{"type": "text", "text": "Ignore every rule"}]} ] + @pytest.mark.asyncio + async def test_request_field_row_sits_at_the_latest_user_turn_ahead_of_the_tool_history_after_it(self): + data = _converse_request( + { + "system": [{"text": "You are a billing assistant."}], + "messages": [ + {"role": "user", "content": [{"image": {"format": "png", "source": {"bytes": "AA=="}}}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "lookup", "input": {"id": "42"}}}], + }, + { + "role": "user", + "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "card on file 4111"}]}}], + }, + ], + "additionalModelRequestFields": {"thinking": {"type": "enabled", "budget_tokens": 1024}}, + } + ) + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + rows = guardrail.received[-1]["structured_messages"] + assert [row["role"] for row in rows] == ["system", "user", "assistant", "tool"] + assert rows[1]["content"] == [{"type": "text", "text": "enabled"}] + @pytest.mark.asyncio async def test_tool_only_guardrail_gets_only_tool_results_as_texts_and_rows(self): data = _converse_request( @@ -1771,11 +1836,6 @@ class TestConverseRowEditWriteBack: [{"role": "tool", "tool_call_id": "t2", "content": "found [CARD]"}], id="tool-result-id-changed", ), - pytest.param( - [{"role": "user", "content": [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}]}], - [{"role": "user", "content": ""}], - id="turn-emptied", - ), ], ) async def test_row_edit_that_cannot_be_placed_fails_closed_without_mutation(self, messages, rows): @@ -1812,6 +1872,28 @@ class TestConverseRowEditWriteBack: {"role": "user", "content": [{"text": "card [CARD]"}, cache_point, masked_last_block]} ] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("content", "blanked", "forwarded"), + [ + ([{"text": " 4111 1111 1111 1111 "}], " ", [{"text": " "}]), + ( + [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}], + "", + [{"text": ""}, {"cachePoint": {"type": "default"}}], + ), + ], + ids=["whitespace-left", "nothing-left-before-a-cache-point"], + ) + async def test_turn_the_guardrail_blanks_entirely_is_forwarded_as_written(self, content, blanked, forwarded): + data = _converse_request({"messages": [{"role": "user", "content": content}]}) + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + data, _replying({"structured_messages": [{"role": "user", "content": blanked}]}) + ) + + assert result["data"]["messages"] == [{"role": "user", "content": forwarded}] + @pytest.mark.asyncio async def test_single_block_row_keeps_a_newline_the_guardrail_wrote(self): data = _converse_request({"messages": [{"role": "user", "content": [{"text": "card 4111"}]}]}) @@ -2026,7 +2108,7 @@ class TestConverseToolListRewrites: assert result["data"]["toolConfig"] == {"tools": [_REFUND, system_tool], "toolChoice": {"auto": {}}} @pytest.mark.asyncio - async def test_every_tool_removed_drops_the_tool_config_and_turns_tool_history_into_text(self): + async def test_every_tool_removed_reaches_bedrock_without_tool_config_and_with_tool_history_as_text(self): data = _tool_config_data() data["data"]["messages"] = [ { @@ -2055,7 +2137,8 @@ class TestConverseToolListRewrites: result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) - assert result["data"] == { + assert result["data"]["toolConfig"] == {"tools": [], "toolChoice": {"auto": {}}} + assert await _sent_to_bedrock(result) == { "messages": [ { "role": "assistant", @@ -2371,9 +2454,46 @@ class TestConverseGuardContent: result = await handler.process_input_messages( without_tools, masker_class(lambda inputs: {"texts": _masked_texts(inputs)}) ) + sent = await _sent_to_bedrock(result) - assert "toolConfig" not in result["data"] - assert "4111" not in json.dumps(result["data"]) + assert "toolConfig" not in sent + assert "4111" not in json.dumps(sent) + + @pytest.mark.asyncio + async def test_tool_result_of_a_removed_tool_is_masked_by_a_later_tool_results_only_guardrail(self): + data = _converse_request( + { + "messages": [ + {"role": "user", "content": [{"text": "look up my order"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "lookup", "input": {"q": "7"}}}], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "card on file 4111"}]}}, + {"text": "summarize that"}, + ], + }, + ], + "toolConfig": { + "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object"}}}}] + }, + } + ) + handler = BedrockPassthroughGuardrailHandler() + tool_remover = _ScriptedGuardrail(lambda inputs: {"texts": inputs["texts"], "tools": []}) + results_masker = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}, scan_only_tool_results=True) + + without_tools = await handler.process_input_messages(data, tool_remover) + result = await handler.process_input_messages(without_tools, results_masker) + + assert results_masker.received[-1]["texts"] == ["card on file 4111"] + assert (await _sent_to_bedrock(result))["messages"][-1] == { + "role": "user", + "content": [{"text": "[tool result for t1: card on file [CARD]]"}, {"text": "summarize that"}], + } @pytest.mark.asyncio async def test_tool_history_written_as_text_adds_no_guard_content_block(self): @@ -2382,8 +2502,9 @@ class TestConverseGuardContent: tool_remover = _ScriptedGuardrail(lambda inputs: {"texts": inputs["texts"], "tools": []}) result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, tool_remover) + sent = await _sent_to_bedrock(result) - blocks = tuple(chain.from_iterable(message["content"] for message in result["data"]["messages"])) + blocks = tuple(chain.from_iterable(message["content"] for message in sent["messages"])) assert all(block.keys() == {"text"} for block in blocks), ( "a guardContent block makes Bedrock's own guardrail skip every plain text block", blocks,