From 9c74725a61fc237ac3726ae532a5c0dba836131b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:08:13 -0700 Subject: [PATCH] fix(mcp): free the held parallel slot when a streamed MCP client hangs up --- .../responses/mcp/chat_completions_handler.py | 16 ++++++++--- .../messages/native/test_mcp_handler.py | 7 ++++- .../mcp/test_chat_completions_handler.py | 27 ++++++++++++++++--- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a43cda98815..c926c57b548 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -338,12 +338,16 @@ async def acompletion_with_mcp( try: return await self._next_chunk() except StopAsyncIteration: - if self.holds_request_slot: - from litellm.proxy.proxy_server import proxy_logging_obj - - await proxy_logging_obj.arelease_max_parallel_requests(self.user_api_key_auth) + await self.release_held_request_slot() raise + async def release_held_request_slot(self) -> None: + if not self.holds_request_slot: + return + from litellm.proxy.proxy_server import proxy_logging_obj + + await proxy_logging_obj.arelease_max_parallel_requests(self.user_api_key_auth) + async def _next_chunk(self) -> ModelResponseStream: # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: @@ -560,6 +564,10 @@ async def acompletion_with_mcp( def __aiter__(self): return self._custom_iterator + async def aclose(self) -> None: + await self._custom_iterator.release_held_request_slot() + await super().aclose() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: diff --git a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py index 0f8bebf3aa1..f3c6ec2f359 100644 --- a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py @@ -488,13 +488,16 @@ async def test_anthropic_messages_with_mcp_keeps_the_requests_parallel_slot_whil return rejected.status_code return 200 - async def run_tool_after_the_first_round_logs(**_: object) -> list[dict[str, str]]: + async def log_the_requests_success() -> None: await limiter.async_log_success_event( kwargs={"litellm_call_id": "messages-owner", "stream": False}, response_obj=None, start_time=None, end_time=datetime.now(), ) + + async def run_tool_after_the_first_round_logs(**_: object) -> list[dict[str, str]]: + await log_the_requests_success() assert await second_request_status() == 429, "A round logged while tools run must not free the request's slot" return [{"tool_call_id": "toolu_1", "result": "ok", "name": "get_weather"}] @@ -530,3 +533,5 @@ async def test_anthropic_messages_with_mcp_keeps_the_requests_parallel_slot_whil ) assert result == responses[1] + await log_the_requests_success() + assert await second_request_status() == 200, "The request's own log after its tools finish must free its slot" diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 8862c1a1483..5cb74535dce 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -26,6 +26,7 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.utils import CustomStreamWrapper @pytest.mark.asyncio @@ -1714,8 +1715,26 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo assert recorder.failed[0]["response_cost"] > 0 +async def _read_the_whole_stream(stream: CustomStreamWrapper) -> str: + return "".join([part.choices[0].delta.content or "" async for part in stream]) + + +async def _hang_up_after_the_first_follow_up_chunk(stream: CustomStreamWrapper) -> str: + parts: Final = stream.__aiter__() + read: Final = (await parts.__anext__(), await parts.__anext__()) + await stream.aclose() + return "".join(part.choices[0].delta.content or "" for part in read) + + @pytest.mark.asyncio -async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the_stream_ends(monkeypatch): +@pytest.mark.parametrize( + ("end_stream", "expected_text"), + [(_read_the_whole_stream, "Sunny."), (_hang_up_after_the_first_follow_up_chunk, "Sunny")], + ids=["stream_ends", "client_hangs_up"], +) +async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the_stream_ends_or_closes( + monkeypatch, end_stream, expected_text +): import contextvars from unittest.mock import MagicMock @@ -1820,8 +1839,8 @@ async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the stream = await acompletion_with_mcp( model="gpt-x", messages=[{"role": "user", "content": "hi"}], tools=tools, stream=True ) - streamed_text: Final = "".join([part.choices[0].delta.content or "" async for part in stream]) + streamed_text: Final = await end_stream(stream) - assert await second_request_status() == 200, "The request's slot must be freed once its stream ends" + assert await second_request_status() == 200, "The request's slot must be freed once its stream ends or closes" - assert streamed_text == "Sunny." + assert streamed_text == expected_text