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>
This commit is contained in:
tombii 2026-03-03 14:59:02 +01:00
parent ac6e1d9fd1
commit 7ea99461dd
2 changed files with 51 additions and 4 deletions

View file

@ -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:

View file

@ -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