From a5dd01d8aca49a94433d4a430958fe369ac18a2f Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:52:23 -0400 Subject: [PATCH 001/220] added bedrock guardrail API exception --- .../guardrail_hooks/bedrock_guardrails.py | 37 ++++++++++++------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8ef188bb23c..02d87a4962d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -82,7 +82,12 @@ def _redact_pii_matches(response_json: dict) -> dict: redacted_response = copy.deepcopy(response_json) # Get assessments from the response - assessments = redacted_response.get("assessments", []) + # 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 @@ -90,13 +95,13 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact PII entities in sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + 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", []) + regexes = sensitive_info_policy.get("regexes") or [] for regex_match in regexes: if "match" in regex_match: regex_match["match"] = "[REDACTED]" @@ -104,12 +109,12 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact custom word matches in word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + 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", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if "match" in managed_word: managed_word["match"] = "[REDACTED]" @@ -682,7 +687,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return False # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED) - assessments = response.get("assessments", []) + # 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. + assessments = response.get("assessments") or [] if not assessments: return False @@ -690,7 +697,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check topic policy topic_policy = assessment.get("topicPolicy") if topic_policy: - topics = topic_policy.get("topics", []) + topics = topic_policy.get("topics") or [] for topic in topics: if topic.get("action") == "BLOCKED": return True @@ -698,7 +705,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check content policy content_policy = assessment.get("contentPolicy") if content_policy: - filters = content_policy.get("filters", []) + filters = content_policy.get("filters") or [] for filter_item in filters: if filter_item.get("action") == "BLOCKED": return True @@ -706,11 +713,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if custom_word.get("action") == "BLOCKED": return True - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if managed_word.get("action") == "BLOCKED": return True @@ -718,12 +725,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] if pii_entities: for pii_entity in pii_entities: if pii_entity.get("action") == "BLOCKED": return True - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] if regexes: for regex in regexes: if regex.get("action") == "BLOCKED": @@ -732,7 +739,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check contextual grounding policy contextual_grounding_policy = assessment.get("contextualGroundingPolicy") if contextual_grounding_policy: - grounding_filters = contextual_grounding_policy.get("filters", []) + grounding_filters = contextual_grounding_policy.get("filters") or [] for grounding_filter in grounding_filters: if grounding_filter.get("action") == "BLOCKED": return True @@ -1391,7 +1398,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Raises: Exception: If content is blocked by Bedrock guardrail """ - texts = inputs.get("texts", []) + # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. + # dict.get("texts", []) would return None if the key exists with a None value. + texts = inputs.get("texts") or [] try: verbose_proxy_logger.debug( f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" From ead822b6985443827d3accb7c0d6cfb7fce40df7 Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:55:36 -0400 Subject: [PATCH 002/220] Added test cases for the null type handling --- .../test_bedrock_guardrails.py | 2972 ++++++++--------- 1 file changed, 1472 insertions(+), 1500 deletions(-) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index cb594c221c4..03fe63b307a 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1,1090 +1,213 @@ -import sys +""" +Unit tests for Bedrock Guardrails +""" +import json import os -import io, asyncio +import sys +from unittest.mock import AsyncMock, MagicMock, patch + import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) -sys.path.insert(0, os.path.abspath("../..")) -import litellm -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.proxy._types import UserAPIKeyAuth -from litellm.caching import DualCache -from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + _redact_pii_matches, +) @pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() +async def test__redact_pii_matches_function(): + """Test the _redact_pii_matches function directly""" - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "I need to cancel my order"}, - { - "role": "user", - "content": "ok, my credit card number is 1234-5678-9012-3456", - }, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print("response after moderation hook", response) - - if response: # Only assert if response is not None - assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert response["messages"][2]["content"] == "I need to cancel my order" - assert ( - response["messages"][3]["content"] - == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking_content_list(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "who is the president of the united states?"}, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print(response) - - if response: # Only assert if response is not None - # Verify that the list content is properly masked - assert isinstance(response["messages"][0]["content"], list) - assert ( - response["messages"][0]["content"][0]["text"] - == "Hello, my phone number is {PHONE}" - ) - assert response["messages"][0]["content"][1]["text"] == "what time is it?" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert ( - response["messages"][2]["content"] - == "who is the president of the united states?" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_messages_api(): - """ - Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "claude-3-5-sonnet-20240620", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "user", "content": "tell me about coffee"}, - ], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="anthropic_messages", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_responses_api(): - """ - Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4.1", - "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee", - "stream": False, - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="responses", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - with pytest.raises(Exception): # Assert that this raises an exception - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Hi I like coffee"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming_no_violation(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_streaming_request_body_mock(): - """Test that the exact request body sent to Bedrock matches expected format when using streaming""" - import json - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - from litellm.types.guardrails import GuardrailEventHooks - - # Create mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - mock_cache = MagicMock(spec=DualCache) - - # Create the guardrail - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - # Mock the assembled response from streaming - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="The capital of Spain is Madrid." - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Mock Bedrock API response - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []} - - # Patch the async_handler.post method to capture the request body - with patch.object(guardrail, "async_handler") as mock_async_handler: - mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) - - # Test data - simulating request data and assembled response - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "what's the capital of spain?"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - # Call the method that should make the Bedrock API request - await guardrail.make_bedrock_api_request( - source="OUTPUT", response=mock_response, request_data=request_data - ) - - # Verify the API call was made - mock_async_handler.post.assert_called_once() - - # Get the request data that was passed - call_args = mock_async_handler.post.call_args - - # The data should be in the 'data' parameter of the prepared request - # We need to parse the JSON from the prepared request body - prepared_request_body = call_args.kwargs.get("data") - - # Parse the JSON body - if isinstance(prepared_request_body, bytes): - actual_body = json.loads(prepared_request_body.decode("utf-8")) - else: - actual_body = json.loads(prepared_request_body) - - # Expected body based on the convert_to_bedrock_format method behavior - expected_body = { - "source": "OUTPUT", - "content": [{"text": {"text": "The capital of Spain is Madrid."}}], - } - - print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) - print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) - - # Assert the request body matches exactly - assert ( - actual_body == expected_body - ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_aws_param_persistence(): - """Test that AWS auth params set on init are used for every request and not popped out.""" - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.guardrails import GuardrailEventHooks - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - aws_access_key_id="test-access-key", - aws_secret_access_key="test-secret-key", - aws_region_name="us-east-1", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - with patch.object( - guardrail, "get_credentials", wraps=guardrail.get_credentials - ) as mock_get_creds: - for i in range(3): - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": f"request {i}"}], - "stream": False, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - # Configure the mock response properly - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.json = MagicMock( - return_value={"action": "NONE", "outputs": []} - ) - mock_post.return_value = mock_response - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - assert mock_get_creds.call_count == 3 - for call in mock_get_creds.call_args_list: - kwargs = call.kwargs - print("used the following kwargs to get credentials=", kwargs) - assert kwargs["aws_access_key_id"] == "test-access-key" - assert kwargs["aws_secret_access_key"] == "test-secret-key" - assert kwargs["aws_region_name"] == "us-east-1" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): - """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - from unittest.mock import MagicMock - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Test 1: ANONYMIZED action should NOT raise exception - anonymized_response: BedrockGuardrailResponse = { + # Test case 1: Response with PII entities + response_with_pii = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ + {"type": "NAME", "match": "John Smith", "action": "BLOCKED"}, { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) - assert should_raise is False, "ANONYMIZED actions should not raise exceptions" - - # Test 2: BLOCKED action should raise exception - blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) - assert should_raise is True, "BLOCKED actions should raise exceptions" - - # Test 3: Mixed actions - should raise if ANY action is BLOCKED - mixed_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" - - # Test 4: NONE action should not raise exception - none_response: BedrockGuardrailResponse = { - "action": "NONE", - "outputs": [], - "assessments": [], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) - assert should_raise is False, "NONE actions should not raise exceptions" - - # Test 5: Test other policy types with BLOCKED actions - content_blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - content_blocked_response - ) - assert ( - should_raise is True - ), "Content policy BLOCKED actions should raise exceptions" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_masking_with_anonymized_response(): - """Test that masking works correctly when guardrail returns ANONYMIZED actions""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - mask_request_content=True, - ) - - # Mock the Bedrock API response with ANONYMIZED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should NOT raise an exception since action is ANONYMIZED - try: - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data with masked content - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE}" - ) - except Exception as e: - pytest.fail( - f"Should not raise exception for ANONYMIZED actions, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): - """Test that masked output from guardrails is used even when masking flags are not enabled""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail WITHOUT masking flags enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - # Note: No mask_request_content=True or mask_response_content=True - ) - - # Mock the Bedrock API response with ANONYMIZED action and masked output - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", + "type": "US_SOCIAL_SECURITY_NUMBER", + "match": "324-12-3212", + "action": "BLOCKED", }, + {"type": "PHONE", "match": "607-456-7890", "action": "BLOCKED"}, + ] + } + } + ], + "outputs": [{"text": "Input blocked by PII policy"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_pii) + + # Verify that PII matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "Phone should be redacted" + + # Verify other fields remain unchanged + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Input blocked by PII policy" + + print("PII redaction function test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_no_pii(): + """Test _redact_pii_matches with response that has no PII""" + + response_no_pii = {"action": "NONE", "assessments": [], "outputs": []} + + # Call the redaction function + redacted_response = _redact_pii_matches(response_no_pii) + + # Should return the same response unchanged + assert redacted_response == response_no_pii + print("No PII redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_empty_assessments(): + """Test _redact_pii_matches with empty assessments""" + + response_empty_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": []}}], + "outputs": [{"text": "Some output"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_empty_assessments) + + # Should return the same response unchanged + assert redacted_response == response_empty_assessments + print("Empty assessments redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_malformed_response(): + """Test _redact_pii_matches with malformed response (should not crash)""" + + # Test with completely malformed response + malformed_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": "not_a_list", # This should cause an exception + } + + # Should not crash and return original response + redacted_response = _redact_pii_matches(malformed_response) + assert redacted_response == malformed_response + + # Test with missing keys + missing_keys_response = { + "action": "GUARDRAIL_INTERVENED" + # Missing assessments key + } + + redacted_response = _redact_pii_matches(missing_keys_response) + assert redacted_response == missing_keys_response + + print("Malformed response redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_multiple_assessments(): + """Test _redact_pii_matches with multiple assessments containing PII""" + + response_multiple_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "EMAIL", - "match": "user@example.com", + "match": "john@example.com", "action": "ANONYMIZED", - }, + } ] } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com", }, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should use the masked output even without masking flags - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Should use the masked content from guardrail output - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE} and email is {EMAIL}" - ) - print("✅ Masked output was applied even without masking flags enabled") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_non_streaming(): - """Test that PII masking is applied to response content in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [ - { - "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - } - ], - "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ { "type": "CREDIT_DEBIT_CARD_NUMBER", "match": "1234-5678-9012-3456", + "action": "BLOCKED", + }, + { + "type": "ADDRESS", + "match": "123 Main St, Anytown USA", "action": "ANONYMIZED", }, + ] + } + }, + ], + "outputs": [{"text": "Multiple PII detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_multiple_assessments) + + # Verify all PII in all assessments are redacted + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" + assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" + assert assessment2_pii[1]["match"] == "[REDACTED]", "Address should be redacted" + + # Verify types remain unchanged + assert assessment1_pii[0]["type"] == "EMAIL" + assert assessment2_pii[0]["type"] == "CREDIT_DEBIT_CARD_NUMBER" + assert assessment2_pii[1]["type"] == "ADDRESS" + + print("Multiple assessments redaction test passed") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_logging_uses_redacted_response(): + """Test that the Bedrock guardrail uses redacted response for logging""" + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with PII + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create a mock response that contains PII - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212", - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your credit card and phone number?"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the post-call success hook - await guardrail.async_post_call_success_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - response=mock_response, - ) - - # Verify that the response content was masked - assert ( - mock_response.choices[0].message.content - == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - ) - print("✓ Non-streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_streaming(): - """Test that PII masking is applied to response content in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED", - }, - { - "type": "US_SSN", - "match": "123-45-6789", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create mock streaming chunks - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="Sure! My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="john@example.com and SSN is " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="123-45-6789"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your email and SSN?"}, - ], - "stream": True, - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the streaming hook - masked_stream = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - - # Collect all chunks from the masked stream - masked_chunks = [] - async for chunk in masked_stream: - masked_chunks.append(chunk) - - # Verify that we got chunks back - assert len(masked_chunks) > 0 - - # Reconstruct the full response from chunks to verify masking - full_content = "" - for chunk in masked_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # Verify that the reconstructed content contains the masked PII - assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content - print("✓ Streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_input_source(): - """Test convert_to_bedrock_format with INPUT source and mock messages""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock messages - mock_messages = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you!"}, - { - "role": "user", - "content": [ - {"type": "text", "text": "What's the weather like?"}, - {"type": "text", "text": "Is it sunny today?"}, - ], - }, - ] - - # Call the method - result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "INPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - expected_content_items = [ - {"text": {"text": "Hello, how are you?"}}, - {"text": {"text": "I'm doing well, thank you!"}}, - {"text": {"text": "What's the weather like?"}}, - {"text": {"text": "Is it sunny today?"}}, - ] - - assert result.get("content") == expected_content_items - print("✅ INPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_output_source(): - """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - import litellm - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock ModelResponse - mock_response = litellm.ModelResponse( - id="test-response-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="This is a test response from the model." - ), - finish_reason="stop", - ), - litellm.Choices( - index=1, - message=litellm.Message( - role="assistant", content="This is a second choice response." - ), - finish_reason="stop", - ), - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Call the method - result = guardrail.convert_to_bedrock_format( - source="OUTPUT", response=mock_response - ) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "OUTPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - should contain both choice contents - expected_content_items = [ - {"text": {"text": "This is a test response from the model."}}, - {"text": {"text": "This is a second choice response."}}, - ] - - assert result.get("content") == expected_content_items - print("✅ OUTPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_post_call_streaming_hook(): - """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock streaming chunks that contain PII - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="john@example.com"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with PII masking - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "My email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", + "match": "+1 412 555 1212", # This should be redacted in logs "action": "ANONYMIZED", } ] @@ -1095,99 +218,82 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): request_data = { "model": "gpt-4o", - "messages": [{"role": "user", "content": "What's your email?"}], - "stream": True, + "messages": [ + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, + ], } - # Track which bedrock API calls were made - bedrock_calls = [] + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None - # 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 - ): - bedrock_calls.append( - { - "source": source, - "messages": messages, - "response": response, - "request_data": request_data, - } - ) - # Return the mock bedrock response - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - return BedrockGuardrailResponse(**mock_bedrock_response.json()) - - # Patch the bedrock API request method + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( - guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request - ): + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" + ) as mock_debug, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: - # Call the streaming hook - result_generator = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), + mock_post.return_value = mock_bedrock_response + + # Call the method that should log the redacted response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), request_data=request_data, ) - # Collect all chunks from the result - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) + # Verify that debug logging was called + mock_debug.assert_called() + + # Get the logged response (second argument to debug call) + logged_calls = mock_debug.call_args_list + bedrock_response_log_call = None + + for call in logged_calls: + args, kwargs = call + if len(args) >= 2 and "Bedrock AI response" in str(args[0]): + bedrock_response_log_call = call + break - # Verify bedrock API calls were made - # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. - # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled - # to avoid redundant validation. Since event_hook=None means all hooks are enabled, - # only OUTPUT validation should be performed in post_call. assert ( - len(bedrock_calls) == 1 - ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" + bedrock_response_log_call is not None + ), "Should have logged Bedrock AI response" - # Verify the OUTPUT call - output_call = bedrock_calls[0] - assert output_call["source"] == "OUTPUT" - assert output_call["response"] is not None - assert output_call["messages"] is None # OUTPUT calls don't need messages + # Extract the logged response data + logged_response = bedrock_response_log_call[0][ + 1 + ] # Second argument to debug call - # Verify that the response content was masked - # The streaming chunks should now contain the masked content - full_content = "" - for chunk in result_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if ( - hasattr(chunk.choices[0], "delta") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # The content should be masked (contains {EMAIL} instead of john@example.com) + # Verify that the logged response has redacted PII assert ( - "{EMAIL}" in full_content - ), f"Expected masked content with {{EMAIL}}, got: {full_content}" - assert ( - "john@example.com" not in full_content - ), f"Original email should be masked, got: {full_content}" - - print( - "✅ Post-call streaming hook test passed - OUTPUT source used for masking" + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "[REDACTED]" ) - print( - f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} " - "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" + + # Verify other fields are preserved + assert logged_response["action"] == "GUARDRAIL_INTERVENED" + assert ( + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["type"] + == "PHONE" ) - print(f"✅ Final masked content: {full_content}") + + print("Bedrock guardrail logging redaction test passed") @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_shows_output_text(): - """Test that BLOCKED actions raise HTTPException with the output text in the detail""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test_bedrock_guardrail_original_response_not_modified(): + """Test that the original response is not modified by redaction, only the logged version""" # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() @@ -1196,408 +302,1274 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock the Bedrock API response with BLOCKED action and output text - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + # Mock the Bedrock API response with PII + original_response_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "this violates litellm corporate guardrail policy"}], + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", # This should NOT be modified in original + "action": "ANONYMIZED", + } ] } } ], } + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = original_response_data + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, ], } - # Patch the async_handler.post method + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: + mock_post.return_value = mock_bedrock_response - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "detail" in exception.__dict__ - - # Check that the detail contains the expected structure - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - - # Verify that the output text from both outputs is included - expected_output_text = "this violates litellm corporate guardrail policy" - assert detail["bedrock_guardrail_response"] == expected_output_text - - print( - "✅ BLOCKED action HTTPException test passed - output text properly included" + # Call the method + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, ) + # Verify that the original response data was not modified + # (The json() method should return the original data) + original_data = mock_bedrock_response.json() + assert ( + original_data["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "+1 412 555 1212" + ) + + # Verify that the returned BedrockGuardrailResponse contains original data + assert ( + result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ + "match" + ] + == "+1 412 555 1212" + ) + + print("Original response not modified test passed") + @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_empty_outputs(): - """Test that BLOCKED actions with empty outputs still raise HTTPException""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test__redact_pii_matches_preserves_non_pii_entities(): + """Test that _redact_pii_matches only affects PII-related entities and preserves other assessment data""" - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock the Bedrock API response with BLOCKED action but no outputs - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + response_with_mixed_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [], # Empty outputs "assessments": [ { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Violent content here"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - - # Check that the detail contains the expected structure with empty output text - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs - - print("✅ BLOCKED action with empty outputs test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - # Mock the Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise HTTPException when disable_exception_on_block=False - with pytest.raises(HTTPException) as exc_info: - await guardrail_default.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "Violated guardrail policy" in str(exception.detail) - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - response = await guardrail_disabled.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data (even though content was blocked) - assert response is not None - print("✅ No exception raised when disable_exception_on_block=True") - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - from fastapi import HTTPException - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Mock streaming chunks that would normally trigger a block - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="Here's how to make explosives: " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="step 1, step 2..."), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Tell me how to make explosives"}], - "stream": True, - } - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise exception during streaming processing - with pytest.raises(HTTPException): - result_generator = ( - guardrail_default.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Try to consume the generator - should raise exception - async for chunk in result_generator: - pass - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - result_generator = ( - guardrail_disabled.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Consume the generator - should succeed without exceptions - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) - - # Should have received chunks back even though content was blocked - assert len(result_chunks) > 0 - print( - "✅ Streaming completed without exception when disable_exception_on_block=True" - ) - - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): - """Test that async_post_call_success_hook skips when there's no output text""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Create a ModelResponse with tool calls (no text content) - # This simulates a response where the LLM is making a tool call - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content=None, # No text content - tool_calls=[ - litellm.utils.ChatCompletionMessageToolCall( - id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", - function=litellm.utils.Function( - name="top_song", arguments='{"sign": "WZPZ"}' - ), - type="function", - ) + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + "confidence": "HIGH", + } ], - ), - finish_reason="tool_calls", - ) + "regexes": [ + { + "name": "custom_pattern", + "match": "some_pattern_match", + "action": "BLOCKED", + } + ], + }, + "contentPolicy": { + "filters": [ + { + "type": "VIOLENCE", + "confidence": "MEDIUM", + "action": "BLOCKED", + } + ] + }, + "topicPolicy": { + "topics": [ + { + "name": "Restricted Topic", + "type": "DENY", + "action": "BLOCKED", + } + ] + }, + } ], - created=1234567890, - model="gpt-4o", - object="chat.completion", + "outputs": [{"text": "Content blocked"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_mixed_data) + + # Verify that PII entity matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" + assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" + assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" + assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" + + # Verify that regex matches are also redacted (updated behavior) + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + + # Verify that other policies are completely unchanged + content_policy = redacted_response["assessments"][0]["contentPolicy"] + assert content_policy["filters"][0]["type"] == "VIOLENCE" + assert content_policy["filters"][0]["confidence"] == "MEDIUM" + + topic_policy = redacted_response["assessments"][0]["topicPolicy"] + assert topic_policy["topics"][0]["name"] == "Restricted Topic" + + # Verify top-level fields are unchanged + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Content blocked" + + print("Preserves non-PII entities test passed") + + +@pytest.mark.asyncio +async def test_pii_redaction_matches_debug_output_format(): + """Test that demonstrates the exact behavior shown in your debug output""" + + # This matches the structure from your debug output + original_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "assessments": [ + { + "invocationMetrics": { + "guardrailCoverage": { + "textCharacters": {"guarded": 84, "total": 84} + }, + "guardrailProcessingLatency": 322, + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "BLOCKED", + "detected": True, + "match": "John Smith", + "type": "NAME", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "324-12-3212", + "type": "US_SOCIAL_SECURITY_NUMBER", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "607-456-7890", + "type": "PHONE", + }, + ] + }, + } + ], + "blockedResponse": "Input blocked by PII policy", + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, + "output": [{"text": "Input blocked by PII policy"}], + "outputs": [{"text": "Input blocked by PII policy"}], + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + } + + # Apply redaction + redacted_response = _redact_pii_matches(original_response) + + # Verify the redacted response matches your expected debug output + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + # All PII matches should be redacted + assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "PHONE should be redacted" + + # But all other fields should be preserved + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert pii_entities[0]["action"] == "BLOCKED" + assert pii_entities[0]["detected"] == True + + # Verify that the original response is unchanged + original_pii_entities = original_response["assessments"][0][ + "sensitiveInformationPolicy" + ]["piiEntities"] + assert ( + original_pii_entities[0]["match"] == "John Smith" + ), "Original should be unchanged" + assert ( + original_pii_entities[1]["match"] == "324-12-3212" + ), "Original should be unchanged" + assert ( + original_pii_entities[2]["match"] == "607-456-7890" + ), "Original should be unchanged" + + # Verify all other metadata is preserved in redacted response + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["actionReason"] == "Guardrail blocked." + assert redacted_response["blockedResponse"] == "Input blocked by PII policy" + assert ( + redacted_response["assessments"][0]["invocationMetrics"][ + "guardrailProcessingLatency" + ] + == 322 ) - data = { + print("PII redaction matches debug output format test passed") + print( + f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" + ) + print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_regex_matches(): + """Test redaction of regex matches in sensitive information policy""" + + response_with_regex = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "regexes": [ + { + "name": "SSN_PATTERN", + "match": "123-45-6789", + "action": "BLOCKED", + }, + { + "name": "CREDIT_CARD_PATTERN", + "match": "4111-1111-1111-1111", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Regex patterns detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_regex) + + # Verify that regex matches are redacted + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + + assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" + assert ( + regexes[1]["match"] == "[REDACTED]" + ), "Credit card regex match should be redacted" + + # Verify other fields are preserved + assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + assert regexes[1]["name"] == "CREDIT_CARD_PATTERN", "Regex name should be preserved" + assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" + + # Verify original response is unchanged + original_regexes = response_with_regex["assessments"][0][ + "sensitiveInformationPolicy" + ]["regexes"] + assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" + assert ( + original_regexes[1]["match"] == "4111-1111-1111-1111" + ), "Original should be unchanged" + + print("Regex matches redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_custom_words(): + """Test redaction of custom word matches in word policy""" + + response_with_custom_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": [ + { + "match": "confidential_data", + "action": "BLOCKED", + }, + { + "match": "secret_information", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Custom words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_custom_words) + + # Verify that custom word matches are redacted + custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] + + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "First custom word match should be redacted" + assert ( + custom_words[1]["match"] == "[REDACTED]" + ), "Second custom word match should be redacted" + + # Verify other fields are preserved + assert ( + custom_words[0]["action"] == "BLOCKED" + ), "Custom word action should be preserved" + assert ( + custom_words[1]["action"] == "ANONYMIZED" + ), "Custom word action should be preserved" + + # Verify original response is unchanged + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ + "customWords" + ] + assert ( + original_custom_words[0]["match"] == "confidential_data" + ), "Original should be unchanged" + assert ( + original_custom_words[1]["match"] == "secret_information" + ), "Original should be unchanged" + + print("Custom words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_managed_words(): + """Test redaction of managed word matches in word policy""" + + response_with_managed_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "managedWordLists": [ + { + "match": "inappropriate_word", + "action": "BLOCKED", + "type": "PROFANITY", + }, + { + "match": "offensive_term", + "action": "ANONYMIZED", + "type": "HATE_SPEECH", + }, + ] + } + } + ], + "outputs": [{"text": "Managed words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_managed_words) + + # Verify that managed word matches are redacted + managed_words = redacted_response["assessments"][0]["wordPolicy"][ + "managedWordLists" + ] + + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "First managed word match should be redacted" + assert ( + managed_words[1]["match"] == "[REDACTED]" + ), "Second managed word match should be redacted" + + # Verify other fields are preserved + assert ( + managed_words[0]["action"] == "BLOCKED" + ), "Managed word action should be preserved" + assert ( + managed_words[0]["type"] == "PROFANITY" + ), "Managed word type should be preserved" + assert ( + managed_words[1]["action"] == "ANONYMIZED" + ), "Managed word action should be preserved" + assert ( + managed_words[1]["type"] == "HATE_SPEECH" + ), "Managed word type should be preserved" + + # Verify original response is unchanged + original_managed_words = response_with_managed_words["assessments"][0][ + "wordPolicy" + ]["managedWordLists"] + assert ( + original_managed_words[0]["match"] == "inappropriate_word" + ), "Original should be unchanged" + assert ( + original_managed_words[1]["match"] == "offensive_term" + ), "Original should be unchanged" + + print("Managed words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_comprehensive_coverage(): + """Test redaction across all supported policy types in a single response""" + + comprehensive_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + } + ], + "regexes": [ + { + "name": "PHONE_PATTERN", + "match": "555-123-4567", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": [ + { + "match": "confidential", + "action": "BLOCKED", + } + ], + "managedWordLists": [ + { + "match": "inappropriate", + "action": "ANONYMIZED", + "type": "PROFANITY", + } + ], + }, + } + ], + "outputs": [{"text": "Multiple policy violations detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(comprehensive_response) + + # Verify all match fields are redacted + assessment = redacted_response["assessments"][0] + + # PII entities + pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] + assert ( + pii_entities[0]["match"] == "[REDACTED]" + ), "PII entity match should be redacted" + + # Regex matches + regexes = assessment["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + + # Custom words + custom_words = assessment["wordPolicy"]["customWords"] + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "Custom word match should be redacted" + + # Managed words + managed_words = assessment["wordPolicy"]["managedWordLists"] + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "Managed word match should be redacted" + + # Verify all other fields are preserved + assert pii_entities[0]["type"] == "EMAIL" + assert regexes[0]["name"] == "PHONE_PATTERN" + assert managed_words[0]["type"] == "PROFANITY" + + # Verify original response is unchanged + original_assessment = comprehensive_response["assessments"][0] + assert ( + original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] + == "user@example.com" + ) + assert ( + original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] + == "555-123-4567" + ) + assert ( + original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + ) + assert ( + original_assessment["wordPolicy"]["managedWordLists"][0]["match"] + == "inappropriate" + ) + + print("Comprehensive coverage redaction test passed") + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set""" + + # Clear any existing environment variable to ensure clean test + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail with custom runtime endpoint + custom_endpoint = "https://custom-bedrock.example.com" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=custom_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method to avoid actual AWS credential loading + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + + print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects AWS_BEDROCK_RUNTIME_ENDPOINT environment variable""" + + custom_endpoint = "https://env-bedrock.example.com" + + # Set the environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) + + # Create guardrail without explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint from environment is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + + print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkeypatch): + """Test that BedrockGuardrail uses default endpoint when no custom endpoint is set""" + + # Ensure no environment variable is set + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail without any custom endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-west-2" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the default endpoint is used + expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected default URL. Got: {prepped_request.url}" + + print(f"Default endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch): + """Test that aws_bedrock_runtime_endpoint parameter takes precedence over environment variable + + This test verifies the corrected behavior where the parameter should take precedence + over the environment variable, consistent with the endpoint_url logic. + """ + + param_endpoint = "https://param-bedrock.example.com" + env_endpoint = "https://env-bedrock.example.com" + + # Set environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_endpoint) + + # Create guardrail with explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=param_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the parameter takes precedence over environment variable + expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + + print(f"Parameter precedence test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): + """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" + # Create a BedrockGuardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the make_bedrock_api_request method + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: + # Test the apply_guardrail method with tool_calls in response + inputs = { + "texts": [], + "tool_calls": [ + { + "id": "call_eFSCWFsyL7MclHYnzKrcQnMK", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location":"São Paulo"}', + }, + } + ], + } + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=None, + ) + + # Verify the result - should succeed without errors + assert guardrailed_inputs is not None + assert "tool_calls" in guardrailed_inputs + assert len(guardrailed_inputs["tool_calls"]) == 1 + assert ( + guardrailed_inputs["tool_calls"][0]["id"] + == "call_eFSCWFsyL7MclHYnzKrcQnMK" + ) + assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" + assert ( + guardrailed_inputs["tool_calls"][0]["function"]["arguments"] + == '{"location":"São Paulo"}' + ) + # Verify that the Bedrock API was NOT called since there's no text to process + mock_api_request.assert_not_called() + print("✅ apply_guardrail with tool_calls test passed - no API call made") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): + """Test that BLOCKED content raises exception even when masking is enabled + + This test verifies the bug fix where previously mask_request_content=True or + mask_response_content=True would bypass all BLOCKED content checks. Now it + properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). + """ + + # Create guardrail with masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, # Masking enabled + mask_response_content=True, # Masking enabled + ) + + # Mock Bedrock response with BLOCKED content (hate speech) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", # Should raise exception + } + ] + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "NAME", + "match": "John Doe", + "action": "ANONYMIZED", # Should be masked + } + ] + }, + } + ], + "outputs": [{"text": "Content blocked due to policy violation"}], + } + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = blocked_response + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Hello"}, + {"role": "user", "content": "Test message with PII and hate speech"}, ], } - mock_user_api_key_dict = UserAPIKeyAuth() + + # Mock AWS-related methods + 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() + ): + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException for BLOCKED content + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + # Verify exception details + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in str(exc_info.value.detail) + + print("✅ BLOCKED content with masking enabled raises exception correctly") - result = await guardrail.async_post_call_success_hook( - data=data, - response=mock_response, - user_api_key_dict=mock_user_api_key_dict, - ) - # If no error is raised and result is None, then the test passes - assert result is None - print("✅ No output text in response test passed") + +# ────────────────────────────────────────────────────────────────────────────── +# Null-safety tests for Bedrock guardrail responses +# +# The Bedrock ApplyGuardrail API can return explicit null/None for list fields +# such as "regexes", "piiEntities", "topics", "filters", "customWords", and +# "managedWordLists" when a particular policy category is present in the +# assessment but has no matches. +# +# Python's dict.get("key", []) returns None (NOT []) when the key exists with +# a None value. The `or []` fallback ensures we always iterate over a list. +# +# Without the fix, iterating over None raises: +# TypeError: 'NoneType' object is not iterable +# which surfaces to callers as: +# openai.InternalServerError: Error code: 500 +# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} +# ────────────────────────────────────────────────────────────────────────────── + + +class TestRedactPiiMatchesNullSafety: + """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" + + @pytest.mark.asyncio + async def test_should_handle_null_regexes_in_sensitive_info_policy(self): + """Bedrock can return regexes: null while piiEntities has data. + + Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no + custom regex patterns configured, so the API returns regexes: null. + """ + response = { + "action": "NONE", + "actionReason": "No action.", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "NONE", + "detected": True, + "match": "joebloggs@gmail.com", + "type": "EMAIL", + } + ], + "regexes": None, # Explicit null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError: 'NoneType' object is not iterable + redacted = _redact_pii_matches(response) + + # PII match should be redacted + pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert pii[0]["match"] == "[REDACTED]" + assert pii[0]["type"] == "EMAIL" + + @pytest.mark.asyncio + async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): + """Bedrock can return piiEntities: null while regexes has data.""" + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, # null from Bedrock API + "regexes": [ + { + "name": "CUSTOM_PATTERN", + "match": "secret-abc-123", + "action": "BLOCKED", + } + ], + }, + } + ], + } + + redacted = _redact_pii_matches(response) + + regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_should_handle_null_custom_words_and_managed_words(self): + """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" + response = { + "action": "NONE", + "assessments": [ + { + "wordPolicy": { + "customWords": None, # null from Bedrock API + "managedWordLists": None, # null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + + # Values should remain None (no crash) + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """Bedrock can return assessments: null.""" + response = { + "action": "NONE", + "assessments": None, # null from Bedrock API + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + assert redacted["assessments"] is None + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists_together(self): + """All sub-list fields are null at the same time — worst-case scenario.""" + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + }, + "topicPolicy": None, + "contentPolicy": None, + "contextualGroundingPolicy": None, + } + ], + } + + # Should not raise any exception + redacted = _redact_pii_matches(response) + assert redacted is not None + + +class TestShouldRaiseGuardrailBlockedExceptionNullSafety: + """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" + + def _create_guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists(self): + """All policy sub-lists are null — should not crash, should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null from Bedrock API + }, + "contentPolicy": { + "filters": None, # null + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": { + "filters": None, # null + }, + } + ], + } + + # No BLOCKED actions found (all lists null) → should return False + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_detect_blocked_despite_other_null_lists(self): + """A mix of null lists and a real BLOCKED action — should still detect it.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null — should not crash + }, + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": None, # entire policy is null + } + ], + } + + # Should return True because contentPolicy has a BLOCKED filter + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """assessments itself is null — should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, # null from Bedrock API + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_handle_null_topics_with_blocked_word_policy(self): + """topics is null but wordPolicy has a BLOCKED customWord.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, + }, + "wordPolicy": { + "customWords": [ + {"match": "badword", "action": "BLOCKED"} + ], + "managedWordLists": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_pii_with_blocked_regex(self): + """piiEntities is null but regexes has a BLOCKED match.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": [ + {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} + ], + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_grounding_filters(self): + """contextualGroundingPolicy.filters is null — should not crash.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contextualGroundingPolicy": { + "filters": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_not_crash_when_action_is_not_intervened(self): + """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" + guardrail = self._create_guardrail() + + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + +class TestApplyGuardrailNullSafety: + """Tests for apply_guardrail handling of null/None texts input.""" + + @pytest.mark.asyncio + async def test_should_handle_none_texts_in_inputs(self): + """inputs[\"texts\"] is explicitly None — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {"texts": None} # Explicit None + + mock_credentials = MagicMock() + + 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() + ): + # With empty texts (from None → []), no Bedrock API call should be made + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + # Should return empty texts without crashing + assert result.get("texts") == [] + # No Bedrock API call should be made for empty input + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_should_handle_missing_texts_key(self): + """inputs has no \"texts\" key at all — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {} # No "texts" key + + mock_credentials = MagicMock() + + 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() + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result.get("texts") == [] + mock_post.assert_not_called() From cf94f4d8b720e63beadd10a28cea8c1787830814 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= Date: Sun, 5 Apr 2026 09:23:32 +0800 Subject: [PATCH 003/220] fix(mcp): is_tool_name_prefixed validates against known server prefixes (#25085) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #25081. is_tool_name_prefixed() checked for the presence of MCP_TOOL_PREFIX_SEPARATOR (default '-') anywhere in the tool name. Any non-MCP tool whose name contains a hyphen (e.g. 'text-to-speech', 'code-review') was silently misclassified as an MCP-prefixed tool. When the semantic tool filter is enabled, these tools would be routed through semantic matching and potentially dropped. Fix: accept an optional known_server_prefixes set. When supplied, the function extracts the candidate prefix (text before the first separator) and checks it against the normalised set of registered server prefixes. Only a genuine match returns True. Without the set, legacy behaviour is preserved for backward compatibility. Updated _get_mcp_server_from_tool_name() to build the prefix set from the live registry and pass it through. 9 new tests. Co-authored-by: d 🔹 <258577966+voidborne-d@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 7 +- .../proxy/_experimental/mcp_server/utils.py | 32 +++++-- .../mcp_server/test_is_tool_name_prefixed.py | 90 +++++++++++++++++++ 3 files changed, 123 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7b87e7e7e61..5363e317ff7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2552,7 +2552,12 @@ class MCPServerManager: return server # If not found and tool name is prefixed, try extracting server name from prefix - if is_tool_name_prefixed(tool_name): + known_prefixes = { + normalize_server_name(get_server_prefix(s)) + for s in self.get_registry().values() + if get_server_prefix(s) + } + if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): ( original_tool_name, server_name_from_prefix, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 8189f212bcb..79942eda54e 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -100,17 +100,39 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: return prefixed_name, "" -def is_tool_name_prefixed(tool_name: str) -> bool: +def is_tool_name_prefixed( + tool_name: str, + known_server_prefixes: Optional[set] = None, +) -> bool: """ - Check if tool name has server prefix + Check if tool name has a known MCP server prefix. + + When ``known_server_prefixes`` is provided the function verifies that the + substring before the first separator is an actual registered server + prefix. Without it the check falls back to the legacy heuristic + (separator present anywhere in the name), which can produce false + positives for non-MCP tools whose names contain hyphens + (e.g. ``text-to-speech``, ``code-review``). Args: - tool_name: Tool name to check + tool_name: Tool name to check. + known_server_prefixes: Optional set of normalised server prefixes + currently registered in the MCP manager. Pass this whenever + the caller has access to the server registry so that the check + is accurate. Returns: - True if tool name is prefixed, False otherwise + True if tool name is prefixed, False otherwise. """ - return MCP_TOOL_PREFIX_SEPARATOR in tool_name + if MCP_TOOL_PREFIX_SEPARATOR not in tool_name: + return False + + if known_server_prefixes is not None: + candidate_prefix = tool_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)[0] + return normalize_server_name(candidate_prefix) in known_server_prefixes + + # Legacy fallback – separator present somewhere in the name. + return True def validate_mcp_server_name( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py new file mode 100644 index 00000000000..d761d9c54cc --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py @@ -0,0 +1,90 @@ +""" +Tests for is_tool_name_prefixed with known_server_prefixes parameter. + +Verifies fix for https://github.com/BerriAI/litellm/issues/25081 +""" + +import pytest + +from litellm.proxy._experimental.mcp_server.utils import is_tool_name_prefixed + + +# --------------------------------------------------------------------------- +# Legacy behaviour (no known_server_prefixes passed) +# --------------------------------------------------------------------------- + + +class TestLegacyBehaviour: + """Without known_server_prefixes the function falls back to heuristic.""" + + def test_plain_name_returns_false(self): + assert is_tool_name_prefixed("get_weather") is False + + def test_hyphenated_name_returns_true_legacy(self): + """Legacy heuristic: any hyphen → True (the bug this issue reports).""" + assert is_tool_name_prefixed("text-to-speech") is True + + def test_prefixed_name_returns_true_legacy(self): + assert is_tool_name_prefixed("myserver-get_weather") is True + + +# --------------------------------------------------------------------------- +# New behaviour (known_server_prefixes supplied) +# --------------------------------------------------------------------------- + + +class TestWithKnownPrefixes: + """When known_server_prefixes is supplied, only real prefixes match.""" + + PREFIXES = {"myserver", "weather_api", "code_tools"} + + def test_known_prefix_returns_true(self): + assert ( + is_tool_name_prefixed( + "myserver-get_weather", known_server_prefixes=self.PREFIXES + ) + is True + ) + + def test_hyphenated_non_mcp_tool_returns_false(self): + """This is the core fix: 'text-to-speech' is NOT an MCP-prefixed tool.""" + assert ( + is_tool_name_prefixed( + "text-to-speech", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_code_review_not_misclassified(self): + assert ( + is_tool_name_prefixed( + "code-review", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_no_separator_returns_false(self): + assert ( + is_tool_name_prefixed( + "simple_tool", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_empty_prefixes_set_rejects_all(self): + """With an empty registry, nothing can be prefixed.""" + assert ( + is_tool_name_prefixed("myserver-get_weather", known_server_prefixes=set()) + is False + ) + + def test_prefix_normalisation(self): + """Server names with spaces are normalised to underscores.""" + prefixes = {"my_server"} + # add_server_prefix_to_name normalises spaces → underscores + assert ( + is_tool_name_prefixed( + "my_server-list_files", known_server_prefixes=prefixes + ) + is True + ) From d6351a3966e5cbbddee1bacd63020bb6aa857614 Mon Sep 17 00:00:00 2001 From: Neha Prasad Date: Sun, 5 Apr 2026 07:09:37 +0530 Subject: [PATCH 004/220] fix(s3_v2): use prepared URL for SigV4-signed S3 requests (#25074) --- litellm/integrations/s3_v2.py | 13 +++--- tests/test_litellm/integrations/test_s3_v2.py | 44 +++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 405bf9698cc..f767f61be87 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -403,9 +403,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) - # Make the request response = await self.async_httpx_client.put( - url, data=json_string, headers=signed_headers + prepped.url, data=json_string, headers=signed_headers ) response.raise_for_status() except Exception as e: @@ -582,8 +581,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_verify is not None else None ) - # Make the request - response = httpx_client.put(url, data=json_string, headers=signed_headers) + response = httpx_client.put( + prepped.url, data=json_string, headers=signed_headers + ) response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error uploading to s3: {str(e)}") @@ -674,8 +674,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) - # Make the request - response = await self.async_httpx_client.get(url, headers=signed_headers) + response = await self.async_httpx_client.get( + prepped.url, headers=signed_headers + ) if response.status_code != 200: verbose_logger.exception( diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index b53c05fa241..943fe4ec37b 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -292,6 +292,50 @@ class TestS3V2UnitTests: assert result == {"downloaded": "data"} + @patch("asyncio.create_task") + @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") + def test_s3_v2_put_url_encodes_spaces_in_object_key( + self, mock_periodic_flush, mock_create_task + ): + import requests + from unittest.mock import AsyncMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + mock_periodic_flush.return_value = None + mock_create_task.return_value = None + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + + s3_object_key = "My Team/2025-09-14/test-key.json" + test_element = s3BatchLoggingElement( + s3_object_key=s3_object_key, + payload={"test": "data"}, + s3_object_download_filename="test-file.json", + ) + + s3_logger = S3Logger( + s3_bucket_name="test-bucket", + s3_endpoint_url="https://s3.amazonaws.com", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + s3_logger.async_httpx_client = AsyncMock() + s3_logger.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger.async_upload_data_to_s3(test_element)) + + call_args = s3_logger.async_httpx_client.put.call_args + assert call_args is not None + actual_url = call_args[0][0] + raw_url = f"https://s3.amazonaws.com/test-bucket/{s3_object_key}" + expected_url = requests.Request("PUT", raw_url).prepare().url + assert actual_url == expected_url + assert " " not in actual_url + @pytest.mark.asyncio async def test_async_log_event_skips_when_standard_logging_object_missing(): """ From e68cfaae0c1238d19a4944efb8af47c41dc949ce Mon Sep 17 00:00:00 2001 From: Christian Reynoso Hunter Date: Sat, 4 Apr 2026 22:40:56 -0300 Subject: [PATCH 005/220] fix(cache): Prevent "multiple values" error in get_cache_key (#20261) ## Problem When `get_cache_key(**kwargs)` is called with kwargs that already contains `preset_cache_key` (which can happen when cache key is recomputed in certain code paths), the call to `_set_preset_cache_key_in_kwargs()` fails with: ``` TypeError: _set_preset_cache_key_in_kwargs() got multiple values for keyword argument 'preset_cache_key' ``` This is because `preset_cache_key` is passed both explicitly: ```python self._set_preset_cache_key_in_kwargs( preset_cache_key=hashed_cache_key, **kwargs ) ``` And implicitly via `**kwargs` unpacking when `kwargs["preset_cache_key"]` exists. ## Solution Filter out `preset_cache_key` from kwargs before passing to `_set_preset_cache_key_in_kwargs()`: ```python kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} self._set_preset_cache_key_in_kwargs( preset_cache_key=hashed_cache_key, **kwargs_for_preset ) ``` ## Testing Added unit tests covering: - kwargs with existing preset_cache_key (the bug case) - kwargs without preset_cache_key (regression test) - Verification that preset_cache_key is correctly set in litellm_params --- litellm/caching/caching.py | 5 +- tests/local_testing/test_cache_preset_key.py | 87 ++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 tests/local_testing/test_cache_preset_key.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 406a4f8c98a..6a68ba8c4d1 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -312,8 +312,11 @@ class Cache: verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError + # when kwargs already contains preset_cache_key from upstream callers + kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} self._set_preset_cache_key_in_kwargs( - preset_cache_key=hashed_cache_key, **kwargs + preset_cache_key=hashed_cache_key, **kwargs_for_preset ) return hashed_cache_key diff --git a/tests/local_testing/test_cache_preset_key.py b/tests/local_testing/test_cache_preset_key.py new file mode 100644 index 00000000000..de0ec05603c --- /dev/null +++ b/tests/local_testing/test_cache_preset_key.py @@ -0,0 +1,87 @@ +""" +Test for preset_cache_key multiple values bug fix. + +This test verifies that get_cache_key doesn't raise TypeError when kwargs +already contains preset_cache_key. + +Issue: When get_cache_key(**kwargs) is called with kwargs containing +preset_cache_key, the call to _set_preset_cache_key_in_kwargs() would fail with: + TypeError: got multiple values for keyword argument 'preset_cache_key' +""" + +import pytest +from unittest.mock import MagicMock, patch + + +class TestPresetCacheKeyFix: + """Tests for the preset_cache_key multiple values fix.""" + + def test_get_cache_key_with_preset_cache_key_in_kwargs(self): + """ + Test that get_cache_key handles kwargs that already contain preset_cache_key. + + This was causing: + TypeError: _set_preset_cache_key_in_kwargs() got multiple values + for keyword argument 'preset_cache_key' + """ + from litellm.caching.caching import Cache + + cache = Cache() + + # Simulate kwargs that already has preset_cache_key (as can happen + # when the cache key is recomputed in certain code paths) + kwargs_with_preset = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "preset_cache_key": "existing_key_12345", # This caused the bug + "litellm_params": {}, + } + + # This should NOT raise TypeError + try: + result = cache.get_cache_key(**kwargs_with_preset) + assert result is not None + assert isinstance(result, str) + except TypeError as e: + if "multiple values for keyword argument" in str(e): + pytest.fail(f"Bug not fixed: {e}") + raise + + def test_get_cache_key_without_preset_cache_key(self): + """Test normal case without preset_cache_key in kwargs still works.""" + from litellm.caching.caching import Cache + + cache = Cache() + + kwargs_normal = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": {}, + } + + result = cache.get_cache_key(**kwargs_normal) + assert result is not None + assert isinstance(result, str) + + def test_preset_cache_key_is_set_in_litellm_params(self): + """Verify that preset_cache_key is correctly set in litellm_params.""" + from litellm.caching.caching import Cache + + cache = Cache() + + litellm_params = {} + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": litellm_params, + } + + result = cache.get_cache_key(**kwargs) + + # The method should set preset_cache_key in litellm_params + assert "preset_cache_key" in litellm_params + assert litellm_params["preset_cache_key"] == result + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From fc75380b88481bff17a0b25d6c7d7cf49e11c361 Mon Sep 17 00:00:00 2001 From: Dmitry Date: Sun, 5 Apr 2026 04:46:43 +0300 Subject: [PATCH 006/220] fix(presidio): use correct text positions in anonymize_text (#24998) * fix(presidio): use correct text positions in anonymize_text (#24160) The Presidio anonymizer endpoint returns items with start/end positions that reference the *anonymized output* text, not the original input. anonymize_text() was applying these positions to the original text, causing garbled output with remnants of un-masked PII data. When output_parse_pii is False, return redacted_text["text"] directly from the anonymizer response instead of manually splicing. When output_parse_pii is True, use analyze_results positions (which correctly reference the original text) to build numbered replacement tokens and the pii_tokens mapping. * address review: remove dead code, fix token numbering order - Remove unused `anon_item_by_entity` dict (Greptile P2) - Number tokens left-to-right ( first in text, not last) - Add assertion for token numbering order in test --- .../guardrails/guardrail_hooks/presidio.py | 106 ++++++----- .../guardrail_hooks/test_presidio.py | 168 ++++++++++++++++++ 2 files changed, 226 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0f4ebbd4880..e048ca21cba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -485,61 +485,71 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): redacted_text = await response.json() - new_text = text if redacted_text is not None: verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - # Process items in reverse order by start position so that - # replacing later spans first does not shift earlier coordinates. - for item in sorted( - redacted_text["items"], key=lambda x: x["start"], reverse=True - ): - start = item["start"] - end = item["end"] - replacement = item["text"] # replacement token - if item["operator"] == "replace" and output_parse_pii is True: - if request_data is None: - verbose_proxy_logger.warning( - "Presidio anonymize_text called without request_data — " - "PII tokens cannot be stored per-request. " - "This may indicate a missing caller update." + + if not output_parse_pii: + # No need to build numbered tokens — just use Presidio's + # already-anonymized text directly. The old code incorrectly + # applied anonymizer item positions (which reference the + # *output* text) to the *original* text, causing offset errors. + for item in redacted_text.get("items", []): + entity_type = item.get("entity_type", None) + if entity_type is not None: + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 ) - request_data = {} - # Store pii_tokens in metadata to avoid leaking to LLM providers. - # Providers like Anthropic reject unknown top-level fields. - if not request_data.get("metadata"): - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] + return redacted_text["text"] - # Append a sequential number to make each token unique - # per request, so unmasking maps back to the correct - # original value. Format: , - # This is LLM-friendly and degrades gracefully if the - # LLM doesn't echo the token verbatim. - seq = len(pii_tokens) + 1 - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" + # output_parse_pii is True — we need sequentially numbered + # tokens and a pii_tokens mapping for later unmasking. + # Use analyze_results positions (which reference the ORIGINAL + # text) instead of anonymizer items (which reference the output). + new_text = text + if request_data is None: + verbose_proxy_logger.warning( + "Presidio anonymize_text called without request_data — " + "PII tokens cannot be stored per-request. " + "This may indicate a missing caller update." + ) + request_data = {} + if not request_data.get("metadata"): + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] - # Use ORIGINAL text (not new_text) since start/end - # reference the original text's coordinates. - pii_tokens[replacement] = text[start:end] + # Assign sequence numbers in forward (left-to-right) order so + # that is the first entity in the text, etc. + sorted_forward = sorted( + analyze_results, key=lambda x: x["start"] + ) + seq_map = {} + for idx, ar in enumerate(sorted_forward, start=1): + seq_map[(ar["start"], ar["end"])] = idx + # Apply replacements in reverse order by start position so + # that replacing later spans first does not shift earlier + # coordinates in the original text. + for ar in reversed(sorted_forward): + start = ar["start"] + end = ar["end"] + entity_type = ar["entity_type"] + replacement = f"<{entity_type}>" + + seq = seq_map[(start, end)] + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" + + pii_tokens[replacement] = text[start:end] new_text = new_text[:start] + replacement + new_text[end:] - entity_type = item.get("entity_type", None) - if entity_type is not None: - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) - # When output_parse_pii is True, new_text contains sequentially - # numbered tokens (e.g. ) that match the keys - # in pii_tokens. Returning redacted_text["text"] (Presidio's - # original output) would send un-numbered tokens to the LLM, - # making unmasking impossible. - # When output_parse_pii is False, new_text == redacted_text["text"] - # because no suffix is appended. + + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 + ) + return new_text else: raise Exception("Invalid anonymizer response: received None") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 32a8c1b1070..38ea42285c1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2230,3 +2230,171 @@ async def test_apply_to_output_streaming_bytes_only_logs_warning(): mock_logger.warning.assert_called_once() warning_msg = mock_logger.warning.call_args[0][0] assert "Output PII masking was skipped" in warning_msg + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_no_parse_pii(): + """ + Regression test for anonymizer offset bug (fixes #24160). + + The Presidio anonymizer returns items with start/end positions that + reference the *anonymized output* text, not the original input text. + When output_parse_pii is False, anonymize_text must return + redacted_text["text"] directly instead of manually splicing the + original text using those positions, which produces garbled output + with remnants of original PII data. + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + # Positions as returned by the analyzer (reference original text) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + # Anonymizer response — positions reference the *anonymized* text + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=False, + masked_entity_count=masked_entity_count, + ) + + expected = "My name is , my email is , phone " + assert result == expected, ( + f"anonymize_text produced garbled output with PII remnants.\n" + f"Expected: {expected!r}\n" + f"Got: {result!r}" + ) + assert masked_entity_count == { + "PERSON": 1, + "EMAIL_ADDRESS": 1, + "PHONE_NUMBER": 1, + } + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_with_parse_pii(): + """ + Regression test for anonymizer offset bug with output_parse_pii=True + (fixes #24160). + + When output_parse_pii is True, anonymize_text must use positions from + analyze_results (which reference the original text) to build numbered + tokens and the pii_tokens mapping, not positions from anonymizer items + (which reference the anonymized output text). + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + output_parse_pii=True, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + request_data = {"metadata": {}} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=True, + masked_entity_count=masked_entity_count, + request_data=request_data, + ) + + # Result must not contain any remnants of original PII + assert "John" not in result + assert "john@example.com" not in result + assert "555-867-5309" not in result + + # pii_tokens must map numbered tokens back to correct original values + pii_tokens = request_data["metadata"]["pii_tokens"] + token_values = set(pii_tokens.values()) + assert "John Smith" in token_values + assert "john@example.com" in token_values + assert "555-867-5309" in token_values + + # Tokens must be numbered in left-to-right order of appearance: + # PERSON (pos 11) → _1, EMAIL_ADDRESS (pos 35) → _2, PHONE_NUMBER (pos 59) → _3 + assert pii_tokens.get("") == "John Smith" + assert pii_tokens.get("") == "john@example.com" + assert pii_tokens.get("") == "555-867-5309" From 23e702dae609aa79682b971967dcf0bc05efab59 Mon Sep 17 00:00:00 2001 From: Bohdan Kulinich Date: Sun, 5 Apr 2026 04:48:39 +0300 Subject: [PATCH 007/220] feat(prometheus): add 7m and 10m latency histogram buckets (#25071) Extend LATENCY_BUCKETS beyond 5 minutes so request/LLM latency metrics can distinguish long runs up to the typical default LLM request timeout. Made-with: Cursor --- litellm/types/integrations/prometheus.py | 2 ++ .../types/test_prometheus_latency_buckets.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 tests/test_litellm/types/test_prometheus_latency_buckets.py diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 0d1501664b9..35b695cd054 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -156,6 +156,8 @@ LATENCY_BUCKETS = ( 180.0, 240.0, 300.0, + 420.0, # 7 minutes + 600.0, # 10 minutes (typical default LLM request timeout) float("inf"), ) diff --git a/tests/test_litellm/types/test_prometheus_latency_buckets.py b/tests/test_litellm/types/test_prometheus_latency_buckets.py new file mode 100644 index 00000000000..85670bb0b74 --- /dev/null +++ b/tests/test_litellm/types/test_prometheus_latency_buckets.py @@ -0,0 +1,17 @@ +"""LATENCY_BUCKETS covers long-running LLM calls (histograms are in seconds).""" + +import math + +from litellm.types.integrations.prometheus import LATENCY_BUCKETS + + +def test_latency_buckets_include_seven_and_ten_minutes(): + """Buckets beyond 5 min so histograms resolve requests up to default LLM timeouts.""" + assert 300.0 in LATENCY_BUCKETS + assert 420.0 in LATENCY_BUCKETS # 7 min + assert 600.0 in LATENCY_BUCKETS # 10 min + assert math.isinf(LATENCY_BUCKETS[-1]) + idx_300 = LATENCY_BUCKETS.index(300.0) + idx_420 = LATENCY_BUCKETS.index(420.0) + idx_600 = LATENCY_BUCKETS.index(600.0) + assert idx_300 < idx_420 < idx_600 From 4ca368923054ef1df73c1f3b815e39716ff2ce71 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 5 Apr 2026 01:15:08 -0700 Subject: [PATCH 008/220] chore: fixes --- .../workflows/run_llm_translation_tests.py | 0 .trivyignore | 12 - ci_cd/.grype.yaml | 36 --- ci_cd/security_scans.sh | 261 ------------------ docs/my-website/.trivyignore | 7 - ui/litellm-dashboard/.trivyignore | 7 - 6 files changed, 323 deletions(-) mode change 100755 => 100644 .github/workflows/run_llm_translation_tests.py delete mode 100644 .trivyignore delete mode 100644 ci_cd/.grype.yaml delete mode 100755 ci_cd/security_scans.sh delete mode 100644 docs/my-website/.trivyignore delete mode 100644 ui/litellm-dashboard/.trivyignore diff --git a/.github/workflows/run_llm_translation_tests.py b/.github/workflows/run_llm_translation_tests.py old mode 100755 new mode 100644 diff --git a/.trivyignore b/.trivyignore deleted file mode 100644 index 0d04ecacdb5..00000000000 --- a/.trivyignore +++ /dev/null @@ -1,12 +0,0 @@ -# LiteLLM Trivy Ignore File -# CVEs listed here are temporarily allowlisted pending fixes - -# Next.js vulnerabilities in UI dashboard (next@14.2.35) -# Allowlisted: 2026-01-31, 7-day fix timeline -# Fix: Upgrade to Next.js 15.5.10+ or 16.1.5+ - -# HIGH: DoS via request deserialization -GHSA-h25m-26qc-wcjf - -# MEDIUM: Image Optimizer DoS -CVE-2025-59471 diff --git a/ci_cd/.grype.yaml b/ci_cd/.grype.yaml deleted file mode 100644 index b9bc9db58f5..00000000000 --- a/ci_cd/.grype.yaml +++ /dev/null @@ -1,36 +0,0 @@ -ignore: - - vulnerability: CVE-2026-22184 - reason: no fixed zlib package is available yet in the Wolfi repositories, so this is ignored temporarily until an upstream release exists - # Wolfi base image: Python 3.13 and Node from apk have no fixed builds in Wolfi yet / not applicable - - vulnerability: CVE-2025-55130 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-59465 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-55131 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-59466 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2026-21637 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-55132 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: GHSA-hx9q-6w63-j58v - reason: orjson dumps recursion; allowlisted - - vulnerability: GHSA-73rr-hh4g-fpgx - reason: diff npm transitive dep; override in package.json, allowlisted - - vulnerability: CVE-2026-0865 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15282 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2026-0672 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15366 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15367 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-11468 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-12781 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2026-1299 - reason: Python 3.13 in Wolfi base; no fixed apk build yet diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh deleted file mode 100755 index 2138fca6cd5..00000000000 --- a/ci_cd/security_scans.sh +++ /dev/null @@ -1,261 +0,0 @@ -#!/bin/bash - -# Security Scans Script for LiteLLM -# This script runs comprehensive security scans including Trivy and Grype - -set -e - -echo "Starting security scans for LiteLLM..." - -# Function to install Trivy and required tools -install_trivy() { - echo "Installing Trivy and required tools..." - TRIVY_VERSION="0.35.0" - sudo apt-get update - sudo apt-get install -y wget jq curl bsdmainutils - wget -qO trivy.deb "https://github.com/aquasecurity/trivy/releases/download/v${TRIVY_VERSION}/trivy_${TRIVY_VERSION}_Linux-64bit.deb" - sudo dpkg -i trivy.deb - rm trivy.deb - echo "Trivy ${TRIVY_VERSION} installed successfully" -} - -# Function to install Grype -install_grype() { - echo "Installing Grype..." - curl -sSfL https://raw.githubusercontent.com/anchore/grype/main/install.sh | sudo sh -s -- -b /usr/local/bin - echo "Grype installed successfully" -} - -# Function to install ggshield -install_ggshield() { - echo "Installing ggshield..." - pip3 install --upgrade pip - pip3 install ggshield - echo "ggshield installed successfully" -} - -# # Function to run secret detection scans -# run_secret_detection() { -# echo "Running secret detection scans..." - -# if ! command -v ggshield &> /dev/null; then -# install_ggshield -# fi - -# # Check if GITGUARDIAN_API_KEY is set (required for CI/CD) -# if [ -z "$GITGUARDIAN_API_KEY" ]; then -# echo "Warning: GITGUARDIAN_API_KEY environment variable is not set." -# echo "ggshield requires a GitGuardian API key to scan for secrets." -# echo "Please set GITGUARDIAN_API_KEY in your CI/CD environment variables." -# exit 1 -# fi - -# echo "Scanning codebase for secrets..." -# echo "Note: Large codebases may take several minutes due to API rate limits (50 requests/minute on free plan)" -# echo "ggshield will automatically handle rate limits and retry as needed." -# echo "Binary files, cache files, and build artifacts are excluded via .gitguardian.yaml" - -# # Use --recursive for directory scanning and auto-confirm if prompted -# # .gitguardian.yaml will automatically exclude binary files, wheel files, etc. -# # GITGUARDIAN_API_KEY environment variable will be used for authentication -# echo y | ggshield secret scan path . --recursive || { -# echo "" -# echo "==========================================" -# echo "ERROR: Secret Detection Failed" -# echo "==========================================" -# echo "ggshield has detected secrets in the codebase." -# echo "Please review discovered secrets above, revoke any actively used secrets" -# echo "from underlying systems and make changes to inject secrets dynamically at runtime." -# echo "" -# echo "For more information, see: https://docs.gitguardian.com/secrets-detection/" -# echo "==========================================" -# echo "" -# exit 1 -# } - -# echo "Secret detection scans completed successfully" -# } - -# Function to run Trivy scans -run_trivy_scans() { - echo "Running Trivy scans..." - - echo "Scanning LiteLLM Docs..." - trivy fs --ignorefile .trivyignore --scanners vuln --dependency-tree --exit-code 1 --severity HIGH,CRITICAL,MEDIUM ./docs/ - - echo "Scanning LiteLLM UI..." - trivy fs --ignorefile .trivyignore --scanners vuln --dependency-tree --exit-code 1 --severity HIGH,CRITICAL,MEDIUM ./ui/ - - echo "Trivy scans completed successfully" -} - -# Function to build and scan Docker images with Grype -run_grype_scans() { - echo "Running Grype scans..." - - # Temporarily add wheel files to .dockerignore for security scans - echo "Temporarily modifying .dockerignore to exclude problematic wheel files..." - cp .dockerignore .dockerignore.backup 2>/dev/null || touch .dockerignore.backup - echo "/*.whl" >> .dockerignore - - # Build and scan Dockerfile.database - echo "Building and scanning Dockerfile.database..." - docker build --no-cache -t litellm-database:latest -f ./docker/Dockerfile.database . - grype litellm-database:latest --config ci_cd/.grype.yaml --fail-on critical - - # Build and scan main Dockerfile - echo "Building and scanning main Dockerfile..." - docker build --no-cache -t litellm:latest . - grype litellm:latest --config ci_cd/.grype.yaml --fail-on critical - - # Restore original .dockerignore - echo "Restoring original .dockerignore..." - mv .dockerignore.backup .dockerignore - - # Scan the locally built LiteLLM image for vulnerabilities with CVSS >= 4.0 - echo "Scanning locally built LiteLLM image for high-severity vulnerabilities..." - echo "Using locally built image: litellm:latest" - - # Allowlist of CVEs to be ignored in failure threshold/reporting - # - CVE-2025-8869: Not applicable on Python >=3.13 (PEP 706 implemented); pip fallback unused; no OS-level fix - # - GHSA-4xh5-x5gv-qwph: GitHub Security Advisory alias for CVE-2025-8869 - # - GHSA-5j98-mcp5-4vw2: glob CLI command injection via -c/--cmd; glob CLI is not used in the litellm runtime image, - # and the vulnerable versions are pulled in only via OS-level/node tooling outside of our application code - ALLOWED_CVES=( - "CVE-2025-8869" - "GHSA-4xh5-x5gv-qwph" - "CVE-2025-8291" # no fix available as of Oct 11, 2025 - "GHSA-5j98-mcp5-4vw2" - "CVE-2025-13836" # Python 3.13 HTTP response reading OOM/DoS - no fix available in base image - "CVE-2025-12084" # Python 3.13 xml.dom.minidom quadratic algorithm - no fix available in base image - "CVE-2025-60876" # BusyBox wget HTTP request splitting - no fix available in Chainguard Wolfi base image - "CVE-2026-0861" # Wolfi glibc still flagged even on 2.42-r5; upstream patched build unavailable yet - "CVE-2010-4756" # glibc glob DoS - awaiting patched Wolfi glibc build - "CVE-2019-1010022" # glibc stack guard bypass - awaiting patched Wolfi glibc build - "CVE-2019-1010023" # glibc ldd remap issue - awaiting patched Wolfi glibc build - "CVE-2019-1010024" # glibc ASLR mitigation bypass - awaiting patched Wolfi glibc build - "CVE-2019-1010025" # glibc pthread heap address leak - awaiting patched Wolfi glibc build - "CVE-2026-22184" # zlib untgz buffer overflow - untgz unused + no fixed Wolfi build yet - "GHSA-58pv-8j8x-9vj2" # jaraco.context path traversal - setuptools vendored only (v5.3.0), not used in application code (using v6.1.0+) - "GHSA-34x7-hfp2-rc4v" # node-tar hardlink path traversal - not applicable, tar CLI not exposed in application code - "GHSA-r6q2-hw4h-h46w" # node-tar not used by application runtime, Linux-only container, not affect by macOS APFS-specific exploit - "GHSA-8rrh-rw8j-w5fx" # wheel is from chainguard and will be handled by then TODO: Remove this after Chainguard updates the wheel - "CVE-2025-59465" # Node only used for Admin UI build/prisma - "CVE-2025-55131" # Node only used for Admin UI build/prisma - "CVE-2025-59466" # Node only used for Admin UI build/prisma - "CVE-2025-55130" # Node only used for Admin UI build/prisma - "CVE-2025-59467" # Node only used for Admin UI build/prisma - "CVE-2026-21637" # Node only used for Admin UI build/prisma - "CVE-2025-55132" # Node only used for Admin UI build/prisma - "GHSA-hx9q-6w63-j58v" # orjson dumps recursion; allowlisted - "CVE-2025-15281" # No fix available yet - "CVE-2026-0865" # No fix available yet - "CVE-2025-15282" # No fix available yet - "CVE-2026-0672" # No fix available yet - "CVE-2025-15366" # No fix available yet - "CVE-2025-15367" # No fix available yet - "CVE-2025-12781" # No fix available yet - "CVE-2025-11468" # No fix available yet - "CVE-2026-1299" # Python 3.13 email module header injection - not applicable, LiteLLM doesn't use BytesGenerator for email serialization - "CVE-2026-0775" # npm cli incorrect permission assignment - no fix available yet, npm is only used at build/prisma-generate time - "GHSA-3ppc-4f35-3m26" # minimatch ReDoS via repeated wildcards - from nodejs_wheel bundled npm, not used in application runtime code - "GHSA-83g3-92jg-28cx" # tar arbitrary file read/write via hardlink - from nodejs_wheel bundled npm, not used in application runtime code - "CVE-2026-2297" # Python 3.13 SourcelessFileLoader audit hook bypass - no fix available in base image - "GHSA-qffp-2rhf-9h96" # tar hardlink path traversal - from nodejs_wheel bundled npm, not used in application runtime code - "CVE-2026-2673" # OpenSSL 3.6.1 TLS 1.3 key exchange group negotiation issue - no fix available yet - "CVE-2026-3644" # Python 3.13 vulnerability - no fix available in base image - "CVE-2026-4224" # Python 3.13 Expat parser stack overflow in ElementDeclHandler - no fix available in base image - ) - - # Build JSON array of allowlisted CVE IDs for jq - ALLOWED_IDS_JSON=$(printf '%s\n' "${ALLOWED_CVES[@]}" | jq -R . | jq -s .) - - echo "Checking for vulnerabilities with CVSS score >= 4.0..." - echo "Allowlisted CVEs (ignored in threshold): ${ALLOWED_CVES[*]}" - echo "" - - # Show all high-severity vulnerabilities for transparency - TOTAL_HIGH_SEVERITY=$(grype litellm:latest -o json | jq -r ' - .matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | .vulnerability.id' | wc -l) - - if [ "$TOTAL_HIGH_SEVERITY" -gt 0 ]; then - echo "Total vulnerabilities found with CVSS >= 4.0: $TOTAL_HIGH_SEVERITY" - echo "" - echo "All high-severity vulnerabilities (including allowlisted):" - grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - ["Package", "Version", "Vulnerability ID", "CVSS Score", "Allowlisted"], - (.matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | [.artifact.name, .artifact.version, .vulnerability.id, .vulnerability.cvss[0].metrics.baseScore, (if (.vulnerability.id as $id | $allow | index($id)) then "YES" else "NO" end)]) - | @tsv' | column -t -s $'\t' - echo "" - fi - - HIGH_SEVERITY_COUNT=$(grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - .matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | select((.vulnerability.id as $id | $allow | index($id) | not)) - | .vulnerability.id' | wc -l) - - if [ "$HIGH_SEVERITY_COUNT" -gt 0 ]; then - echo "" - echo "==========================================" - echo "ERROR: Security Scan Failed" - echo "==========================================" - echo "Found $HIGH_SEVERITY_COUNT non-allowlisted vulnerabilities with CVSS score >= 4.0 in litellm:latest" - echo "" - echo "These vulnerabilities are NOT in the allowlist and must be addressed." - echo "Current allowlisted CVEs: ${ALLOWED_CVES[*]}" - echo "" - echo "Detailed vulnerability report:" - echo "" - grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - ["Package", "Version", "Vulnerability ID", "CVSS Score", "Severity", "Fix Version", "Description"], - (.matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | select((.vulnerability.id as $id | $allow | index($id) | not)) - | [.artifact.name, .artifact.version, .vulnerability.id, .vulnerability.cvss[0].metrics.baseScore, .vulnerability.severity, (.vulnerability.fix.versions[0] // "No fix available"), .vulnerability.description]) - | @tsv' | column -t -s $'\t' - echo "" - echo "==========================================" - echo "Action Required:" - echo "==========================================" - echo "1. If a fix is available, update the package to the fixed version" - echo "2. If the vulnerability is not applicable or has no fix:" - echo " - Add the CVE/GHSA ID to ALLOWED_CVES array in ci_cd/security_scans.sh" - echo " - Add a comment explaining why it's safe to ignore" - echo "" - echo "Note: Some vulnerabilities may have multiple IDs (CVE-XXXX and GHSA-XXXX)." - echo "Add all relevant IDs to the allowlist if they refer to the same issue." - echo "==========================================" - echo "" - exit 1 - else - echo "No high-severity vulnerabilities (CVSS >= 4.0) found in litellm:latest" - fi - - echo "Grype scans completed successfully" -} - -# Main execution -main() { - echo "Installing security scanning tools..." - install_trivy - install_grype - - # echo "Running secret detection scans..." - # run_secret_detection - - echo "Running filesystem vulnerability scans..." - run_trivy_scans - - echo "Running Docker image vulnerability scans..." - run_grype_scans - - echo "All security scans completed successfully!" -} - -# Execute main function -main "$@" diff --git a/docs/my-website/.trivyignore b/docs/my-website/.trivyignore deleted file mode 100644 index 977504f2670..00000000000 --- a/docs/my-website/.trivyignore +++ /dev/null @@ -1,7 +0,0 @@ -# js-yaml CVE-2025-64718 -# This vulnerability is not applicable because we've forced js-yaml to version 4.1.1 -# via npm overrides in package.json. Trivy incorrectly reports this based on -# dependency requirements in the lockfile, but the actual installed version is 4.1.1. -# Verified with: npm list js-yaml -CVE-2025-64718 - diff --git a/ui/litellm-dashboard/.trivyignore b/ui/litellm-dashboard/.trivyignore deleted file mode 100644 index 977504f2670..00000000000 --- a/ui/litellm-dashboard/.trivyignore +++ /dev/null @@ -1,7 +0,0 @@ -# js-yaml CVE-2025-64718 -# This vulnerability is not applicable because we've forced js-yaml to version 4.1.1 -# via npm overrides in package.json. Trivy incorrectly reports this based on -# dependency requirements in the lockfile, but the actual installed version is 4.1.1. -# Verified with: npm list js-yaml -CVE-2025-64718 - From f233520c44061a07e3c43ba34ba6f79c73cc0bae Mon Sep 17 00:00:00 2001 From: Hendrik Jaks Date: Mon, 6 Apr 2026 21:13:06 +0300 Subject: [PATCH 009/220] fix(ui): resolve login redirect loop when reverse proxy adds HttpOnly to cookies (#23532) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ui): resolve login redirect loop when reverse proxy adds HttpOnly to cookies When LiteLLM is behind nginx-ingress or similar with security-hardened configs, the reverse proxy adds HttpOnly to all Set-Cookie headers. This makes the JWT token unreadable by JavaScript, causing an infinite login redirect loop. Fix by returning the JWT token in the /v2/login response body so the frontend can set a JS-accessible cookie directly. Fixes #19663 Co-Authored-By: Claude Opus 4.6 * fix: address Greptile review feedback - Add window guard to setTokenCookie for SSR consistency with clearTokenCookies - Add SSR test for window undefined case - Add code comment explaining why JWT is included in response body Co-Authored-By: Claude Opus 4.6 * fix: address second round of Greptile review feedback - Add loginCall integration tests verifying setTokenCookie is called with token and skipped when absent (backward-compatibility path) - Use encodeURIComponent/decodeURIComponent in setTokenCookie/getCookie for defense-in-depth against non-standard token formats Co-Authored-By: Claude Opus 4.6 * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(ui): use sessionStorage instead of cookie for login token storage Replace setTokenCookie (which is a no-op when reverse proxy adds HttpOnly) with storeLoginToken using sessionStorage. Add sessionStorage fallback to getCookie so the token is found even when the cookie is HttpOnly. Also handle '=' in cookie values with .slice(1).join("=") and clear sessionStorage on logout. Co-Authored-By: Claude Opus 4.6 * fix(ui): use shared getCookie in page.tsx and user_dashboard.tsx Replace local getCookie functions in page.tsx and user_dashboard.tsx with the shared one from cookieUtils that has the sessionStorage fallback. Without this, the HttpOnly cookie fix was incomplete — page.tsx (the dashboard entry point) could not read the token, causing the redirect loop to persist. Also scope the sessionStorage fallback to the "token" key only, and clear sessionStorage in page.tsx deleteCookie. Co-Authored-By: Claude Opus 4.6 * fix(ui): scope deleteCookie sessionStorage cleanup to token key only Also document the sessionStorage cross-tab trade-off: per-tab scope means users behind an HttpOnly proxy must log in once per tab, but this is intentional to avoid localStorage XSS exposure. Co-Authored-By: Claude Opus 4.6 * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * style: remove stray double blank line in user_dashboard.tsx Co-Authored-By: Claude Opus 4.6 * fix(ui): guard storeLoginToken against empty/whitespace-only tokens Co-Authored-By: Claude Opus 4.6 * fix(ui): preserve sessionStorage token across beforeunload clear The existing beforeunload handler calls sessionStorage.clear() to flush cached UI data on page refresh. This also wiped the token stored by storeLoginToken, re-introducing the redirect loop after any page refresh in the HttpOnly proxy scenario. Now the token is saved and restored across the clear. Co-Authored-By: Claude Opus 4.6 * fix(ui): set JS-accessible cookie at /ui path as HttpOnly workaround sessionStorage alone is unreliable. Also set the token via document.cookie at path=/ui — nginx only adds HttpOnly to server-set Set-Cookie headers, so a JS-set cookie is always readable. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(ui): use dynamic cookie path based on server_root_path Hardcoded path=/ui breaks when LiteLLM is deployed with a custom server_root_path. Now derives the cookie path from serverRootPath so it works at /ui, /myapp/ui, etc. Also reuse clearTokenCookies() in deleteCookie() to avoid duplication. Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(ui): remove circular dependency in cookieUtils.ts Derive the UI cookie path from window.location.pathname instead of importing serverRootPath from networking.tsx. This breaks the cookieUtils → networking → cookieUtils cycle that could cause serverRootPath to be undefined under certain bundler configurations. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(ui): harden getUiCookiePath regex and add missing tests - Use regex /\/ui(?=\/|$)/ to match "/ui" only as a full path segment, preventing false matches on paths like "/my-ui-tool/login". - Add unit tests for storeLoginToken empty/whitespace guard and cookie-at-/ui-path behavior. Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Black formatting in audit_logs.py Co-Authored-By: Claude Opus 4.6 (1M context) * fix CI: formatting, test params, remove token from login JSON Co-Authored-By: Claude Opus 4.6 (1M context) * fix: reformat with Black 23.x to match CI Co-Authored-By: Claude Opus 4.6 (1M context) * fix: keep token in login JSON body for UI storeLoginToken flow Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use storeLoginToken in exchangeLoginCode, add credentials include Co-Authored-By: Claude Opus 4.6 (1M context) * revert: remove unrelated changes from HttpOnly cookie fix branch Reset files not related to the login cookie fix back to main: - prometheus.py, bedrock converse, guardrail handler - auth_checks.py, reset_budget_job.py, audit_logs.py - test_user_api_key_auth.py Co-Authored-By: Claude Opus 4.6 (1M context) * Revert "revert: remove unrelated changes from HttpOnly cookie fix branch" This reverts commit 0684a1e27521ada35cf8a4afbdff7aeaa87ff41d. * Revert "fix: use storeLoginToken in exchangeLoginCode, add credentials include" This reverts commit 866405f443eb2004ee56ed2f52c9047593d8cc6e. * Revert "fix: keep token in login JSON body for UI storeLoginToken flow" This reverts commit 086c41640c5749399f11aba298321d1d16e10a92. * Revert "fix: reformat with Black 23.x to match CI" This reverts commit b2c3334c888a9dbb60fd94cce12f43e8f3aa7e82. * Revert "fix CI: formatting, test params, remove token from login JSON" This reverts commit 2905d47bd4013b81fbaa01cf2c7402d51360f1e6. --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +- ui/litellm-dashboard/src/app/page.tsx | 16 +--- .../src/components/networking.test.ts | 33 ++++++++ .../src/components/networking.tsx | 6 +- .../src/components/user_dashboard.tsx | 12 ++- .../src/utils/cookieUtils.test.ts | 73 ++++++++++++++++- ui/litellm-dashboard/src/utils/cookieUtils.ts | 81 ++++++++++++++++++- 7 files changed, 198 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a2..accf669f2e3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11506,8 +11506,11 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" + # Token is included in the response body so the UI can set a JS-accessible + # cookie even when a reverse proxy (e.g. nginx-ingress) adds HttpOnly to the + # server-set cookie, which would otherwise cause an infinite login redirect. json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui}, + content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd41..73c27c00724 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -42,6 +42,7 @@ import ToolPoliciesView from "@/components/ToolPoliciesView"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; import { formatUserRole, isAdminRole } from "@/utils/roles"; @@ -51,21 +52,12 @@ import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; -function getCookie(name: string) { - // Safer cookie read + decoding; handles '=' inside values - const match = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - if (!match) return null; - const value = match.slice(name.length + 1); - try { - return decodeURIComponent(value); - } catch { - return value; - } -} - function deleteCookie(name: string, path = "/") { // Best-effort client-side clear (works for non-HttpOnly cookies without Domain) document.cookie = `${name}=; Max-Age=0; Path=${path}`; + if (name === "token") { + clearTokenCookies(); + } } interface ProxySettings { diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index c57dcb97eb9..3c107fa586f 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -5,6 +5,7 @@ import * as Networking from "./networking"; vi.mock("@/utils/cookieUtils", () => ({ clearTokenCookies: vi.fn(), getCookie: vi.fn(), + storeLoginToken: vi.fn(), })); vi.mock("./molecules/notifications_manager", () => ({ @@ -79,6 +80,38 @@ describe("networking - expired session handling", () => { }); }); +describe("loginCall - storeLoginToken integration", () => { + const originalFetch = global.fetch; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it("calls storeLoginToken when response includes token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success", token: "my-jwt" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).toHaveBeenCalledWith("my-jwt"); + }); + + it("does not call storeLoginToken when response has no token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).not.toHaveBeenCalled(); + }); +}); + describe("daily activity helpers", () => { const startTime = new Date("2025-02-12T00:00:00.000Z"); const endTime = new Date("2025-02-19T00:00:00.000Z"); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 28f8d308de7..16f35605877 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -69,7 +69,7 @@ export const getInProductNudgesCall = async (accessToken: string) => { * Helper file for calls being made to proxy */ import MessageManager from "@/components/molecules/message_manager"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, storeLoginToken } from "@/utils/cookieUtils"; import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types"; import { Team } from "./key_team_helpers/key_list"; import { UserInfo } from "./view_users/types"; @@ -9255,14 +9255,14 @@ export const loginCall = async (username: string, password: string, useV3?: bool const exchangeData: LoginResponse = await exchangeResponse.json(); if (exchangeData.token) { - document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`; + storeLoginToken(exchangeData.token); } return exchangeData; } // Backwards compatibility: v2 or old v3 returns token directly if (data.token) { - document.cookie = `token=${data.token}; path=/; SameSite=Lax`; + storeLoginToken(data.token); } return data; diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index f97d8ffab04..90eac56540d 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,5 +1,5 @@ "use client"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { Col, Grid } from "@tremor/react"; import { Typography } from "antd"; import { jwtDecode } from "jwt-decode"; @@ -35,12 +35,6 @@ export type UserInfo = { spend: number; }; -function getCookie(name: string) { - console.log("COOKIES", document.cookie); - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; -} - interface UserDashboardProps { userID: string | null; userRole: string | null; @@ -103,7 +97,11 @@ const UserDashboard: React.FC = ({ // They are only cleared on logout useEffect(() => { const handleBeforeUnload = () => { + const token = sessionStorage.getItem("token"); sessionStorage.clear(); + if (token) { + sessionStorage.setItem("token", token); + } }; window.addEventListener("beforeunload", handleBeforeUnload); return () => window.removeEventListener("beforeunload", handleBeforeUnload); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts index 8b066e6a8ea..c7bd27a6a85 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -1,11 +1,12 @@ import { describe, it, expect, beforeEach, vi } from "vitest"; -import { clearTokenCookies, getCookie } from "./cookieUtils"; +import { clearTokenCookies, getCookie, storeLoginToken } from "./cookieUtils"; describe("cookieUtils", () => { beforeEach(() => { document.cookie.split(";").forEach((c) => { document.cookie = c.replace(/^ +/, "").replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); }); + sessionStorage.clear(); vi.spyOn(console, "log").mockImplementation(() => {}); }); @@ -116,6 +117,55 @@ describe("cookieUtils", () => { vi.restoreAllMocks(); }); + + it("should clear sessionStorage token", () => { + sessionStorage.setItem("token", "stored-token"); + clearTokenCookies(); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + }); + + describe("storeLoginToken", () => { + it("should store the token in sessionStorage", () => { + storeLoginToken("my-jwt-token"); + expect(sessionStorage.getItem("token")).toBe("my-jwt-token"); + }); + + it("should overwrite an existing token in sessionStorage", () => { + storeLoginToken("old-token"); + expect(sessionStorage.getItem("token")).toBe("old-token"); + + storeLoginToken("new-token"); + expect(sessionStorage.getItem("token")).toBe("new-token"); + }); + + it("should not throw when window is undefined (server-side rendering)", () => { + const originalWindow = global.window; + delete (global as any).window; + + expect(() => storeLoginToken("token")).not.toThrow(); + + global.window = originalWindow; + }); + + it("should not store empty string token", () => { + storeLoginToken(""); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should not store whitespace-only token", () => { + storeLoginToken(" "); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should set a JS-accessible cookie at /ui path", () => { + const cookieSpy = vi.spyOn(document, "cookie", "set"); + storeLoginToken("my-jwt-token"); + expect(cookieSpy).toHaveBeenCalledWith( + expect.stringContaining("path=/ui") + ); + vi.restoreAllMocks(); + }); }); describe("getCookie", () => { @@ -141,5 +191,26 @@ describe("cookieUtils", () => { expect(getCookie("token")).toBe("token-value"); expect(getCookie("other")).toBe("other-value"); }); + + it("should handle values containing '=' characters", () => { + document.cookie = "token=abc=def=ghi; path=/"; + expect(getCookie("token")).toBe("abc=def=ghi"); + }); + + it("should fall back to sessionStorage when cookie is not found", () => { + sessionStorage.setItem("token", "session-stored-jwt"); + expect(getCookie("token")).toBe("session-stored-jwt"); + }); + + it("should prefer cookie over sessionStorage", () => { + document.cookie = "token=cookie-value; path=/"; + sessionStorage.setItem("token", "session-value"); + expect(getCookie("token")).toBe("cookie-value"); + }); + + it("should not fall back to sessionStorage for non-token keys", () => { + sessionStorage.setItem("other", "other-value"); + expect(getCookie("other")).toBeNull(); + }); }); }); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts index 01add36542c..b4493744ad4 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -2,6 +2,23 @@ * Utility functions for managing cookies */ +/** + * Returns the cookie path for the UI. + * Derives the path from window.location.pathname so it works when + * LiteLLM is deployed behind a subpath (e.g. /myapp/ui instead of /ui). + * No imports from networking.tsx to avoid circular dependencies. + */ +function getUiCookiePath(): string { + if (typeof window === "undefined") return "/ui"; + // Match "/ui" only as a full path segment (followed by "/" or end of string) + // to avoid false matches like "/my-ui-tool/login" → "/my-ui". + const match = window.location.pathname.match(/\/ui(?=\/|$)/); + if (match && match.index !== undefined) { + return window.location.pathname.substring(0, match.index + 3); + } + return "/ui"; +} + /** * Clears the token cookie from both root and /ui paths */ @@ -16,7 +33,8 @@ export function clearTokenCookies() { // Clear with various combinations of path and SameSite // Include current path in case of custom server root path const currentPath = window.location.pathname; - const paths = ["/", "/ui"]; + const uiCookiePath = getUiCookiePath(); + const paths = ["/", uiCookiePath]; // Add the current path directory if it's different from root and /ui if (currentPath && currentPath !== "/" && !currentPath.startsWith("/ui")) { @@ -43,7 +61,45 @@ export function clearTokenCookies() { }); }); - console.log("After clearing cookies:", document.cookie); + try { + sessionStorage.removeItem("token"); + } catch { + // sessionStorage may be unavailable + } + +} + +/** + * Stores the login token so the UI can read it even when a reverse proxy + * (e.g. nginx-ingress) adds HttpOnly to the server-set cookie. + * + * Strategy: + * 1. Set a JS-accessible cookie at path "/ui". Because nginx only modifies + * server-set Set-Cookie headers, a cookie created via document.cookie will + * never carry HttpOnly. Using path "/ui" avoids colliding with the + * server-set HttpOnly cookie at path "/". + * 2. Also store in sessionStorage as a secondary fallback. + */ +export function storeLoginToken(token: string) { + if (typeof window === "undefined") return; + if (!token || !token.trim()) return; + + // 1. JS-accessible cookie at /ui — survives same-tab navigations and + // is readable by getCookie() via document.cookie. + try { + const secure = window.location.protocol === "https:" ? "; Secure" : ""; + const cookiePath = getUiCookiePath(); + document.cookie = `token=${encodeURIComponent(token)}; path=${cookiePath}; SameSite=Lax${secure}`; + } catch { + // cookie setting may fail in restrictive environments + } + + // 2. sessionStorage backup + try { + sessionStorage.setItem("token", token); + } catch { + // sessionStorage may be unavailable (e.g. private browsing quota exceeded) + } } /** @@ -53,6 +109,23 @@ export function clearTokenCookies() { */ export function getCookie(name: string) { if (typeof document === "undefined") return null; - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; + const row = document.cookie.split("; ").find((r) => r.startsWith(name + "=")); + if (row) { + const raw = row.split("=").slice(1).join("="); + try { + return decodeURIComponent(raw); + } catch { + return raw; + } + } + // Fallback to sessionStorage — covers the case where a reverse proxy + // added HttpOnly to the server-set cookie, making it invisible to JS. + if (name === "token" && typeof window !== "undefined") { + try { + return sessionStorage.getItem(name); + } catch { + return null; + } + } + return null; } From 168b0a05c41bd39bc63322312cefd92d0ac7cbbc Mon Sep 17 00:00:00 2001 From: kothamah Date: Tue, 7 Apr 2026 16:25:03 -0400 Subject: [PATCH 010/220] added changes based on the feedback --- .../test_bedrock_guardrails.py | 3062 +++++++++-------- .../test_bedrock_guardrails.py | 476 +++ 2 files changed, 2085 insertions(+), 1453 deletions(-) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 03fe63b307a..146d16c242a 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1,201 +1,567 @@ -""" -Unit tests for Bedrock Guardrails -""" -import json -import os import sys -from unittest.mock import AsyncMock, MagicMock, patch - +import os +import io, asyncio import pytest -from fastapi import HTTPException -sys.path.insert(0, os.path.abspath("../../../../../..")) - -from litellm.proxy._types import UserAPIKeyAuth +sys.path.insert(0, os.path.abspath("../..")) +import litellm from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, _redact_pii_matches, ) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching import DualCache +from unittest.mock import MagicMock, AsyncMock, patch @pytest.mark.asyncio -async def test__redact_pii_matches_function(): - """Test the _redact_pii_matches function directly""" +async def test_bedrock_guardrails_pii_masking(): + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() - # Test case 1: Response with PII entities - response_with_pii = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ + guardrail = BedrockGuardrail( + guardrailIdentifier="wf0hkdb5x07f", + guardrailVersion="DRAFT", + ) + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, + {"role": "assistant", "content": "Hello, how can I help you today?"}, + {"role": "user", "content": "I need to cancel my order"}, { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "John Smith", "action": "BLOCKED"}, - { - "type": "US_SOCIAL_SECURITY_NUMBER", - "match": "324-12-3212", - "action": "BLOCKED", - }, - {"type": "PHONE", "match": "607-456-7890", "action": "BLOCKED"}, - ] - } - } - ], - "outputs": [{"text": "Input blocked by PII policy"}], - } - - # Call the redaction function - redacted_response = _redact_pii_matches(response_with_pii) - - # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - - assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" - assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" - assert pii_entities[2]["match"] == "[REDACTED]", "Phone should be redacted" - - # Verify other fields remain unchanged - assert pii_entities[0]["type"] == "NAME" - assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" - assert pii_entities[2]["type"] == "PHONE" - assert redacted_response["action"] == "GUARDRAIL_INTERVENED" - assert redacted_response["outputs"][0]["text"] == "Input blocked by PII policy" - - print("PII redaction function test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_no_pii(): - """Test _redact_pii_matches with response that has no PII""" - - response_no_pii = {"action": "NONE", "assessments": [], "outputs": []} - - # Call the redaction function - redacted_response = _redact_pii_matches(response_no_pii) - - # Should return the same response unchanged - assert redacted_response == response_no_pii - print("No PII redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_empty_assessments(): - """Test _redact_pii_matches with empty assessments""" - - response_empty_assessments = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [{"sensitiveInformationPolicy": {"piiEntities": []}}], - "outputs": [{"text": "Some output"}], - } - - # Call the redaction function - redacted_response = _redact_pii_matches(response_empty_assessments) - - # Should return the same response unchanged - assert redacted_response == response_empty_assessments - print("Empty assessments redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_malformed_response(): - """Test _redact_pii_matches with malformed response (should not crash)""" - - # Test with completely malformed response - malformed_response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": "not_a_list", # This should cause an exception - } - - # Should not crash and return original response - redacted_response = _redact_pii_matches(malformed_response) - assert redacted_response == malformed_response - - # Test with missing keys - missing_keys_response = { - "action": "GUARDRAIL_INTERVENED" - # Missing assessments key - } - - redacted_response = _redact_pii_matches(missing_keys_response) - assert redacted_response == missing_keys_response - - print("Malformed response redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_multiple_assessments(): - """Test _redact_pii_matches with multiple assessments containing PII""" - - response_multiple_assessments = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED", - } - ] - } - }, - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "CREDIT_DEBIT_CARD_NUMBER", - "match": "1234-5678-9012-3456", - "action": "BLOCKED", - }, - { - "type": "ADDRESS", - "match": "123 Main St, Anytown USA", - "action": "ANONYMIZED", - }, - ] - } + "role": "user", + "content": "ok, my credit card number is 1234-5678-9012-3456", }, ], - "outputs": [{"text": "Multiple PII detected"}], } - # Call the redaction function - redacted_response = _redact_pii_matches(response_multiple_assessments) + response = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + print("response after moderation hook", response) - # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] - - assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" - assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" - assert assessment2_pii[1]["match"] == "[REDACTED]", "Address should be redacted" - - # Verify types remain unchanged - assert assessment1_pii[0]["type"] == "EMAIL" - assert assessment2_pii[0]["type"] == "CREDIT_DEBIT_CARD_NUMBER" - assert assessment2_pii[1]["type"] == "ADDRESS" - - print("Multiple assessments redaction test passed") + if response: # Only assert if response is not None + assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" + assert response["messages"][1]["content"] == "Hello, how can I help you today?" + assert response["messages"][2]["content"] == "I need to cancel my order" + assert ( + response["messages"][3]["content"] + == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" + ) @pytest.mark.asyncio -async def test_bedrock_guardrail_logging_uses_redacted_response(): - """Test that the Bedrock guardrail uses redacted response for logging""" +async def test_bedrock_guardrails_pii_masking_content_list(): + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="wf0hkdb5x07f", + guardrailVersion="DRAFT", + ) + + request_data = { + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Hello, my phone number is +1 412 555 1212", + }, + {"type": "text", "text": "what time is it?"}, + ], + }, + {"role": "assistant", "content": "Hello, how can I help you today?"}, + {"role": "user", "content": "who is the president of the united states?"}, + ], + } + + response = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + print(response) + + if response: # Only assert if response is not None + # Verify that the list content is properly masked + assert isinstance(response["messages"][0]["content"], list) + assert ( + response["messages"][0]["content"][0]["text"] + == "Hello, my phone number is {PHONE}" + ) + assert response["messages"][0]["content"][1]["text"] == "what time is it?" + assert response["messages"][1]["content"] == "Hello, how can I help you today?" + assert ( + response["messages"][2]["content"] + == "who is the president of the united states?" + ) + + +@pytest.mark.asyncio +async def test_bedrock_guardrails_block_messages_api(): + """ + Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. + """ + from fastapi import HTTPException # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() + guardrail = BedrockGuardrail( + guardrailIdentifier="ff6ujrregl1q", + guardrailVersion="DRAFT", + ) + + request_data = { + "model": "claude-3-5-sonnet-20240620", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Hello, my phone number is +1 412 555 1212", + }, + {"type": "text", "text": "what time is it?"}, + ], + }, + {"role": "user", "content": "tell me about coffee"}, + ], + } + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="anthropic_messages", + cache=MagicMock(spec=DualCache), + ) + + exception = exc_info.value + assert exception.status_code == 400 + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + assert ( + detail["bedrock_guardrail_response"] + == "Sorry, the model cannot answer this question. coffee guardrail applied " + ) + + +@pytest.mark.asyncio +async def test_bedrock_guardrails_block_responses_api(): + """ + Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. + """ + from fastapi import HTTPException + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="ff6ujrregl1q", + guardrailVersion="DRAFT", + ) + + request_data = { + "model": "gpt-4.1", + "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee", + "stream": False, + } + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="responses", + cache=MagicMock(spec=DualCache), + ) + + exception = exc_info.value + assert exception.status_code == 400 + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + assert ( + detail["bedrock_guardrail_response"] + == "Sorry, the model cannot answer this question. coffee guardrail applied " + ) + + +@pytest.mark.asyncio +async def test_bedrock_guardrails_with_streaming(): + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + # Create proper mock objects + mock_user_api_key_cache = MagicMock(spec=DualCache) + mock_user_api_key_dict = UserAPIKeyAuth() + + with pytest.raises(Exception): # Assert that this raises an exception + proxy_logging_obj = ProxyLogging( + user_api_key_cache=mock_user_api_key_cache, + premium_user=True, + ) + + guardrail = BedrockGuardrail( + guardrailIdentifier="ff6ujrregl1q", + guardrailVersion="DRAFT", + supported_event_hooks=[GuardrailEventHooks.post_call], + guardrail_name="bedrock-post-guard", + ) + + litellm.callbacks.append(guardrail) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hi I like coffee"}], + "stream": True, + "metadata": {"guardrails": ["bedrock-post-guard"]}, + } + + response = await litellm.acompletion( + **request_data, + ) + + response = proxy_logging_obj.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=response, + request_data=request_data, + ) + + async for chunk in response: + print(chunk) + + +@pytest.mark.asyncio +async def test_bedrock_guardrails_with_streaming_no_violation(): + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + # Create proper mock objects + mock_user_api_key_cache = MagicMock(spec=DualCache) + mock_user_api_key_dict = UserAPIKeyAuth() + + proxy_logging_obj = ProxyLogging( + user_api_key_cache=mock_user_api_key_cache, + premium_user=True, + ) + + guardrail = BedrockGuardrail( + guardrailIdentifier="ff6ujrregl1q", + guardrailVersion="DRAFT", + supported_event_hooks=[GuardrailEventHooks.post_call], + guardrail_name="bedrock-post-guard", + ) + + litellm.callbacks.append(guardrail) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "metadata": {"guardrails": ["bedrock-post-guard"]}, + } + + response = await litellm.acompletion( + **request_data, + ) + + response = proxy_logging_obj.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=response, + request_data=request_data, + ) + + async for chunk in response: + print(chunk) + + +@pytest.mark.asyncio +async def test_bedrock_guardrails_streaming_request_body_mock(): + """Test that the exact request body sent to Bedrock matches expected format when using streaming""" + import json + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.caching import DualCache + from litellm.types.guardrails import GuardrailEventHooks + + # Create mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + mock_cache = MagicMock(spec=DualCache) + + # Create the guardrail + guardrail = BedrockGuardrail( + guardrailIdentifier="wf0hkdb5x07f", + guardrailVersion="DRAFT", + supported_event_hooks=[GuardrailEventHooks.post_call], + guardrail_name="bedrock-post-guard", + ) + + # Mock the assembled response from streaming + mock_response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", content="The capital of Spain is Madrid." + ), + finish_reason="stop", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + ) + + # Mock Bedrock API response + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []} + + # Patch the async_handler.post method to capture the request body + with patch.object(guardrail, "async_handler") as mock_async_handler: + mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) + + # Test data - simulating request data and assembled response + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "what's the capital of spain?"}], + "stream": True, + "metadata": {"guardrails": ["bedrock-post-guard"]}, + } + + # Call the method that should make the Bedrock API request + await guardrail.make_bedrock_api_request( + source="OUTPUT", response=mock_response, request_data=request_data + ) + + # Verify the API call was made + mock_async_handler.post.assert_called_once() + + # Get the request data that was passed + call_args = mock_async_handler.post.call_args + + # The data should be in the 'data' parameter of the prepared request + # We need to parse the JSON from the prepared request body + prepared_request_body = call_args.kwargs.get("data") + + # Parse the JSON body + if isinstance(prepared_request_body, bytes): + actual_body = json.loads(prepared_request_body.decode("utf-8")) + else: + actual_body = json.loads(prepared_request_body) + + # Expected body based on the convert_to_bedrock_format method behavior + expected_body = { + "source": "OUTPUT", + "content": [{"text": {"text": "The capital of Spain is Madrid."}}], + } + + print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) + print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) + + # Assert the request body matches exactly + assert ( + actual_body == expected_body + ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_aws_param_persistence(): + """Test that AWS auth params set on init are used for every request and not popped out.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = BedrockGuardrail( + guardrailIdentifier="wf0hkdb5x07f", + guardrailVersion="DRAFT", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + supported_event_hooks=[GuardrailEventHooks.post_call], + guardrail_name="bedrock-post-guard", + ) + + with patch.object( + guardrail, "get_credentials", wraps=guardrail.get_credentials + ) as mock_get_creds: + for i in range(3): + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": f"request {i}"}], + "stream": False, + "metadata": {"guardrails": ["bedrock-post-guard"]}, + } + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + # Configure the mock response properly + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = MagicMock( + return_value={"action": "NONE", "outputs": []} + ) + mock_post.return_value = mock_response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + assert mock_get_creds.call_count == 3 + for call in mock_get_creds.call_args_list: + kwargs = call.kwargs + print("used the following kwargs to get credentials=", kwargs) + assert kwargs["aws_access_key_id"] == "test-access-key" + assert kwargs["aws_secret_access_key"] == "test-secret-key" + assert kwargs["aws_region_name"] == "us-east-1" + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): + """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" + from unittest.mock import MagicMock + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrailResponse, + ) + guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock the Bedrock API response with PII + # Test 1: ANONYMIZED action should NOT raise exception + anonymized_response: BedrockGuardrailResponse = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + anonymized_response + ) + assert should_raise is False, "ANONYMIZED actions should not raise exceptions" + + # Test 2: BLOCKED action should raise exception + blocked_response: BedrockGuardrailResponse = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) + assert should_raise is True, "BLOCKED actions should raise exceptions" + + # Test 3: Mixed actions - should raise if ANY action is BLOCKED + mixed_response: BedrockGuardrailResponse = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + }, + "topicPolicy": { + "topics": [ + {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} + ] + }, + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) + assert ( + should_raise is True + ), "Mixed actions with any BLOCKED should raise exceptions" + + # Test 4: NONE action should not raise exception + none_response: BedrockGuardrailResponse = { + "action": "NONE", + "outputs": [], + "assessments": [], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) + assert should_raise is False, "NONE actions should not raise exceptions" + + # Test 5: Test other policy types with BLOCKED actions + content_blocked_response: BedrockGuardrailResponse = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "contentPolicy": { + "filters": [ + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + content_blocked_response + ) + assert ( + should_raise is True + ), "Content policy BLOCKED actions should raise exceptions" + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_masking_with_anonymized_response(): + """Test that masking works correctly when guardrail returns ANONYMIZED actions""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.caching import DualCache + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, + ) + + # Mock the Bedrock API response with ANONYMIZED action mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { @@ -207,7 +573,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): "piiEntities": [ { "type": "PHONE", - "match": "+1 412 555 1212", # This should be redacted in logs + "match": "+1 412 555 1212", "action": "ANONYMIZED", } ] @@ -223,77 +589,608 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): ], } - # Mock AWS credentials to avoid credential loading issues in CI - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Mock AWS-related methods to ensure test runs without external dependencies + # Patch the async_handler.post method with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request: - + ) as mock_post: mock_post.return_value = mock_bedrock_response - # Call the method that should log the redacted response - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - # Verify that debug logging was called - mock_debug.assert_called() - - # Get the logged response (second argument to debug call) - logged_calls = mock_debug.call_args_list - bedrock_response_log_call = None - - for call in logged_calls: - args, kwargs = call - if len(args) >= 2 and "Bedrock AI response" in str(args[0]): - bedrock_response_log_call = call - break - - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" - - # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call - - # Verify that the logged response has redacted PII - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" - ) - - # Verify other fields are preserved - assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) - - print("Bedrock guardrail logging redaction test passed") + # This should NOT raise an exception since action is ANONYMIZED + try: + response = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + # Should succeed and return data with masked content + assert response is not None + assert ( + response["messages"][0]["content"] + == "Hello, my phone number is {PHONE}" + ) + except Exception as e: + pytest.fail( + f"Should not raise exception for ANONYMIZED actions, but got: {e}" + ) @pytest.mark.asyncio -async def test_bedrock_guardrail_original_response_not_modified(): - """Test that the original response is not modified by redaction, only the logged version""" +async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): + """Test that masked output from guardrails is used even when masking flags are not enabled""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail WITHOUT masking flags enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + # Note: No mask_request_content=True or mask_response_content=True + ) + + # Mock the Bedrock API response with ANONYMIZED action and masked output + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + }, + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + }, + ] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com", + }, + ], + } + + # Patch the async_handler.post method + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # This should use the masked output even without masking flags + response = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + + # Should use the masked content from guardrail output + assert response is not None + assert ( + response["messages"][0]["content"] + == "Hello, my phone number is {PHONE} and email is {EMAIL}" + ) + print("✅ Masked output was applied even without masking flags enabled") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_response_pii_masking_non_streaming(): + """Test that PII masking is applied to response content in non-streaming scenarios""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail with response masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + ) + + # Mock the Bedrock API response with ANONYMIZED PII + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [ + { + "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" + } + ], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "CREDIT_DEBIT_CARD_NUMBER", + "match": "1234-5678-9012-3456", + "action": "ANONYMIZED", + }, + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + }, + ] + } + } + ], + } + + # Create a mock response that contains PII + mock_response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212", + ), + finish_reason="stop", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + ) + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "What's your credit card and phone number?"}, + ], + } + + # Patch the async_handler.post method + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Call the post-call success hook + await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + response=mock_response, + ) + + # Verify that the response content was masked + assert ( + mock_response.choices[0].message.content + == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" + ) + print("✓ Non-streaming response PII masking test passed") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_response_pii_masking_streaming(): + """Test that PII masking is applied to response content in streaming scenarios""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail with response masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + ) + + # Mock the Bedrock API response with ANONYMIZED PII + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "john@example.com", + "action": "ANONYMIZED", + }, + { + "type": "US_SSN", + "match": "123-45-6789", + "action": "ANONYMIZED", + }, + ] + } + } + ], + } + + # Create mock streaming chunks + async def mock_streaming_response(): + chunks = [ + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta(content="Sure! My email is "), + finish_reason=None, + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta( + content="john@example.com and SSN is " + ), + finish_reason=None, + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta(content="123-45-6789"), + finish_reason="stop", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ] + for chunk in chunks: + yield chunk + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "What's your email and SSN?"}, + ], + "stream": True, + } + + # Patch the async_handler.post method + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Call the streaming hook + masked_stream = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) + + # Collect all chunks from the masked stream + masked_chunks = [] + async for chunk in masked_stream: + masked_chunks.append(chunk) + + # Verify that we got chunks back + assert len(masked_chunks) > 0 + + # Reconstruct the full response from chunks to verify masking + full_content = "" + for chunk in masked_chunks: + if hasattr(chunk, "choices") and chunk.choices: + if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): + full_content += chunk.choices[0].delta.content + + # Verify that the reconstructed content contains the masked PII + assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content + print("✓ Streaming response PII masking test passed") + + +@pytest.mark.asyncio +async def test_convert_to_bedrock_format_input_source(): + """Test convert_to_bedrock_format with INPUT source and mock messages""" + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockRequest, + ) + from unittest.mock import patch + + # Create the guardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock messages + mock_messages = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you!"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "What's the weather like?"}, + {"type": "text", "text": "Is it sunny today?"}, + ], + }, + ] + + # Call the method + result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) + + # Verify the result structure + assert isinstance(result, dict) + assert result.get("source") == "INPUT" + assert "content" in result + assert isinstance(result.get("content"), list) + + # Verify content items + expected_content_items = [ + {"text": {"text": "Hello, how are you?"}}, + {"text": {"text": "I'm doing well, thank you!"}}, + {"text": {"text": "What's the weather like?"}}, + {"text": {"text": "Is it sunny today?"}}, + ] + + assert result.get("content") == expected_content_items + print("✅ INPUT source test passed - result:", result) + + +@pytest.mark.asyncio +async def test_convert_to_bedrock_format_output_source(): + """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockRequest, + ) + import litellm + from unittest.mock import patch + + # Create the guardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock ModelResponse + mock_response = litellm.ModelResponse( + id="test-response-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", content="This is a test response from the model." + ), + finish_reason="stop", + ), + litellm.Choices( + index=1, + message=litellm.Message( + role="assistant", content="This is a second choice response." + ), + finish_reason="stop", + ), + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + ) + + # Call the method + result = guardrail.convert_to_bedrock_format( + source="OUTPUT", response=mock_response + ) + + # Verify the result structure + assert isinstance(result, dict) + assert result.get("source") == "OUTPUT" + assert "content" in result + assert isinstance(result.get("content"), list) + + # Verify content items - should contain both choice contents + expected_content_items = [ + {"text": {"text": "This is a test response from the model."}}, + {"text": {"text": "This is a second choice response."}}, + ] + + assert result.get("content") == expected_content_items + print("✅ OUTPUT source test passed - result:", result) + + +@pytest.mark.asyncio +async def test_convert_to_bedrock_format_post_call_streaming_hook(): + """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + import litellm + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock streaming chunks that contain PII + async def mock_streaming_response(): + chunks = [ + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta(content="My email is "), + finish_reason=None, + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta(content="john@example.com"), + finish_reason="stop", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ] + for chunk in chunks: + yield chunk + + # Mock Bedrock API response with PII masking + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "My email is {EMAIL}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "john@example.com", + "action": "ANONYMIZED", + } + ] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "What's your email?"}], + "stream": True, + } + + # Track which bedrock API calls were made + bedrock_calls = [] + + # 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 + ): + bedrock_calls.append( + { + "source": source, + "messages": messages, + "response": response, + "request_data": request_data, + } + ) + # Return the mock bedrock response + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrailResponse, + ) + + return BedrockGuardrailResponse(**mock_bedrock_response.json()) + + # Patch the bedrock API request method + with patch.object( + guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request + ): + + # Call the streaming hook + result_generator = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) + + # Collect all chunks from the result + result_chunks = [] + async for chunk in result_generator: + result_chunks.append(chunk) + + # Verify bedrock API calls were made + # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. + # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled + # to avoid redundant validation. Since event_hook=None means all hooks are enabled, + # only OUTPUT validation should be performed in post_call. + assert ( + len(bedrock_calls) == 1 + ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" + + # Verify the OUTPUT call + output_call = bedrock_calls[0] + assert output_call["source"] == "OUTPUT" + assert output_call["response"] is not None + assert output_call["messages"] is None # OUTPUT calls don't need messages + + # Verify that the response content was masked + # The streaming chunks should now contain the masked content + full_content = "" + for chunk in result_chunks: + if hasattr(chunk, "choices") and chunk.choices: + if ( + hasattr(chunk.choices[0], "delta") + and chunk.choices[0].delta.content + ): + full_content += chunk.choices[0].delta.content + + # The content should be masked (contains {EMAIL} instead of john@example.com) + assert ( + "{EMAIL}" in full_content + ), f"Expected masked content with {{EMAIL}}, got: {full_content}" + assert ( + "john@example.com" not in full_content + ), f"Original email should be masked, got: {full_content}" + + print( + "✅ Post-call streaming hook test passed - OUTPUT source used for masking" + ) + print( + f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} " + "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" + ) + print(f"✅ Final masked content: {full_content}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_action_shows_output_text(): + """Test that BLOCKED actions raise HTTPException with the output text in the detail""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() @@ -302,1274 +1199,533 @@ async def test_bedrock_guardrail_original_response_not_modified(): guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock the Bedrock API response with PII - original_response_data = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", # This should NOT be modified in original - "action": "ANONYMIZED", - } - ] - } - } - ], - } - + # Mock the Bedrock API response with BLOCKED action and output text mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = original_response_data - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - ], - } - - # Mock AWS credentials to avoid credential loading issues in CI - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Mock AWS-related methods to ensure test runs without external dependencies - 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") - ) as mock_load_creds, patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request: - - mock_post.return_value = mock_bedrock_response - - # Call the method - result = await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - # Verify that the original response data was not modified - # (The json() method should return the original data) - original_data = mock_bedrock_response.json() - assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "+1 412 555 1212" - ) - - # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) - - print("Original response not modified test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_preserves_non_pii_entities(): - """Test that _redact_pii_matches only affects PII-related entities and preserves other assessment data""" - - response_with_mixed_data = { + mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "this violates litellm corporate guardrail policy"}], "assessments": [ { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "user@example.com", - "action": "ANONYMIZED", - "confidence": "HIGH", - } - ], - "regexes": [ - { - "name": "custom_pattern", - "match": "some_pattern_match", - "action": "BLOCKED", - } - ], - }, - "contentPolicy": { - "filters": [ - { - "type": "VIOLENCE", - "confidence": "MEDIUM", - "action": "BLOCKED", - } - ] - }, "topicPolicy": { "topics": [ - { - "name": "Restricted Topic", - "type": "DENY", - "action": "BLOCKED", - } - ] - }, - } - ], - "outputs": [{"text": "Content blocked"}], - } - - # Call the redaction function - redacted_response = _redact_pii_matches(response_with_mixed_data) - - # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" - assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" - assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" - assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" - - # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] - assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" - assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" - assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" - - # Verify that other policies are completely unchanged - content_policy = redacted_response["assessments"][0]["contentPolicy"] - assert content_policy["filters"][0]["type"] == "VIOLENCE" - assert content_policy["filters"][0]["confidence"] == "MEDIUM" - - topic_policy = redacted_response["assessments"][0]["topicPolicy"] - assert topic_policy["topics"][0]["name"] == "Restricted Topic" - - # Verify top-level fields are unchanged - assert redacted_response["action"] == "GUARDRAIL_INTERVENED" - assert redacted_response["outputs"][0]["text"] == "Content blocked" - - print("Preserves non-PII entities test passed") - - -@pytest.mark.asyncio -async def test_pii_redaction_matches_debug_output_format(): - """Test that demonstrates the exact behavior shown in your debug output""" - - # This matches the structure from your debug output - original_response = { - "action": "GUARDRAIL_INTERVENED", - "actionReason": "Guardrail blocked.", - "assessments": [ - { - "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, - "guardrailProcessingLatency": 322, - "usage": { - "contentPolicyImageUnits": 0, - "contentPolicyUnits": 0, - "contextualGroundingPolicyUnits": 0, - "sensitiveInformationPolicyFreeUnits": 0, - "sensitiveInformationPolicyUnits": 1, - "topicPolicyUnits": 0, - "wordPolicyUnits": 0, - }, - }, - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "action": "BLOCKED", - "detected": True, - "match": "John Smith", - "type": "NAME", - }, - { - "action": "BLOCKED", - "detected": True, - "match": "324-12-3212", - "type": "US_SOCIAL_SECURITY_NUMBER", - }, - { - "action": "BLOCKED", - "detected": True, - "match": "607-456-7890", - "type": "PHONE", - }, - ] - }, - } - ], - "blockedResponse": "Input blocked by PII policy", - "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, - "output": [{"text": "Input blocked by PII policy"}], - "outputs": [{"text": "Input blocked by PII policy"}], - "usage": { - "contentPolicyImageUnits": 0, - "contentPolicyUnits": 0, - "contextualGroundingPolicyUnits": 0, - "sensitiveInformationPolicyFreeUnits": 0, - "sensitiveInformationPolicyUnits": 1, - "topicPolicyUnits": 0, - "wordPolicyUnits": 0, - }, - } - - # Apply redaction - redacted_response = _redact_pii_matches(original_response) - - # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - - # All PII matches should be redacted - assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" - assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" - assert pii_entities[2]["match"] == "[REDACTED]", "PHONE should be redacted" - - # But all other fields should be preserved - assert pii_entities[0]["type"] == "NAME" - assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" - assert pii_entities[2]["type"] == "PHONE" - assert pii_entities[0]["action"] == "BLOCKED" - assert pii_entities[0]["detected"] == True - - # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" - - # Verify all other metadata is preserved in redacted response - assert redacted_response["action"] == "GUARDRAIL_INTERVENED" - assert redacted_response["actionReason"] == "Guardrail blocked." - assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) - - print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) - print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_with_regex_matches(): - """Test redaction of regex matches in sensitive information policy""" - - response_with_regex = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "regexes": [ - { - "name": "SSN_PATTERN", - "match": "123-45-6789", - "action": "BLOCKED", - }, - { - "name": "CREDIT_CARD_PATTERN", - "match": "4111-1111-1111-1111", - "action": "ANONYMIZED", - }, + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} ] } } ], - "outputs": [{"text": "Regex patterns detected"}], } - # Call the redaction function - redacted_response = _redact_pii_matches(response_with_regex) - - # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] - - assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" - - # Verify other fields are preserved - assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" - assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" - assert regexes[1]["name"] == "CREDIT_CARD_PATTERN", "Regex name should be preserved" - assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" - - # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] - assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" - - print("Regex matches redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_with_custom_words(): - """Test redaction of custom word matches in word policy""" - - response_with_custom_words = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "wordPolicy": { - "customWords": [ - { - "match": "confidential_data", - "action": "BLOCKED", - }, - { - "match": "secret_information", - "action": "ANONYMIZED", - }, - ] - } - } + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Tell me how to make explosives"}, ], - "outputs": [{"text": "Custom words detected"}], } - # Call the redaction function - redacted_response = _redact_pii_matches(response_with_custom_words) - - # Verify that custom word matches are redacted - custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" - - # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" - - # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" - - print("Custom words redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_with_managed_words(): - """Test redaction of managed word matches in word policy""" - - response_with_managed_words = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "wordPolicy": { - "managedWordLists": [ - { - "match": "inappropriate_word", - "action": "BLOCKED", - "type": "PROFANITY", - }, - { - "match": "offensive_term", - "action": "ANONYMIZED", - "type": "HATE_SPEECH", - }, - ] - } - } - ], - "outputs": [{"text": "Managed words detected"}], - } - - # Call the redaction function - redacted_response = _redact_pii_matches(response_with_managed_words) - - # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] - - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" - - # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" - - # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" - - print("Managed words redaction test passed") - - -@pytest.mark.asyncio -async def test__redact_pii_matches_comprehensive_coverage(): - """Test redaction across all supported policy types in a single response""" - - comprehensive_response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "user@example.com", - "action": "ANONYMIZED", - } - ], - "regexes": [ - { - "name": "PHONE_PATTERN", - "match": "555-123-4567", - "action": "BLOCKED", - } - ], - }, - "wordPolicy": { - "customWords": [ - { - "match": "confidential", - "action": "BLOCKED", - } - ], - "managedWordLists": [ - { - "match": "inappropriate", - "action": "ANONYMIZED", - "type": "PROFANITY", - } - ], - }, - } - ], - "outputs": [{"text": "Multiple policy violations detected"}], - } - - # Call the redaction function - redacted_response = _redact_pii_matches(comprehensive_response) - - # Verify all match fields are redacted - assessment = redacted_response["assessments"][0] - - # PII entities - pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" - - # Regex matches - regexes = assessment["sensitiveInformationPolicy"]["regexes"] - assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" - - # Custom words - custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" - - # Managed words - managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" - - # Verify all other fields are preserved - assert pii_entities[0]["type"] == "EMAIL" - assert regexes[0]["name"] == "PHONE_PATTERN" - assert managed_words[0]["type"] == "PROFANITY" - - # Verify original response is unchanged - original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) - - print("Comprehensive coverage redaction test passed") - -@pytest.mark.asyncio -async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): - """Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set""" - - # Clear any existing environment variable to ensure clean test - monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) - - # Create guardrail with custom runtime endpoint - custom_endpoint = "https://custom-bedrock.example.com" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - aws_bedrock_runtime_endpoint=custom_endpoint, - ) - - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Test data - data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - optional_params = {} - aws_region_name = "us-east-1" - - # Mock the _load_credentials method to avoid actual AWS credential loading + # Patch the async_handler.post method with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): - # Call _prepare_request which internally calls get_runtime_endpoint - prepped_request = guardrail._prepare_request( - credentials=mock_credentials, - data=data, - optional_params=optional_params, - aws_region_name=aws_region_name, + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # This should raise HTTPException due to BLOCKED action + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 + assert "detail" in exception.__dict__ + + # Check that the detail contains the expected structure + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + + # Verify that the output text from both outputs is included + expected_output_text = "this violates litellm corporate guardrail policy" + assert detail["bedrock_guardrail_response"] == expected_output_text + + print( + "✅ BLOCKED action HTTPException test passed - output text properly included" ) - # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" - - print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") - @pytest.mark.asyncio -async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): - """Test that BedrockGuardrail respects AWS_BEDROCK_RUNTIME_ENDPOINT environment variable""" +async def test_bedrock_guardrail_blocked_action_empty_outputs(): + """Test that BLOCKED actions with empty outputs still raise HTTPException""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException - custom_endpoint = "https://env-bedrock.example.com" + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() - # Set the environment variable - monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) - - # Create guardrail without explicit aws_bedrock_runtime_endpoint guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Test data - data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - optional_params = {} - aws_region_name = "us-east-1" - - # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): - # Call _prepare_request which internally calls get_runtime_endpoint - prepped_request = guardrail._prepare_request( - credentials=mock_credentials, - data=data, - optional_params=optional_params, - aws_region_name=aws_region_name, - ) - - # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" - - print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkeypatch): - """Test that BedrockGuardrail uses default endpoint when no custom endpoint is set""" - - # Ensure no environment variable is set - monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) - - # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Test data - data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - optional_params = {} - aws_region_name = "us-west-2" - - # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): - # Call _prepare_request which internally calls get_runtime_endpoint - prepped_request = guardrail._prepare_request( - credentials=mock_credentials, - data=data, - optional_params=optional_params, - aws_region_name=aws_region_name, - ) - - # Verify that the default endpoint is used - expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" - - print(f"Default endpoint test passed. URL: {prepped_request.url}") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch): - """Test that aws_bedrock_runtime_endpoint parameter takes precedence over environment variable - - This test verifies the corrected behavior where the parameter should take precedence - over the environment variable, consistent with the endpoint_url logic. - """ - - param_endpoint = "https://param-bedrock.example.com" - env_endpoint = "https://env-bedrock.example.com" - - # Set environment variable - monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_endpoint) - - # Create guardrail with explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - aws_bedrock_runtime_endpoint=param_endpoint, - ) - - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - # Test data - data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - optional_params = {} - aws_region_name = "us-east-1" - - # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): - # Call _prepare_request which internally calls get_runtime_endpoint - prepped_request = guardrail._prepare_request( - credentials=mock_credentials, - data=data, - optional_params=optional_params, - aws_region_name=aws_region_name, - ) - - # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" - - print(f"Parameter precedence test passed. URL: {prepped_request.url}") - - -@pytest.mark.asyncio -async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): - """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" - # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: - # Test the apply_guardrail method with tool_calls in response - inputs = { - "texts": [], - "tool_calls": [ - { - "id": "call_eFSCWFsyL7MclHYnzKrcQnMK", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"location":"São Paulo"}', - }, - } - ], - } - - guardrailed_inputs = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="response", - logging_obj=None, - ) - - # Verify the result - should succeed without errors - assert guardrailed_inputs is not None - assert "tool_calls" in guardrailed_inputs - assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] - == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) - assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) - # Verify that the Bedrock API was NOT called since there's no text to process - mock_api_request.assert_not_called() - print("✅ apply_guardrail with tool_calls test passed - no API call made") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): - """Test that BLOCKED content raises exception even when masking is enabled - - This test verifies the bug fix where previously mask_request_content=True or - mask_response_content=True would bypass all BLOCKED content checks. Now it - properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). - """ - - # Create guardrail with masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - mask_request_content=True, # Masking enabled - mask_response_content=True, # Masking enabled - ) - - # Mock Bedrock response with BLOCKED content (hate speech) - blocked_response = { + # Mock the Bedrock API response with BLOCKED action but no outputs + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", + "outputs": [], # Empty outputs "assessments": [ { "contentPolicy": { "filters": [ - { - "type": "HATE", - "confidence": "HIGH", - "action": "BLOCKED", # Should raise exception - } + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} ] - }, - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "NAME", - "match": "John Doe", - "action": "ANONYMIZED", # Should be masked - } - ] - }, + } } ], - "outputs": [{"text": "Content blocked due to policy violation"}], } - - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = blocked_response - - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Test message with PII and hate speech"}, + {"role": "user", "content": "Violent content here"}, ], } - - # Mock AWS-related methods + + # Patch the async_handler.post method 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() - ): + ) as mock_post: mock_post.return_value = mock_bedrock_response - - # Should raise HTTPException for BLOCKED content + + # This should raise HTTPException due to BLOCKED action with pytest.raises(HTTPException) as exc_info: - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - # Verify exception details - assert exc_info.value.status_code == 400 - assert "Violated guardrail policy" in str(exc_info.value.detail) - - print("✅ BLOCKED content with masking enabled raises exception correctly") - - -# ────────────────────────────────────────────────────────────────────────────── -# Null-safety tests for Bedrock guardrail responses -# -# The Bedrock ApplyGuardrail API can return explicit null/None for list fields -# such as "regexes", "piiEntities", "topics", "filters", "customWords", and -# "managedWordLists" when a particular policy category is present in the -# assessment but has no matches. -# -# Python's dict.get("key", []) returns None (NOT []) when the key exists with -# a None value. The `or []` fallback ensures we always iterate over a list. -# -# Without the fix, iterating over None raises: -# TypeError: 'NoneType' object is not iterable -# which surfaces to callers as: -# openai.InternalServerError: Error code: 500 -# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} -# ────────────────────────────────────────────────────────────────────────────── - - -class TestRedactPiiMatchesNullSafety: - """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" - - @pytest.mark.asyncio - async def test_should_handle_null_regexes_in_sensitive_info_policy(self): - """Bedrock can return regexes: null while piiEntities has data. - - Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no - custom regex patterns configured, so the API returns regexes: null. - """ - response = { - "action": "NONE", - "actionReason": "No action.", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "action": "NONE", - "detected": True, - "match": "joebloggs@gmail.com", - "type": "EMAIL", - } - ], - "regexes": None, # Explicit null from Bedrock API - }, - } - ], - } - - # Should not raise TypeError: 'NoneType' object is not iterable - redacted = _redact_pii_matches(response) - - # PII match should be redacted - pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] - assert pii[0]["match"] == "[REDACTED]" - assert pii[0]["type"] == "EMAIL" - - @pytest.mark.asyncio - async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): - """Bedrock can return piiEntities: null while regexes has data.""" - response = { - "action": "NONE", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": None, # null from Bedrock API - "regexes": [ - { - "name": "CUSTOM_PATTERN", - "match": "secret-abc-123", - "action": "BLOCKED", - } - ], - }, - } - ], - } - - redacted = _redact_pii_matches(response) - - regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] - assert regexes[0]["match"] == "[REDACTED]" - - @pytest.mark.asyncio - async def test_should_handle_null_custom_words_and_managed_words(self): - """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" - response = { - "action": "NONE", - "assessments": [ - { - "wordPolicy": { - "customWords": None, # null from Bedrock API - "managedWordLists": None, # null from Bedrock API - }, - } - ], - } - - # Should not raise TypeError - redacted = _redact_pii_matches(response) - - # Values should remain None (no crash) - assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None - assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None - - @pytest.mark.asyncio - async def test_should_handle_null_assessments_list(self): - """Bedrock can return assessments: null.""" - response = { - "action": "NONE", - "assessments": None, # null from Bedrock API - } - - # Should not raise TypeError - redacted = _redact_pii_matches(response) - assert redacted["assessments"] is None - - @pytest.mark.asyncio - async def test_should_handle_all_null_policy_sub_lists_together(self): - """All sub-list fields are null at the same time — worst-case scenario.""" - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": None, - "regexes": None, - }, - "wordPolicy": { - "customWords": None, - "managedWordLists": None, - }, - "topicPolicy": None, - "contentPolicy": None, - "contextualGroundingPolicy": None, - } - ], - } - - # Should not raise any exception - redacted = _redact_pii_matches(response) - assert redacted is not None - - -class TestShouldRaiseGuardrailBlockedExceptionNullSafety: - """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" - - def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - @pytest.mark.asyncio - async def test_should_handle_all_null_policy_sub_lists(self): - """All policy sub-lists are null — should not crash, should return False.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": None, # null from Bedrock API - }, - "contentPolicy": { - "filters": None, # null - }, - "wordPolicy": { - "customWords": None, # null - "managedWordLists": None, # null - }, - "sensitiveInformationPolicy": { - "piiEntities": None, # null - "regexes": None, # null - }, - "contextualGroundingPolicy": { - "filters": None, # null - }, - } - ], - } - - # No BLOCKED actions found (all lists null) → should return False - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is False - - @pytest.mark.asyncio - async def test_should_detect_blocked_despite_other_null_lists(self): - """A mix of null lists and a real BLOCKED action — should still detect it.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": None, # null — should not crash - }, - "contentPolicy": { - "filters": [ - { - "type": "HATE", - "confidence": "HIGH", - "action": "BLOCKED", - } - ], - }, - "wordPolicy": { - "customWords": None, # null - "managedWordLists": None, # null - }, - "sensitiveInformationPolicy": { - "piiEntities": None, # null - "regexes": None, # null - }, - "contextualGroundingPolicy": None, # entire policy is null - } - ], - } - - # Should return True because contentPolicy has a BLOCKED filter - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is True - - @pytest.mark.asyncio - async def test_should_handle_null_assessments_list(self): - """assessments itself is null — should return False.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": None, # null from Bedrock API - } - - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is False - - @pytest.mark.asyncio - async def test_should_handle_null_topics_with_blocked_word_policy(self): - """topics is null but wordPolicy has a BLOCKED customWord.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": None, - }, - "wordPolicy": { - "customWords": [ - {"match": "badword", "action": "BLOCKED"} - ], - "managedWordLists": None, - }, - } - ], - } - - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is True - - @pytest.mark.asyncio - async def test_should_handle_null_pii_with_blocked_regex(self): - """piiEntities is null but regexes has a BLOCKED match.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], - }, - } - ], - } - - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is True - - @pytest.mark.asyncio - async def test_should_handle_null_grounding_filters(self): - """contextualGroundingPolicy.filters is null — should not crash.""" - guardrail = self._create_guardrail() - - response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "contextualGroundingPolicy": { - "filters": None, - }, - } - ], - } - - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is False - - @pytest.mark.asyncio - async def test_should_not_crash_when_action_is_not_intervened(self): - """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" - guardrail = self._create_guardrail() - - response = { - "action": "NONE", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": None, - "regexes": None, - }, - } - ], - } - - result = guardrail._should_raise_guardrail_blocked_exception(response) - assert result is False - - -class TestApplyGuardrailNullSafety: - """Tests for apply_guardrail handling of null/None texts input.""" - - @pytest.mark.asyncio - async def test_should_handle_none_texts_in_inputs(self): - """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - inputs = {"texts": None} # Explicit None - - mock_credentials = MagicMock() - - 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() - ): - # With empty texts (from None → []), no Bedrock API call should be made - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", ) - # Should return empty texts without crashing - assert result.get("texts") == [] - # No Bedrock API call should be made for empty input - mock_post.assert_not_called() + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 - @pytest.mark.asyncio - async def test_should_handle_missing_texts_key(self): - """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + # Check that the detail contains the expected structure with empty output text + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs - inputs = {} # No "texts" key + print("✅ BLOCKED action with empty outputs test passed") - mock_credentials = MagicMock() - 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() - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", +@pytest.mark.asyncio +async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): + """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Test 1: disable_exception_on_block=False (default) - should raise exception + guardrail_default = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + # Mock the Bedrock API response with BLOCKED action + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Tell me how to make explosives"}, + ], + } + + # Patch the async_handler.post method + with patch.object( + guardrail_default.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException when disable_exception_on_block=False + with pytest.raises(HTTPException) as exc_info: + await guardrail_default.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", ) - assert result.get("texts") == [] - mock_post.assert_not_called() + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 + assert "Violated guardrail policy" in str(exception.detail) + + # Test 2: disable_exception_on_block=True - should NOT raise exception + guardrail_disabled = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=True, + ) + + with patch.object( + guardrail_disabled.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Should NOT raise exception when disable_exception_on_block=True + try: + response = await guardrail_disabled.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + # Should succeed and return data (even though content was blocked) + assert response is not None + print("✅ No exception raised when disable_exception_on_block=True") + except Exception as e: + pytest.fail( + f"Should not raise exception when disable_exception_on_block=True, but got: {e}" + ) + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_disable_exception_on_block_streaming(): + """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + from fastapi import HTTPException + import litellm + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Mock streaming chunks that would normally trigger a block + async def mock_streaming_response(): + chunks = [ + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta( + content="Here's how to make explosives: " + ), + finish_reason=None, + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ModelResponseStream( + id="test-id", + choices=[ + litellm.utils.StreamingChoices( + index=0, + delta=litellm.utils.Delta(content="step 1, step 2..."), + finish_reason="stop", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion.chunk", + ), + ] + for chunk in chunks: + yield chunk + + # Mock Bedrock API response with BLOCKED action + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "contentPolicy": { + "filters": [ + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} + ] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Tell me how to make explosives"}], + "stream": True, + } + + # Test 1: disable_exception_on_block=False (default) - should raise exception + guardrail_default = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + with patch.object( + guardrail_default.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Should raise exception during streaming processing + with pytest.raises(HTTPException): + result_generator = ( + guardrail_default.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) + ) + + # Try to consume the generator - should raise exception + async for chunk in result_generator: + pass + + # Test 2: disable_exception_on_block=True - should NOT raise exception + guardrail_disabled = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=True, + ) + + with patch.object( + guardrail_disabled.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_bedrock_response + + # Should NOT raise exception when disable_exception_on_block=True + try: + result_generator = ( + guardrail_disabled.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) + ) + + # Consume the generator - should succeed without exceptions + result_chunks = [] + async for chunk in result_generator: + result_chunks.append(chunk) + + # Should have received chunks back even though content was blocked + assert len(result_chunks) > 0 + print( + "✅ Streaming completed without exception when disable_exception_on_block=True" + ) + + except Exception as e: + pytest.fail( + f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}" + ) + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): + """Test that async_post_call_success_hook skips when there's no output text""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + import litellm + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Create a ModelResponse with tool calls (no text content) + # This simulates a response where the LLM is making a tool call + mock_response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content=None, # No text content + tool_calls=[ + litellm.utils.ChatCompletionMessageToolCall( + id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", + function=litellm.utils.Function( + name="top_song", arguments='{"sign": "WZPZ"}' + ), + type="function", + ) + ], + ), + finish_reason="tool_calls", + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + ) + + data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Hello"}, + ], + } + mock_user_api_key_dict = UserAPIKeyAuth() + + result = await guardrail.async_post_call_success_hook( + data=data, + response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + # If no error is raised and result is None, then the test passes + assert result is None + print("✅ No output text in response test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_null_list_fields(): + """Test that explicit null values from Bedrock API are handled correctly. + + The Bedrock API can return explicit JSON null for list fields like + piiEntities, regexes, customWords, managedWordLists. This would cause + TypeError: 'NoneType' object is not iterable if not handled. + """ + # Test 1: null piiEntities and regexes + response_with_null_pii = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + } + } + ], + } + redacted = _redact_pii_matches(response_with_null_pii) + assert redacted is not None + assert redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] is None + assert redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] is None + + # Test 2: null customWords and managedWordLists + response_with_null_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + } + } + ], + } + redacted = _redact_pii_matches(response_with_null_words) + assert redacted is not None + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + # Test 3: null assessments at top level + response_with_null_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, + } + redacted = _redact_pii_matches(response_with_null_assessments) + assert redacted is not None + + +@pytest.mark.asyncio +async def test__redact_pii_matches_malformed_response(): + """Test _redact_pii_matches with malformed response (should not crash)""" + + # Test with completely malformed response + malformed_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": "not_a_list", + } + redacted_response = _redact_pii_matches(malformed_response) + assert redacted_response == malformed_response + + # Test with missing keys + missing_keys_response = { + "action": "GUARDRAIL_INTERVENED", + } + redacted_response = _redact_pii_matches(missing_keys_response) + assert redacted_response == missing_keys_response + + +@pytest.mark.asyncio +async def test_should_raise_guardrail_blocked_exception_null_fields(): + """Test that _should_raise_guardrail_blocked_exception handles null list fields. + + Validates the or [] null-safety pattern works for all policy fields + in _should_raise_guardrail_blocked_exception. + """ + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Test with null assessments + response_null_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_assessments) is False + + # Test with null topics in topicPolicy + response_null_topics = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"topicPolicy": {"topics": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_topics) is False + + # Test with null filters in contentPolicy + response_null_filters = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"contentPolicy": {"filters": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_filters) is False + + # Test with null customWords and managedWordLists in wordPolicy + response_null_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"wordPolicy": {"customWords": None, "managedWordLists": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_words) is False + + # Test with null piiEntities and regexes in sensitiveInformationPolicy + response_null_pii = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": None, "regexes": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_pii) is False + + # Test with null filters in contextualGroundingPolicy + response_null_grounding = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"contextualGroundingPolicy": {"filters": None}}], + } + assert guardrail._should_raise_guardrail_blocked_exception(response_null_grounding) is False 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 84d320a0a27..68aaabafed8 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 @@ -1189,3 +1189,479 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): print("✅ BLOCKED content with masking enabled raises exception correctly") + +# ────────────────────────────────────────────────────────────────────────────── +# Null-safety tests for Bedrock guardrail responses +# +# The Bedrock ApplyGuardrail API can return explicit null/None for list fields +# such as "regexes", "piiEntities", "topics", "filters", "customWords", and +# "managedWordLists" when a particular policy category is present in the +# assessment but has no matches. +# +# Python's dict.get("key", []) returns None (NOT []) when the key exists with +# a None value. The `or []` fallback ensures we always iterate over a list. +# +# Without the fix, iterating over None raises: +# TypeError: 'NoneType' object is not iterable +# which surfaces to callers as: +# openai.InternalServerError: Error code: 500 +# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} +# ────────────────────────────────────────────────────────────────────────────── + + +class TestRedactPiiMatchesNullSafety: + """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" + + @pytest.mark.asyncio + async def test_should_handle_null_regexes_in_sensitive_info_policy(self): + """Bedrock can return regexes: null while piiEntities has data. + + Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no + custom regex patterns configured, so the API returns regexes: null. + """ + response = { + "action": "NONE", + "actionReason": "No action.", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "NONE", + "detected": True, + "match": "joebloggs@gmail.com", + "type": "EMAIL", + } + ], + "regexes": None, # Explicit null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError: 'NoneType' object is not iterable + redacted = _redact_pii_matches(response) + + # PII match should be redacted + pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert pii[0]["match"] == "[REDACTED]" + assert pii[0]["type"] == "EMAIL" + + @pytest.mark.asyncio + async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): + """Bedrock can return piiEntities: null while regexes has data.""" + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, # null from Bedrock API + "regexes": [ + { + "name": "CUSTOM_PATTERN", + "match": "secret-abc-123", + "action": "BLOCKED", + } + ], + }, + } + ], + } + + redacted = _redact_pii_matches(response) + + regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_should_handle_null_custom_words_and_managed_words(self): + """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" + response = { + "action": "NONE", + "assessments": [ + { + "wordPolicy": { + "customWords": None, # null from Bedrock API + "managedWordLists": None, # null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + + # Values should remain None (no crash) + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """Bedrock can return assessments: null.""" + response = { + "action": "NONE", + "assessments": None, # null from Bedrock API + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + assert redacted["assessments"] is None + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists_together(self): + """All sub-list fields are null at the same time — worst-case scenario.""" + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + }, + "topicPolicy": None, + "contentPolicy": None, + "contextualGroundingPolicy": None, + } + ], + } + + # Should not raise any exception + redacted = _redact_pii_matches(response) + assert redacted is not None + + +class TestShouldRaiseGuardrailBlockedExceptionNullSafety: + """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" + + def _create_guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists(self): + """All policy sub-lists are null — should not crash, should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null from Bedrock API + }, + "contentPolicy": { + "filters": None, # null + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": { + "filters": None, # null + }, + } + ], + } + + # No BLOCKED actions found (all lists null) → should return False + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_detect_blocked_despite_other_null_lists(self): + """A mix of null lists and a real BLOCKED action — should still detect it.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null — should not crash + }, + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": None, # entire policy is null + } + ], + } + + # Should return True because contentPolicy has a BLOCKED filter + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """assessments itself is null — should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, # null from Bedrock API + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_handle_null_topics_with_blocked_word_policy(self): + """topics is null but wordPolicy has a BLOCKED customWord.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, + }, + "wordPolicy": { + "customWords": [ + {"match": "badword", "action": "BLOCKED"} + ], + "managedWordLists": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_pii_with_blocked_regex(self): + """piiEntities is null but regexes has a BLOCKED match.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": [ + {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} + ], + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_grounding_filters(self): + """contextualGroundingPolicy.filters is null — should not crash.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contextualGroundingPolicy": { + "filters": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_not_crash_when_action_is_not_intervened(self): + """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" + guardrail = self._create_guardrail() + + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + +class TestApplyGuardrailNullSafety: + """Tests for apply_guardrail handling of null/None texts input.""" + + @pytest.mark.asyncio + async def test_should_handle_none_texts_in_inputs(self): + """inputs[\"texts\"] is explicitly None — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {"texts": None} # Explicit None + + mock_credentials = MagicMock() + + 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() + ): + # With empty texts (from None → []), no Bedrock API call should be made + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + # Should return empty texts without crashing + assert result.get("texts") == [] + # No Bedrock API call should be made for empty input + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_should_handle_missing_texts_key(self): + """inputs has no \"texts\" key at all — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {} # No "texts" key + + mock_credentials = MagicMock() + + 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() + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result.get("texts") == [] + mock_post.assert_not_called() + + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): + """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Test 1: ANONYMIZED action should NOT raise exception + anonymized_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + anonymized_response + ) + assert should_raise is False, "ANONYMIZED actions should not raise exceptions" + + # Test 2: BLOCKED action should raise exception + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) + assert should_raise is True, "BLOCKED actions should raise exceptions" + + # Test 3: Mixed actions - should raise if ANY action is BLOCKED + mixed_response = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + }, + "topicPolicy": { + "topics": [ + {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} + ] + }, + } + ], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) + assert ( + should_raise is True + ), "Mixed actions with any BLOCKED should raise exceptions" + + # Test 4: NONE action should not raise exception + none_response = { + "action": "NONE", + "outputs": [], + "assessments": [], + } + + should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) + assert should_raise is False, "NONE action should not raise exceptions" + + print("✅ BLOCKED vs ANONYMIZED actions test passed") From a02ec3bfa0d72fcf674c86828974bd6d5853c5c4 Mon Sep 17 00:00:00 2001 From: Lucas Song Date: Tue, 7 Apr 2026 23:06:13 -0700 Subject: [PATCH 011/220] fix(ui): delete policy attachments via controlled modal Replace static Modal.confirm with DeleteResourceModal so attachment delete reliably triggers the API call. Add a regression test covering the confirm->delete flow. Made-with: Cursor --- .../src/components/policies/index.tsx | 70 ++++--- .../policies/policies_panel.test.tsx | 185 ++++++++++++++++++ 2 files changed, 232 insertions(+), 23 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/policies/policies_panel.test.tsx diff --git a/ui/litellm-dashboard/src/components/policies/index.tsx b/ui/litellm-dashboard/src/components/policies/index.tsx index f47b8d78b9d..00ad1874842 100644 --- a/ui/litellm-dashboard/src/components/policies/index.tsx +++ b/ui/litellm-dashboard/src/components/policies/index.tsx @@ -1,8 +1,8 @@ import React, { useState, useEffect, useCallback } from "react"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; -import { Modal, Alert } from "antd"; +import { Alert } from "antd"; import MessageManager from "@/components/molecules/message_manager"; -import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { isAdminRole } from "@/utils/roles"; import PolicyTable from "./policy_table"; import PolicyInfoView from "./policy_info"; @@ -57,6 +57,9 @@ const PoliciesPanel: React.FC = ({ const [isDeleting, setIsDeleting] = useState(false); const [policyToDelete, setPolicyToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); + const [isDeletingAttachment, setIsDeletingAttachment] = useState(false); + const [attachmentToDelete, setAttachmentToDelete] = useState(null); + const [isDeleteAttachmentModalOpen, setIsDeleteAttachmentModalOpen] = useState(false); const [isGuardrailSelectionModalOpen, setIsGuardrailSelectionModalOpen] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState(null); const [existingGuardrailNames, setExistingGuardrailNames] = useState>(new Set()); @@ -166,26 +169,32 @@ const PoliciesPanel: React.FC = ({ setPolicyToDelete(null); }; - const handleDeleteAttachment = (attachmentId: string) => { - Modal.confirm({ - title: "Delete Attachment", - icon: , - content: "Are you sure you want to delete this attachment? This action cannot be undone.", - okText: "Delete", - okType: "danger", - cancelText: "Cancel", - onOk: async () => { - if (!accessToken) return; - try { - await deletePolicyAttachmentCall(accessToken, attachmentId); - MessageManager.success("Attachment deleted successfully"); - fetchAttachments(); - } catch (error) { - console.error("Error deleting attachment:", error); - MessageManager.error("Failed to delete attachment"); - } - }, - }); + const handleDeleteAttachmentClick = (attachmentId: string, policyName?: string) => { + const attachment = attachmentsList.find((a) => a.attachment_id === attachmentId) || null; + setAttachmentToDelete(attachment); + setIsDeleteAttachmentModalOpen(true); + }; + + const handleAttachmentDeleteCancel = () => { + setIsDeleteAttachmentModalOpen(false); + setAttachmentToDelete(null); + }; + + const handleAttachmentDeleteConfirm = async () => { + if (!attachmentToDelete || !accessToken) return; + setIsDeletingAttachment(true); + try { + await deletePolicyAttachmentCall(accessToken, attachmentToDelete.attachment_id); + MessageManager.success("Attachment deleted successfully"); + await fetchAttachments(); + } catch (error) { + console.error("Error deleting attachment:", error); + MessageManager.error("Failed to delete attachment"); + } finally { + setIsDeletingAttachment(false); + setIsDeleteAttachmentModalOpen(false); + setAttachmentToDelete(null); + } }; const handleAttachmentSuccess = () => { @@ -579,7 +588,7 @@ const PoliciesPanel: React.FC = ({ @@ -600,6 +609,21 @@ const PoliciesPanel: React.FC = ({ + + { diff --git a/ui/litellm-dashboard/src/components/policies/policies_panel.test.tsx b/ui/litellm-dashboard/src/components/policies/policies_panel.test.tsx new file mode 100644 index 00000000000..ea1fdd9aa61 --- /dev/null +++ b/ui/litellm-dashboard/src/components/policies/policies_panel.test.tsx @@ -0,0 +1,185 @@ +import React from "react"; +import { screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import PoliciesPanel from "./index"; + +/** + * Ant Design's static Modal.confirm often does not run onOk in the real app (React 18+). + * In jsdom it may still run; we mock confirm as a no-op so the test fails until the panel + * uses a controlled DeleteResourceModal instead of Modal.confirm. + */ +vi.mock("antd", async (importOriginal) => { + const mod = await importOriginal(); + return { + ...mod, + Modal: Object.assign(mod.Modal, { + confirm: vi.fn(), + }), + }; +}); + +const EXPECTED_ATTACHMENT_ID = "att-11111111-2222-3333-4444-555555555555" as const; + +const networkingMocks = vi.hoisted(() => ({ + deletePolicyAttachmentCall: vi.fn().mockResolvedValue(undefined), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), + getPolicyAttachmentsList: vi.fn().mockResolvedValue({ + attachments: [ + { + attachment_id: "att-11111111-2222-3333-4444-555555555555", + policy_name: "test-policy", + scope: null, + teams: [], + keys: [], + models: [], + tags: [], + }, + ], + }), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + getPolicyInfo: vi.fn().mockResolvedValue({}), + deletePolicyCall: vi.fn().mockResolvedValue(undefined), + createPolicyCall: vi.fn(), + updatePolicyCall: vi.fn(), + createPolicyAttachmentCall: vi.fn(), + createGuardrailCall: vi.fn(), + enrichPolicyTemplate: vi.fn(), +})); + +vi.mock("../networking", () => ({ + ...networkingMocks, +})); + +vi.mock("./impact_popover", () => ({ + default: () =>