diff --git a/litellm/llms/anthropic/messages/native/mcp_handler.py b/litellm/llms/anthropic/messages/native/mcp_handler.py index 8fe18d79f24..215d1324c61 100644 --- a/litellm/llms/anthropic/messages/native/mcp_handler.py +++ b/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(context.litellm_call_id): + 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)) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6877d6c0118..c0ccd067b5b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2351,6 +2351,7 @@ class ProxyBaseLLMRequestProcessing: rate_limited_model: str | None = None, skip_guardrails: bool = False, logging_auth: UserAPIKeyAuth | None = None, + litellm_call_id: str | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2487,10 +2488,10 @@ class ProxyBaseLLMRequestProcessing: ) ) - litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) - self.data["litellm_call_id"] = litellm_call_id - bind_telemetry_call_id(litellm_call_id) - DDSpanTagger.tag_call_id(litellm_call_id) + call_id: Final = litellm_call_id or resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + self.data["litellm_call_id"] = call_id + bind_telemetry_call_id(call_id) + DDSpanTagger.tag_call_id(call_id) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, requested_model=aliased_model, @@ -2634,6 +2635,7 @@ class ProxyBaseLLMRequestProcessing: else None ) pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: return await self.common_processing_pre_call_logic( @@ -2651,6 +2653,7 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, + litellm_call_id=litellm_call_id, ) except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data @@ -2707,6 +2710,7 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, rate_limited_model=original_model, + litellm_call_id=litellm_call_id, ) except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index d174d9fecb1..8bc973ce364 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -616,7 +616,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if "model" not in data: return None - claim_request_stash_for_data(data) + claim_request_stash_for_data(data, call_type) model: Final = data["model"] priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index e44d9304ac1..7d70f1532da 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/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 @@ -713,6 +713,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 @@ -801,10 +802,28 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: +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 + return stash + + +@contextmanager +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: + 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") - if isinstance(owner_call_id, str): + if isinstance(owner_call_id, str) and call_type != CallTypes.pass_through.value: stash.owner_litellm_call_id = owner_call_id return stash @@ -2117,7 +2136,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 @@ -4052,7 +4075,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - stash: Final = claim_request_stash_for_data(data) + stash: Final = claim_request_stash_for_data(data, call_type) io_token_descriptors: Final = [ d for d in descriptors if d["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) ] @@ -4281,7 +4304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") - stash: Final = claim_request_stash_for_data(data) + stash: Final = claim_request_stash_for_data(data, call_type) reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None if reuses_request_slot: stash.parallel_slot_reused_at = self._get_current_time() @@ -5787,7 +5810,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @with_service_target("rate_limits") async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -5810,7 +5833,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): that usage instead of refunding it. """ try: - stash: Final = get_request_stash() + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(request_data)) if stash is None: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) @@ -5851,7 +5874,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ptu_ops: Final = ( self._build_ptu_failure_settlement_ops( stash, - request_data, # pyright: ignore[reportUnknownArgumentType] # request_data is a bare dict + request_data, tpm_actual, ) if reserved_tokens > 0 diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..4bf0dcd4f0f 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/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(context.litellm_call_id) is not None def __aiter__(self): return self @@ -332,7 +334,21 @@ 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: + 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: # Get the iterator from the stream wrapper @@ -548,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/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c01d07fb5c3..f3f3867b671 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -742,6 +742,13 @@ class LiteLLM_Proxy_MCP_Handler: tool_arguments, tool_call_id, ) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call) + tool_logging_call_id = ( + f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" if litellm_call_id else str(uuid.uuid4()) + ) + logging_request_data = { + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": tool_logging_call_id, + } if not tool_name: verbose_logger.warning("Tool call missing name: %s", tool_call) @@ -772,9 +779,6 @@ class LiteLLM_Proxy_MCP_Handler: }, } ] - tool_logging_call_id = ( - f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" if litellm_call_id else str(uuid.uuid4()) - ) logging_metadata: dict[str, object] = { "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, 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 73c2edbce19..f3c6ec2f359 100644 --- a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py @@ -454,3 +454,84 @@ 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 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"}] + + 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] + 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/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index e056e723a04..bb00eeb83c1 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -41,6 +41,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _request_stash, get_or_create_request_stash, get_request_stash, + request_slot_held_by_tool_loop, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, @@ -4271,6 +4272,146 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() assert await in_flight() == 1 +@pytest.mark.asyncio +async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reservation(monkeypatch) -> None: + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + api_key: Final = hash_token("sk-mcp-auto-execute-tool-failure") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1, tpm_limit=10_000) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + tokens_key: Final = handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens") + owner_call_id: Final = "mcp-tool-failure-owner" + owner_data: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "weather in Paris"}], + "max_tokens": 50, + "litellm_call_id": owner_call_id, + } + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) + + async def reserved_tokens() -> int: + return int(await local_cache.async_get_cache(key=tokens_key) or 0) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + await pre_call(owner_data, CallTypes.acompletion.value) + await pre_call( + {"model": "MCP: get_weather", "messages": [{"role": "user", "content": "Tool: get_weather"}]}, + CallTypes.call_mcp_tool.value, + ) + reserved_before_tool_failure: Final = await reserved_tokens() + + await handler.async_post_call_failure_hook( + request_data={"call_type": CallTypes.call_mcp_tool.value, "litellm_call_id": f"{owner_call_id}_call_1"}, + original_exception=HTTPException(status_code=504, detail="MCP tool call timed out after 3.0s"), + user_api_key_dict=user_api_key_dict, + ) + + assert reserved_before_tool_failure > 0 + assert await reserved_tokens() == reserved_before_tool_failure + with pytest.raises(HTTPException) as rejected: + await contextvars.Context().run( + asyncio.create_task, pre_call({"model": "gpt-4o-mini", "input": "next"}, CallTypes.aresponses.value) + ) + assert rejected.value.status_code == 429 + + await handler.async_post_call_failure_hook( + request_data=owner_data, + original_exception=Exception("request failed"), + user_api_key_dict=user_api_key_dict, + ) + + assert await in_flight() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream_round", "stream_round"]) +async def test_round_logged_while_a_tool_loop_holds_the_slot_keeps_it_until_the_request_ends(stream: bool) -> None: + now: Final = datetime(2026, 10, 28, 12, 0, 0) + api_key: Final = hash_token(f"sk-tool-loop-hold-{stream}") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), time_provider=lambda: now + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + owner_call_id: Final = f"tool-loop-hold-owner-{stream}" + + async def pre_call(data: dict[str, object]) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=CallTypes.anthropic_messages.value, + ) + + async def next_request_status() -> int: + try: + await contextvars.Context().run(asyncio.create_task, pre_call({"model": "claude-x", "messages": []})) + except HTTPException as rejected: + return rejected.status_code + return 200 + + await pre_call( + {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}], "litellm_call_id": owner_call_id} + ) + 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"), + user_api_key_dict=user_api_key_dict, + ) + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": stream}, + response_obj=None, + start_time=None, + end_time=now - timedelta(milliseconds=1), + ) + + assert await next_request_status() == 429 + + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + + assert await next_request_status() == 200 + + +@pytest.mark.asyncio +async def test_failed_pass_through_request_frees_its_slot_when_the_client_body_carries_a_litellm_call_id() -> None: + api_key: Final = hash_token("sk-pass-through-client-call-id") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + + async def pre_call(data: dict[str, object]) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=CallTypes.pass_through.value, + ) + + await pre_call({"model": "gpt-4o-mini", "max_completion_tokens": -5, "litellm_call_id": "client-chosen-id"}) + await handler.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "litellm_call_id": "proxy-assigned-id"}, + original_exception=HTTPException(status_code=400, detail="max_completion_tokens must be positive"), + user_api_key_dict=user_api_key_dict, + ) + + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + await contextvars.Context().run(asyncio.create_task, pre_call({"model": "gpt-4o-mini"})) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("stream", "in_flight_after_first_round_logs"), diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..046b38fd1ba 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -8092,6 +8092,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) + proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=limiter.async_post_call_failure_hook) guardrails_by_group = model_guardrails or {} router = litellm.Router( model_list=[ @@ -8259,6 +8260,42 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["litellm_logging_obj"].model == primary_model assert processor.data["litellm_call_id"] + @pytest.mark.asyncio + async def test_failed_request_served_by_a_fallback_frees_its_parallel_slot(self, monkeypatch: pytest.MonkeyPatch): + """/completions reports a failure with the endpoint's own request body, the one the + rate-limited first attempt ran on, so the fallback attempt that took the slot has to carry + the same call id or the slot never comes back.""" + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = ProxyUserAPIKeyAuth(api_key="hashed-key", max_parallel_requests=1) + rig: Final = self._v3_limiter_rig( + monkeypatch, + key, + [{primary_model: [fallback_model]}], + untagged_rate_limited_models=frozenset({primary_model}), + ) + + async def request_that_fails_after_falling_back() -> None: + client_data: Final[dict[str, object]] = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + } + await self._pre_call(client_data, key, rig) + await rig[0].post_call_failure_hook( + user_api_key_dict=key, + original_exception=Exception("upstream failed"), + request_data=client_data, + ) + + await contextvars.Context().run(asyncio.create_task, request_that_fails_after_falling_back()) + _, (next_data, _) = await contextvars.Context().run( + asyncio.create_task, + self._pre_call({"model": fallback_model, "messages": [{"role": "user", "content": "next"}]}, key, rig), + ) + + assert next_data["model"] == fallback_model + assert rig[3] == [primary_model, fallback_model, fallback_model] + HARD_PER_MODEL_LIMITS: Final = MappingProxyType({"disable_fallbacks_on_per_model_rate_limits": True}) @pytest.mark.parametrize("cap_owner", ["key", "team"]) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..b09106f5b81 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 @@ -1712,3 +1713,145 @@ 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 + + +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 +@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"], +) +@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, stream_call_id: str, status_after_stream: int +): + 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": stream_call_id, "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, + litellm_call_id=stream_call_id, + ) + streamed_text: Final = await end_stream(stream) + + 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 diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 229d5cbe676..7600e7bea3b 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -33,7 +33,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.responses.main import OutputFunctionToolCall -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, ModelResponse class _DummyMCPResult: @@ -419,6 +419,27 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" +@pytest.mark.asyncio +async def test_execute_tool_calls_reports_an_unmapped_tool_failure_under_the_tool_calls_own_id(monkeypatch): + call_tool_mock: Final = _setup_mcp_call_environment(monkeypatch) + post_call_failure_hook: Final = _setup_proxy_logging(monkeypatch) + + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={}, + tool_calls=[{"id": "call-unmapped", "function": {"name": "unknown_tool", "arguments": "{}"}}], + user_api_key_auth=types.SimpleNamespace(api_key="test_key", user_id="test_user"), + litellm_call_id="cid", + ) + + assert results[0]["result"].startswith("Error executing tool") + call_tool_mock.assert_not_awaited() + post_call_failure_hook.assert_awaited_once() + assert post_call_failure_hook.await_args is not None + reported_request: Final = post_call_failure_hook.await_args.kwargs["request_data"] + assert reported_request["litellm_call_id"] == "cid_call-unmapped" + assert reported_request["call_type"] == CallTypes.call_mcp_tool.value + + @pytest.mark.asyncio async def test_execute_tool_calls_logs_each_tool_call_under_its_own_id_on_the_request_trace( monkeypatch,