diff --git a/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py b/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py index 620dcc4cc49..756cc38f7bb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py +++ b/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 # ------------------------------------------------------------------ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py index 1e06724768a..decdebe5f9b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py +++ b/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()