From 90c8031dd76c5565c25b2adfe301a4ce715a6964 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 29 Aug 2026 14:32:35 -0700
Subject: [PATCH] fix(policy_engine): fail closed on content filter category
 MASK steps for streaming pipelines

---
 .../litellm_content_filter/content_filter.py  |  2 ++
 .../content_filter/test_content_filter.py     | 20 +++++++++++++++++
 .../proxy_logging/test_guardrail_pipeline.py  | 22 +++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
index bd31882841e..85eb50c78e7 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
@@ -1952,6 +1952,8 @@ class ContentFilterGuardrail(CustomGuardrail):
             super().rewrites_streamed_output()
             or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns)
             or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values())
+            or any(action == ContentFilterAction.MASK for _, _, action in self.category_keywords.values())
+            or any(action == ContentFilterAction.MASK for _, _, action in self.always_block_category_keywords.values())
         )
 
     async def async_post_call_streaming_iterator_hook(
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
index ffbedfa43ff..73020fe3e6f 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
@@ -3104,3 +3104,23 @@ class TestRewritesStreamedOutput:
         )
 
         assert guardrail.rewrites_streamed_output() is True
+
+    @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)])
+    def test_category_keywords_follow_the_category_action(self, action, expected):
+        guardrail = ContentFilterGuardrail(
+            guardrail_name="cf",
+            categories=[{"category": "bias_gender", "enabled": True, "action": action}],
+        )
+
+        assert guardrail.category_keywords and not guardrail.always_block_category_keywords
+        assert guardrail.rewrites_streamed_output() is expected
+
+    @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)])
+    def test_always_block_category_keywords_follow_the_category_action(self, action, expected):
+        guardrail = ContentFilterGuardrail(
+            guardrail_name="cf",
+            categories=[{"category": "age_discrimination", "enabled": True, "action": action}],
+        )
+
+        assert guardrail.always_block_category_keywords and not guardrail.category_keywords
+        assert guardrail.rewrites_streamed_output() is expected
diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
index ee0a9a10172..895a986f348 100644
--- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
+++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
@@ -1493,6 +1493,28 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas
     assert "a MASK action" in info.value.detail["error"]["message"]
 
 
+@pytest.mark.asyncio
+async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    guardrail = ContentFilterGuardrail(
+        guardrail_name="gr-post",
+        event_hook=GuardrailEventHooks.post_call,
+        categories=[{"category": "bias_gender", "enabled": True, "action": "MASK"}],
+    )
+    monkeypatch.setattr(litellm, "callbacks", [guardrail])
+    data = _post_call_pipeline_data(stream=True)
+    user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions")
+
+    with pytest.raises(HTTPException) as info:
+        await proxy_logging.pre_call_hook(
+            user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
+        )
+
+    assert info.value.status_code == 400
+    assert info.value.detail["error"]["guardrails"] == ("gr-post",)
+
+
 @pytest.mark.asyncio
 async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation(
     proxy_logging, make_user_api_key_auth, monkeypatch