fix(proxy): hold the request's parallel slot through MCP tool loops on streamed chat and messages

litellm/llms/anthropic/messages/native/mcp_handler.py

@@ -82,6 +82,7 @@ async def anthropic_messages_with_mcp(
     from litellm.mcp_client.tools import (
         transform_mcp_tool_to_anthropic_tool,
     )
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import request_slot_held_by_tool_loop
     from litellm.responses.mcp.litellm_proxy_mcp_handler import (
         LiteLLM_Proxy_MCP_Handler,
     )
@@ -145,52 +146,53 @@ async def anthropic_messages_with_mcp(
         else base_call_args
     )
     working_messages: Sequence[Mapping[str, object]] = tuple(messages)
-    response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
-        messages=list(working_messages), stream=False, **loop_call_args
-    )
-
-    for _ in range(MAX_MCP_TOOL_USE_ITERATIONS):
-        if _get_stop_reason(response) != "tool_use":
-            break
-
-        tool_use_blocks = _extract_tool_use_blocks(response)
-        if not tool_use_blocks:
-            break
-
-        tool_results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(
-            tool_server_map=tool_server_map,
-            served_tools=deduplicated_mcp_tools,
-            tool_calls=list(tool_use_blocks),
-            user_api_key_auth=context.user_api_key_auth,
-            mcp_auth_header=context.mcp_auth_header,
-            mcp_server_auth_headers=context.mcp_server_auth_headers,
-            oauth2_headers=context.oauth2_headers,
-            raw_headers=context.raw_headers,
-            litellm_call_id=context.litellm_call_id,
-            litellm_trace_id=context.litellm_trace_id,
-            request_tags=list(context.request_tags) if context.request_tags else None,
-            guardrail_context=context.guardrail_context,
-        )
-
-        # Every tool call was skipped, so there is nothing to feed back; a
-        # tool_result message with empty content is rejected by Anthropic.
-        if not tool_results:
-            break
-
-        working_messages = (
-            *working_messages,
-            {"role": "assistant", "content": list(_get_response_content(response))},
-            _build_tool_result_message(tool_results),
-        )
-        response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
+    with request_slot_held_by_tool_loop():
+        response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
             messages=list(working_messages), stream=False, **loop_call_args
         )
