diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py index 2ae9986ce94..cfea31cb655 100644 --- a/litellm/litellm_core_utils/api_route_to_call_types.py +++ b/litellm/litellm_core_utils/api_route_to_call_types.py @@ -67,3 +67,31 @@ def get_routes_for_call_type(call_type: CallTypes) -> list: if call_type in types: routes.append(route) return routes + + +OPENAI_CHUNK_STREAMING_CALL_TYPES = frozenset( + { + CallTypes.completion, + CallTypes.acompletion, + CallTypes.text_completion, + CallTypes.atext_completion, + } +) + + +def route_streams_openai_format_chunks(route: Optional[str]) -> bool: + """ + Whether streaming chunks emitted by `route` are OpenAI-format + ``ModelResponseStream`` objects. + + Routes such as ``/v1/messages`` or ``/v1/responses`` relay provider-native + payloads (Anthropic SSE bytes, Responses API events), which hooks written + against ``ModelResponseStream`` cannot read. Unknown routes are treated as + OpenAI-format to preserve existing behavior. + """ + if route is None: + return True + call_types = get_call_types_for_route(route) + if call_types is None: + return True + return any(call_type in OPENAI_CHUNK_STREAMING_CALL_TYPES for call_type in call_types) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 924189fed4b..e6e864a8e13 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -105,6 +105,9 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert +from litellm.litellm_core_utils.api_route_to_call_types import ( + route_streams_openai_format_chunks, +) from litellm.litellm_core_utils.core_helpers import coerce_token_limit from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -2669,6 +2672,8 @@ class ProxyLogging: current_response = response + streams_openai_chunks = route_streams_openai_format_chunks(user_api_key_dict.request_route) + for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): if ( @@ -2676,7 +2681,13 @@ class ProxyLogging: is not True ): continue - if kind == "override": + effective_kind = ProxyLogging._resolve_iterator_override_kind( + resolved_callback=resolved_callback, + kind=kind, + streams_openai_chunks=streams_openai_chunks, + request_route=user_api_key_dict.request_route, + ) + if effective_kind == "override": current_response = self._wrap_streaming_iterator_with_enrichment( resolved_callback, resolved_callback.async_post_call_streaming_iterator_hook( @@ -2686,7 +2697,7 @@ class ProxyLogging: ), ) else: - # kind == "apply_guardrail": route through unified_guardrail + # effective_kind == "apply_guardrail": route through unified_guardrail request_data["guardrail_to_apply"] = resolved_callback current_response = self._wrap_streaming_iterator_with_enrichment( resolved_callback, @@ -2707,6 +2718,35 @@ class ProxyLogging: # we reach this point the metadata is fully populated. ProxyLogging._fire_deferred_stream_logging(request_data) + @staticmethod + def _resolve_iterator_override_kind( + resolved_callback: Any, + kind: str, + streams_openai_chunks: bool, + request_route: Optional[str], + ) -> str: + """ + Pick how a callback participates in the streaming post-call chain. + + A callback's own ``async_post_call_streaming_iterator_hook`` is written + against OpenAI-format ``ModelResponseStream`` chunks, so it cannot scan + the provider-native payloads relayed by routes like ``/v1/messages`` + (raw Anthropic SSE bytes). On those routes prefer the translation-aware + ``apply_guardrail`` path when the callback provides one; otherwise warn, + since the guardrail will not see the response. + """ + if kind != "override" or streams_openai_chunks: + return kind + if "apply_guardrail" in type(resolved_callback).__dict__: + return "apply_guardrail" + verbose_proxy_logger.warning( + "Guardrail %s only implements async_post_call_streaming_iterator_hook, which expects " + "OpenAI-format chunks; it cannot scan the streaming response on %s", + getattr(resolved_callback, "guardrail_name", None) or type(resolved_callback).__name__, + request_route, + ) + return kind + @staticmethod def _fire_deferred_stream_logging(request_data: dict) -> None: """ diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 65d3c3c8079..82b892f1923 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -430,3 +430,148 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log data={}, user_api_key_dict=make_user_api_key_auth(), response=response ) assert out == {} + +# --------------------------------------------------------------------------- +# non-OpenAI streaming surfaces (/v1/messages) -- issue #35257 +# --------------------------------------------------------------------------- + + +def _anthropic_sse(event_type: str, data: Dict[str, Any]) -> bytes: + import json + + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +async def _anthropic_stream(text: str): + yield _anthropic_sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ) + yield _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + yield _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + yield _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}) + yield _anthropic_sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 3}}, + ) + yield _anthropic_sse("message_stop", {"type": "message_stop"}) + + +def _content_filter_guardrail(): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + + guardrail = ContentFilterGuardrail( + guardrail_name="output-filter", + event_hook="post_call", + blocked_words=[{"keyword": "zebra", "action": "BLOCK"}], + default_on=True, + ) + # scan every chunk so the block decision lands before the offending delta + # is forwarded, matching the /chat/completions withholding behavior + guardrail.streaming_sampling_rate = 1 + return guardrail + + +async def _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, text: str): + request_data: Dict[str, Any] = { + "metadata": {"guardrails": ["output-filter"]}, + "messages": [{"role": "user", "content": "hi"}], + } + collected: List[Any] = [] + blocked = False + try: + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=_anthropic_stream(text), + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + request_data=request_data, + ): + collected.append(chunk) + except HTTPException: + blocked = True + raw = b"".join(c if isinstance(c, bytes) else str(c).encode() for c in collected).decode() + statuses = [ + entry.get("guardrail_status") + for entry in request_data.get("litellm_metadata", {}).get("standard_logging_guardrail_information", []) + ] + return raw, blocked, statuses + + +@pytest.mark.asyncio +async def test_v1_messages_streaming_guardrail_blocks_instead_of_leaking( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A guardrail with its own ModelResponseStream-shaped iterator hook must not + be handed raw Anthropic SSE bytes: on /v1/messages the translation-aware + apply_guardrail path runs, so blocked content is withheld and the violation + is recorded rather than silently logged as a successful scan.""" + monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail()]) + + raw, blocked, statuses = await _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, "the zebra runs") + + assert "zebra" not in raw + assert blocked is True + assert "guardrail_intervened" in statuses + + +@pytest.mark.asyncio +async def test_v1_messages_streaming_guardrail_passes_clean_content( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail()]) + + raw, blocked, statuses = await _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, "the horse runs") + + assert "the horse runs" in raw + assert blocked is False + assert "guardrail_intervened" not in statuses + + +@pytest.mark.asyncio +async def test_chat_completions_streaming_still_uses_callback_iterator_hook( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """The reroute is scoped to routes that stream provider-native chunks; the + OpenAI chat path must keep using the callback's own iterator hook.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _Guardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, **kwargs): # type: ignore[override] + async for chunk in kwargs["response"]: + yield f"{chunk}*" + + async def apply_guardrail(self, **kwargs): # type: ignore[override] + raise AssertionError("apply_guardrail should not run on /chat/completions") + + monkeypatch.setattr(litellm, "callbacks", [_Guardrail(guardrail_name="g", event_hook="post_call", default_on=True)]) + + async def gen(): + for ch in ("a", "b"): + yield ch + + out: List[str] = [] + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=gen(), + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + request_data={"metadata": {"guardrails": ["g"]}}, + ): + out.append(chunk) + assert out == ["a*", "b*"]