diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 3d2ab334402..e749b167c75 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -106,8 +106,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # so should_run_guardrail returns True for both pre_call and post_call. if (self.output_parse_pii or self.apply_to_output) and not logging_only: current_hook = self.event_hook - if isinstance(current_hook, str) and current_hook == "pre_call": - self.event_hook = ["pre_call", "post_call"] + if isinstance(current_hook, str) and current_hook != "post_call": + self.event_hook = [current_hook, "post_call"] elif isinstance(current_hook, list) and "post_call" not in current_hook: self.event_hook = current_hook + ["post_call"] self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = ( 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 e336f2f1bf8..55137176a60 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1711,3 +1711,34 @@ async def test_pii_tokens_in_metadata_used_for_unmasking(): ) assert response.choices[0].message.content == "Hello John, how can I help you?" + + +@pytest.mark.parametrize( + "initial_hook", + ["pre_call", "during_call", "pre_mcp_call"], +) +def test_event_hook_auto_expansion_for_all_string_hooks(initial_hook): + """ + Regression test: when output_parse_pii is True, the guardrail must add + 'post_call' to event_hook regardless of the initial string hook value, + not just when it's 'pre_call'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=initial_hook, + ) + assert isinstance(guardrail.event_hook, list) + assert initial_hook in guardrail.event_hook + assert "post_call" in guardrail.event_hook + + +def test_event_hook_no_expansion_when_already_post_call(): + """post_call alone should stay as-is — no expansion needed.""" + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook="post_call", + ) + # Should remain a string "post_call", not expanded to a list + assert guardrail.event_hook == "post_call"