fix(guardrails): scan privaite extra_body, gate the agent prompt, refuse token IDs under block_entities

litellm/proxy/guardrails/guardrail_hooks/privaite/engine.py

@@ -46,6 +46,8 @@ class PrivaiteEngine(Protocol):
 
     async def process_request(self, messages: Sequence[object]) -> tuple[Sequence[object], RestoreMapping]: ...
 
+    async def gate_document(self, value: object) -> None: ...
+
     async def process_response(self, text: str, mapping: RestoreMapping) -> str: ...
 
 

litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py

@@ -6,12 +6,17 @@ the reversible map in the request metadata, and the post-call hooks consume it t
 the real values in the response, streaming included. The request surface scanned matches
 the PrivAiTe core: chat `messages`, every Responses `input` item shape and
 `prompt.variables`, the `/v1/completions` `prompt` and `suffix`, and the chat
-`prediction.content` and `web_search_options.user_location` fields. Responses API
+`prediction.content` and `web_search_options.user_location` fields. Every string leaf
+of `extra_body` is scanned too, since LiteLLM merges it over the request it sends.
+Other top-level fields are not scanned (`user`, `metadata`, tool definitions, and any
+field LiteLLM does not know, which it forwards to the provider as it is). Responses API
 streaming restore is not implemented yet.
 
-If `block_entities` is configured, a request containing any of those PII types is
-rejected with a 400 before anything is forwarded. The error names the types, never the
-values.
+If `block_entities` is configured, a request with any of those PII types in a scanned
+field is rejected with a 400 before anything is forwarded. The error names the types,
+never the values. The gate also reads the caller's own prompt (Responses `instructions`,
+Anthropic `system`), which is relayed verbatim, and refuses a prompt or input sent as
+token IDs, which it could not check.
 
 There is no detection or masking logic here, it all lives in the `privaite` package,
 which is imported lazily so this module stays importable for the guardrail registry
