fix(proxy): close common streaming responses on exit - #28536
Conversation
|
|
Greptile SummaryThis PR adds a shielded
Confidence Score: 4/5The change is a targeted cleanup fix that adds a The cleanup itself is correct and the tests cover the main exit scenarios. Two gaps worth noting: the shielded cancel scope has no timeout, so a hung
|
| Filename | Overview |
|---|---|
| litellm/proxy/common_request_processing.py | Adds a shielded finally block to async_streaming_data_generator that calls response.aclose(), preventing upstream stream leaks on early client disconnect; anyio added as import. No timeout bounds the shielded cleanup. |
| tests/test_litellm/test_common_request_processing_streaming_cleanup.py | New test file with 4 mock-only tests covering early exit, normal completion, mid-stream error, and swallowed close-error; does not exercise the task-cancellation path the shield is designed to protect. |
Reviews (1): Last reviewed commit: "test(proxy): cover streaming response cl..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| with anyio.CancelScope(shield=True): | ||
| if hasattr(response, "aclose"): | ||
| try: | ||
| # This generator owns final cleanup of the raw upstream stream; | ||
| # iterator hooks should transform/drain it without closing it. | ||
| await response.aclose() | ||
| except BaseException as e: | ||
| verbose_proxy_logger.debug( | ||
| "async_streaming_data_generator: error closing response stream: %s", | ||
| e, | ||
| ) |
There was a problem hiding this comment.
aclose()
anyio.CancelScope(shield=True) prevents the cleanup from being cancelled, but there's no deadline/timeout. If response.aclose() blocks (e.g., a misbehaving proxy or hung upstream TCP connection), this shielded scope will hold the task open indefinitely after the client has already disconnected. Adding a deadline to the scope — e.g., anyio.CancelScope(shield=True, deadline=anyio.current_time() + 5) — would bound the worst-case hang to a fixed window and still leave enough time for normal close operations to complete.
| @pytest.mark.asyncio | ||
| async def test_async_streaming_data_generator_closes_response_on_early_exit(): | ||
| mock_response = MagicMock() | ||
| mock_response.aclose = AsyncMock() | ||
| mock_proxy_logging_obj = _mock_proxy_logging_obj() | ||
|
|
||
| async def mock_streaming_iterator(*args, **kwargs): | ||
| yield {"content": "hello"} | ||
| yield {"content": " world"} | ||
|
|
||
| mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( | ||
| mock_streaming_iterator | ||
| ) | ||
|
|
||
| generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( | ||
| response=mock_response, | ||
| user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), | ||
| request_data={"model": "gpt-3.5-turbo"}, | ||
| proxy_logging_obj=mock_proxy_logging_obj, | ||
| serialize_chunk=lambda chunk: str(chunk), | ||
| serialize_error=lambda proxy_exception: str(proxy_exception.to_dict()), | ||
| ) | ||
|
|
||
| first_chunk = await generator.__anext__() | ||
| assert first_chunk == "{'content': 'hello'}" | ||
|
|
||
| await generator.aclose() | ||
|
|
||
| mock_response.aclose.assert_awaited_once() |
There was a problem hiding this comment.
GeneratorExit, not task cancellation
The shield in the implementation exists specifically to protect cleanup from asyncio task cancellation (i.e., the client-disconnect path). The test exercises generator.aclose(), which raises GeneratorExit inside the generator — a different code path than an asyncio.CancelledError delivered to the running task. Because GeneratorExit reaches the finally block regardless of shielding, the test passes even if the CancelScope(shield=True) were removed. A complementary test that cancels the task mid-stream (via asyncio.create_task + task.cancel()) would give confidence that the shield actually works under the condition it's guarding against.
|
🤖 litellm-agent: This PR is currently BLOCKED from merge. Score: 4/5 ❌ Why blocked:
Details: Score docked for: 1 unresolved reviewer concern (greptile). Fix the issues above and push an update — the bot will re-review automatically.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. |
2 participants
Relevant issues
Related to #25776.
Linear ticket
N/A
Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
Targeted tests pass locally:
.venv/bin/python -m pytest tests/test_litellm/test_common_request_processing_streaming_cleanup.py # 5 passedLint on changed files passes locally:
uv run --no-sync ruff check litellm/proxy/common_request_processing.py tests/test_litellm/test_common_request_processing_streaming_cleanup.py # All checks passedType
🐛 Bug Fix
✅ Test
Changes
response.aclose()fromProxyBaseLLMRequestProcessing.async_streaming_data_generatorin a shieldedfinallyblock.