From 7ea99461dd06bf472b21efa9e5fc409172d9b384 Mon Sep 17 00:00:00 2001
From: tombii <github@tombii.com>
Date: Tue, 3 Mar 2026 14:59:02 +0100
Subject: [PATCH] fix(proxy): guard against None return in
 async_post_call_streaming_iterator_hook

Fixes "TypeError: 'async for' requires an object with __aiter__ method, got NoneType"
error that occurs when a user-defined custom callback has a poorly implemented
async_post_call_streaming_iterator_hook that returns None instead of an async generator.

Changes:
- Fix type(callback).__dict__ -> type(_callback).__dict__ to check the resolved instance
- Add None check after calling hook to prevent current_response from becoming None
- Add regression test for the NoneType crash scenario

The fix ensures that if a callback returns None (e.g., a sync method instead of
async generator), the streaming chain continues with the previous valid response
iterator instead of crashing.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 litellm/proxy/utils.py                        | 12 ++++--
 ...async_post_call_streaming_iterator_hook.py | 43 +++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

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