From 87ceb3486b3ccda7a748e96081f86354540b13ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:50:46 +0000 Subject: [PATCH] fix(proxy): close guardrail bypass on bedrock invoke passthrough routes Pre-call extraction and post-call output processing only handled Converse shapes, so /bedrock/model/{modelId}/invoke and invoke-with-response-stream returned unguarded. An authenticated caller could move blocked content into an InvokeModel payload and skip the key/team guardrail entirely. Non-Converse Bedrock routes now fall back to the generic passthrough handler, which scans the full request and response payloads so blocking guardrails still run, matching how other passthrough providers are guarded. --- .../guardrail_translation/handler.py | 35 +++++++++++--- .../guardrail_translation/test_handler.py | 48 +++++++++++++++++-- 2 files changed, 74 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 021b39927b..f86bd9a7a3 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -22,6 +22,19 @@ def _is_converse_endpoint(endpoint: str) -> bool: return bool(parts) and parts[-1] in _CONVERSE_ACTIONS +def _generic_passthrough_handler() -> BaseTranslation: + """ + Fallback for non-Converse Bedrock routes (e.g. invoke). The generic + handler scans the full request/response payload so blocking guardrails + still run, matching how other passthrough providers are guarded. + """ + from litellm.llms.pass_through.guardrail_translation.handler import ( + PassThroughEndpointHandler, + ) + + return PassThroughEndpointHandler() + + def _extract_converse_texts( body: dict, skip_system: bool, @@ -226,14 +239,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): endpoint = data.get("endpoint", "") body = data.get("data") - if not _is_converse_endpoint(endpoint) or not isinstance(body, dict): - verbose_proxy_logger.debug( - "BedrockPassthroughGuardrailHandler: skipping non-converse endpoint %s", - endpoint, + if not _is_converse_endpoint(endpoint): + return await _generic_passthrough_handler().process_input_messages( + data=data, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, ) - return data - if not isinstance(body.get("messages"), list): + if not isinstance(body, dict) or not isinstance(body.get("messages"), list): return data skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply) @@ -270,6 +283,16 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): user_api_key_dict: Optional[Any] = None, request_data: Optional[dict] = None, ) -> Any: + endpoint = (request_data or {}).get("endpoint", "") + if endpoint and not _is_converse_endpoint(endpoint): + return await _generic_passthrough_handler().process_output_response( + response=response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if not isinstance(response, dict): return response diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 447a8af979..51c16ac8fd 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -274,14 +274,38 @@ class TestBedrockPassthroughGuardrailHandlerInput: assert tool_result["content"][0]["text"] == "[REDACTED]" @pytest.mark.asyncio - async def test_non_converse_endpoint_skips_apply_guardrail(self): + async def test_non_converse_endpoint_scans_full_payload(self): + """Invoke routes must not bypass guardrails: the full request payload is + scanned so blocking guardrails still see user-controlled text.""" handler = BedrockPassthroughGuardrailHandler() - data = _converse_data(endpoint="model/anthropic.claude-3-sonnet/invoke") + data = { + "endpoint": "model/anthropic.claude-3-sonnet/invoke", + "custom_llm_provider": "bedrock", + "model": "anthropic.claude-3-sonnet", + "data": {"messages": [{"role": "user", "content": "blocked invoke text"}]}, + } guardrail = _make_guardrail({"texts": []}) await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - guardrail.apply_guardrail.assert_not_called() + guardrail.apply_guardrail.assert_called_once() + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "blocked invoke text" in sent_texts[0] + + @pytest.mark.asyncio + async def test_non_converse_endpoint_blocking_propagates(self): + handler = BedrockPassthroughGuardrailHandler() + data = { + "endpoint": "model/anthropic.claude-3-sonnet/invoke-with-response-stream", + "custom_llm_provider": "bedrock", + "data": {"prompt": "blocked"}, + } + guardrail = MagicMock() + guardrail.guardrail_name = "block-guard" + guardrail.apply_guardrail = AsyncMock(side_effect=HTTPException(status_code=400, detail="Blocked")) + + with pytest.raises(HTTPException): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) @pytest.mark.asyncio async def test_missing_messages_field_skips(self): @@ -366,6 +390,24 @@ class TestBedrockPassthroughGuardrailHandlerOutput: assert result == {"stopReason": "end_turn"} guardrail.apply_guardrail.assert_not_called() + @pytest.mark.asyncio + async def test_non_converse_response_scanned_via_generic_handler(self): + """Invoke responses are not Converse-shaped; they must still be scanned + through the generic passthrough handler rather than skipped.""" + handler = BedrockPassthroughGuardrailHandler() + response = {"completion": "blocked model output"} + guardrail = _make_guardrail({"texts": []}) + + await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data={"endpoint": "model/anthropic.claude-3-sonnet/invoke"}, + ) + + guardrail.apply_guardrail.assert_called_once() + sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] + assert "blocked model output" in sent_texts[0] + def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: import json