From 828cd42207c995f2418be18ecfb3942ad98724aa Mon Sep 17 00:00:00 2001 From: davida-ps Date: Tue, 21 Apr 2026 17:08:20 +0300 Subject: [PATCH] Fix Prompt Security redaction persistence and role writeback (#25895) * fix(prompt-security): enforce redaction before model call * fix(prompt-security): preserve original roles in modify writeback * chore(prompt-security): add compatibility hook flag and typing overloads --- .../prompt_security/prompt_security.py | 137 +++++- .../test_prompt_security_guardrails.py | 437 ++++++++++++++++++ 2 files changed, 566 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index b3e761869b0..30f33841305 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -1,7 +1,7 @@ import asyncio import base64 import os -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Type +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Type, Union, overload from fastapi import HTTPException @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -56,6 +57,24 @@ class PromptSecurityGuardrail(CustomGuardrail): else: self.check_tool_results = check_tool_results + # Optional compatibility control: + # - True (default): remap during_call -> pre_call + post_call for enforced modify behavior + # - False: preserve configured event_hook exactly + expand_during_call_hooks = kwargs.pop("expand_during_call_hooks", True) + if isinstance(expand_during_call_hooks, str): + expand_during_call_hooks = expand_during_call_hooks.lower() in ( + "true", + "1", + "yes", + ) + + # Prompt Security supports request and response checks for "during_call" mode. + # In LiteLLM, during_call hooks execute in parallel and cannot safely apply modify actions. + # Translate during_call to pre_call + post_call so modifications are actually enforced. + event_hook = kwargs.get("event_hook") + if expand_during_call_hooks: + kwargs["event_hook"] = self._expand_event_hooks_for_prompt_security(event_hook) + if not self.api_key or not self.api_base: msg = ( "Couldn't get Prompt Security api base or key, " @@ -70,6 +89,58 @@ class PromptSecurityGuardrail(CustomGuardrail): super().__init__(**kwargs) + @staticmethod + def _expand_event_hooks_for_prompt_security( + event_hook: Optional[ + Union[ + GuardrailEventHooks, + List[GuardrailEventHooks], + str, + List[str], + ] + ], + ) -> Optional[Union[str, List[str], GuardrailEventHooks, List[GuardrailEventHooks]]]: + """Expand during_call to pre_call + post_call for Prompt Security.""" + if event_hook is None: + return None + + during_call = GuardrailEventHooks.during_call.value + pre_call = GuardrailEventHooks.pre_call.value + post_call = GuardrailEventHooks.post_call.value + + if isinstance(event_hook, GuardrailEventHooks): + if event_hook == GuardrailEventHooks.during_call: + return [pre_call, post_call] + return event_hook + + if isinstance(event_hook, str): + if event_hook == during_call: + return [pre_call, post_call] + return event_hook + + normalized_hooks: List[str] = [] + for hook in event_hook: + if isinstance(hook, GuardrailEventHooks): + normalized_hooks.append(hook.value) + else: + normalized_hooks.append(hook) + + if during_call not in normalized_hooks: + return normalized_hooks + + expanded_hooks = [hook for hook in normalized_hooks if hook != during_call] + if pre_call not in expanded_hooks: + expanded_hooks.append(pre_call) + if post_call not in expanded_hooks: + expanded_hooks.append(post_call) + + verbose_proxy_logger.debug( + "Prompt Security Guardrail: Expanded event_hook=%s to %s", + event_hook, + expanded_hooks, + ) + return expanded_hooks + @log_guardrail_information async def apply_guardrail( self, @@ -167,7 +238,9 @@ class PromptSecurityGuardrail(CustomGuardrail): await self._process_standalone_images(images, user_api_key_alias) # Filter messages by role for the API call - filtered_messages = self.filter_messages_by_role(messages) + filtered_messages, filtered_message_indexes = self.filter_messages_by_role( + messages, include_original_indices=True + ) if not filtered_messages: verbose_proxy_logger.debug( @@ -218,11 +291,42 @@ class PromptSecurityGuardrail(CustomGuardrail): + ", ".join(violations), ) elif action == "modify": - # Extract modified texts from modified_messages modified_messages = result.get("modified_messages", []) - modified_texts = self._extract_texts_from_messages(modified_messages) - if modified_texts: - inputs["texts"] = modified_texts + if modified_messages: + # If Prompt Security returned a message per scanned input message, map + # those back to the original message list so we preserve full context. + if len(modified_messages) == len(filtered_messages): + for modified_idx, modified_message in enumerate(modified_messages): + if not isinstance(modified_message, dict): + continue + original_idx = filtered_message_indexes[modified_idx] + original_message = messages[original_idx] + if not isinstance(original_message, dict): + messages[original_idx] = modified_message + continue + + # Preserve the original role for in-flight request messages. + # Prompt Security may transform non-standard roles to "other" + # for scanning, but "other" should never be sent to model providers. + merged_message = {**original_message, **modified_message} + merged_message["role"] = original_message.get( + "role", merged_message.get("role") + ) + messages[original_idx] = merged_message + + inputs["texts"] = self._extract_texts_from_messages(messages) + if structured_messages: + inputs["structured_messages"] = messages + else: + verbose_proxy_logger.warning( + "Prompt Security Guardrail: modified_messages length (%d) did not match scanned messages length (%d); " + "falling back to modified text extraction only.", + len(modified_messages), + len(filtered_messages), + ) + modified_texts = self._extract_texts_from_messages(modified_messages) + if modified_texts: + inputs["texts"] = modified_texts return inputs @@ -623,7 +727,19 @@ class PromptSecurityGuardrail(CustomGuardrail): return processed_messages - def filter_messages_by_role(self, messages: list) -> list: + @overload + def filter_messages_by_role( + self, messages: list, include_original_indices: Literal[True] + ) -> Tuple[list, List[int]]: ... + + @overload + def filter_messages_by_role( + self, messages: list, include_original_indices: Literal[False] = False + ) -> list: ... + + def filter_messages_by_role( + self, messages: list, include_original_indices: bool = False + ) -> Union[list, Tuple[list, List[int]]]: """Filter messages to only include standard OpenAI/Anthropic roles. Behavior depends on check_tool_results flag: @@ -634,13 +750,15 @@ class PromptSecurityGuardrail(CustomGuardrail): """ supported_roles = ["system", "user", "assistant"] filtered_messages = [] + filtered_message_indexes: List[int] = [] transformed_count = 0 filtered_count = 0 - for message in messages: + for message_idx, message in enumerate(messages): role = message.get("role", "") if role in supported_roles: filtered_messages.append(message) + filtered_message_indexes.append(message_idx) else: if self.check_tool_results: transformed_message = { @@ -652,6 +770,7 @@ class PromptSecurityGuardrail(CustomGuardrail): }, } filtered_messages.append(transformed_message) + filtered_message_indexes.append(message_idx) transformed_count += 1 verbose_proxy_logger.debug( "Prompt Security Guardrail: Transformed message from role '%s' to 'other'", @@ -678,6 +797,8 @@ class PromptSecurityGuardrail(CustomGuardrail): len(filtered_messages), ) + if include_original_indices: + return filtered_messages, filtered_message_indexes return filtered_messages def _build_headers(self, user_api_key_alias: Optional[str] = None) -> dict: diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index f35d64b89e3..fe151305713 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -11,6 +11,7 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im PromptSecurityGuardrailMissingSecrets, PromptSecurityGuardrail, ) +from litellm.types.guardrails import GuardrailEventHooks sys.path.insert( 0, os.path.abspath("../..") @@ -77,6 +78,129 @@ def test_prompt_security_guard_config_no_api_key(): ) +def test_prompt_security_during_call_expands_to_pre_and_post_hooks(): + """Test that during_call mode is expanded to pre_call + post_call for Prompt Security.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", event_hook="during_call", default_on=True + ) + + assert isinstance(guardrail.event_hook, list) + assert set(guardrail.event_hook) == {"pre_call", "post_call"} + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.pre_call) is True + ) + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.post_call) is True + ) + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.during_call) is False + ) + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +def test_prompt_security_during_call_expands_to_pre_and_post_hooks_for_enum_input(): + """Test enum input expands during_call to pre_call + post_call.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook=GuardrailEventHooks.during_call, + default_on=True, + ) + + assert isinstance(guardrail.event_hook, list) + assert set(guardrail.event_hook) == {"pre_call", "post_call"} + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +def test_prompt_security_during_call_expands_to_pre_and_post_hooks_for_list_input(): + """Test list input with during_call expands to include pre_call + post_call.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], + default_on=True, + ) + + assert isinstance(guardrail.event_hook, list) + assert set(guardrail.event_hook) == {"pre_call", "post_call"} + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +def test_prompt_security_event_hook_list_without_during_call_is_preserved(): + """Test list input without during_call remains unchanged.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + default_on=True, + ) + + assert guardrail.event_hook == ["pre_call", "post_call"] + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +def test_prompt_security_during_call_not_expanded_when_flag_disabled(): + """Test compatibility flag can keep during_call behavior unchanged.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="during_call", + default_on=True, + expand_during_call_hooks=False, + ) + + assert guardrail.event_hook == "during_call" + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.during_call) is True + ) + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.pre_call) is False + ) + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.post_call) is False + ) + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +def test_prompt_security_during_call_not_expanded_when_flag_disabled_string_value(): + """Test string config values for compatibility flag are respected.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="during_call", + default_on=True, + expand_during_call_hooks="false", + ) + + assert guardrail.event_hook == "during_call" + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + @pytest.mark.asyncio async def test_apply_guardrail_block_request(): """Test that apply_guardrail blocks malicious prompts""" @@ -534,6 +658,319 @@ async def test_user_api_key_alias_forwarding(): del os.environ["PROMPT_SECURITY_API_BASE"] +@pytest.mark.asyncio +async def test_apply_guardrail_modify_request_preserves_filtered_message_alignment(): + """Test modify action preserves full message alignment when tool/function messages are filtered out.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + check_tool_results=False, + ) + + request_data = { + "messages": [ + {"role": "system", "content": "System context"}, + {"role": "user", "content": "my id is 228230355"}, + {"role": "tool", "content": "tool output should remain"}, + {"role": "assistant", "content": "Acknowledged"}, + ] + } + + inputs = { + "texts": [ + "System context", + "my id is 228230355", + "tool output should remain", + "Acknowledged", + ], + "structured_messages": request_data["messages"], + } + + # Prompt Security receives only system/user/assistant (tool filtered out), + # but we still need the returned texts aligned with the original messages. + modified_messages = [ + {"role": "system", "content": "System context"}, + {"role": "user", "content": "my id is [REDACTED]"}, + {"role": "assistant", "content": "Acknowledged"}, + ] + + mock_response = Response( + json={ + "result": { + "prompt": {"action": "modify", "modified_messages": modified_messages} + } + }, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == [ + "System context", + "my id is [REDACTED]", + "tool output should remain", + "Acknowledged", + ] + assert result["structured_messages"][0]["content"] == "System context" + assert result["structured_messages"][1]["content"] == "my id is [REDACTED]" + assert result["structured_messages"][2]["content"] == "tool output should remain" + assert result["structured_messages"][2]["role"] == "tool" + assert result["structured_messages"][3]["content"] == "Acknowledged" + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_modify_request_preserves_original_tool_role_when_checking_tool_results(): + """Test modify action does not leak Prompt Security's temporary role='other' to model messages.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + check_tool_results=True, + ) + + request_data = { + "messages": [ + {"role": "user", "content": "Summarize this tool output"}, + {"role": "tool", "content": "customer id is 228230355"}, + ] + } + inputs = { + "texts": ["Summarize this tool output", "customer id is 228230355"], + "structured_messages": request_data["messages"], + } + + # Prompt Security sees the tool message transformed to role='other'. + modified_messages = [ + {"role": "user", "content": "Summarize this tool output"}, + {"role": "other", "content": "customer id is [REDACTED]"}, + ] + + mock_response = Response( + json={ + "result": { + "prompt": {"action": "modify", "modified_messages": modified_messages} + } + }, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["Summarize this tool output", "customer id is [REDACTED]"] + assert result["structured_messages"][1]["role"] == "tool" + assert result["structured_messages"][1]["content"] == "customer id is [REDACTED]" + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_modify_request_length_mismatch_falls_back_to_modified_texts(): + """Test mismatch between scanned and modified message lengths uses modified text fallback.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + + request_data = { + "messages": [ + {"role": "user", "content": "my id is 228230355"}, + {"role": "assistant", "content": "Acknowledged"}, + ] + } + inputs = { + "texts": ["my id is 228230355", "Acknowledged"], + "structured_messages": request_data["messages"], + } + + mock_response = Response( + json={ + "result": { + "prompt": { + "action": "modify", + # Scanned length is 2 but modified length is 1. + "modified_messages": [ + {"role": "user", "content": "my id is [REDACTED]"} + ], + } + } + }, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["my id is [REDACTED]"] + assert result["structured_messages"] == request_data["messages"] + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + +@pytest.mark.asyncio +async def test_simulated_two_turn_conversation_redacts_id_before_model_memory(): + """Simulate the reported two-turn leak pattern and ensure redaction persists.""" + os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" + os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="during_call", + default_on=True, + ) + + secret_id = "228230355" + redacted_id = "[REDACTED_ID_IL_ID_NUMBER_1]" + + class MockMemoryLLM: + def __init__(self): + self.memory = "" + + def complete(self, messages): + for msg in messages: + if msg.get("role") != "user": + continue + content = msg.get("content") + if not isinstance(content, str): + continue + if secret_id in content: + self.memory = secret_id + elif redacted_id in content: + self.memory = redacted_id + + latest_user = next( + ( + m.get("content") + for m in reversed(messages) + if m.get("role") == "user" and isinstance(m.get("content"), str) + ), + "", + ) + if "echo it" in latest_user.lower(): + return self.memory + return f"I understand your ID is {self.memory}." + + async def mock_prompt_security_post(*args, **kwargs): + payload_messages = kwargs.get("json", {}).get("messages", []) + modified_messages = [] + modified = False + + for message in payload_messages: + if not isinstance(message, dict): + modified_messages.append(message) + continue + content = message.get("content") + if isinstance(content, str): + new_content = content.replace(secret_id, redacted_id) + if new_content != content: + modified = True + modified_messages.append({**message, "content": new_content}) + else: + modified_messages.append(message) + + prompt_result = {"action": "allow"} + if modified: + prompt_result = {"action": "modify", "modified_messages": modified_messages} + + mock_response = Response( + json={"result": {"prompt": prompt_result}}, + status_code=200, + request=Request( + method="POST", url="https://test.prompt.security/api/protect" + ), + ) + mock_response.raise_for_status = lambda: None + return mock_response + + async def apply_pre_call_guardrail(messages): + request_data = {"messages": messages} + if ( + guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) + is not True + ): + return messages + + inputs = { + "texts": [m["content"] for m in messages if isinstance(m.get("content"), str)], + "structured_messages": messages, + } + guardrailed_inputs = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + return guardrailed_inputs.get("structured_messages", messages) + + llm = MockMemoryLLM() + + with patch.object( + guardrail.async_handler, "post", side_effect=mock_prompt_security_post + ): + # Turn 1: user provides sensitive ID + turn_1_messages = [{"role": "user", "content": f"my id is {secret_id}"}] + guarded_turn_1_messages = await apply_pre_call_guardrail(turn_1_messages) + + assert guardrail.should_run_guardrail({}, GuardrailEventHooks.pre_call) is True + assert ( + guardrail.should_run_guardrail({}, GuardrailEventHooks.during_call) is False + ) + + assert secret_id not in guarded_turn_1_messages[0]["content"] + assert redacted_id in guarded_turn_1_messages[0]["content"] + + assistant_turn_1 = llm.complete(guarded_turn_1_messages) + assert secret_id not in assistant_turn_1 + + # Turn 2: user asks the model to repeat + turn_2_messages = guarded_turn_1_messages + [ + {"role": "assistant", "content": assistant_turn_1}, + {"role": "user", "content": "echo it"}, + ] + guarded_turn_2_messages = await apply_pre_call_guardrail(turn_2_messages) + assistant_turn_2 = llm.complete(guarded_turn_2_messages) + + assert assistant_turn_2 == redacted_id + assert secret_id not in assistant_turn_2 + + del os.environ["PROMPT_SECURITY_API_KEY"] + del os.environ["PROMPT_SECURITY_API_BASE"] + + @pytest.mark.asyncio async def test_role_filtering(): """Test that tool/function messages are filtered out by default"""