From 11b31dc37a5f936a41dcb4ad5b2ce2b75cc37eb8 Mon Sep 17 00:00:00 2001
From: Aishwary Dongre <dongreaishwary@gmail.com>
Date: Thu, 15 Oct 2026 10:27:05 +0530
Subject: [PATCH] fix(proxy): preserve PII-redacted prompt in image generation
 requests

When a pre-call guardrail (e.g. Presidio) redacts data['prompt'] for
/v1/images/generations, the endpoint was unconditionally rebuilding
the prompt from messages, discarding the redaction. The upstream model
received the original unredacted text.

This preserves the redacted prompt by only rebuilding from messages
if the guardrail didn't modify it.

Fixes #46473
---
 litellm/proxy/image_endpoints/endpoints.py    |  9 ++-
 .../hooks/test_image_generation_guardrails.py | 60 +++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py
index 28fb9baf2aa..e236d24c77f 100644
--- a/litellm/proxy/image_endpoints/endpoints.py
+++ b/litellm/proxy/image_endpoints/endpoints.py
@@ -145,9 +145,12 @@ async def image_generation(
             user_api_key_dict=user_api_key_dict, data=data, call_type="image_generation"
         )
 
-        messages: Final = data.get("messages")
-        if isinstance(messages, list) and messages:
-            data["prompt"] = get_str_from_messages(messages)
+        # Guardrails may redact data["prompt"] directly. Only rebuild from messages if prompt wasn't modified.
+        prompt_after_hook: Final = data.get("prompt")
+        if prompt_after_hook == prompt_value:
+            messages: Final = data.get("messages")
+            if isinstance(messages, list) and messages:
+                data["prompt"] = get_str_from_messages(messages)
         data.pop("messages", None)
 
         ## ROUTE TO CORRECT ENDPOINT ##
diff --git a/tests/unit/proxy/hooks/test_image_generation_guardrails.py b/tests/unit/proxy/hooks/test_image_generation_guardrails.py
index fd8299b07ec..5702d40f2ae 100644
--- a/tests/unit/proxy/hooks/test_image_generation_guardrails.py
+++ b/tests/unit/proxy/hooks/test_image_generation_guardrails.py
@@ -292,3 +292,63 @@ async def test_non_default_guardrail_skipped_for_image_generation():
     assert (
         guardrail.called is False
     ), "Opt-in guardrail should not fire without explicit request"
+
+
+# ---------------------------------------------------------------------------
+# 6. Pre-call PII guardrail redaction is preserved for image generation
+# ---------------------------------------------------------------------------
+
+
+class MaskingGuardrail(CustomGuardrail):
+    """Guardrail that masks PII in the prompt (pre-call hook)."""
+
+    def __init__(self):
+        super().__init__(
+            guardrail_name="masking_guardrail",
+            default_on=True,
+            event_hook=GuardrailEventHooks.pre_call,
+        )
+
+    async def async_pre_call_hook(
+        self,
+        data: dict,
+        user_api_key_dict: UserAPIKeyAuth,
+        call_type: str,
+    ) -> dict:
+        # Simulate PII redaction: replace SSN pattern with [REDACTED]
+        prompt = data.get("prompt", "")
+        if "123-45-6789" in prompt:
+            data["prompt"] = prompt.replace("123-45-6789", "[REDACTED_SSN]")
+        return data
+
+
+@pytest.mark.asyncio
+async def test_pre_call_pii_redaction_preserved_for_image_generation():
+    """
+    Verify that when a pre-call guardrail redacts data["prompt"], the redacted
+    value is preserved and not overwritten by rebuilding from messages.
+    
+    This tests the fix for issue #46473 where the image generation endpoint
+    was rebuilding prompt from messages after the guardrail ran, which
+    discarded the redacted value.
+    """
+    guardrail = MaskingGuardrail()
+    original_prompt = "Generate an image for user SSN: 123-45-6789"
+    
+    with patch("litellm.callbacks", [guardrail]):
+        proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
+        
+        data = {"model": "dall-e-3", "prompt": original_prompt}
+        user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
+        
+        # Run pre_call_hook (this is what the endpoint does)
+        result = await proxy_logging.pre_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            data=data,
+            call_type="image_generation",
+        )
+    
+    # The prompt should be redacted
+    assert result["prompt"] == "Generate an image for user SSN: [REDACTED_SSN]"
+    # Original prompt should NOT appear in the result
+    assert "123-45-6789" not in result["prompt"]