diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 015dcd9b5db..12817902765 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -51,6 +51,17 @@ def test_has_streaming_callbacks_detects_guardrails(monkeypatch): assert ProxyLogging.has_streaming_callbacks() is True +def test_custom_guardrail_streaming_call_type_capability_is_opt_in(): + from litellm.types.utils import CallTypes + + class _ProtocolAwareGuardrail(CustomGuardrail): + supported_streaming_call_types = frozenset({CallTypes.anthropic_messages}) + + assert CustomGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is False + assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is True + assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.aresponses) is False + + @pytest.mark.asyncio async def test_post_call_response_headers_hook_returns_early_without_callbacks( monkeypatch, @@ -427,6 +438,46 @@ async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(mon assert delivered == [] +@pytest.mark.asyncio +async def test_protocol_aware_guardrail_keeps_own_anthropic_streaming_iterator(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + first_chunk = _sse_bytes("message_start", {"type": "message_start"}) + + class _ProtocolAwareGuardrail(CustomGuardrail): + supported_streaming_call_types = frozenset({CallTypes.anthropic_messages}) + + def __init__(self): + super().__init__(guardrail_name="protocol-aware", event_hook="post_call", default_on=True) + self.own_iterator_called = False + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + self.own_iterator_called = True + async for item in response: + yield item + + async def response_stream(): + yield first_chunk + raise AssertionError("upstream requested another chunk before the first chunk was emitted") + + guardrail = _ProtocolAwareGuardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + guarded_stream = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/messages"), + response=response_stream(), + request_data={"model": "claude-sonnet-5", "metadata": {}}, + ) + + assert await anext(guarded_stream) == first_chunk + assert guardrail.own_iterator_called is True + await guarded_stream.aclose() + + @pytest.mark.asyncio async def test_post_call_stream_masking_guardrail_keeps_own_iterator_on_anthropic(monkeypatch): """