From 6fae4b3c3977edcddfeb9e0080f91718ae5c77d6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 2 Sep 2026 13:32:21 -0700
Subject: [PATCH] fix(guardrails): keep the presidio output masker from
 unmasking after an in-memory update

---
 .../proxy/guardrails/guardrail_hooks/presidio.py    |  2 ++
 .../guardrails/guardrail_hooks/test_presidio.py     | 12 ++++++++++++
 .../proxy/guardrails/test_guardrail_registry.py     | 13 +++++++++++--
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
index da51a905ae3..70ea21320ee 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
@@ -1633,6 +1633,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         Update the guardrails litellm params in memory
         """
         super().update_in_memory_litellm_params(litellm_params)
+        if self.apply_to_output:
+            self.output_parse_pii = False
         if litellm_params.pii_entities_config:
             self.pii_entities_config = litellm_params.pii_entities_config
         if litellm_params.presidio_score_thresholds:
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
index fcf940afd0d..84f7611c0c0 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
@@ -3129,6 +3129,18 @@ def test_update_in_memory_applies_analyze_chunk_size():
     assert guardrail.presidio_analyze_chunk_size_bytes == 99_000
 
 
+def test_update_in_memory_keeps_output_masker_from_unmasking():
+    masker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, apply_to_output=True, output_parse_pii=False)
+    unmasker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True)
+    params = LitellmParams(guardrail="presidio", mode="pre_call", output_parse_pii=True)
+
+    masker.update_in_memory_litellm_params(params)
+    unmasker.update_in_memory_litellm_params(params)
+
+    assert (masker.apply_to_output, masker.output_parse_pii) == (True, False)
+    assert (unmasker.apply_to_output, unmasker.output_parse_pii) == (False, True)
+
+
 def test_merge_drops_truncated_same_type_fragment_from_overlap():
     """A boundary entity seen truncated by chunk 1 and whole by chunk 2 must
     merge to the single full span; keeping both overlapping spans corrupts the
diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
index 5cbdef5f92f..24742e1bac2 100644
--- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
+++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
@@ -566,7 +566,14 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st
     try:
         handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"}))
         tracked = _presidio_callbacks_in(litellm.callbacks)
-        roles_before = [(callback.apply_to_output, callback.event_hook) for callback in tracked]
+        roles_before = [
+            (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked
+        ]
+        assert roles_before == [
+            (False, True, [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]),
+            (False, True, GuardrailEventHooks.post_call),
+            (True, False, GuardrailEventHooks.post_call),
+        ]
 
         updated = Guardrail(
             guardrail_id=PRESIDIO_SIBLINGS_GID,
@@ -585,7 +592,9 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st
         handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated)
 
         assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3
-        assert [(callback.apply_to_output, callback.event_hook) for callback in tracked] == roles_before
+        assert [
+            (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked
+        ] == roles_before
         assert _presidio_callbacks_in(litellm.callbacks) == tracked
     finally:
         for cb_list, snapshot in zip(lists, snapshots):