From b51a8d87c7e24828ce4dfff7e981753814bf117f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:50:25 -0700 Subject: [PATCH] fix(proxy): only the request's own MCP tool loop holds its max_parallel_requests slot A silent_model shadow copies the request's context but has its own call id, so its MCP stream could set or clear the request's tool-loop hold and free the slot early. The hold now checks the call id against the slot's owner. --- .../anthropic/messages/native/mcp_handler.py | 2 +- .../hooks/parallel_request_limiter_v3.py | 8 ++++---- .../responses/mcp/chat_completions_handler.py | 2 +- .../hooks/test_parallel_request_limiter_v3.py | 4 +++- .../mcp/test_chat_completions_handler.py | 19 +++++++++++++++---- 5 files changed, 24 insertions(+), 11 deletions(-) diff --git a/litellm/llms/anthropic/messages/native/mcp_handler.py b/litellm/llms/anthropic/messages/native/mcp_handler.py index c60da9416e3..215d1324c61 100644 --- a/litellm/llms/anthropic/messages/native/mcp_handler.py +++ b/litellm/llms/anthropic/messages/native/mcp_handler.py @@ -146,7 +146,7 @@ async def anthropic_messages_with_mcp( else base_call_args ) working_messages: Sequence[Mapping[str, object]] = tuple(messages) - with request_slot_held_by_tool_loop(): + with request_slot_held_by_tool_loop(context.litellm_call_id): response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( messages=list(working_messages), stream=False, **loop_call_args ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 63491dff9d0..fb50412370b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -786,8 +786,8 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def hold_request_slot_for_tool_loop() -> RequestRateLimiterStash | None: - stash: Final = get_request_stash() +def hold_request_slot_for_tool_loop(litellm_call_id: str | None) -> RequestRateLimiterStash | None: + stash: Final = get_request_stash_for_call(litellm_call_id) if stash is None or stash.parallel_slot is None: return None stash.parallel_slot_held_by_tool_loop = True @@ -795,8 +795,8 @@ def hold_request_slot_for_tool_loop() -> RequestRateLimiterStash | None: @contextmanager -def request_slot_held_by_tool_loop() -> Generator[None, None, None]: - held_stash: Final = hold_request_slot_for_tool_loop() +def request_slot_held_by_tool_loop(litellm_call_id: str | None) -> Generator[None, None, None]: + held_stash: Final = hold_request_slot_for_tool_loop(litellm_call_id) try: yield finally: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index c926c57b548..4bf0dcd4f0f 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -259,7 +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 + self.holds_request_slot = hold_request_slot_for_tool_loop(context.litellm_call_id) is not None def __aiter__(self): return self diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index c559d5a2cd0..bb00eeb83c1 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4363,7 +4363,9 @@ async def test_round_logged_while_a_tool_loop_holds_the_slot_keeps_it_until_the_ await pre_call( {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": owner_call_id} ) - with request_slot_held_by_tool_loop(): + with request_slot_held_by_tool_loop(owner_call_id): + with request_slot_held_by_tool_loop(f"{owner_call_id}-silent-shadow"): + pass await handler.async_post_call_failure_hook( request_data={"call_type": CallTypes.call_mcp_tool.value, "litellm_call_id": f"{owner_call_id}_toolu_1"}, original_exception=KeyError("lookup_order"), diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 5cb74535dce..b09106f5b81 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1732,8 +1732,13 @@ async def _hang_up_after_the_first_follow_up_chunk(stream: CustomStreamWrapper) [(_read_the_whole_stream, "Sunny."), (_hang_up_after_the_first_follow_up_chunk, "Sunny")], ids=["stream_ends", "client_hangs_up"], ) +@pytest.mark.parametrize( + ("stream_call_id", "status_after_stream"), + [("chat-owner", 200), ("silent-shadow", 429)], + ids=["requests_own_stream", "silent_model_shadow_stream"], +) async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the_stream_ends_or_closes( - monkeypatch, end_stream, expected_text + monkeypatch, end_stream, expected_text, stream_call_id: str, status_after_stream: int ): import contextvars from unittest.mock import MagicMock @@ -1795,7 +1800,7 @@ async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the async def run_tool_after_the_first_round_logs(**_): await limiter.async_log_success_event( - kwargs={"litellm_call_id": "chat-owner", "stream": True}, + kwargs={"litellm_call_id": stream_call_id, "stream": True}, response_obj=None, start_time=None, end_time=datetime.now(), @@ -1837,10 +1842,16 @@ async def test_streamed_mcp_auto_exec_keeps_the_requests_parallel_slot_until_the 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 + model="gpt-x", + messages=[{"role": "user", "content": "hi"}], + tools=tools, + stream=True, + litellm_call_id=stream_call_id, ) streamed_text: Final = await end_stream(stream) - assert await second_request_status() == 200, "The request's slot must be freed once its stream ends or closes" + assert ( + await second_request_status() == status_after_stream + ), "Only the request's own stream may free its slot, once that stream ends or closes" assert streamed_text == expected_text