From 40c26178fae3b8a174fa8188a671571cd353b1d6 Mon Sep 17 00:00:00 2001
From: Ninad Phalak <ninadphalak@gmail.com>
Date: Sat, 26 Sep 2026 09:31:05 -0500
Subject: [PATCH] fix(guardrails): redact llm_shield_proxy predicted outputs
 and output schemas

`prediction.content` is the caller's own draft of the reply, so it is
redacted into the caller vault and restored with the reply. The
descriptions in a structured-output schema (Chat
response_format.json_schema, Responses text.format) are application
authored like tool schemas, so they go to the non-restorable vault.
---
 .../llm_shield_proxy/llm_shield_proxy.py      | 24 +++++++++++++++-
 .../guardrail_hooks/test_llm_shield_proxy.py  | 28 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
index ebafb10831a..01770f83141 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
@@ -332,6 +332,27 @@ def _collect_schema_descriptions(schema: object, privileged: _SlotSink) -> None:
             pending.extend((value, depth + 1) for value in node if isinstance(value, (dict, list)))
 
 
+def _collect_output_contracts(data: MutableRequest, slots: _SlotSink, privileged: _SlotSink) -> None:
+    """Text the caller sends to shape the reply rather than to prompt it.
+
+    A predicted output (`prediction.content`) is the caller's own draft of the answer, so
+    it goes with their text: the model largely repeats it, and it has to come back. A
+    structured-output schema -- Chat `response_format.json_schema`, Responses
+    `text.format` -- is application-authored like a tool schema, so its descriptions go
+    to the privileged sink, and its names and types stay as sent.
+    """
+    prediction: Final = data.get("prediction")
+    if isinstance(prediction, dict):
+        _collect(prediction, "content", slots)
+        _collect_text_parts(prediction, "content", slots)
+    response_format: Final = data.get("response_format")
+    if isinstance(response_format, dict):
+        _collect_schema_descriptions(response_format.get("json_schema"), privileged)
+    text_options: Final = data.get("text")
+    if isinstance(text_options, dict):
+        _collect_schema_descriptions(text_options.get("format"), privileged)
+
+
 def _collect_end_user_ids(data: MutableRequest, privileged: _SlotSink) -> None:
     """`user` and `safety_identifier` are forwarded to the provider and often hold an email.
 
@@ -857,7 +878,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
 
         The split exists because the response is restored against one vault only.
         Server-authored spans -- system and developer turns, Anthropic's top-level
-        `system`, the Responses API `instructions`, tool definitions -- go into a
+        `system`, the Responses API `instructions`, tool and output schemas -- go into a
         vault nothing is ever restored against, so a caller who gets the model to
         echo one of their placeholders back receives the placeholder, not the value
         behind it. End-user identifiers go there too: nothing in a reply needs them.
@@ -880,6 +901,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
         _collect_prompt(data, slots)
         _collect_system(data, privileged)
         _collect_tool_definitions(data, privileged)
+        _collect_output_contracts(data, slots, privileged)
         _collect_end_user_ids(data, privileged)
         return tuple(slots), tuple(privileged)
 
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py
index a24b3845f9b..64768b39070 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py
@@ -829,6 +829,34 @@ class TestVaultIsolation:
                 },
                 id="anthropic-schema-description",
             ),
+            pytest.param(
+                {
+                    "messages": [{"role": "user", "content": "U"}],
+                    "response_format": {
+                        "type": "json_schema",
+                        "json_schema": {"name": "n", "description": "S", "schema": {"type": "object"}},
+                    },
+                },
+                id="chat-response-format",
+            ),
+            pytest.param(
+                {
+                    "input": "U",
+                    "text": {
+                        "format": {
+                            "type": "json_schema",
+                            "name": "n",
+                            "schema": {"properties": {"a": {"description": "S"}}},
+                        }
+                    },
+                },
+                id="responses-text-format",
+            ),
+            pytest.param({"prediction": {"type": "content", "content": "U"}, "instructions": "S"}, id="prediction"),
+            pytest.param(
+                {"prediction": {"type": "content", "content": [{"type": "text", "text": "U"}]}, "instructions": "S"},
+                id="prediction-parts",
+            ),
             pytest.param({"messages": [{"role": "user", "content": "U"}], "user": "S"}, id="end-user-id"),
             pytest.param({"input": "U", "safety_identifier": "S"}, id="safety-identifier"),
         ],