feat(atr-guardrail): add async_post_call_streaming_hook

Addresses the veria-ai streaming-bypass finding. Scans the aggregated
streamed response after stream completion using LiteLLM's existing
post-call streaming surface; per-chunk scanning would emit false
negatives for split-across-chunk attack patterns, so we wait for the
aggregated text.

Three new tests covering the streaming hook: block-on-match,
pass-when-no-match, and no-op-on-empty-response.

Signed-off-by: Adam Lin <adam@agentthreatrule.org>

litellm/proxy/guardrails/guardrail_hooks/atr/atr.py

@@ -220,6 +220,42 @@ class ATRGuardrail(CustomGuardrail):
             )
         return response
 
+    @log_guardrail_information
+    async def async_post_call_streaming_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: str,
+    ) -> Any:
+        """
+        Scan the aggregated streamed response after stream completion.
+
+        ATR rules match against complete content (a regex over a full
+        response). Per-chunk scanning would emit false negatives (the
+        attack pattern split across two chunks never appears in either)
+        and inconsistent false positives. LiteLLM aggregates the streamed
+        response before this hook fires, so we get a uniform policy
+        whether the caller opts into streaming or not.
+
+        Known limitation (documented for honesty rather than fixed): an
+        attacker who streams a long-running response specifically to
+        inject content that is acted on mid-stream is out of scope. That
+        requires per-chunk inspection with a stateful aggregator and a
+        semantic gate, not a regex catalog.
+        """
+        if response is None or len(response) == 0:
+            return response
+
+        matches = self._scan(response, event_type="llm_output")
+        if matches:
+            import json
+
+            error_detail = {
+                "error": "Streamed response blocked by ATR guardrail",
+                "matched_rules": [self._summarize_match(m) for m in matches],
+            }
+            return f"data: {json.dumps({'error': error_detail})}\n\n"
+        return response
+
     # ------------------------------------------------------------------
     # Internals
     # ------------------------------------------------------------------

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py

@@ -466,3 +466,95 @@ def test_scan_unknown_severity_treated_conservatively(fake_pyatr, tmp_path):
     matches = guard._scan("some content", event_type="llm_input")
     assert len(matches) == 1
     assert matches[0].rule_id == "ATR-601"
+
+
+@pytest.mark.asyncio
+async def test_post_call_streaming_blocks_on_match(fake_pyatr, tmp_path):
+    """Streaming hook returns SSE error frame when aggregated response matches."""
+    import json as _json
+
+    _, engine = fake_pyatr
+    ATRGuardrail, _, _ = _import_guardrail()
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rules_dir = tmp_path / "rules"
+    rules_dir.mkdir()
+
+    engine.evaluate.return_value = [
+        MagicMock(rule_id="ATR-700", title="Stream leak", severity="critical")
+    ]
+
+    guard = ATRGuardrail(
+        rules_path=str(rules_dir),
+        severity_threshold="high",
+        guardrail_name="atr-test",
+        event_hook="post_call",
+        default_on=True,
+    )
+
+    result = await guard.async_post_call_streaming_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        response="here is your API key: sk-abc123",
+    )
+
+    assert isinstance(result, str)
+    assert result.startswith("data: ")
+    payload = _json.loads(result[len("data: ") :].strip())
+    assert payload["error"]["error"] == "Streamed response blocked by ATR guardrail"
+    assert payload["error"]["matched_rules"][0]["rule_id"] == "ATR-700"
+
+
+@pytest.mark.asyncio
+async def test_post_call_streaming_passes_when_no_match(fake_pyatr, tmp_path):
+    """Streaming hook returns the response unchanged when no rules fire."""
+    _, engine = fake_pyatr
+    ATRGuardrail, _, _ = _import_guardrail()
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rules_dir = tmp_path / "rules"
+    rules_dir.mkdir()
+
+    engine.evaluate.return_value = []
+
+    guard = ATRGuardrail(
+        rules_path=str(rules_dir),
+        severity_threshold="high",
+        guardrail_name="atr-test",
+        event_hook="post_call",
+        default_on=True,
+    )
+
+    aggregated = "Sure, here is the summary you asked for."
+    result = await guard.async_post_call_streaming_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        response=aggregated,
+    )
+
+    assert result == aggregated
+
+
+@pytest.mark.asyncio
+async def test_post_call_streaming_passes_empty_response(fake_pyatr, tmp_path):
+    """Streaming hook is a no-op when the aggregated response is empty."""
+    _, engine = fake_pyatr
+    ATRGuardrail, _, _ = _import_guardrail()
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rules_dir = tmp_path / "rules"
+    rules_dir.mkdir()
+
+    guard = ATRGuardrail(
+        rules_path=str(rules_dir),
+        severity_threshold="high",
+        guardrail_name="atr-test",
+        event_hook="post_call",
+        default_on=True,
+    )
+
+    result = await guard.async_post_call_streaming_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        response="",
+    )
+
+    assert result == ""
+    engine.evaluate.assert_not_called()