fix(mcp): free the held parallel slot when a streamed MCP client hangs up

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:

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"

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