diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5e0d5336aa9..091214d78ab 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2109,16 +2109,18 @@ class ProxyLogging: ): if ( "async_post_call_streaming_iterator_hook" - in type(callback).__dict__ + in type(_callback).__dict__ ): - current_response = ( + _new_response = ( _callback.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=current_response, request_data=request_data, ) ) - elif "apply_guardrail" in type(callback).__dict__: + if _new_response is not None: + current_response = _new_response + elif "apply_guardrail" in type(_callback).__dict__: request_data["guardrail_to_apply"] = callback current_response = ( unified_guardrail.async_post_call_streaming_iterator_hook( @@ -2128,13 +2130,15 @@ class ProxyLogging: ) ) else: - current_response = ( + _new_response = ( _callback.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=current_response, request_data=request_data, ) ) + if _new_response is not None: + current_response = _new_response # Actually iterate through the chained async generator and yield chunks async for chunk in current_response: diff --git a/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py b/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py index 50c6a580f91..71156075da4 100644 --- a/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py +++ b/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py @@ -192,3 +192,46 @@ async def test_streaming_hook_propagates_callback_errors(): with pytest.raises(RuntimeError, match="Callback failed!"): async for _ in result: pass + + +@pytest.mark.asyncio +async def test_no_double_strip_on_second_call(): + """Regression test: callback returning None should not break the streaming chain. + + This tests the fix for: 'async for' requires an object with __aiter__ method, got NoneType + Caused by a user-defined callback whose async_post_call_streaming_iterator_hook + is a regular sync method (not an async generator) that returns None. + """ + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + class NoneReturningCallback(CustomLogger): + """Simulates a badly-implemented callback that returns None instead of an async generator.""" + + def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ): + # Regular sync method, not an async generator — returns None implicitly + pass + + bad_callback = NoneReturningCallback() + + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + request_data = {"model": "gpt-4", "messages": []} + + with patch.object(litellm, "callbacks", [bad_callback]): + result = proxy_logging.async_post_call_streaming_iterator_hook( + response=mock_streaming_response(), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + # Should not raise TypeError about NoneType + collected_chunks = [] + async for chunk in result: + collected_chunks.append(chunk) + + # All 4 original chunks should pass through unmodified + assert len(collected_chunks) == 4