-    else:
-        verbose_logger.warning(
-            "MCP tool loop hit its %s iteration cap for model %s; returning the last response",
-            MAX_MCP_TOOL_USE_ITERATIONS,
-            model,
-        )
+
+        for _ in range(MAX_MCP_TOOL_USE_ITERATIONS):
+            if _get_stop_reason(response) != "tool_use":
+                break
+
+            tool_use_blocks = _extract_tool_use_blocks(response)
+            if not tool_use_blocks:
+                break
+
+            tool_results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls(
+                tool_server_map=tool_server_map,
+                served_tools=deduplicated_mcp_tools,
+                tool_calls=list(tool_use_blocks),
+                user_api_key_auth=context.user_api_key_auth,
+                mcp_auth_header=context.mcp_auth_header,
+                mcp_server_auth_headers=context.mcp_server_auth_headers,
+                oauth2_headers=context.oauth2_headers,
+                raw_headers=context.raw_headers,
+                litellm_call_id=context.litellm_call_id,
+                litellm_trace_id=context.litellm_trace_id,
+                request_tags=list(context.request_tags) if context.request_tags else None,
+                guardrail_context=context.guardrail_context,
+            )
+
+            # Every tool call was skipped, so there is nothing to feed back; a
+            # tool_result message with empty content is rejected by Anthropic.
+            if not tool_results:
+                break
+
+            working_messages = (
+                *working_messages,
+                {"role": "assistant", "content": list(_get_response_content(response))},
+                _build_tool_result_message(tool_results),
+            )
+            response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
+                messages=list(working_messages), stream=False, **loop_call_args
+            )
+        else:
+            verbose_logger.warning(
+                "MCP tool loop hit its %s iteration cap for model %s; returning the last response",
+                MAX_MCP_TOOL_USE_ITERATIONS,
+                model,
+            )
 
     client_response: Final = (
         await with_safeguard_results(response, request_safeguards_evaluator(kwargs, working_messages))

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -11,8 +11,8 @@ import json
 import logging
 import os
 import uuid
-from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
-from contextlib import asynccontextmanager
+from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence, Set
+from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
@@ -697,6 +697,7 @@ class RequestRateLimiterStash:
     rate_limit_response: RateLimitResponse | None = None
     parallel_slot: ParallelSlotAcquisition | None = None
     parallel_slot_reused_at: datetime | None = None
+    parallel_slot_held_by_tool_loop: bool = False
     parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
     reserved_tokens: int = 0
     reserved_model: RateLimitedModel | None = None
@@ -785,6 +786,24 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
     return stash
 
 
+def hold_request_slot_for_tool_loop() -> RequestRateLimiterStash | None:
+    stash: Final = get_request_stash()
+    if stash is None or stash.parallel_slot is None:
+        return None
+    stash.parallel_slot_held_by_tool_loop = True
+    return stash
+
+
+@contextmanager
+def request_slot_held_by_tool_loop() -> Generator[None, None, None]:
+    held_stash: Final = hold_request_slot_for_tool_loop()
+    try:
+        yield
+    finally:
+        if held_stash is not None:
+            held_stash.parallel_slot_held_by_tool_loop = False
+
+
 def claim_request_stash_for_data(data: dict, call_type: str | None) -> RequestRateLimiterStash:
     stash: Final = get_or_create_request_stash()
     owner_call_id: Final = data.get("litellm_call_id")
@@ -2051,7 +2070,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             return
         async with stash.parallel_slot_release_lock:
             acquisition: Final = stash.parallel_slot
-            if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at):
+            if (
+                acquisition is None
+                or (in_logging_callback and stash.parallel_slot_held_by_tool_loop)
+                or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at)
+            ):
                 return
             deferred: Final = in_logging_callback and await self._defer_parallel_slot_release(
                 acquisition, parent_otel_span
@@ -5713,15 +5736,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
             if popped is not None:
                 await self.batch_enqueued_token_store.refund(reservation=popped, litellm_parent_otel_span=span)
 
-    def _mark_request_slot_reused_by_failed_tool_call(self, request_data: Mapping[str, object]) -> None:
-        request_stash: Final = get_request_stash()
-        if (
-            request_stash is not None
-            and request_stash.parallel_slot is not None
-            and request_data.get("call_type") == CallTypes.call_mcp_tool.value
-        ):
-            request_stash.parallel_slot_reused_at = self._get_current_time()
-
     @with_service_target("rate_limits")
     async def async_post_call_failure_hook(
         self,
@@ -5750,7 +5764,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         try:
             stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(request_data))
             if stash is None:
-                self._mark_request_slot_reused_by_failed_tool_call(request_data)
                 return
             await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
 

litellm/responses/mcp/chat_completions_handler.py

@@ -215,6 +215,7 @@ async def acompletion_with_mcp(
 
         # Create a custom async generator that collects chunks and handles tool execution
         from litellm.main import stream_chunk_builder
+        from litellm.proxy.hooks.parallel_request_limiter_v3 import hold_request_slot_for_tool_loop
         from litellm.types.utils import ModelResponseStream
 
         class MCPStreamingIterator:
@@ -258,6 +259,7 @@ async def acompletion_with_mcp(
                 self.follow_up_stream = None
                 self.follow_up_iterator = None
                 self.follow_up_exhausted = False
+                self.holds_request_slot = hold_request_slot_for_tool_loop() is not None
 
             def __aiter__(self):
                 return self
@@ -332,7 +334,17 @@ async def acompletion_with_mcp(
                         "Error draining inner MCP stream after final chunk; spend logging may be incomplete"
                     )
 
-            async def __anext__(self):
+            async def __anext__(self) -> ModelResponseStream:
+                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)
+                    raise
+
+            async def _next_chunk(self) -> ModelResponseStream:
                 # Phase 1: Collect and yield initial stream chunks
                 if not self.stream_exhausted:
                     # Get the iterator from the stream wrapper

tests/unit/llms/anthropic/messages/native/test_mcp_handler.py

@@ -454,3 +454,79 @@ async def test_anthropic_messages_with_mcp_loads_a_filtered_routes_beta_map_off_
 @pytest.mark.asyncio
 async def test_anthropic_messages_with_mcp_never_loads_the_beta_map_for_first_party_anthropic():
     assert await _mcp_safeguards_turn_recording_beta_map_loads("anthropic", "claude-sonnet-4-6") == ()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_with_mcp_keeps_the_requests_parallel_slot_while_its_tools_run():
+    import asyncio
+    import contextvars
+    from datetime import datetime
+
+    from fastapi import HTTPException
+
+    from litellm.caching.caching import DualCache
+    from litellm.llms.anthropic.messages.native import mcp_handler
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
+    from litellm.proxy.utils import InternalUsageCache, hash_token
+    from litellm.responses.mcp.request_context import MCPRequestContext
+    from litellm.types.utils import CallTypes
+
+    local_cache: Final = DualCache()
+    limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache))
+    key: Final = UserAPIKeyAuth(api_key=hash_token("sk-messages-tool-loop"), max_parallel_requests=1)
+
+    async def admit(data: dict[str, object]) -> None:
+        await limiter.async_pre_call_hook(
+            user_api_key_dict=key, cache=local_cache, data=data, call_type=CallTypes.anthropic_messages.value
+        )
+
+    async def second_request_status() -> int:
+        try:
+            await contextvars.Context().run(asyncio.create_task, admit({"model": "claude-x", "messages": []}))
+        except HTTPException as rejected:
+            return rejected.status_code
+        return 200
+
+    async def run_tool_after_the_first_round_logs(**_: object) -> list[dict[str, str]]:
+        await limiter.async_log_success_event(
+            kwargs={"litellm_call_id": "messages-owner", "stream": False},
+            response_obj=None,
+            start_time=None,
+            end_time=datetime.now(),
+        )
+        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"}]
+
+    responses = [
+        {
+            "stop_reason": "tool_use",
+            "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}],
+        },
+        {"stop_reason": "end_turn", "content": [{"type": "text", "text": "done"}]},
+    ]
+    handler_class: Final = import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler
+
+    await admit(
+        {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": "messages-owner"}
+    )
+    with (
+        patch.object(MCPRequestContext, "resolve", return_value=MCPRequestContext(user_api_key_auth="auth")),
+        patch.object(
+            handler_class,
+            "process_mcp_tools_without_openai_transform",
+            new=AsyncMock(return_value=([], {"get_weather": "weather"})),
+        ),
+        patch.object(
+            handler_class, "execute_tool_calls", new=AsyncMock(side_effect=run_tool_after_the_first_round_logs)
+        ),
+        patch("litellm.anthropic_messages", new=AsyncMock(side_effect=responses)),
+    ):
+        result = await mcp_handler.anthropic_messages_with_mcp(
+            max_tokens=100,
+            messages=[{"role": "user", "content": "hi"}],
+            model="claude-sonnet-4-5",
+            tools=[MCP_REFERENCE],
+        )
+
+    assert result == responses[1]

tests/unit/responses/mcp/test_chat_completions_handler.py

@@ -1712,3 +1712,116 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo
     assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19)
     assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE))
     assert recorder.failed[0]["response_cost"] > 0
