From ec735074a28502e9235773dfb3c835857233c601 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 22 Apr 2026 23:00:32 +0300 Subject: [PATCH 1/3] fix(proxy): reapply Bedrock guardrail spend logging (#25854) Restore guardrail spend/UI event_type wiring, request_data on streaming OUTPUT paths, and centralized match redaction after the upstream revert. Made-with: Cursor --- litellm/integrations/custom_guardrail.py | 12 + litellm/litellm_core_utils/core_helpers.py | 36 +++ .../guardrail_hooks/bedrock_guardrails.py | 140 +++++---- litellm/proxy/utils.py | 12 +- .../integrations/test_custom_guardrail.py | 47 +++ .../litellm_core_utils/test_core_helpers.py | 35 +++ .../test_bedrock_guardrails.py | 273 +++++++++++++++++- 7 files changed, 486 insertions(+), 69 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index abf010e0d6..b1bf3483a9 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -2,6 +2,7 @@ from datetime import datetime from typing import ( TYPE_CHECKING, Any, + ClassVar, Dict, List, Literal, @@ -12,6 +13,7 @@ from typing import ( ) from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.types.guardrails import ( @@ -81,6 +83,9 @@ class ModifyResponseException(Exception): class CustomGuardrail(CustomLogger): + # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. + use_native_during_call_hook: ClassVar[bool] = False + def __init__( self, guardrail_name: Optional[str] = None, @@ -637,6 +642,13 @@ class CustomGuardrail(CustomLogger): if isinstance(item, dict): item.pop("secret_fields", None) + # Default-safe behavior: never persist raw matched spans in standard + # guardrail logging payloads (single shared implementation; Bedrock hooks pass + # raw provider JSON so redaction is not duplicated upstream). + clean_guardrail_response = redact_nested_match_and_regex_keys( + clean_guardrail_response + ) + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 22006be21a..07239a6886 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,5 +1,6 @@ # What is this? ## Helper utilities +import copy from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union import httpx @@ -435,3 +436,38 @@ def filter_internal_params( # Filter out internal parameters return {k: v for k, v in data.items() if k not in internal_params} + + +def redact_nested_match_and_regex_keys( + payload: Union[dict, List[Any], str, None], +) -> Union[dict, List[Any], str, None]: + """ + Deep-copy `payload` and replace every `match` / `regex` string field with + "[REDACTED]" anywhere in nested dict/list structures. + + Used for guardrail spend/compliance logging so raw spans are not persisted. + """ + if payload is None or isinstance(payload, str): + return payload + try: + redacted: Union[dict, List[Any], str, None] = copy.deepcopy(payload) + except Exception: + return payload + + def _walk(node: Any) -> None: + if isinstance(node, dict): + if "match" in node: + node["match"] = "[REDACTED]" + if "regex" in node: + node["regex"] = "[REDACTED]" + for value in node.values(): + _walk(value) + elif isinstance(node, list): + for item in node: + _walk(item) + + try: + _walk(redacted) + except Exception: + return payload + return redacted diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 77b2f466f2..8bfe5027b7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -5,7 +5,6 @@ # +-------------------------------------------------------------+ # Thank you users! We ❤️ you! - Krrish & Ishaan -import copy import os import sys @@ -18,6 +17,7 @@ from typing import ( TYPE_CHECKING, Any, AsyncGenerator, + ClassVar, Dict, List, Literal, @@ -33,6 +33,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.caching import DualCache from litellm.exceptions import GuardrailInterventionNormalStringError from litellm.integrations.custom_guardrail import CustomGuardrail @@ -79,56 +80,33 @@ class GuardrailMessageFilterResult(NamedTuple): def _redact_pii_matches(response_json: dict) -> dict: - try: - # Create a deep copy to avoid modifying the original response - redacted_response = copy.deepcopy(response_json) + """ + Redact match-like fields from a Bedrock ApplyGuardrail JSON payload. - # Get assessments from the response - # NOTE: We use `.get("key") or []` instead of `.get("key", [])` because - # the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null). - # In Python, dict.get("key", []) returns None (not []) when the key exists - # with a None/null value. The `or []` ensures we always get an iterable, - # preventing "TypeError: 'NoneType' object is not iterable". - assessments = redacted_response.get("assessments") or [] - if not assessments: - return redacted_response + Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend + logging). Kept as a Bedrock-module entry point for existing unit tests. + """ + redacted = redact_nested_match_and_regex_keys(response_json) + return redacted if isinstance(redacted, dict) else response_json - for assessment in assessments: - # Redact PII entities in sensitive information policy - sensitive_info_policy = assessment.get("sensitiveInformationPolicy") - if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities") or [] - for pii_entity in pii_entities: - if "match" in pii_entity: - pii_entity["match"] = "[REDACTED]" - # Redact regex matches - regexes = sensitive_info_policy.get("regexes") or [] - for regex_match in regexes: - if "match" in regex_match: - regex_match["match"] = "[REDACTED]" +def _redact_assessment_match_fields(assessments: List[dict]) -> List[dict]: + """ + Redact sensitive match-like fields from blocked assessment summaries. - # Redact custom word matches in word policy - word_policy = assessment.get("wordPolicy") - if word_policy: - custom_words = word_policy.get("customWords") or [] - for custom_word in custom_words: - if "match" in custom_word: - custom_word["match"] = "[REDACTED]" - - managed_words = word_policy.get("managedWordLists") or [] - for managed_word in managed_words: - if "match" in managed_word: - managed_word["match"] = "[REDACTED]" - - return redacted_response - except Exception as e: - # We do not want to fail in any case so this is just a warning - verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e)) - return response_json + This is used for customer-visible error payloads (HTTPException.detail) where + we want to preserve policy/type/action metadata without echoing raw matched + content. + """ + redacted = redact_nested_match_and_regex_keys(assessments) + return redacted if isinstance(redacted, list) else assessments class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): + # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise + # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL. + use_native_during_call_hook: ClassVar[bool] = True + def __init__( self, guardrailIdentifier: Optional[str] = None, @@ -419,6 +397,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): messages: Optional[List[AllMessageValues]] = None, response: Optional[Union[Any, litellm.ModelResponse]] = None, request_data: Optional[dict] = None, + logging_event_type: Optional[GuardrailEventHooks] = None, ) -> BedrockGuardrailResponse: from datetime import datetime @@ -456,11 +435,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - event_type = ( - GuardrailEventHooks.pre_call - if source == "INPUT" - else GuardrailEventHooks.post_call - ) + # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API + # body, which must not be confused with the proxy hook (pre_call / during_call / + # post_call). When omitted, keep legacy mapping for backward compatibility. + if logging_event_type is not None: + event_type = logging_event_type + else: + event_type = ( + GuardrailEventHooks.pre_call + if source == "INPUT" + else GuardrailEventHooks.post_call + ) try: httpx_response = await self.async_handler.post( @@ -515,9 +500,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### # Add guardrail information to request trace ######################################################### + _json_response = httpx_response.json() + # Raw Bedrock JSON is passed here; match/regex redaction runs once inside + # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, - guardrail_json_response=httpx_response.json(), + guardrail_json_response=_json_response, request_data=request_data or {}, guardrail_status=self._get_bedrock_guardrail_response_status( response=httpx_response @@ -530,9 +518,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### if httpx_response.status_code == 200: # check if the response was flagged - _json_response = httpx_response.json() - redacted_response = _redact_pii_matches(_json_response) - verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response) + verbose_proxy_logger.debug( + "Bedrock AI response : %s", + redact_nested_match_and_regex_keys(_json_response), + ) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception( bedrock_guardrail_response @@ -809,7 +798,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): assessments = self._extract_blocked_assessments(response) if assessments: - detail["assessments"] = assessments + detail["assessments"] = _redact_assessment_match_fields(assessments) return HTTPException(status_code=400, detail=detail) @@ -831,8 +820,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return False # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED) - # NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API. - # See _redact_pii_matches() for detailed explanation of the null safety pattern. + # NOTE: Use `.get("k") or []` not `.get("k", [])` — Bedrock can return explicit + # JSON null; dict.get("k", []) then yields None, and `for x in None` raises. assessments = response.get("assessments") or [] if not assessments: return False @@ -952,7 +941,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) try: bedrock_guardrail_response = await self.make_bedrock_api_request( - source="INPUT", messages=filtered_messages, request_data=data + source="INPUT", + messages=filtered_messages, + request_data=data, + logging_event_type=GuardrailEventHooks.pre_call, ) except GuardrailInterventionNormalStringError as e: bedrock_guardrail_response = e.message @@ -1024,7 +1016,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) try: bedrock_guardrail_response = await self.make_bedrock_api_request( - source="INPUT", messages=filtered_messages, request_data=data + source="INPUT", + messages=filtered_messages, + request_data=data, + logging_event_type=GuardrailEventHooks.during_call, ) except GuardrailInterventionNormalStringError as e: bedrock_guardrail_response = e.message @@ -1128,9 +1123,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="INPUT", messages=input_messages, request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) output_task = self.make_bedrock_api_request( - source="OUTPUT", response=response, request_data=data + source="OUTPUT", + response=response, + request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) # Execute both requests in parallel @@ -1144,7 +1143,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) try: output_content_bedrock = await self.make_bedrock_api_request( - source="OUTPUT", response=response, request_data=data + source="OUTPUT", + response=response, + request_data=data, + logging_event_type=GuardrailEventHooks.post_call, ) except GuardrailInterventionNormalStringError as e: output_content_bedrock = e.message @@ -1271,9 +1273,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="INPUT", messages=input_messages, request_data=request_data, + logging_event_type=GuardrailEventHooks.post_call, ) # Only input messages output_task = self.make_bedrock_api_request( - source="OUTPUT", response=assembled_model_response + source="OUTPUT", + response=assembled_model_response, + request_data=request_data, + logging_event_type=GuardrailEventHooks.post_call, ) # Only response # Execute both requests in parallel @@ -1287,7 +1293,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) try: output_guardrail_response = await self.make_bedrock_api_request( - source="OUTPUT", response=assembled_model_response + source="OUTPUT", + response=assembled_model_response, + request_data=request_data, + logging_event_type=GuardrailEventHooks.post_call, ) except GuardrailInterventionNormalStringError as e: output_guardrail_response = e.message @@ -1564,6 +1573,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Bedrock will throw an error if there is no text to process if filtered_messages: + _log_hook = ( + GuardrailEventHooks.pre_call + if input_type == "request" + else GuardrailEventHooks.post_call + ) # Map the abstract input_type to the Bedrock source parameter. # "request" -> INPUT (scan user-supplied content) # "response" -> OUTPUT (scan model-generated content) @@ -1594,12 +1608,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source="OUTPUT", response=synthetic_response, request_data=request_data, + logging_event_type=_log_hook, ) else: bedrock_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=request_data, + logging_event_type=_log_hook, ) # Apply any masking that was applied by the guardrail diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f21a729f55..561f8e5c55 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -940,7 +940,11 @@ class ProxyLogging: Result from the guardrail execution """ # Use unified_guardrail if callback has apply_guardrail method - use_unified = "apply_guardrail" in type(callback).__dict__ + has_apply_guardrail = "apply_guardrail" in type(callback).__dict__ + use_unified = has_apply_guardrail and not ( + hook_type == "during_call" + and getattr(callback, "use_native_during_call_hook", False) + ) if use_unified: data["guardrail_to_apply"] = callback @@ -1537,10 +1541,12 @@ class ProxyLogging: else: user_api_key_auth_dict = user_api_key_dict # Add task to list for parallel execution - if ( + use_unified_during = ( "apply_guardrail" in type(callback).__dict__ and user_api_key_dict is not None - ): + and not getattr(callback, "use_native_during_call_hook", False) + ) + if use_unified_during: data["guardrail_to_apply"] = callback guardrail_task = self._run_guardrail_task_with_enrichment( callback, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 0c904e9df5..d09c4ac2c3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1055,3 +1055,50 @@ class TestTracingFieldsPopulation: assert slg["classification"] == classification assert slg["detection_method"] == "llm-judge" assert slg["confidence_score"] == 0.94 + + +class TestCustomGuardrailSpendLogMatchRedaction: + """Guardrail JSON persisted via standard_logging must not contain raw match spans.""" + + def test_add_standard_logging_redacts_nested_match(self): + cg = CustomGuardrail(guardrail_name="test-rail") + raw = { + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + {"type": "NAME", "match": "GG", "action": "BLOCKED"} + ] + } + } + ] + } + request_data: dict = {"metadata": {}} + cg.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=raw, + request_data=request_data, + guardrail_status="guardrail_intervened", + ) + slg = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert ( + slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "[REDACTED]" + ) + assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ + "match" + ] == "GG" + + def test_add_standard_logging_redacts_regex_field(self): + cg = CustomGuardrail(guardrail_name="test-rail") + raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]} + request_data: dict = {"metadata": {}} + cg.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=raw, + request_data=request_data, + guardrail_status="success", + ) + slg = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]" + assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}" diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 6f95a8b603..aa5ce5fa6a 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -6,6 +6,7 @@ from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, map_finish_reason, reconstruct_model_name, + redact_nested_match_and_regex_keys, ) @@ -158,3 +159,37 @@ class TestFinishReasonMapOutputsAreValid: f"Mapped value '{openai_reason}' (from '{provider_reason}') " f"is not a valid OpenAI finish reason" ) + + +class TestRedactNestedMatchAndRegexKeys: + def test_redacts_match_and_regex_recursively(self): + payload = { + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + {"type": "NAME", "match": "secret-name", "action": "BLOCKED"} + ] + }, + "wordPolicy": { + "customWords": [{"match": "badword", "action": "BLOCKED"}] + }, + } + ], + "regex": "should-redact-key-named-regex", + } + out = redact_nested_match_and_regex_keys(payload) + assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ + "match" + ] == "[REDACTED]" + assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ( + "[REDACTED]" + ) + assert out["regex"] == "[REDACTED]" + assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][ + 0 + ]["match"] == "secret-name" + + def test_passes_through_none_and_str(self): + assert redact_nested_match_and_regex_keys(None) is None + assert redact_nested_match_and_regex_keys("plain") == "plain" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 7d454eb6fe..fef984d704 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -12,11 +12,15 @@ from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../../../..")) +import litellm +from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, _redact_pii_matches, ) +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ModelResponse @@ -106,10 +110,12 @@ async def test__redact_pii_matches_malformed_response(): # Test with completely malformed response malformed_response = { "action": "GUARDRAIL_INTERVENED", - "assessments": "not_a_list", # This should cause an exception + # Wrong type for assessments; redact_nested_match_and_regex_keys walks dict + # values and skips non-dict/list nodes, so this must not raise. + "assessments": "not_a_list", } - # Should not crash and return original response + # Should not crash (deep copy + walk skips the string value under assessments) redacted_response = _redact_pii_matches(malformed_response) assert redacted_response == malformed_response @@ -188,7 +194,7 @@ async def test__redact_pii_matches_multiple_assessments(): @pytest.mark.asyncio async def test_bedrock_guardrail_logging_uses_redacted_response(): - """Test that the Bedrock guardrail uses redacted response for logging""" + """Debug logs and standard_logging payloads must not include raw match values.""" # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() @@ -295,6 +301,14 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): == "PHONE" ) + slg_list = request_data["metadata"]["standard_logging_guardrail_information"] + assert ( + slg_list[0]["guardrail_response"]["assessments"][0][ + "sensitiveInformationPolicy" + ]["piiEntities"][0]["match"] + == "[REDACTED]" + ) + print("Bedrock guardrail logging redaction test passed") @@ -1751,6 +1765,124 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): print("\u2705 BLOCKED vs ANONYMIZED actions test passed") +# --------------------------------------------------------------------------- +# Spend logs: guardrail_mode (pre/during/post) vs Bedrock INPUT/OUTPUT +# --------------------------------------------------------------------------- + + +def test_bedrock_guardrail_uses_native_during_call_hook(): + """during_call must use async_moderation_hook, not unified apply_guardrail(input=request).""" + assert BedrockGuardrail.use_native_during_call_hook is True + + +@pytest.mark.asyncio +async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): + """ + Spend/UI use event_type from the proxy hook, not Bedrock's INPUT/OUTPUT alone. + When logging_event_type is set, it must be forwarded to standard guardrail logging. + When omitted, INPUT maps to pre_call (legacy). + """ + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + {"type": "NAME", "match": "GG", "action": "BLOCKED"} + ] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + } + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log: + mock_post.return_value = mock_bedrock_response + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + logging_event_type=GuardrailEventHooks.during_call, + ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call + # Raw Bedrock JSON is forwarded; redaction runs once in + # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. + assert ( + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ + "sensitiveInformationPolicy" + ]["piiEntities"][0]["match"] + == "GG" + ) + + mock_log.reset_mock() + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.pre_call + + +@pytest.mark.asyncio +async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): + """ + Bedrock sets use_native_during_call_hook so ProxyLogging runs the real + async_moderation_hook (unified apply_guardrail would log INPUT as pre_call). + """ + cache = DualCache() + proxy_logging = ProxyLogging(user_api_key_cache=cache) + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-during-test", + guardrailIdentifier="gid", + guardrailVersion="1", + event_hook=GuardrailEventHooks.during_call, + default_on=True, + ) + mock_mod = AsyncMock(return_value=None) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + try: + litellm.callbacks = [guardrail] + with patch.object(guardrail, "async_moderation_hook", new=mock_mod): + await proxy_logging.during_call_hook( + data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "test"}], + }, + user_api_key_dict=UserAPIKeyAuth( + api_key="test_key", user_id="test_user" + ), + call_type="completion", + ) + finally: + litellm.callbacks = original_callbacks + + mock_mod.assert_awaited_once() + + # --------------------------------------------------------------------------- # L3: _extract_blocked_assessments + _get_http_exception_for_blocked_guardrail # Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error. @@ -1766,7 +1898,7 @@ def _make_guardrail() -> BedrockGuardrail: def test_extract_blocked_assessments_pii_entity(): - """L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term.""" + """L3: PII entity match (BLOCKED) is surfaced with category, type, and match.""" g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", @@ -1877,6 +2009,7 @@ def test_get_http_exception_includes_assessments_and_identifier(): assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" assert exc.detail["assessments"][0]["matches"][0]["type"] == "NAME" + assert exc.detail["assessments"][0]["matches"][0]["match"] == "[REDACTED]" def test_get_http_exception_no_blocked_assessments_omits_field(): @@ -1899,3 +2032,135 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): assert isinstance(exc, HTTPException) assert "assessments" not in exc.detail assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" + + +@pytest.mark.asyncio +async def test_streaming_post_call_parallel_output_passes_request_data_to_make_bedrock(): + """ + async_post_call_streaming_iterator_hook must pass request_data into OUTPUT + make_bedrock_api_request so spend/standard_logging attaches to the real request + (Greptile: previously OUTPUT used request_data=None / ephemeral {}). + """ + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"stream_guardrail_logging": True}, + } + guardrail = BedrockGuardrail( + guardrail_name="bedrock-stream-reqdata", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + mock_chunks = [ + litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="Hi", role="assistant"), + finish_reason=None, + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ), + litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="!", role="assistant"), + finish_reason="stop", + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ), + ] + + async def mock_stream(): + for c in mock_chunks: + yield c + + minimal = {"action": "NONE", "assessments": [], "outputs": []} + with patch.object( + guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) + ) as mock_make: + out = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data=request_data, + ): + out.append(chunk) + + assert len(out) >= 1 + output_calls = [ + c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" + ] + assert len(output_calls) == 1 + assert output_calls[0].kwargs.get("request_data") is request_data + assert ( + output_calls[0].kwargs.get("logging_event_type") + == GuardrailEventHooks.post_call + ) + input_calls = [ + c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" + ] + assert len(input_calls) == 1 + assert input_calls[0].kwargs.get("request_data") is request_data + + +@pytest.mark.asyncio +async def test_streaming_post_call_output_only_path_passes_request_data_to_make_bedrock(): + """When INPUT validation is skipped (pre/during already ran), OUTPUT still gets request_data.""" + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + } + guardrail = BedrockGuardrail( + guardrail_name="bedrock-stream-out-only", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.during_call, + default_on=True, + ) + mock_chunks = [ + litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="x", role="assistant"), + finish_reason="stop", + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ), + ] + + async def mock_stream(): + for c in mock_chunks: + yield c + + minimal = {"action": "NONE", "assessments": [], "outputs": []} + with patch.object( + guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) + ) as mock_make: + async for _ in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data=request_data, + ): + pass + + assert mock_make.call_count == 1 + c = mock_make.call_args + assert c.kwargs.get("source") == "OUTPUT" + assert c.kwargs.get("request_data") is request_data From 9577d87158d7d3969a50e8daf0e1d1adbff6aab9 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 22 Apr 2026 23:22:35 +0300 Subject: [PATCH 2/3] fix(proxy): guardrail header dedupe, mypy during_call, test mock kwargs - Dedupe names in add_guardrail_to_applied_guardrails_header (matches policies). - Inline unified during_call condition so mypy narrows UserAPIKeyAuth. - Extend bedrock guardrails test mock for logging_event_type. Made-with: Cursor --- litellm/proxy/common_utils/callback_utils.py | 3 ++- litellm/proxy/utils.py | 5 ++--- tests/guardrails_tests/test_bedrock_guardrails.py | 8 +++++++- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index e31c76dcac..7ddd722a80 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -433,7 +433,8 @@ def add_guardrail_to_applied_guardrails_header( return _metadata = request_data.get("metadata", None) or {} if "applied_guardrails" in _metadata: - _metadata["applied_guardrails"].append(guardrail_name) + if guardrail_name not in _metadata["applied_guardrails"]: + _metadata["applied_guardrails"].append(guardrail_name) else: _metadata["applied_guardrails"] = [guardrail_name] # Ensure metadata is set back to request_data (important when metadata didn't exist) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 561f8e5c55..712853a33c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1541,12 +1541,11 @@ class ProxyLogging: else: user_api_key_auth_dict = user_api_key_dict # Add task to list for parallel execution - use_unified_during = ( + if ( "apply_guardrail" in type(callback).__dict__ and user_api_key_dict is not None and not getattr(callback, "use_native_during_call_hook", False) - ) - if use_unified_during: + ): data["guardrail_to_apply"] = callback guardrail_task = self._run_guardrail_task_with_enrichment( callback, diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 7eaac60bf2..5435721620 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1107,7 +1107,12 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): # Mock the make_bedrock_api_request method to track calls async def mock_make_bedrock_api_request( - source, messages=None, response=None, request_data=None + source, + messages=None, + response=None, + request_data=None, + logging_event_type=None, + **kwargs, ): bedrock_calls.append( { @@ -1115,6 +1120,7 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): "messages": messages, "response": response, "request_data": request_data, + "logging_event_type": logging_event_type, } ) # Return the mock bedrock response From 3df9780c0286ff89a5d94be078b62db0b15cf895 Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 23 Apr 2026 00:12:30 +0300 Subject: [PATCH 3/3] fix(core_helpers): make redact_nested_match_and_regex_keys iterative Replace recursive `_walk` helper with a stack-based traversal so the recursive_detector CI check passes without adding to the ignore list, and avoid Python recursion limits on deeply nested payloads. Made-with: Cursor --- litellm/litellm_core_utils/core_helpers.py | 30 ++++++++++++---------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 07239a6886..b7a8b6f9ad 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -454,20 +454,24 @@ def redact_nested_match_and_regex_keys( except Exception: return payload - def _walk(node: Any) -> None: - if isinstance(node, dict): - if "match" in node: - node["match"] = "[REDACTED]" - if "regex" in node: - node["regex"] = "[REDACTED]" - for value in node.values(): - _walk(value) - elif isinstance(node, list): - for item in node: - _walk(item) - + # Iterative traversal; `seen` guards against cyclic refs preserved by deepcopy. try: - _walk(redacted) + seen: set = set() + stack: List[Any] = [redacted] + while stack: + node = stack.pop() + node_id = id(node) + if node_id in seen: + continue + seen.add(node_id) + if isinstance(node, dict): + if "match" in node: + node["match"] = "[REDACTED]" + if "regex" in node: + node["regex"] = "[REDACTED]" + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) except Exception: return payload return redacted