@@ -41,11 +46,12 @@ from .engine import (
 )
 from .payload import JsonObject, as_array, as_object, read_field, read_list, write_field
 from .request_walk import (
+    agent_prompts,
     aux_field_slots,
-    has_aux_fields,
+    carries_request_text,
     prompt_variable_slots,
-    prompt_variables,
     responses_input_slots,
+    tokenized_fields,
 )
 from .restore import (
     StreamRestorer,
@@ -102,6 +108,20 @@ def _pop_restore_map(request: JsonObject) -> tuple[tuple[str, str], ...]:
     return tuple((fake, original) for fake, original in pairs if isinstance(original, str))
 
 
+def _refuse_tokenized(data: JsonObject) -> None:
+    """With block_entities set, a prompt or input sent as token IDs is refused: no
+    detector reads token IDs, so the gate could not check it. Without block_entities it
+    passes unscanned, as in the PrivAiTe core."""
+    fields: Final = tokenized_fields(data)
+    if fields:
+        from fastapi import HTTPException
+
+        raise HTTPException(
+            status_code=400,
+            detail={"error": f"block_entities cannot check a tokenized {', '.join(fields)}; send it as text"},
+        )
+
+
 def _pseudo_message(text: str) -> JsonObject:
     """A string leaf wrapped as a message, so it crosses the engine's single choke point like chat content."""
     return {"role": "user", "content": text}
@@ -194,17 +214,19 @@ class PrivaiteGuardrail(CustomGuardrail):
         # metadata is caller controlled, so an incoming privaite_map is never trusted:
         # the post-call hooks may only restore from a map this pre-call pass produced.
         _drop_restore_map(data)
-        if (
-            not data.get("messages")
-            and not data.get("input")
-            and prompt_variables(data) is None
-            and not has_aux_fields(data)
-        ):
+        prompts: Final = agent_prompts(data) if self.block_entities else ()
+        if not prompts and not carries_request_text(data):
             return data
+        if self.block_entities:
+            _refuse_tokenized(data)
 
         engine: Final = await self._engine_for(self._languages())
         blocked: Final = blocked_error()
         try:
+            # The agent prompt is relayed verbatim, but a blocked type in it still
+            # rejects the request, before the caller's dict is touched.
+            for prompt in prompts:
+                await engine.gate_document(prompt)
             mapping: Final = await self._anonymize_request(data, engine)
         except blocked as exc:
             from fastapi import HTTPException

litellm/proxy/guardrails/guardrail_hooks/privaite/request_walk.py

@@ -63,6 +63,10 @@ GENERIC_TEXT_FIELDS: Final = ("output", "arguments", "input", "text", "reason")
 """Text-bearing fields of an unknown Responses item. Every one is scanned, since an
 mcp_call carries both `arguments` and `output`."""
 
+AGENT_PROMPT_FIELDS: Final = ("instructions", "system")
+"""The caller's own prompt: Responses `instructions`, Anthropic `system`. Relayed
+verbatim, as the PrivAiTe gateway does, but read by the block gate."""
+
 
 def leaf(text: object, write: Write) -> Iterator[Slot]:
     if isinstance(text, str) and text:
@@ -208,6 +212,8 @@ def has_aux_fields(data: JsonObject) -> bool:
     web_search: Final = as_object(data.get("web_search_options"))
     if web_search is not None and "user_location" in web_search:
         return True
+    if data.get("extra_body"):
+        return True
     suffix: Final = data.get("suffix")
     if isinstance(suffix, str) and suffix:
         return True
@@ -217,14 +223,21 @@ def has_aux_fields(data: JsonObject) -> bool:
 
 def aux_field_slots(data: JsonObject) -> Iterator[Slot]:
     """Request text outside messages and input that LiteLLM forwards verbatim: chat
-    `prediction.content`, `web_search_options.user_location`, and the completions `suffix`
-    and `prompt`. A tokenized (integer array) prompt passes through unscanned."""
+    `prediction.content`, `web_search_options.user_location`, every string leaf of
+    `extra_body`, and the completions `suffix` and `prompt`. A tokenized (integer array)
+    prompt passes through unscanned.
+
+    LiteLLM merges `extra_body` over the request it sends, so `extra_body.messages`
+    replaced the scrubbed messages on the wire. Its keys are provider specific
+    (`contents` for Gemini, say), so it is walked whole rather than by field name."""
     prediction: Final = as_object(data.get("prediction"))
     if prediction is not None and prediction.get("content") is not None:
         yield from data_slots(prediction["content"], field_writer(prediction, "content"))
     web_search: Final = as_object(data.get("web_search_options"))
     if web_search is not None and web_search.get("user_location") is not None:
         yield from data_slots(web_search["user_location"], field_writer(web_search, "user_location"))
+    if data.get("extra_body") is not None:
+        yield from data_slots(data["extra_body"], top_level_writer(data, "extra_body"))
     yield from leaf(data.get("suffix"), top_level_writer(data, "suffix"))
     prompt: Final = data.get("prompt")
     if isinstance(prompt, str):
@@ -235,3 +248,28 @@ def aux_field_slots(data: JsonObject) -> Iterator[Slot]:
         return
     for index, part in enumerate(parts):
         yield from leaf(part, item_writer(parts, index))
+
+
+def carries_request_text(data: JsonObject) -> bool:
+    """Whether the request has a field the anonymization pass reads."""
+    return bool(data.get("messages") or data.get("input")) or prompt_variables(data) is not None or has_aux_fields(data)
+
+
+def agent_prompts(data: JsonObject) -> tuple[object, ...]:
+    """The non-empty agent prompt fields, for the block gate."""
+    return tuple(data[name] for name in AGENT_PROMPT_FIELDS if data.get(name))
+
+
+def _is_token_list(value: object) -> bool:
+    items: Final = as_array(value) or ()
+    return len(items) > 0 and all(type(item) is int for item in items)
+
+
+def tokenized_fields(data: JsonObject) -> tuple[str, ...]:
+    """The completions `prompt` or embeddings `input` sent as token IDs: a list of
+    integers, or a batch holding one. No detector reads token IDs."""
+    return tuple(
+        name
+        for name in ("prompt", "input")
+        if _is_token_list(data.get(name)) or any(_is_token_list(item) for item in as_array(data.get(name)) or ())
+    )

tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py

@@ -103,6 +103,16 @@ class _FakeEngine:
                 raise _FakePIIBlockedError(hit)
         return anonymized, mapping
 
+    async def gate_document(self, value):
+        """The block gate without the rewrite, like PIIEngine.gate_document."""
+        self.gated = [*getattr(self, "gated", []), value]
+        blocked = set(getattr(self.config, "block_entities", None) or [])
+        mapping = _FakeMapping()
+        self._scrub(json.dumps(value), mapping)
+        detected = {fake.strip("<>").rsplit("_", 1)[0] for fake in mapping.get_all_fakes()}
+        if detected & blocked:
+            raise _FakePIIBlockedError(detected & blocked)
+
     async def process_response(self, text, _mapping):
         for fake, original in _FAKES.items():
             text = text.replace(fake, original)
@@ -1458,3 +1468,106 @@ async def test_streaming_restores_copies_and_leaves_the_chunks_litellm_reassembl
     assert [chunk.choices[0].delta.content for chunk in source] == ["Hi <PERS", "ON_1>"], (
         "the chunks LiteLLM keeps to rebuild the logged and cached response must keep the placeholder"
     )
+
+
+@pytest.mark.asyncio
+async def test_pre_call_scrubs_extra_body_in_place():
+    # LiteLLM merges extra_body over the request it sends, so extra_body.messages
+    # would otherwise replace the scrubbed messages on the wire.
+    gr = _make_guardrail()
+    extra_body = {"messages": [{"role": "user", "content": "Email marie@acme.com"}], "top_k": 5}
+    data = {"messages": [{"role": "user", "content": "hi"}], "extra_body": extra_body}
+    data["proxy_server_request"] = {"body": dict(data)}
+    out = await gr.async_pre_call_hook(None, None, data, "completion")
+
+    assert out["extra_body"] is extra_body
+    assert extra_body == {"messages": [{"role": "user", "content": "Email <EMAIL_ADDRESS_1>"}], "top_k": 5}
+    assert data["proxy_server_request"]["body"]["extra_body"]["messages"][0]["content"] == "Email <EMAIL_ADDRESS_1>"
+    assert out["metadata"]["privaite_map"] == {"<EMAIL_ADDRESS_1>": "marie@acme.com"}
+
+
+@pytest.mark.asyncio
+async def test_pre_call_block_gate_reads_extra_body():
+    from fastapi import HTTPException
+
+    gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"])
+    data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "extra_body": {"messages": [{"role": "user", "content": "marie@acme.com"}]},
+    }
+    with pytest.raises(HTTPException) as ei:
+        await gr.async_pre_call_hook(None, None, data, "completion")
+    assert ei.value.status_code == 400
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "data",
+    [
+        {"input": "hi", "instructions": "The user is marie@acme.com."},
+        {"system": [{"type": "text", "text": "The user is marie@acme.com."}]},
+    ],
+)
+async def test_pre_call_block_gate_reads_the_agent_prompt(data):
+    from fastapi import HTTPException
+
+    gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"])
+    with pytest.raises(HTTPException) as ei:
+        await gr.async_pre_call_hook(None, None, data, "aresponses")
+    assert ei.value.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_pre_call_agent_prompt_is_relayed_verbatim():
+    gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"])
+    data = {"input": "hi", "instructions": "The user is Marie Dupont."}
+    out = await gr.async_pre_call_hook(None, None, data, "aresponses")
+    assert out["instructions"] == "The user is Marie Dupont."
+    assert gr._engine.gated == ["The user is Marie Dupont."]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_agent_prompt_is_not_read_without_block_entities():
+    gr = _make_guardrail()
+    data = {"input": "hi", "instructions": "The user is marie@acme.com."}
+    out = await gr.async_pre_call_hook(None, None, data, "aresponses")
+    assert out["instructions"] == "The user is marie@acme.com."
+    assert getattr(gr._engine, "gated", []) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("data", "field"),
+    [
+        ({"prompt": [9906, 1917]}, "prompt"),
+        ({"prompt": [[9906], [1917]]}, "prompt"),
+        ({"prompt": ["hi", [9906]]}, "prompt"),
+        ({"input": [9906, 1917]}, "input"),
+    ],
+)
+async def test_pre_call_block_gate_refuses_token_ids(data, field):
+    # No detector reads token IDs, so the gate cannot check them.
+    from fastapi import HTTPException
+
+    gr = _make_guardrail(block_entities=["EMAIL_ADDRESS"])
+    with pytest.raises(HTTPException) as ei:
+        await gr.async_pre_call_hook(None, None, data, "atext_completion")
+    assert ei.value.status_code == 400
+    assert f"tokenized {field}" in ei.value.detail["error"]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_token_ids_pass_without_block_entities():
+    gr = _make_guardrail()
+    out = await gr.async_pre_call_hook(None, None, {"prompt": [9906, 1917]}, "atext_completion")
+    assert out["prompt"] == [9906, 1917]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_extra_body_only_request_is_scanned():
+    # No top-level messages at all: the early return must not skip the scan.
+    gr = _make_guardrail()
+    extra_body = {"messages": [{"role": "user", "content": "Email marie@acme.com"}]}
+    out = await gr.async_pre_call_hook(None, None, {"extra_body": extra_body}, "completion")
+    assert extra_body["messages"][0]["content"] == "Email <EMAIL_ADDRESS_1>"
+    assert out["metadata"]["privaite_map"] == {"<EMAIL_ADDRESS_1>": "marie@acme.com"}