+
+
+@pytest.mark.asyncio
+async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the_stream_ends(monkeypatch):
+    import contextvars
+    from unittest.mock import MagicMock
+
+    from fastapi import HTTPException
+
+    from litellm.caching.caching import DualCache
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
+    from litellm.proxy.utils import InternalUsageCache, hash_token
+    from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices
+    from litellm.utils import CustomStreamWrapper
+
+    local_cache: Final = DualCache()
+    limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache))
+    key: Final = UserAPIKeyAuth(api_key=hash_token("sk-chat-stream-tool-loop"), max_parallel_requests=1)
+    tools: Final = [{"type": "function", "function": {"name": "get_weather"}}]
+    tool_call: Final = {"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}
+
+    async def admit(data: dict[str, object]) -> None:
+        await limiter.async_pre_call_hook(
+            user_api_key_dict=key, cache=local_cache, data=data, call_type=CallTypes.acompletion.value
+        )
+
+    async def second_request_status() -> int:
+        try:
+            await contextvars.Context().run(asyncio.create_task, admit({"model": "gpt-x", "messages": []}))
+        except HTTPException as rejected:
+            return rejected.status_code
+        return 200
+
+    def chunk(content: str, finish_reason: str | None) -> ModelResponseStream:
+        return ModelResponseStream(
+            id="stream",
+            model="gpt-x",
+            created=0,
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(index=0, delta=Delta(content=content, role="assistant"), finish_reason=finish_reason)
+            ],
+        )
+
+    logging_obj: Final = MagicMock()
+    logging_obj.model_call_details = {}
+
+    class FixedStream(CustomStreamWrapper):
+        def __init__(self, chunks: tuple[ModelResponseStream, ...]):
+            super().__init__(completion_stream=None, model="gpt-x", logging_obj=logging_obj)
+            self.pending = iter(chunks)
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            try:
+                return next(self.pending)
+            except StopIteration:
+                raise StopAsyncIteration
+
+    async def run_tool_after_the_first_round_logs(**_):
+        await limiter.async_log_success_event(
+            kwargs={"litellm_call_id": "chat-owner", "stream": True},
+            response_obj=None,
+            start_time=None,
+            end_time=datetime.now(),
+        )
+        assert await second_request_status() == 429, "The first round's stream log must not free the request's slot"
+        return [{"tool_call_id": "call-1", "result": "sunny"}]
+
+    acompletion: Final = AsyncMock(
+        side_effect=[FixedStream((chunk("", "tool_calls"),)), FixedStream((chunk("Sunny", None), chunk(".", "stop")))]
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.proxy_logging_obj",
+        types.SimpleNamespace(arelease_max_parallel_requests=limiter.async_release_max_parallel_requests_on_disconnect),
+    )
+    monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "should_use_litellm_mcp_gateway", staticmethod(lambda tools: True))
+    monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "_parse_mcp_tools", staticmethod(lambda tools: (tools, [])))
+    monkeypatch.setattr(
+        LiteLLM_Proxy_MCP_Handler,
+        "process_mcp_tools_without_openai_transform",
+        AsyncMock(return_value=(tools, {"get_weather": "weather"})),
+    )
+    monkeypatch.setattr(
+        LiteLLM_Proxy_MCP_Handler, "transform_mcp_tools_to_openai", staticmethod(lambda *_, **__: tools)
+    )
+    monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "should_auto_execute_tools", staticmethod(lambda **_: True))
+    monkeypatch.setattr(
+        LiteLLM_Proxy_MCP_Handler, "extract_tool_calls_from_chat_response", staticmethod(lambda **_: [tool_call])
+    )
+    monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", run_tool_after_the_first_round_logs)
+    monkeypatch.setattr(
+        LiteLLM_Proxy_MCP_Handler,
+        "create_follow_up_messages_for_chat",
+        staticmethod(lambda **_: [{"role": "tool", "tool_call_id": "call-1", "content": "sunny"}]),
+    )
+    monkeypatch.setattr(
+        ResponsesAPIRequestUtils, "extract_mcp_headers_from_request", staticmethod(lambda **_: (None, None, None, None))
+    )
+
+    await admit({"model": "gpt-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": "chat-owner"})
+    with patch("litellm.acompletion", acompletion):
+        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])
+
+        assert await second_request_status() == 200, "The request's slot must be freed once its stream ends"
+
+    assert streamed_text == "Sunny."