fix(guardrails): restore privaite responses on a copy

The post-call hooks wrote the restored originals into the objects LiteLLM handed them. LiteLLM keeps those objects, the response for its cache write and success logs and every stream chunk in the wrapper's chunk list, which it reassembles into the cached and logged response at the end of the stream. A cached reply could then carry one caller's values under a key built from the anonymized request, where another caller's identical-looking request would hit them

Both hooks now restore a deep copy and return or yield that, so what LiteLLM keeps holds placeholders only. The proxy adopts a guardrail's returned response, so callers still get the restored values

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

@@ -21,6 +21,7 @@ without it. See https://github.com/crp4222/PrivAiTe.
 from __future__ import annotations
 
 import asyncio
+import copy
 import itertools
 from collections.abc import AsyncGenerator, AsyncIterable, Iterable, Iterator, Sequence
 from typing import TYPE_CHECKING, Final, TypeAlias
@@ -228,6 +229,13 @@ class PrivaiteGuardrail(CustomGuardrail):
         user_api_key_dict: UserAPIKeyAuth,
         response: object,
     ) -> object:
+        """Restores the original values in a copy of the response.
+
+        LiteLLM keeps the object it hands this hook for its response cache and its
+        success logs, and a cache write may serialize it after this hook ran. Restoring it
+        in place would put this caller's values in the cache under a key built from the
+        anonymized request, where another caller's identical-looking request hits them.
+        """
         if not self.deanonymize:
             return response
         pairs: Final = _pop_restore_map(data)
@@ -235,12 +243,13 @@ class PrivaiteGuardrail(CustomGuardrail):
             return response
         mapping: Final = restore_mapping(pairs)
         engine: Final = await self._engine_for(self._languages())
-        for choice in read_list(response, "choices"):
+        restored: Final = copy.deepcopy(response)
+        for choice in read_list(restored, "choices"):
             message = read_field(choice, "message")
             if message is not None:
                 await restore_message(message, engine, mapping)
-        await restore_responses_output(response, engine, mapping)
-        return response
+        await restore_responses_output(restored, engine, mapping)
+        return restored
 
     async def async_post_call_streaming_iterator_hook(
         self,
@@ -248,6 +257,12 @@ class PrivaiteGuardrail(CustomGuardrail):
         response: AsyncIterable[ModelResponseStream],
         request_data: JsonObject,
     ) -> AsyncGenerator[ModelResponseStream, None]:
+        """Restores the original values in copies of the chunks.
+
+        The stream wrapper keeps every chunk it yields and rebuilds the whole response from
+        them for the cache and the success logs once the stream ends, so the chunks
+        themselves keep their placeholders, as in the non-streaming hook.
+        """
         pairs: Final = _pop_restore_map(request_data)
         if not self.deanonymize or not pairs:
             async for chunk in response:
@@ -260,8 +275,9 @@ class PrivaiteGuardrail(CustomGuardrail):
             restore_buffer_factory(),
         )
         async for chunk in response:
-            restorer.restore_chunk(chunk)
-            yield chunk
+            restored = copy.deepcopy(chunk)
+            restorer.restore_chunk(restored)
+            yield restored
         for drained in restorer.drain():
             yield drained
 

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

@@ -1424,3 +1424,37 @@ async def test_engine_for_a_privaite_without_pii_engine_recommends_the_install(m
 
     with pytest.raises(RuntimeError, match=r"pip install 'privaite>=0\.7\.2'"):
         await _make_guardrail()._engine_for(["en"])
+
+
+@pytest.mark.asyncio
+async def test_post_call_restores_a_copy_and_leaves_the_response_litellm_caches_and_logs():
+    gr = _make_guardrail()
+    message = types.SimpleNamespace(content="Sending to <PERSON_1>", tool_calls=None, function_call=None)
+    response = types.SimpleNamespace(choices=[types.SimpleNamespace(message=message)])
+
+    out = await gr.async_post_call_success_hook({"metadata": {"privaite_map": _FAKES}}, None, response)
+
+    assert out.choices[0].message.content == "Sending to Marie Dupont"
+    assert message.content == "Sending to <PERSON_1>", "the object LiteLLM caches and logs must keep the placeholder"
+
+
+@pytest.mark.asyncio
+async def test_streaming_restores_copies_and_leaves_the_chunks_litellm_reassembles():
+    gr = _make_guardrail()
+    source = [
+        _choice_chunk(_bare_delta(content="Hi <PERS")),
+        _choice_chunk(_bare_delta(content="ON_1>"), finish="stop"),
+    ]
+
+    async def _source():
+        for chunk in source:
+            yield chunk
+
+    chunks = await _collect(
+        gr.async_post_call_streaming_iterator_hook(None, _source(), {"metadata": {"privaite_map": _FAKES}})
+    )
+
+    assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hi Marie Dupont"
+    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"
+    )