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

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:

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"""