From 62ca98a3f7b5eea61616d7765acebe79a74316f3 Mon Sep 17 00:00:00 2001
From: vivekvar-dl <domainluther1234@gmail.com>
Date: Sat, 8 Aug 2026 13:57:22 +0530
Subject: [PATCH] fix(responses): propagate aclose through
 MCPEnhancedStreamingIterator and type test helpers

Greptile review: the MCP wrapper stores its upstream-owning iterator in
base_iterator, so the inherited response-based aclose was a no-op there.
Delegate to the base iterator's aclose when it has one. Also add the
missing parameter and return annotations on the new passthrough test
helpers
---
 .../responses/mcp/mcp_streaming_iterator.py   |  5 ++++
 .../test_streaming_handler.py                 |  6 ++---
 .../mcp/test_mcp_streaming_iterator.py        | 26 +++++++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py
index e4cc36de06c..e4cf858b49b 100644
--- a/litellm/responses/mcp/mcp_streaming_iterator.py
+++ b/litellm/responses/mcp/mcp_streaming_iterator.py
@@ -326,6 +326,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
         self._error_event_emitted = False
         self._last_sequence_number = 0
 
+    async def aclose(self) -> None:
+        base_aclose: Final = getattr(self.base_iterator, "aclose", None)
+        if base_aclose is not None:
+            await base_aclose()
+
     def _extract_mcp_headers_from_params(self) -> None:
         """Extract MCP headers from original request params to pass to tool calls"""
 
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
index 7648491e81c..f03f4312a1a 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py
@@ -3,7 +3,7 @@ when the stream ends, including on client disconnect (generator aclose), so
 the provider connection is released and backends like vLLM stop generating.
 """
 
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Coroutine
 from datetime import datetime, timezone
 from typing import Final
 from unittest.mock import MagicMock, patch
@@ -51,7 +51,7 @@ def _logging_obj_stub() -> MagicMock:
     return logging_obj
 
 
-def _chunk_processor(response: httpx.Response):
+def _chunk_processor(response: httpx.Response) -> AsyncIterator[bytes]:
     return PassThroughStreamingHandler.chunk_processor(
         response=response,
         request_body={"model": "claude-sonnet-4-5"},
@@ -63,7 +63,7 @@ def _chunk_processor(response: httpx.Response):
     )
 
 
-def _drop_scheduled_coroutine(async_coroutine) -> None:
+def _drop_scheduled_coroutine(async_coroutine: Coroutine[object, object, object]) -> None:
     async_coroutine.close()
 
 
diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py
index 24edf12fffe..8ec38fd4a8b 100644
--- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py
+++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py
@@ -257,3 +257,29 @@ async def test_initial_call_failure_is_stashed_for_eager_reraise(monkeypatch):
 
     assert iterator._initial_creation_error is not None
     assert "initial boom" in str(iterator._initial_creation_error)
+
+
+@pytest.mark.asyncio
+async def test_aclose_closes_base_iterator():
+    class _ClosableStream:
+        def __init__(self):
+            self.aclose_called = False
+
+        async def aclose(self) -> None:
+            self.aclose_called = True
+
+    iterator = _make_iterator([])
+    closable = _ClosableStream()
+    iterator.base_iterator = closable
+
+    await iterator.aclose()
+
+    assert closable.aclose_called
+
+
+@pytest.mark.asyncio
+async def test_aclose_is_a_noop_without_base_iterator():
+    iterator = _make_iterator([])
+    iterator.base_iterator = None
+
+    await iterator.aclose()