From 220a2261d7b63b6666f7dcf0222792bb7da02e4b Mon Sep 17 00:00:00 2001 From: mayank Date: Wed, 4 Nov 2026 19:46:14 +0530 Subject: [PATCH 1/3] fix(guardrails): treat an incomplete Model Armor scan as an API error Model Armor can answer HTTP 200 with sanitizationResult.invocationResult set to PARTIAL or FAILURE, meaning some filters did not run. The guardrail read that as a clean scan. Raise ModelArmorAPIError for those results so the hooks apply fail_on_error. A partial scan that already matched a filter is still returned, so the caller blocks it instead of failing open. Fixes #49910 --- .../model_armor/model_armor.py | 26 ++++- .../guardrail_hooks/test_model_armor.py | 110 ++++++++++++++++++ 2 files changed, 134 insertions(+), 2 deletions(-) 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 b9871b1e276..713c2d5a0d8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -108,6 +108,7 @@ class _Filter(TypedDict): @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _SanitizationResult(TypedDict): + invocationResult: NotRequired[ReadOnly[str]] filterResults: NotRequired[ReadOnly[Mapping[str, _Filter] | Sequence[_Filter] | str | int | float | bool | None]] @@ -154,6 +155,9 @@ class ModelArmorAPIError(Exception): self.detail = detail +# invocationResult values that mean Model Armor did not finish scanning, although it answered HTTP 200 +_INCOMPLETE_INVOCATION_RESULTS: Final = frozenset({"PARTIAL", "FAILURE"}) + _SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"}) @@ -419,8 +423,26 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): json_response: Final = _RESPONSE_JSON.validate_python(response.json()) if isinstance(json_response, _AwaitableJSON): - return _ModelArmorResponseRoot.model_validate(await json_response).root - return _ModelArmorResponseRoot.model_validate(json_response).root + armor_response: Final = _ModelArmorResponseRoot.model_validate(await json_response).root + else: + armor_response = _ModelArmorResponseRoot.model_validate(json_response).root + self._raise_if_scan_incomplete(armor_response) + return armor_response + + def _raise_if_scan_incomplete(self, armor_response: _ModelArmorResponse) -> None: + """A 200 whose invocationResult is PARTIAL or FAILURE means some filters did not run. + + Accepting it would treat an unfinished scan as a clean one. It is raised like any other API + failure, so every hook applies fail_on_error. A scan that already matched a filter is still + returned, so the caller blocks instead of failing open on content that was flagged. + """ + result: Final = armor_response.get("sanitizationResult", {}).get("invocationResult") + if result not in _INCOMPLETE_INVOCATION_RESULTS: + return + if self._should_block_content(armor_response): + return + verbose_proxy_logger.error("Model Armor scan did not complete - invocationResult: %s", result) + raise ModelArmorAPIError(f"Model Armor scan incomplete (invocationResult={result})") def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict: """ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..f00f310241d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -658,6 +658,116 @@ async def test_model_armor_api_failure_raises_sanitized_error(): assert "Internal Server Error" not in str(exc_info.value.detail) +def _armor_guardrail(**kwargs) -> ModelArmorGuardrail: + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + **kwargs, + ) + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + guardrail.should_run_guardrail = Mock(return_value=True) + return guardrail + + +def _armor_http_response(sanitization_result: dict) -> Mock: + response = Mock() + response.status_code = 200 + response.text = json.dumps({"sanitizationResult": sanitization_result}) + response.json = Mock(return_value={"sanitizationResult": sanitization_result}) + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invocation_result", ["PARTIAL", "FAILURE"]) +async def test_model_armor_incomplete_scan_raises_api_error(invocation_result: str): + """A 200 whose invocationResult is PARTIAL or FAILURE is an unfinished scan, not a clean one.""" + guardrail = _armor_guardrail() + response = _armor_http_response( + {"invocationResult": invocation_result, "filterMatchState": "NO_MATCH_FOUND", "filterResults": {}} + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content="test content", source="user_prompt") + + assert exc_info.value.detail == f"Model Armor scan incomplete (invocationResult={invocation_result})" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sanitization_result", + [ + {"invocationResult": "SUCCESS", "filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}, + {"invocationResult": "INVOCATION_RESULT_UNSPECIFIED", "filterResults": {}}, + {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}, + ], +) +async def test_model_armor_complete_or_unreported_scan_is_accepted(sanitization_result: dict): + """SUCCESS, an unspecified result and a response without the field are all treated as before.""" + guardrail = _armor_guardrail() + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=_armor_http_response(sanitization_result)) + ): + result = await guardrail.make_model_armor_request(content="test content", source="user_prompt") + + assert result["sanitizationResult"] == sanitization_result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_incomplete_scan_honors_fail_on_error(fail_on_error: bool): + """An unfinished scan blocks when fail_on_error is true and passes through when fail-open is configured.""" + guardrail = _armor_guardrail(fail_on_error=fail_on_error) + response = _armor_http_response({"invocationResult": "PARTIAL", "filterResults": {}}) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "synthetic input"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + if fail_on_error: + with pytest.raises(ModelArmorAPIError): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + else: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + assert result is request_data + + +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_match_still_blocks_when_fail_open(): + """A partial scan that already found a match must block, even with fail_on_error disabled.""" + guardrail = _armor_guardrail(fail_on_error=False) + response = _armor_http_response( + { + "invocationResult": "PARTIAL", + "filterMatchState": "MATCH_FOUND", + "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}, + } + ) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "synthetic input"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio @pytest.mark.parametrize("sanitize", [True, False]) async def test_model_armor_error_output_sanitization(sanitize: bool): From 0fa346d761d85eb49170fa2f31faf430f443ff1f Mon Sep 17 00:00:00 2001 From: mayank Date: Wed, 4 Nov 2026 20:29:03 +0530 Subject: [PATCH 2/3] fix(guardrails): do not let a masked match hide an incomplete Model Armor scan A de-identification match is masked rather than blocked when mask_request_content is on, so counting it as a block let a partial scan through while the other filters had not run. Only a match that blocks even with masking allowed now returns the response instead of raising. --- .../model_armor/model_armor.py | 6 +++-- .../guardrail_hooks/test_model_armor.py | 25 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) 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 713c2d5a0d8..b005c068856 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -434,12 +434,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Accepting it would treat an unfinished scan as a clean one. It is raised like any other API failure, so every hook applies fail_on_error. A scan that already matched a filter is still - returned, so the caller blocks instead of failing open on content that was flagged. + returned, so the caller blocks instead of failing open on content that was flagged. Only a + match that blocks even when masking is enabled counts: a de-identification match is masked + rather than blocked, so it must not hide the filters that did not run. """ result: Final = armor_response.get("sanitizationResult", {}).get("invocationResult") if result not in _INCOMPLETE_INVOCATION_RESULTS: return - if self._should_block_content(armor_response): + if self._should_block_content(armor_response, allow_sanitization=True): return verbose_proxy_logger.error("Model Armor scan did not complete - invocationResult: %s", result) raise ModelArmorAPIError(f"Model Armor scan incomplete (invocationResult={result})") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index f00f310241d..da3f674cacd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -768,6 +768,31 @@ async def test_model_armor_partial_scan_with_match_still_blocks_when_fail_open() assert exc_info.value.status_code == 400 +@pytest.mark.asyncio +@pytest.mark.parametrize("mask_request_content", [True, False]) +async def test_model_armor_partial_scan_with_only_a_deidentify_match_is_an_error(mask_request_content: bool): + """A de-identification match is masked, not blocked, so it must not hide the filters that did not run.""" + guardrail = _armor_guardrail(mask_request_content=mask_request_content) + response = _armor_http_response( + { + "invocationResult": "PARTIAL", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}} + } + } + }, + } + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content="call me at 555-0100", source="user_prompt") + + assert exc_info.value.detail == "Model Armor scan incomplete (invocationResult=PARTIAL)" + + @pytest.mark.asyncio @pytest.mark.parametrize("sanitize", [True, False]) async def test_model_armor_error_output_sanitization(sanitize: bool): From b0c8f4836e7a1a2049e001ce271a7d31040c80d2 Mon Sep 17 00:00:00 2001 From: mayank Date: Wed, 4 Nov 2026 22:59:09 +0530 Subject: [PATCH 3/3] fix(guardrails): judge an incomplete Model Armor scan by what the caller does with the match A de-identification match is blocked for files and for text with masking off, and masked only for text with masking on. The incomplete-scan check assumed masking was always on, so with fail_on_error=False a blocking match became an error that the caller swallowed. Pass the caller's masking setting into the check. --- .../model_armor/model_armor.py | 20 ++++-- .../guardrail_hooks/test_model_armor.py | 65 ++++++++++++++----- 2 files changed, 63 insertions(+), 22 deletions(-) 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 b005c068856..ae7611975ea 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -426,22 +426,30 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): armor_response: Final = _ModelArmorResponseRoot.model_validate(await json_response).root else: armor_response = _ModelArmorResponseRoot.model_validate(json_response).root - self._raise_if_scan_incomplete(armor_response) + # What the caller can do with a de-identification match: files and text with masking off are + # blocked on it, only masked text is rewritten and let through. + if file_bytes is not None: + masks_match = False + elif source == "user_prompt": + masks_match = self.mask_request_content + else: + masks_match = self.mask_response_content + self._raise_if_scan_incomplete(armor_response, allow_sanitization=masks_match) return armor_response - def _raise_if_scan_incomplete(self, armor_response: _ModelArmorResponse) -> None: + def _raise_if_scan_incomplete(self, armor_response: _ModelArmorResponse, allow_sanitization: bool) -> None: """A 200 whose invocationResult is PARTIAL or FAILURE means some filters did not run. Accepting it would treat an unfinished scan as a clean one. It is raised like any other API failure, so every hook applies fail_on_error. A scan that already matched a filter is still - returned, so the caller blocks instead of failing open on content that was flagged. Only a - match that blocks even when masking is enabled counts: a de-identification match is masked - rather than blocked, so it must not hide the filters that did not run. + returned, so the caller blocks instead of failing open on content that was flagged. A + de-identification match counts as one only where the caller blocks on it (files, or text with + masking off); where it is masked instead, it must not hide the filters that did not run. """ result: Final = armor_response.get("sanitizationResult", {}).get("invocationResult") if result not in _INCOMPLETE_INVOCATION_RESULTS: return - if self._should_block_content(armor_response, allow_sanitization=True): + if self._should_block_content(armor_response, allow_sanitization=allow_sanitization): return verbose_proxy_logger.error("Model Armor scan did not complete - invocationResult: %s", result) raise ModelArmorAPIError(f"Model Armor scan incomplete (invocationResult={result})") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index da3f674cacd..af0c2d4f883 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -768,23 +768,23 @@ async def test_model_armor_partial_scan_with_match_still_blocks_when_fail_open() assert exc_info.value.status_code == 400 -@pytest.mark.asyncio -@pytest.mark.parametrize("mask_request_content", [True, False]) -async def test_model_armor_partial_scan_with_only_a_deidentify_match_is_an_error(mask_request_content: bool): - """A de-identification match is masked, not blocked, so it must not hide the filters that did not run.""" - guardrail = _armor_guardrail(mask_request_content=mask_request_content) - response = _armor_http_response( - { - "invocationResult": "PARTIAL", - "filterResults": { - "sdp": { - "sdpFilterResult": { - "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}} - } - } - }, +_PARTIAL_DEIDENTIFY_ONLY = { + "invocationResult": "PARTIAL", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}} + } } - ) + }, +} + + +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_masked_deidentify_match_is_an_error(): + """Where a de-identification match is masked, not blocked, it must not hide the filters that did not run.""" + guardrail = _armor_guardrail(mask_request_content=True) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): with pytest.raises(ModelArmorAPIError) as exc_info: @@ -793,6 +793,39 @@ async def test_model_armor_partial_scan_with_only_a_deidentify_match_is_an_error assert exc_info.value.detail == "Model Armor scan incomplete (invocationResult=PARTIAL)" +@pytest.mark.asyncio +async def test_model_armor_partial_scan_with_blocking_deidentify_match_still_blocks_when_fail_open(): + """With masking off the caller blocks on a de-identification match, so an incomplete scan must not turn it + into an error that fail_on_error=False then swallows.""" + guardrail = _armor_guardrail(fail_on_error=False, mask_request_content=False) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) + request_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "call me at 555-0100"}], "metadata": {}} + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_model_armor_partial_file_scan_with_deidentify_match_is_returned_even_with_masking_on(): + """Files cannot be masked, so the caller blocks on the match whatever the masking setting is.""" + guardrail = _armor_guardrail(mask_request_content=True) + response = _armor_http_response(_PARTIAL_DEIDENTIFY_ONLY) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=response)): + result = await guardrail.make_model_armor_request(file_bytes=b"call me at 555-0100", file_type="PLAINTEXT_UTF8") + + assert result["sanitizationResult"]["invocationResult"] == "PARTIAL" + assert guardrail._should_block_content(result, allow_sanitization=False) + + @pytest.mark.asyncio @pytest.mark.parametrize("sanitize", [True, False]) async def test_model_armor_error_output_sanitization(sanitize: bool):