fix(responses): close the MCP round's upstream stream on client disconnect

MCPEnhancedStreamingIterator bypasses the base constructor and swaps
base_iterator on every tool-execution round, so the inherited aclose found no
self.response and silently left the live provider stream generating.

Override it to close the current round's iterator.
This commit is contained in:
Mohammad Ali Farhan 2026-08-08 00:47:47 +05:30
parent 3e869f3a05
commit 867fd64efb
2 changed files with 62 additions and 0 deletions

View file

@ -326,6 +326,18 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self._error_event_emitted = False
self._last_sequence_number = 0
async def aclose(self) -> None:
"""
Close the provider stream of the current tool-execution round.
This class bypasses ``BaseResponsesAPIStreamingIterator.__init__`` and swaps
``base_iterator`` on every round, so the inherited close finds no
``self.response`` and would leave the live upstream generating on disconnect.
"""
base_iterator: Final = self.base_iterator
if isinstance(base_iterator, BaseResponsesAPIStreamingIterator):
await base_iterator.aclose()
def _extract_mcp_headers_from_params(self) -> None:
"""Extract MCP headers from original request params to pass to tool calls"""

View file

@ -3,6 +3,7 @@ import types
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from mcp.types import CallToolResult, TextContent
@ -11,6 +12,7 @@ from litellm.responses.mcp.mcp_streaming_iterator import (
MAX_MCP_TOOL_CALL_ROUNDS,
MCPEnhancedStreamingIterator,
)
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamEvents
# `litellm.__init__` re-exports a function named `responses`, which shadows the
@ -257,3 +259,51 @@ 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)
class _RecordingByteStream(httpx.AsyncByteStream):
def __init__(self):
self.aclose_calls = 0
async def __aiter__(self):
yield b""
async def aclose(self) -> None:
self.aclose_calls += 1
def _responses_iterator_over(stream: _RecordingByteStream) -> ResponsesAPIStreamingIterator:
return ResponsesAPIStreamingIterator(
response=httpx.Response(200, stream=stream),
model="gpt-4o-mini",
responses_api_provider_config=MagicMock(),
logging_obj=MagicMock(model_call_details={"litellm_params": {}}),
custom_llm_provider="openai",
)
@pytest.mark.asyncio
async def test_aclose_closes_the_current_round_provider_stream():
"""The MCP iterator swaps `base_iterator` per tool-execution round and holds no
`self.response`, so the inherited close is a no-op. Without an override, a client
disconnect leaves the round's provider stream generating (#36123)."""
stream = _RecordingByteStream()
iterator = _make_iterator([])
iterator.base_iterator = _responses_iterator_over(stream)
await iterator.aclose()
assert stream.aclose_calls == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
"base_iterator",
[None, ResponsesAPIResponse(id="resp-1", created_at=0, output=[])],
ids=["no_iterator_yet", "non_streaming_response"],
)
async def test_aclose_is_a_noop_when_no_provider_stream_is_live(base_iterator):
iterator = _make_iterator([])
iterator.base_iterator = base_iterator
await iterator.aclose()