fix(policy_engine): fail closed on content filter category MASK steps for streaming pipelines

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(

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

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