From ca207140ee88b2450dcbfd1313f2cd83b124ed1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:12:57 -0700 Subject: [PATCH] fix(rate-limit): keep one call id across proxy fallbacks and hold the slot when a tool call fails unserved A proxy-level router fallback now reuses the request's call id, so a failure reported with the endpoint's original body still frees the slot the fallback attempt took. A /v1/messages tool call that fails before any stash owner matches now marks the request slot as reused, so a round logged before the failure no longer frees it early. --- litellm/proxy/common_request_processing.py | 10 ++-- .../hooks/parallel_request_limiter_v3.py | 10 ++++ .../mcp/litellm_proxy_mcp_handler.py | 5 +- .../hooks/test_parallel_request_limiter_v3.py | 51 +++++++++++++++++++ .../proxy/test_common_request_processing.py | 37 ++++++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 6 ++- 6 files changed, 113 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..a9f1fe3e4be 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2305,6 +2305,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 @@ -2435,9 +2436,9 @@ class ProxyBaseLLMRequestProcessing: aliased_model: Final = self.data.get("model") - 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) + 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(self.data.get("litellm_call_id")) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, @@ -2583,6 +2584,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( @@ -2600,6 +2602,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 @@ -2656,6 +2659,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/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5d54e57338a..1386dd2973b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5713,6 +5713,15 @@ 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, @@ -5741,6 +5750,7 @@ 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) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 6d5b634b0ec..f3f3867b671 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -745,7 +745,10 @@ 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_request_data = {"litellm_call_id": tool_logging_call_id} + 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) 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 941aff9328b..f391193af7f 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4055,6 +4055,57 @@ async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reserva assert await in_flight() == 0 +@pytest.mark.asyncio +async def test_round_logged_before_a_tool_call_fails_unserved_keeps_the_requests_parallel_slot() -> None: + now: Final = datetime(2026, 10, 28, 12, 0, 0) + api_key: Final = hash_token("sk-messages-unserved-tool") + 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 = "messages-unserved-tool-owner" + + 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 log_round(ended_at: datetime) -> None: + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": False}, + response_obj=None, + start_time=None, + end_time=ended_at, + ) + + 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} + ) + 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 log_round(now - timedelta(milliseconds=1)) + + assert await next_request_status() == 429 + + await log_round(now + timedelta(milliseconds=1)) + + 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") diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 52816dc50a8..d0bb1723dbe 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -7929,6 +7929,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=[ @@ -8096,6 +8097,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_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 825694dbc9f..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: @@ -435,7 +435,9 @@ async def test_execute_tool_calls_reports_an_unmapped_tool_failure_under_the_too call_tool_mock.assert_not_awaited() post_call_failure_hook.assert_awaited_once() assert post_call_failure_hook.await_args is not None - assert post_call_failure_hook.await_args.kwargs["request_data"]["litellm_call_id"] == "cid_call-unmapped" + 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