From d041c0ffe17c71469af5163ffe3cca2087f5b0c9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:58:07 -0700 Subject: [PATCH 1/6] fix(rate-limit): keep a request's parallel slot when its auto-executed MCP tool call fails --- .../hooks/parallel_request_limiter_v3.py | 6 +- .../hooks/test_parallel_request_limiter_v3.py | 65 +++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..de958d532e0 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5716,7 +5716,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, @@ -5739,7 +5739,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) @@ -5780,7 +5780,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/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index e056e723a04..4daf98fe6bf 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4271,6 +4271,71 @@ 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", "in_flight_after_first_round_logs"), From d259ea474d362497f519bdf21ae6bbff81d4c8ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:18:41 -0700 Subject: [PATCH 2/6] fix(rate-limit): keep the slot owner consistent for unmapped MCP tools and pass-through bodies A tool call that failed before its logging data was built, such as a tool name the gateway does not serve, reached the failure hook with no call id and still freed its parent's slot. The tool call's own id is now set right after the call is read. A pass-through body that carried its own litellm_call_id made that value the slot's owner, so the failure hook, which sees the proxy's id, skipped the release and the slot stayed taken for an hour. The limiter no longer takes an owner from a pass-through body. --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 2 +- .../hooks/parallel_request_limiter_v3.py | 8 +- .../mcp/litellm_proxy_mcp_handler.py | 7 +- .../hooks/test_parallel_request_limiter_v3.py | 1108 +++++------------ .../mcp/test_litellm_proxy_mcp_handler.py | 19 + 5 files changed, 370 insertions(+), 774 deletions(-) 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 de958d532e0..5d54e57338a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -785,10 +785,10 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: +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 @@ -3984,7 +3984,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) ] @@ -4212,7 +4212,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() diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c01d07fb5c3..6d5b634b0ec 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -742,6 +742,10 @@ 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 = {"litellm_call_id": tool_logging_call_id} if not tool_name: verbose_logger.warning("Tool call missing name: %s", tool_call) @@ -772,9 +776,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/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 4daf98fe6bf..941aff9328b 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -97,14 +97,10 @@ def _isolated_request_stash(): (0.5, 50, 500), ], ) -def test_api_key_descriptor_applies_budget_throttle( - throttle_pct, expected_rpm, expected_tpm -): +def test_api_key_descriptor_applies_budget_throttle(throttle_pct, expected_rpm, expected_tpm): """The api_key rate-limit descriptor scales the key's configured TPM/RPM by the request-scoped budget_throttle_pct, leaving the configured limits intact.""" - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-throttle"), rpm_limit=100, @@ -176,8 +172,11 @@ async def test_mcp_description_does_not_change_admission_or_reserved_tokens( @pytest.mark.parametrize("itpm_limit,otpm_limit", [(64, 4096), (4096, 64), (4096, 4096)]) @pytest.mark.parametrize("arguments_rewritten", [False, True]) async def test_mcp_description_preserves_project_input_and_output_reservations( - description: str | None, itpm_limit: int, otpm_limit: int, - arguments_rewritten: bool, monkeypatch: pytest.MonkeyPatch + description: str | None, + itpm_limit: int, + otpm_limit: int, + arguments_rewritten: bool, + monkeypatch: pytest.MonkeyPatch, ) -> None: cache: Final = DualCache() handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) @@ -315,18 +314,12 @@ async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -408,18 +401,12 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -451,9 +438,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller new_window_value = await local_cache.async_get_cache(key=window_key) new_counter_value = await local_cache.async_get_cache(key=counter_key) - assert ( - new_window_value == window_value - ), "Window value should not change within window" + assert new_window_value == window_value, "Window value should not change within window" assert new_counter_value == 2, "Counter should be 2 after second request" # Wait for window to expire @@ -484,9 +469,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller ) @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio -async def test_normal_router_call_tpm_v3( - monkeypatch, rate_limit_object, time_controller -): +async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_controller): """ Test normal router call with parallel request limiter v3 for TPM rate limiting """ @@ -562,18 +545,12 @@ async def test_normal_router_call_tpm_v3( window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -596,9 +573,7 @@ async def test_normal_router_call_tpm_v3( return None value = get_value_for_key(rate_limit_object, user_api_key_dict, "azure-model") - counter_key = parallel_request_handler.create_rate_limit_keys( - rate_limit_object, value, "tokens" - ) + counter_key = parallel_request_handler.create_rate_limit_keys(rate_limit_object, value, "tokens") # First request should succeed. Include messages + a tight max_tokens so # the atomic reserve_tpm_tokens path populates the :tokens counter with a @@ -611,12 +586,8 @@ async def test_normal_router_call_tpm_v3( "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5, } - expected_reservation = parallel_request_handler._estimate_tokens_for_request( - data=pre_call_data - ) - assert ( - expected_reservation < 10 - ), "Test premise: reservation must fit under tpm_limit=10" + expected_reservation = parallel_request_handler._estimate_tokens_for_request(data=pre_call_data) + assert expected_reservation < 10, "Test premise: reservation must fit under tpm_limit=10" await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -644,15 +615,11 @@ async def test_normal_router_call_tpm_v3( counter_value = await local_cache.async_get_cache(key=counter_key) print(f"local_cache: {local_cache.in_memory_cache.cache_dict}") - assert ( - counter_value is not None - ), f"Counter value should be stored in cache for {counter_key}" + assert counter_value is not None, f"Counter value should be stored in cache for {counter_key}" # Manually increment the token counter to simulate token usage from previous call # This simulates what would happen after a successful call - await local_cache.async_increment_cache( - key=counter_key, value=15, ttl=2 - ) # Use up most of our 10 token limit + await local_cache.async_increment_cache(key=counter_key, value=15, ttl=2) # Use up most of our 10 token limit # Make another request to test rate limiting - this should fail as we've consumed tokens with pytest.raises(HTTPException) as exc_info: @@ -700,17 +667,13 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ _api_key = hash_token(_api_key) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=100) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method directly since it imports general_settings internally def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response with different token counts mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) @@ -759,9 +722,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ) # Verify that the correct token count was used based on the rate limit type - assert ( - len(captured_operations) == 1 - ), "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" + assert len(captured_operations) == 1, ( + "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" + ) tpm_operation = None for op in captured_operations: @@ -778,9 +741,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ "total": mock_usage.total_tokens, # 50 } - assert ( - tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] - ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( + f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + ) @pytest.mark.parametrize( @@ -797,9 +760,7 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ], ) @pytest.mark.asyncio -async def test_async_log_success_event_counts_non_chat_response_tokens( - monkeypatch, response_obj -): +async def test_async_log_success_event_counts_non_chat_response_tokens(monkeypatch, response_obj): """ Embedding and text completion responses must increment the TPM counter, not just chat completion ModelResponse objects. @@ -807,12 +768,8 @@ async def test_async_log_success_event_counts_non_chat_response_tokens( monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-98765") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: "total" - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") mock_kwargs = { "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, @@ -838,9 +795,7 @@ async def test_async_log_success_event_counts_non_chat_response_tokens( end_time=datetime.now(), ) - tpm_operation = next( - (op for op in captured_operations if op["key"].endswith(":tokens")), None - ) + tpm_operation = next((op for op in captured_operations if op["key"].endswith(":tokens")), None) assert tpm_operation is not None, "Should have a TPM increment operation" assert tpm_operation["increment_value"] == 50 @@ -856,9 +811,7 @@ async def test_async_log_failure_event_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b"]) @@ -908,29 +861,18 @@ async def test_failure_event_without_acquired_slot_does_not_release_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - await _seed_max_parallel_requests_slots( - local_cache, counter_key, ["slot-a", "slot-b", "slot-c"] - ) + await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b", "slot-c"]) await handler.async_log_failure_event( - kwargs={ - "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} - }, + kwargs={"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}}, response_obj=None, start_time=None, end_time=None, ) - assert ( - handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) - == 3 - ) + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 3 @pytest.mark.asyncio @@ -981,9 +923,7 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3(): rejected requests that should have been admitted after a release. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) _api_key = hash_token("sk-98765") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) @@ -1036,9 +976,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): and an over-limit script result maps to a 429 without occupying a slot. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) _api_key = hash_token("sk-98765") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -1063,9 +1001,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): stashed_slot_id = stashed_acquisition["slot_id"] assert isinstance(stashed_slot_id, str) and stashed_slot_id assert stashed_acquisition["counter_keys"] == [counter_key] - assert captured_calls == [ - ([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id]) - ] + assert captured_calls == [([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id])] assert ( await handler.internal_usage_cache.async_get_cache( key=counter_key, litellm_parent_otel_span=None, local_only=True @@ -1112,9 +1048,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -1143,9 +1077,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - not should_rate_limit_called - ), "should_rate_limit should not be called when no rate limits are configured" + assert not should_rate_limit_called, "should_rate_limit should not be called when no rate limits are configured" # Test 2: API key rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1161,9 +1093,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when API key rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when API key rate limits are configured" # Test 3: User rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1180,9 +1110,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when user rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when user rate limits are configured" # Test 4: Team rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1199,9 +1127,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when team rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when team rate limits are configured" # Test 5: End user rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1218,9 +1144,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when end user rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when end user rate limits are configured" # Test 6: Max parallel requests configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1236,9 +1160,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when max parallel requests are configured" + assert should_rate_limit_called, "should_rate_limit should be called when max parallel requests are configured" @pytest.mark.asyncio @@ -1254,9 +1176,7 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -1272,9 +1192,7 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): should_rate_limit_called = False user_api_key_dict_with_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={ - "model_tpm_limit": {"gpt-4": 1000} - }, # Rate limit for gpt-4, not gpt-3.5-turbo + metadata={"model_tpm_limit": {"gpt-4": 1000}}, # Rate limit for gpt-4, not gpt-3.5-turbo ) await parallel_request_handler.async_pre_call_hook( @@ -1284,17 +1202,15 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert ( - not should_rate_limit_called - ), "should_rate_limit should not be called when model-specific limits don't match requested model" + assert not should_rate_limit_called, ( + "should_rate_limit should not be called when model-specific limits don't match requested model" + ) # Test 2: Model-specific rate limits configured for requested model - SHOULD be called should_rate_limit_called = False user_api_key_dict_with_matching_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={ - "model_tpm_limit": {"gpt-3.5-turbo": 1000} - }, # Rate limit for requested model + metadata={"model_tpm_limit": {"gpt-3.5-turbo": 1000}}, # Rate limit for requested model ) await parallel_request_handler.async_pre_call_hook( @@ -1304,9 +1220,9 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when model-specific limits match requested model" + assert should_rate_limit_called, ( + "should_rate_limit should be called when model-specific limits match requested model" + ) @pytest.mark.asyncio @@ -1333,9 +1249,7 @@ async def test_tpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1393,15 +1307,11 @@ async def test_tpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert ( - model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" - ), "Api-Key value should combine api_key and model" - assert ( - model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit - ), "Api-Key RPM limit should be set" - assert ( - model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit - ), "Api-Key TPM limit should be set" + assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( + "Api-Key value should combine api_key and model" + ) + assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" + assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1428,9 +1338,7 @@ async def test_rpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1488,15 +1396,11 @@ async def test_rpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert ( - model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" - ), "Api-Key value should combine api_key and model" - assert ( - model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit - ), "Api-Key RPM limit should be set" - assert ( - model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit - ), "Api-Key TPM limit should be set" + assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( + "Api-Key value should combine api_key and model" + ) + assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" + assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1518,9 +1422,7 @@ async def test_team_member_rate_limits_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1552,18 +1454,12 @@ async def test_team_member_rate_limits_v3(): team_member_descriptor = descriptor break - assert ( - team_member_descriptor is not None - ), "Team member descriptor should be present" - assert ( - team_member_descriptor["value"] == f"{_team_id}:{_user_id}" - ), "Team member value should combine team_id and user_id" - assert ( - team_member_descriptor["rate_limit"]["requests_per_unit"] == 10 - ), "Team member RPM limit should be set" - assert ( - team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000 - ), "Team member TPM limit should be set" + assert team_member_descriptor is not None, "Team member descriptor should be present" + assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", ( + "Team member value should combine team_id and user_id" + ) + assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set" + assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set" @pytest.mark.asyncio @@ -1586,9 +1482,7 @@ async def test_team_member_rate_limits_v3_raises_429_when_over_limit(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = None @@ -1664,9 +1558,7 @@ async def test_dynamic_rate_limiting_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if limits are enforced captured_descriptors = [] @@ -1699,9 +1591,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert ( - api_key_descriptor["rate_limit"]["requests_per_unit"] is None - ), "RPM limit should be None when dynamic mode and no failures" + assert api_key_descriptor["rate_limit"]["requests_per_unit"] is None, ( + "RPM limit should be None when dynamic mode and no failures" + ) # Test 2: With failures - rate limits SHOULD be enforced (rpm_limit should be set) async def mock_check_with_failures(*args, **kwargs): @@ -1725,9 +1617,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert ( - api_key_descriptor["rate_limit"]["requests_per_unit"] == 2 - ), "RPM limit should be enforced when dynamic mode and failures detected" + assert api_key_descriptor["rate_limit"]["requests_per_unit"] == 2, ( + "RPM limit should be enforced when dynamic mode and failures detected" + ) @pytest.mark.asyncio @@ -1738,44 +1630,30 @@ async def test_async_increment_tokens_fallback_behavior(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the token_increment_script to None to simulate unavailable script parallel_request_handler.token_increment_script = None # Mock the fallback method fallback_called = False - original_method = ( - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline - ) + original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline async def mock_fallback(*args, **kwargs): nonlocal fallback_called fallback_called = True return await original_method(*args, **kwargs) - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( - mock_fallback - ) + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback # Test operations - pipeline_operations = [ - RedisPipelineIncrementOperation( - key="test_fallback_key", increment_value=10.0, ttl=60 - ) - ] + pipeline_operations = [RedisPipelineIncrementOperation(key="test_fallback_key", increment_value=10.0, ttl=60)] # Execute increment - await parallel_request_handler.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations - ) + await parallel_request_handler.async_increment_tokens_with_ttl_preservation(pipeline_operations=pipeline_operations) # Verify fallback was called - assert ( - fallback_called - ), "Fallback method should be called when Lua script is not available" + assert fallback_called, "Fallback method should be called when Lua script is not available" # Redis Cluster Compatibility Tests @@ -1786,9 +1664,7 @@ def test_group_keys_by_hash_tag_regular_redis(): For regular Redis, all keys should be grouped together under a single group. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Test keys with different hash tags test_keys = [ @@ -1808,9 +1684,7 @@ def test_group_keys_by_hash_tag_regular_redis(): # Verify all keys are in single group for regular Redis assert len(groups) == 1, f"Expected 1 group for regular Redis, got {len(groups)}" assert "all_keys" in groups, "Expected 'all_keys' group for regular Redis" - assert set(groups["all_keys"]) == set( - test_keys - ), "All keys should be in single group" + assert set(groups["all_keys"]) == set(test_keys), "All keys should be in single group" def test_group_keys_by_hash_tag_redis_cluster(): @@ -1822,9 +1696,7 @@ def test_group_keys_by_hash_tag_redis_cluster(): from unittest.mock import patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -1844,17 +1716,13 @@ def test_group_keys_by_hash_tag_redis_cluster(): # All group keys should start with "slot_" for group_key in groups.keys(): - assert group_key.startswith( - "slot_" - ), f"Group key {group_key} should start with 'slot_'" + assert group_key.startswith("slot_"), f"Group key {group_key} should start with 'slot_'" # Verify all original keys are present across groups all_grouped_keys = [] for group_keys in groups.values(): all_grouped_keys.extend(group_keys) - assert set(all_grouped_keys) == set( - test_keys - ), "All keys should be present in groups" + assert set(all_grouped_keys) == set(test_keys), "All keys should be present in groups" def test_keyslot_for_redis_cluster(): @@ -1862,9 +1730,7 @@ def test_keyslot_for_redis_cluster(): Test the keyslot calculation for Redis cluster. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Test basic key slot1 = handler.keyslot_for_redis_cluster("user:1000") @@ -1892,18 +1758,14 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): from unittest.mock import AsyncMock, patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script that simulates Redis cluster slot conflict mock_script = AsyncMock() mock_script.side_effect = [ - Exception( - "EVALSHA - all keys must map to the same key slot" - ), # First group fails + Exception("EVALSHA - all keys must map to the same key slot"), # First group fails [1234, 2], # Second group succeeds ] handler.batch_rate_limiter_script = mock_script @@ -1920,9 +1782,7 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): ] # Execute the method - results = await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=test_keys, now_int=1234 - ) + results = await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=test_keys, now_int=1234) assert results == [1234, 1, 1234, 2] @@ -1946,9 +1806,7 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): # Should have processed all keys (some might be duplicated due to fallback) unique_processed_keys = set(all_processed_keys) - assert ( - len(unique_processed_keys) >= 2 - ), "Should have processed at least some keys" + assert len(unique_processed_keys) >= 2, "Should have processed at least some keys" @pytest.mark.asyncio @@ -1974,9 +1832,7 @@ async def test_multiple_rate_limits_per_descriptor(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to return a response with multiple statuses where one hits the limit # This simulates the case where we have more statuses than descriptors due to multiple rate limit types @@ -2054,9 +1910,7 @@ async def test_missing_descriptor_fallback(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to return a status with descriptor_key that doesn't match descriptors async def mock_should_rate_limit(descriptors, **kwargs): @@ -2100,9 +1954,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): This verifies the change from 'output' to 'total' as the default value. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return empty dict (no token_rate_limit_type set) import litellm.proxy.proxy_server as proxy_server @@ -2112,9 +1964,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): try: result = parallel_request_handler.get_rate_limit_type() - assert ( - result == "total" - ), f"Default rate limit type should be 'total', got '{result}'" + assert result == "total", f"Default rate limit type should be 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -2125,23 +1975,17 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): Test that get_rate_limit_type falls back to 'total' when an invalid value is specified. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return an invalid token_rate_limit_type import litellm.proxy.proxy_server as proxy_server original_settings = getattr(proxy_server, "general_settings", {}) - monkeypatch.setattr( - proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"} - ) + monkeypatch.setattr(proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"}) try: result = parallel_request_handler.get_rate_limit_type() - assert ( - result == "total" - ), f"Invalid rate limit type should fall back to 'total', got '{result}'" + assert result == "total", f"Invalid rate limit type should fall back to 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -2155,9 +1999,7 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): ], ) @pytest.mark.asyncio -async def test_async_log_success_event_with_dict_usage( - monkeypatch, token_rate_limit_type, expected_field -): +async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_limit_type, expected_field): """ Test that async_log_success_event correctly handles usage as a dict (Responses API format). @@ -2169,17 +2011,13 @@ async def test_async_log_success_event_with_dict_usage( _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response object with usage as a dict (Responses API format) from litellm.types.utils import BaseLiteLLMOpenAIResponseObject @@ -2242,9 +2080,9 @@ async def test_async_log_success_event_with_dict_usage( "total": 60, # total_tokens } - assert ( - tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] - ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( + f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + ) @pytest.mark.asyncio @@ -2259,17 +2097,13 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return "output" - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response object with usage as a dict missing some fields mock_response = MagicMock() @@ -2280,9 +2114,7 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc } from litellm.types.utils import BaseLiteLLMOpenAIResponseObject - mock_response.__class__ = type( - "MockResponse", (BaseLiteLLMOpenAIResponseObject,), {} - ) + mock_response.__class__ = type("MockResponse", (BaseLiteLLMOpenAIResponseObject,), {}) # Create mock kwargs for the success event mock_kwargs = { @@ -2338,9 +2170,7 @@ async def test_execute_token_increment_script_cluster_compatibility(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -2378,18 +2208,14 @@ async def test_execute_token_increment_script_cluster_compatibility(): "{api_key:sk-123}:max_parallel_requests", "{user:user-456}:tokens", } - assert ( - set(all_processed_keys) == expected_keys - ), "All operation keys should be processed" + assert set(all_processed_keys) == expected_keys, "All operation keys should be processed" # Verify args structure is correct for each call for call_args in call_args_list: keys = call_args[1]["keys"] args = call_args[1]["args"] # Each key should have 2 args (increment_value, ttl) - assert ( - len(args) == len(keys) * 2 - ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" + assert len(args) == len(keys) * 2, f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" @pytest.mark.asyncio @@ -2412,9 +2238,7 @@ async def test_agent_level_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2479,9 +2303,7 @@ async def test_agent_session_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2548,9 +2370,7 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2583,8 +2403,7 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): # should_rate_limit should not have been called (no agent-level limits, only session limits # but no session_id) assert captured_descriptors is None, ( - "No descriptors should be created when agent has only session limits " - "but no session_id in request" + "No descriptors should be created when agent has only session limits but no session_id in request" ) @@ -2607,9 +2426,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2650,9 +2467,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): agent_descriptor = d break - assert ( - agent_descriptor is not None - ), "Agent descriptor should be created from metadata agent_id" + assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" assert agent_descriptor["value"] == _agent_id assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 @@ -2678,9 +2493,7 @@ async def test_agent_both_agent_and_session_rate_limits(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2747,16 +2560,12 @@ async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): _session_id = "sess_tpm_test" local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) def mock_get_rate_limit_type(): return "total" - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) mock_response = ModelResponse( @@ -2867,18 +2676,12 @@ async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -3016,9 +2819,7 @@ async def test_project_model_rate_limits_enforced_v3(): """ _api_key = hash_token("sk-project-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3046,13 +2847,9 @@ async def test_project_model_rate_limits_enforced_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert ( - "model_per_project" in descriptor_keys - ), f"Expected model_per_project descriptor, got: {descriptor_keys}" + assert "model_per_project" in descriptor_keys, f"Expected model_per_project descriptor, got: {descriptor_keys}" - model_per_project = next( - d for d in captured_descriptors if d["key"] == "model_per_project" - ) + model_per_project = next(d for d in captured_descriptors if d["key"] == "model_per_project") assert model_per_project["value"] == "proj-abc123:gpt-4" assert model_per_project["rate_limit"]["requests_per_unit"] == 5 assert model_per_project["rate_limit"]["tokens_per_unit"] == 1000 @@ -3063,9 +2860,7 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): """Project model limits should not trigger for a model not in project_metadata.""" _api_key = hash_token("sk-project-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3092,9 +2887,9 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert ( - "model_per_project" not in descriptor_keys - ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" + assert "model_per_project" not in descriptor_keys, ( + f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" + ) @pytest.mark.asyncio @@ -3105,9 +2900,7 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): """ _api_key = hash_token("sk-project-io-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3138,12 +2931,8 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): assert "model_per_project_otpm" in descriptor_keys assert "model_per_project" not in descriptor_keys - itpm_descriptor = next( - d for d in captured_descriptors if d["key"] == "model_per_project_itpm" - ) - otpm_descriptor = next( - d for d in captured_descriptors if d["key"] == "model_per_project_otpm" - ) + itpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_itpm") + otpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_otpm") assert itpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" assert itpm_descriptor["rate_limit"]["tokens_per_unit"] == 20000000 assert otpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" @@ -3155,9 +2944,7 @@ async def test_project_model_itpm_otpm_limits_not_triggered_for_other_model_v3() """Split project limits must not apply to an unrelated model.""" _api_key = hash_token("sk-project-io-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3192,9 +2979,7 @@ async def test_project_model_itpm_and_tpm_limits_coexist_v3(): """Combined project TPM and split ITPM/OTPM limits are enforced together.""" _api_key = hash_token("sk-project-io-test-3") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3236,9 +3021,7 @@ async def test_enforce_project_io_token_quota_for_frame_blocks_over_limit_otpm() limit and reject once a frame's estimated output tokens exceed it.""" _api_key = hash_token("sk-ws-frame-otpm") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, project_id="proj-mantle-ws", @@ -3270,9 +3053,7 @@ async def test_enforce_project_io_token_quota_for_frame_noop_without_project_lim the per-frame check (no descriptors to reserve against).""" _api_key = hash_token("sk-ws-frame-no-limits") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) await handler.enforce_project_io_token_quota_for_frame( @@ -3407,17 +3188,13 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-refund-lifecycle") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys( - key="api_key", value=_api_key, rate_limit_type="tokens" - ) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3434,28 +3211,18 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): reserved = get_request_stash().reserved_tokens assert reserved > 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 kwargs = {"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}} - await handler.async_log_failure_event( - kwargs=kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 assert get_request_stash().reservation_released is True - await handler.async_log_failure_event( - kwargs=kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 @pytest.mark.asyncio @@ -3499,9 +3266,7 @@ async def test_pre_call_hook_ignores_caller_supplied_stash_values(): async def spy_increment_pipeline(increment_list, **kwargs): refund_calls.append(increment_list) - handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( - spy_increment_pipeline - ) + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = spy_increment_pipeline await handler.async_post_call_failure_hook( request_data=data, @@ -3527,17 +3292,13 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-nested-guard") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys( - key="api_key", value=_api_key, rate_limit_type="tokens" - ) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3562,33 +3323,23 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) "litellm_call_id": "nested-guardrail-call", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_success_event( - kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None - ) - await handler.async_log_failure_event( - kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_success_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_failure_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) assert stash.parallel_slot is not None assert stash.reservation_released is False - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved owner_kwargs = { "litellm_call_id": "owner-call-id", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_failure_event( - kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None) assert stash.parallel_slot is None assert stash.reservation_released is True - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 @@ -3602,9 +3353,7 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): do not thread ``litellm_call_id`` into their logging kwargs. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) unclaimed = get_or_create_request_stash() unclaimed.reserved_tokens = 42 @@ -3616,9 +3365,7 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): ) assert unclaimed.reservation_released is True - claimed = RequestRateLimiterStash( - owner_litellm_call_id="owner-1", reserved_tokens=42 - ) + claimed = RequestRateLimiterStash(owner_litellm_call_id="owner-1", reserved_tokens=42) _request_stash.set(claimed) await handler.async_log_failure_event( kwargs={"standard_logging_object": {}}, @@ -3801,9 +3548,7 @@ async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usag def _make_mcp_handler() -> tuple[PROXY_MaxParallelRequestsHandler_v3, DualCache]: local_cache: Final = DualCache() - handler: Final = PROXY_MaxParallelRequestsHandler_v3( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache)) return handler, local_cache @@ -3827,11 +3572,7 @@ async def test_mcp_per_key_rate_limit_uses_trusted_server_alias_v3() -> None: with pytest.raises(ProxyRateLimitError, match="mcp_per_key"): await handler.enforce_mcp_server_rate_limits(user_api_key_dict, server) - assert all( - value == 0 - for key, value in local_cache.in_memory_cache.cache_dict.items() - if key.endswith(":tokens") - ) + assert all(value == 0 for key, value in local_cache.in_memory_cache.cache_dict.items() if key.endswith(":tokens")) @pytest.mark.asyncio @@ -4146,9 +3887,7 @@ def test_get_key_mcp_rpm_limit_precedence(): _TEST_SLOT_ID = "slot-disconnect-test" -async def _seed_max_parallel_requests_slots( - dual_cache: DualCache, counter_key: str, slot_ids: List[str] -) -> None: +async def _seed_max_parallel_requests_slots(dual_cache: DualCache, counter_key: str, slot_ids: List[str]) -> None: await dual_cache.async_set_cache( key=counter_key, value={slot_id: time.time() for slot_id in slot_ids}, @@ -4160,9 +3899,7 @@ async def _seed_max_parallel_requests_slots( async def test_completed_responses_post_call_releases_parallel_slot() -> None: api_key = hash_token("sk-responses-post-call") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) data = { "model": "gpt-4o-mini", @@ -4177,9 +3914,7 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: data=data, call_type="aresponses", ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 await handler.async_post_call_success_hook( data=data, @@ -4193,9 +3928,7 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: status="completed", ), ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 await handler.async_log_success_event( kwargs={"litellm_call_id": data["litellm_call_id"]}, @@ -4203,26 +3936,20 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 @pytest.mark.asyncio async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: api_key: Final = hash_token("sk-mcp-auto-execute") local_cache: Final = DualCache() - handler: Final = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + 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" owner_call_id: Final = "mcp-auto-execute-owner" async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) + return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) def mcp_tool_call() -> dict[str, object]: return { @@ -4239,9 +3966,7 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() call_type=call_type, ) - async def pre_call_from_another_request( - data: dict[str, object], call_type: str - ) -> None: + async def pre_call_from_another_request(data: dict[str, object], call_type: str) -> None: await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) await pre_call( @@ -4253,9 +3978,7 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() assert await in_flight() == 1 with pytest.raises(HTTPException): - await pre_call_from_another_request( - mcp_tool_call(), CallTypes.call_mcp_tool.value - ) + await pre_call_from_another_request(mcp_tool_call(), CallTypes.call_mcp_tool.value) await handler.async_log_success_event( kwargs={"litellm_call_id": owner_call_id}, @@ -4276,9 +3999,7 @@ async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reserva 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) - ) + 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") @@ -4291,9 +4012,7 @@ async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reserva } async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) + 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) @@ -4336,6 +4055,33 @@ 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_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"), @@ -4357,9 +4103,7 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request owner_call_id: Final = f"mcp-late-round-log-{stream}" async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) + return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) async def pre_call(data: dict[str, object], call_type: str) -> None: await handler.async_pre_call_hook( @@ -4420,9 +4164,7 @@ async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kep messages: Final = [{"role": "user", "content": "add 2 and 3"}] async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) + return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) async def pre_call(data: dict[str, object], call_type: str) -> None: await handler.async_pre_call_hook( @@ -4492,9 +4234,7 @@ async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kep async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: api_key: Final = hash_token("sk-mcp-reuse-headers") local_cache: Final = DualCache() - handler: Final = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) data: Final[dict[str, object]] = { "model": "gpt-4o-mini", @@ -4524,9 +4264,7 @@ async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_hea ) response._hidden_params = {} - await handler.async_post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response - ) + await handler.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) headers: Final = response._hidden_params["additional_headers"] assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 @@ -4542,18 +4280,14 @@ async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redi api_key = hash_token("sk-concurrent-release") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) call_id = "concurrent-release-owner" parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" release_started = asyncio.Event() allow_redis_failure = asyncio.Event() - async def failing_release( - keys: Sequence[str], args: Sequence[object] - ) -> list[int]: + async def failing_release(keys: Sequence[str], args: Sequence[object]) -> list[int]: release_started.set() await allow_redis_failure.wait() raise ConnectionError("redis unavailable") @@ -4618,9 +4352,7 @@ async def _build_seeded_limiter(): """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") cache = DualCache() - limiter = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(cache) - ) + limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) counter_key = f"{{api_key:{api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(cache, counter_key, [_TEST_SLOT_ID]) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) @@ -4655,16 +4387,12 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 get_or_create_request_stash().parallel_slot = ParallelSlotAcquisition( slot_id=_TEST_SLOT_ID, @@ -4673,9 +4401,7 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert get_request_stash().parallel_slot is None - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -4688,9 +4414,7 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) @@ -4701,9 +4425,7 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): await handler.async_release_max_parallel_requests_on_disconnect( UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=None) ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -4719,9 +4441,7 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4732,18 +4452,14 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 await handler.async_post_call_failure_hook( request_data=admitted_data, original_exception=Exception("guardrail rejected the request"), user_api_key_dict=user_api_key_dict, ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_log_failure_event( kwargs={ @@ -4753,9 +4469,7 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4774,9 +4488,7 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4788,23 +4500,17 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 await handler.async_log_success_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, }, - response_obj=ModelResponse( - usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) - ), + response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), start_time=datetime.now(), end_time=datetime.now(), ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4821,9 +4527,7 @@ _SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1) @pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"]) @pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) @pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) -async def test_sse_keepalive_releases_parallel_capacity( - after_ping: bool, fails: bool, preadmitted: bool -) -> None: +async def test_sse_keepalive_releases_parallel_capacity(after_ping: bool, fails: bool, preadmitted: bool) -> None: from fastapi.responses import StreamingResponse from litellm.proxy.common_request_processing import open_sse_before_first_byte @@ -4842,9 +4546,7 @@ async def test_sse_keepalive_releases_parallel_capacity( upstream_answers.set() async def admit() -> None: - await handler.async_pre_call_hook( - user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" - ) + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion") async def failure(exc: Exception) -> None: await handler.async_post_call_failure_hook( @@ -4958,9 +4660,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" descriptors = [ { @@ -4979,9 +4679,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): handler.parallel_count_script = fake_count response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) - assert captured_calls == [ - ([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS]) - ] + assert captured_calls == [([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS])] assert response["overall_code"] == "OK" assert response["statuses"] == [ { @@ -4998,9 +4696,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): raise ConnectionError("redis unavailable") handler.parallel_count_script = failing_count - await _seed_max_parallel_requests_slots( - local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"] - ) + await _seed_max_parallel_requests_slots(local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"]) response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) assert response["overall_code"] == "OVER_LIMIT" assert response["statuses"][0]["rate_limit_type"] == "max_parallel_requests" @@ -5015,9 +4711,7 @@ async def test_redis_release_script_updates_local_mirror_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" captured_calls = [] @@ -5055,12 +4749,8 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) - user_api_key_dict = UserAPIKeyAuth( - api_key=_api_key, max_parallel_requests=5, tpm_limit=100 - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5, tpm_limit=100) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None, **_kwargs): @@ -5087,9 +4777,7 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): call_type="", ) assert exc_info.value.status_code == 429 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -5104,9 +4792,7 @@ async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -5156,9 +4842,7 @@ async def test_release_max_parallel_requests_on_disconnect_noop_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_release_max_parallel_requests_on_disconnect( @@ -5189,9 +4873,7 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 1 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 1 proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter @@ -5222,9 +4904,7 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.parametrize("disconnect", ["cancel", "aclose"]) @@ -5271,14 +4951,10 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect else: await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( - saved_hook - ) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -5296,9 +4972,7 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): from litellm.integrations.custom_logger import CustomLogger class _PassthroughIteratorOverride(CustomLogger): - async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict, response, request_data - ): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): async for chunk in response: yield chunk @@ -5326,14 +5000,10 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): await gen.__anext__() await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( - saved_hook - ) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -5341,18 +5011,14 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): def test_tpm_reservation_enabled_by_default(monkeypatch): """Upfront TPM reservation is on unless explicitly disabled via env.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) assert handler.tpm_reservation_enabled is True @pytest.mark.parametrize("value", ["false", "False", "FALSE"]) def test_tpm_reservation_disabled_via_env(monkeypatch, value): monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", value) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) assert handler.tpm_reservation_enabled is False @@ -5364,9 +5030,7 @@ async def test_pre_call_hook_reserves_tpm_when_enabled(monkeypatch): only the reservation path owns it. """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -5406,9 +5070,7 @@ async def test_pre_call_hook_skips_reservation_when_disabled(monkeypatch): pre-v1.82 post-call accounting behavior. """ monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -5456,9 +5118,7 @@ async def test_per_tag_rate_limit_independent_counters_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def call(tag: str) -> None: await handler.async_pre_call_hook( @@ -5492,9 +5152,7 @@ async def test_per_tag_descriptor_creation_v3(): api_key=_api_key, metadata={"tag_rpm_limit": {"cell-1": 5}}, ) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -5518,9 +5176,7 @@ async def test_per_tag_descriptor_absent_without_config_v3(): api_key=hash_token("sk-no-tag"), rpm_limit=10, ) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -5551,9 +5207,7 @@ async def test_per_tag_untagged_request_governed_by_key_limit_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def call(metadata: dict) -> None: await handler.async_pre_call_hook( @@ -5750,9 +5404,7 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) tpm_limit=10000, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Real pre-call: populates data and stashes the response into metadata # so the success callback can find it via litellm_params.metadata. @@ -5809,27 +5461,16 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) call_type="acompletion", ) - additional_headers = ( - mock_kwargs["standard_logging_object"] - .get("hidden_params", {}) - .get("additional_headers", {}) - ) + additional_headers = mock_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) # api_key-scoped remaining/limit values are the baseline every request # emits and must always reach the SLP. - remaining_keys = [ - k for k in additional_headers if "-remaining-" in k - ] - assert ( - remaining_keys - ), f"streaming success must populate remaining values, got {additional_headers!r}" + remaining_keys = [k for k in additional_headers if "-remaining-" in k] + assert remaining_keys, f"streaming success must populate remaining values, got {additional_headers!r}" limit_keys = [k for k in additional_headers if "-limit-" in k] assert limit_keys, "streaming success must also populate limit values" - assert ( - additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99 - ), ( - "api_key remaining requests should reflect the just-consumed slot;" - f" got {additional_headers!r}" + assert additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99, ( + f"api_key remaining requests should reflect the just-consumed slot; got {additional_headers!r}" ) @@ -5849,9 +5490,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def _noop_increment(increment_list, **_): return True @@ -5901,9 +5540,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} - assert ( - additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99 - ), f"got {additional_headers!r}" + assert additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99, f"got {additional_headers!r}" assert additional_headers.get("x-ratelimit-model_per_key-limit-requests") == 100 # response._hidden_params is also updated for late readers. @@ -5920,9 +5557,7 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): """ monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-stream-no-mirror") - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) async def _noop_increment(increment_list, **_): return True @@ -5964,9 +5599,7 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} ratelimit_keys = [k for k in additional_headers if k.startswith("x-ratelimit-")] - assert ( - not ratelimit_keys - ), f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" + assert not ratelimit_keys, f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" @pytest.mark.asyncio @@ -5986,9 +5619,7 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def _noop_increment(increment_list, **_): return True @@ -6023,15 +5654,11 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): user_api_key_dict=user_api_key_dict, response=non_stream_response, ) - non_stream_headers = non_stream_response._hidden_params.get( - "additional_headers", {} - ) + non_stream_headers = non_stream_response._hidden_params.get("additional_headers", {}) # Streaming path: async_logging_hook mirrors into standard_logging_object. stream_kwargs: Dict[str, Any] = { - "standard_logging_object": { - "metadata": {"user_api_key_hash": _api_key} - }, + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, "litellm_params": {"metadata": data["metadata"]}, "model": "gpt-4o-mini", } @@ -6048,18 +5675,13 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): result=stream_response, call_type="acompletion", ) - stream_slp_headers = ( - stream_kwargs["standard_logging_object"] - .get("hidden_params", {}) - .get("additional_headers", {}) - ) + stream_slp_headers = stream_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) def _rl_only(headers: Dict[str, Any]) -> Dict[str, Any]: return {k: v for k, v in headers.items() if k.startswith("x-ratelimit-")} assert _rl_only(stream_slp_headers) == _rl_only(non_stream_headers), ( - f"streaming={_rl_only(stream_slp_headers)}" - f" non_streaming={_rl_only(non_stream_headers)}" + f"streaming={_rl_only(stream_slp_headers)} non_streaming={_rl_only(non_stream_headers)}" ) assert "x-ratelimit-model_per_key-remaining-requests" in stream_slp_headers @@ -6075,12 +5697,8 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-passthrough") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: "total" - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") captured_operations = [] @@ -6116,9 +5734,7 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey @pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"]) @pytest.mark.asyncio -async def test_aggregate_only_usage_charges_tpm_under_every_limit_type( - monkeypatch, rate_limit_type -): +async def test_aggregate_only_usage_charges_tpm_under_every_limit_type(monkeypatch, rate_limit_type): """ A pass-through target reports one total for the whole request and cannot split it into prompt/completion. Reading a split out of it yields 0, which @@ -6128,12 +5744,8 @@ async def test_aggregate_only_usage_charges_tpm_under_every_limit_type( monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-aggregate-only") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type) captured_operations = [] @@ -6168,9 +5780,7 @@ async def test_split_usage_still_respects_the_configured_limit_type(monkeypatch) monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-split-usage") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "output") captured_operations = [] @@ -6212,9 +5822,7 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): """ from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptor = RateLimitDescriptor( key="model_saturation_check", value="tpm-only-model", @@ -6229,15 +5837,9 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): assert under_limit["overall_code"] == "OK" assert [s["rate_limit_type"] for s in under_limit["statuses"]] == ["tokens"] - counter_key = handler.create_rate_limit_keys( - "model_saturation_check", "tpm-only-model", "tokens" - ) + counter_key = handler.create_rate_limit_keys("model_saturation_check", "tpm-only-model", "tokens") await handler.async_increment_tokens_with_ttl_preservation( - pipeline_operations=[ - RedisPipelineIncrementOperation( - key=counter_key, increment_value=100, ttl=60 - ) - ], + pipeline_operations=[RedisPipelineIncrementOperation(key=counter_key, increment_value=100, ttl=60)], ) at_limit = await handler.atomic_check_and_increment_by_n( @@ -6282,9 +5884,7 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): its descriptors or an exhausted RPM budget would double-enforce here.""" from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptor = RateLimitDescriptor( key="api_key", value="reserve-test-key", @@ -6296,8 +5896,7 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): estimated_tokens=10, ) assert response["overall_code"] == "OK", ( - "an exhausted requests budget (limit 0) must not block the token " - f"reservation pass, got: {response}" + f"an exhausted requests budget (limit 0) must not block the token reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] @@ -6548,9 +6147,7 @@ async def test_estimated_output_tokens_resolution_precedence( """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-{expected_output_estimate}-{tier}"), tpm_limit=1_000_000, @@ -6573,9 +6170,7 @@ async def test_request_max_tokens_outranks_configured_estimate(monkeypatch): """An explicit request-level max_tokens stays the top of the precedence order.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-explicit-max-tokens"), tpm_limit=1_000_000, @@ -6597,9 +6192,7 @@ async def test_configured_estimate_does_not_apply_to_embeddings(monkeypatch): """Embeddings generate no output, so a declared output estimate must not be reserved.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-embeddings"), tpm_limit=1_000_000, @@ -6626,9 +6219,7 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) configured = UserAPIKeyAuth( api_key=hash_token("sk-estimate-contentless-configured"), tpm_limit=1_000_000, @@ -6640,17 +6231,9 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): ) assert ( - await _reserved_tokens_for( - handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []} - ) - == 2002 - ) - assert ( - await _reserved_tokens_for( - handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []} - ) - == 1 + await _reserved_tokens_for(handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []}) == 2002 ) + assert await _reserved_tokens_for(handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []}) == 1 @pytest.mark.asyncio @@ -6667,9 +6250,7 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} raised_reserved = await _reserved_tokens_for( handler, @@ -6721,9 +6302,7 @@ async def test_one_malformed_estimate_field_does_not_discard_the_other(monkeypat """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) broken_map = await _reserved_tokens_for( handler, @@ -6772,9 +6351,7 @@ async def test_declared_estimate_over_the_tpm_budget_is_honored_and_explained(mo """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-over-budget-{declared}"), tpm_limit=5000, @@ -6813,9 +6390,7 @@ async def test_a_key_that_declared_nothing_is_never_blamed_for_a_declaration(mon """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): with pytest.raises(HTTPException) as exc_info: @@ -6843,9 +6418,7 @@ async def test_declared_estimate_inside_the_tpm_budget_is_not_explained(monkeypa """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): await handler.async_pre_call_hook( @@ -6878,9 +6451,7 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo async def admitted(metadata): local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-overrun-{metadata}"), tpm_limit=8000, @@ -6908,9 +6479,7 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo @pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"]) def test_internal_call_origin_success_ops_are_skipped(origin, rate_limit_type): """Foreground compaction charges the same scopes as ordinary caller traffic.""" - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) response = ModelResponse( id="internal-origin-tpm", object="chat.completion", @@ -6962,15 +6531,10 @@ def test_conflicting_token_limits_reserve_the_larger_declared_budget(): A request declaring max_tokens=1 alongside max_completion_tokens=10000 previously reserved one output token while the provider stayed free to emit ten thousand. """ - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) bodies = _conflicting_budget_bodies() - reserved = { - label: handler._estimate_tokens_for_request(data=body) - for label, body in bodies.items() - } + reserved = {label: handler._estimate_tokens_for_request(data=body) for label, body in bodies.items()} assert reserved["both"] == reserved["only_large"] assert reserved["both"] > reserved["only_small"] @@ -6984,9 +6548,7 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): output floor, so dropping it from the estimate under-reserves and reopens the same TPM bypass that reading both spellings was meant to close. """ - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]} reserved = handler._estimate_tokens_for_request(data={**base, "max_tokens": declared}) @@ -6999,12 +6561,8 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): """The pre-call hook must refuse a request whose larger declared budget exceeds the TPM limit.""" local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) - user_api_key_dict = UserAPIKeyAuth( - api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[] - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[]) bodies = _conflicting_budget_bodies() await handler.async_pre_call_hook( @@ -7031,7 +6589,9 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): def _enqueued_test_handler() -> _PROXY_MaxParallelRequestsHandler: - return PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60))) + return PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)) + ) def _batch_response(batch_id: str, status: str): @@ -7415,9 +6975,7 @@ async def test_redis_batch_rate_limiter_returns_values_in_caller_order(): "{end_user:u28551}:requests", ] - results: Final = await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=keys_to_fetch, now_int=now - ) + results: Final = await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=keys_to_fetch, now_int=now) assert results == [now, 1, now, 2, now, 1] @@ -7528,9 +7086,7 @@ async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired( acquired = get_or_create_request_stash().parallel_slot assert acquired is not None - await handler.async_post_call_failure_hook( - request_data=data, original_exception=exc.value, user_api_key_dict=auth - ) + await handler.async_post_call_failure_hook(request_data=data, original_exception=exc.value, user_api_key_dict=auth) assert redis.released_slots == [(list(acquired["counter_keys"]), [acquired["slot_id"]])] assert get_or_create_request_stash().parallel_slot is None @@ -7715,12 +7271,11 @@ async def test_batch_increment_refunds_pipelined_groups_declared_after_the_one_t ) @pytest.mark.parametrize("request_kind", ["count", "generation"]) @pytest.mark.asyncio -async def test_request_capacity_enforces_shared_rpm_scopes( - limits, request_data, counter_scope, request_kind -): +async def test_request_capacity_enforces_shared_rpm_scopes(limits, request_data, counter_scope, request_kind): cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) auth = UserAPIKeyAuth(api_key=hash_token("sk-count-rpm"), **limits) + async def request(): if request_kind == "generation": await handler.async_pre_call_hook( @@ -7744,11 +7299,15 @@ async def test_request_capacity_enforces_shared_rpm_scopes( async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch): import litellm.proxy.proxy_server as proxy_server - router = Router(model_list=[{ - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - }]) + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, + } + ] + ) monkeypatch.setattr(proxy_server, "llm_router", router) handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) auth = UserAPIKeyAuth( @@ -8005,7 +7564,9 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li choices=[], ) kwargs = { - "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-pool"), "user_api_key_team_id": "t"}}, + "standard_logging_object": { + "metadata": {"user_api_key_hash": hash_token("sk-pool"), "user_api_key_team_id": "t"} + }, "litellm_params": { "metadata": { "model_group": "test-model", @@ -8016,10 +7577,14 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li "model": "test-model", } - ops = handler._build_success_event_pipeline_operations(kwargs=kwargs, response_obj=response, rate_limit_type="output") + ops = handler._build_success_event_pipeline_operations( + kwargs=kwargs, response_obj=response, rate_limit_type="output" + ) charged_keys = {op["key"] for op in ops} - assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys + assert ( + handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys + ) team_pool_key = handler.create_rate_limit_keys("model_per_team", "t:test-model", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool @@ -8073,9 +7638,7 @@ def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_t def test_per_model_rate_limit_error_carries_descriptor_key() -> None: - handler: Final = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) over_limit: Final[RateLimitResponse] = { "overall_code": "OVER_LIMIT", "statuses": [ @@ -8093,9 +7656,7 @@ def test_per_model_rate_limit_error_carries_descriptor_key() -> None: with pytest.raises(ProxyRateLimitError) as exc_info: handler._handle_rate_limit_error( response=over_limit, - descriptors=[ - {"key": "model_per_key", "value": "gpt-4o-mini", "rate_limit": None} - ], + descriptors=[{"key": "model_per_key", "value": "gpt-4o-mini", "rate_limit": None}], requested_model="gpt-4o-mini", ) @@ -8106,8 +7667,12 @@ def _resolve_alias_to_target(model: str) -> str | None: return "target" if model == "alias" else None -async def _rpm_request(handler: _PROXY_MaxParallelRequestsHandler, cache: DualCache, auth: UserAPIKeyAuth, model: str) -> None: - await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data={"model": model}, call_type="acompletion") +async def _rpm_request( + handler: _PROXY_MaxParallelRequestsHandler, cache: DualCache, auth: UserAPIKeyAuth, model: str +) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data={"model": model}, call_type="acompletion" + ) @pytest.mark.asyncio @@ -8234,7 +7799,10 @@ def test_success_tpm_accounting_charges_the_alias_target_bucket( ) charged_keys: Final = {op["key"] for op in ops} - assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-alias-tpm')}:target", "tokens") in charged_keys + assert ( + handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-alias-tpm')}:target", "tokens") + in charged_keys + ) assert not any(":alias" in key for key in charged_keys) team_pool_key: Final = handler.create_rate_limit_keys("model_per_team", "t:target", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool @@ -9127,11 +8695,17 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( from litellm.types.agents import AgentResponse expensive: Final = AgentResponse( - agent_id="expensive", agent_name="Expensive", agent_card_params={}, rpm_limit=1, + agent_id="expensive", + agent_name="Expensive", + agent_card_params={}, + rpm_limit=1, litellm_params={"cost_per_query": 0.25}, ) cheap: Final = AgentResponse( - agent_id="cheap", agent_name="Cheap", agent_card_params={}, rpm_limit=100, + agent_id="cheap", + agent_name="Cheap", + agent_card_params={}, + rpm_limit=100, litellm_params={"cost_per_query": 0.01}, ) registry: Final = agent_registry.AgentRegistry() @@ -9145,7 +8719,9 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( monkeypatch.setattr(proxy_server, "prisma_client", database) auth: Final = UserAPIKeyAuth(agent_id="caller") auth.managed_agent_policy = AgentResponse( - agent_id="caller", agent_name="Caller", agent_card_params={}, + agent_id="caller", + agent_name="Caller", + agent_card_params={}, object_permission={"object_permission_id": "both-targets", "agents": ["expensive", "cheap"]}, ) body: Final = {"model": "a2a/cheap"} 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..825694dbc9f 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -419,6 +419,25 @@ 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 + assert post_call_failure_hook.await_args.kwargs["request_data"]["litellm_call_id"] == "cid_call-unmapped" + + @pytest.mark.asyncio async def test_execute_tool_calls_logs_each_tool_call_under_its_own_id_on_the_request_trace( monkeypatch, 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 3/6] 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 From c37ce7861db60db876359fa873866992c551991f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:51:33 -0700 Subject: [PATCH 4/6] fix(proxy): hold the request's parallel slot through MCP tool loops on streamed chat and messages --- .../anthropic/messages/native/mcp_handler.py | 90 +- .../hooks/parallel_request_limiter_v3.py | 39 +- .../responses/mcp/chat_completions_handler.py | 14 +- .../messages/native/test_mcp_handler.py | 76 ++ .../hooks/test_parallel_request_limiter_v3.py | 1111 ++++++++++++----- .../mcp/test_chat_completions_handler.py | 113 ++ 6 files changed, 1053 insertions(+), 390 deletions(-) diff --git a/litellm/llms/anthropic/messages/native/mcp_handler.py b/litellm/llms/anthropic/messages/native/mcp_handler.py index 8fe18d79f24..c60da9416e3 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(): + 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/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1386dd2973b..63491dff9d0 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 @@ -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) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..a43cda98815 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() 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 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..0f8bebf3aa1 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,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] 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 f391193af7f..c559d5a2cd0 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, @@ -97,10 +98,14 @@ def _isolated_request_stash(): (0.5, 50, 500), ], ) -def test_api_key_descriptor_applies_budget_throttle(throttle_pct, expected_rpm, expected_tpm): +def test_api_key_descriptor_applies_budget_throttle( + throttle_pct, expected_rpm, expected_tpm +): """The api_key rate-limit descriptor scales the key's configured TPM/RPM by the request-scoped budget_throttle_pct, leaving the configured limits intact.""" - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-throttle"), rpm_limit=100, @@ -172,11 +177,8 @@ async def test_mcp_description_does_not_change_admission_or_reserved_tokens( @pytest.mark.parametrize("itpm_limit,otpm_limit", [(64, 4096), (4096, 64), (4096, 4096)]) @pytest.mark.parametrize("arguments_rewritten", [False, True]) async def test_mcp_description_preserves_project_input_and_output_reservations( - description: str | None, - itpm_limit: int, - otpm_limit: int, - arguments_rewritten: bool, - monkeypatch: pytest.MonkeyPatch, + description: str | None, itpm_limit: int, otpm_limit: int, + arguments_rewritten: bool, monkeypatch: pytest.MonkeyPatch ) -> None: cache: Final = DualCache() handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) @@ -314,12 +316,18 @@ async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) results.append(now) results.append(new_counter) return results @@ -401,12 +409,18 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) results.append(now) results.append(new_counter) return results @@ -438,7 +452,9 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller new_window_value = await local_cache.async_get_cache(key=window_key) new_counter_value = await local_cache.async_get_cache(key=counter_key) - assert new_window_value == window_value, "Window value should not change within window" + assert ( + new_window_value == window_value + ), "Window value should not change within window" assert new_counter_value == 2, "Counter should be 2 after second request" # Wait for window to expire @@ -469,7 +485,9 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller ) @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio -async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_controller): +async def test_normal_router_call_tpm_v3( + monkeypatch, rate_limit_object, time_controller +): """ Test normal router call with parallel request limiter v3 for TPM rate limiting """ @@ -545,12 +563,18 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) results.append(now) results.append(new_counter) return results @@ -573,7 +597,9 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co return None value = get_value_for_key(rate_limit_object, user_api_key_dict, "azure-model") - counter_key = parallel_request_handler.create_rate_limit_keys(rate_limit_object, value, "tokens") + counter_key = parallel_request_handler.create_rate_limit_keys( + rate_limit_object, value, "tokens" + ) # First request should succeed. Include messages + a tight max_tokens so # the atomic reserve_tpm_tokens path populates the :tokens counter with a @@ -586,8 +612,12 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5, } - expected_reservation = parallel_request_handler._estimate_tokens_for_request(data=pre_call_data) - assert expected_reservation < 10, "Test premise: reservation must fit under tpm_limit=10" + expected_reservation = parallel_request_handler._estimate_tokens_for_request( + data=pre_call_data + ) + assert ( + expected_reservation < 10 + ), "Test premise: reservation must fit under tpm_limit=10" await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -615,11 +645,15 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co counter_value = await local_cache.async_get_cache(key=counter_key) print(f"local_cache: {local_cache.in_memory_cache.cache_dict}") - assert counter_value is not None, f"Counter value should be stored in cache for {counter_key}" + assert ( + counter_value is not None + ), f"Counter value should be stored in cache for {counter_key}" # Manually increment the token counter to simulate token usage from previous call # This simulates what would happen after a successful call - await local_cache.async_increment_cache(key=counter_key, value=15, ttl=2) # Use up most of our 10 token limit + await local_cache.async_increment_cache( + key=counter_key, value=15, ttl=2 + ) # Use up most of our 10 token limit # Make another request to test rate limiting - this should fail as we've consumed tokens with pytest.raises(HTTPException) as exc_info: @@ -667,13 +701,17 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ _api_key = hash_token(_api_key) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=100) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock the get_rate_limit_type method directly since it imports general_settings internally def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) # Create a mock response with different token counts mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) @@ -722,9 +760,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ) # Verify that the correct token count was used based on the rate limit type - assert len(captured_operations) == 1, ( - "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" - ) + assert ( + len(captured_operations) == 1 + ), "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" tpm_operation = None for op in captured_operations: @@ -741,9 +779,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ "total": mock_usage.total_tokens, # 50 } - assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( - f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" - ) + assert ( + tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] + ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" @pytest.mark.parametrize( @@ -760,7 +798,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ], ) @pytest.mark.asyncio -async def test_async_log_success_event_counts_non_chat_response_tokens(monkeypatch, response_obj): +async def test_async_log_success_event_counts_non_chat_response_tokens( + monkeypatch, response_obj +): """ Embedding and text completion responses must increment the TPM counter, not just chat completion ModelResponse objects. @@ -768,8 +808,12 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(monkeypat monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-98765") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", lambda: "total" + ) mock_kwargs = { "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, @@ -795,7 +839,9 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(monkeypat end_time=datetime.now(), ) - tpm_operation = next((op for op in captured_operations if op["key"].endswith(":tokens")), None) + tpm_operation = next( + (op for op in captured_operations if op["key"].endswith(":tokens")), None + ) assert tpm_operation is not None, "Should have a TPM increment operation" assert tpm_operation["increment_value"] == 50 @@ -811,7 +857,9 @@ async def test_async_log_failure_event_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b"]) @@ -861,18 +909,29 @@ async def test_failure_event_without_acquired_slot_does_not_release_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b", "slot-c"]) + await _seed_max_parallel_requests_slots( + local_cache, counter_key, ["slot-a", "slot-b", "slot-c"] + ) await handler.async_log_failure_event( - kwargs={"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}}, + kwargs={ + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} + }, response_obj=None, start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 3 + assert ( + handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) + == 3 + ) @pytest.mark.asyncio @@ -923,7 +982,9 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3(): rejected requests that should have been admitted after a release. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) _api_key = hash_token("sk-98765") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) @@ -976,7 +1037,9 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): and an over-limit script result maps to a 429 without occupying a slot. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) _api_key = hash_token("sk-98765") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -1001,7 +1064,9 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): stashed_slot_id = stashed_acquisition["slot_id"] assert isinstance(stashed_slot_id, str) and stashed_slot_id assert stashed_acquisition["counter_keys"] == [counter_key] - assert captured_calls == [([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id])] + assert captured_calls == [ + ([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id]) + ] assert ( await handler.internal_usage_cache.async_get_cache( key=counter_key, litellm_parent_otel_span=None, local_only=True @@ -1048,7 +1113,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -1077,7 +1144,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert not should_rate_limit_called, "should_rate_limit should not be called when no rate limits are configured" + assert ( + not should_rate_limit_called + ), "should_rate_limit should not be called when no rate limits are configured" # Test 2: API key rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1093,7 +1162,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert should_rate_limit_called, "should_rate_limit should be called when API key rate limits are configured" + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when API key rate limits are configured" # Test 3: User rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1110,7 +1181,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert should_rate_limit_called, "should_rate_limit should be called when user rate limits are configured" + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when user rate limits are configured" # Test 4: Team rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1127,7 +1200,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert should_rate_limit_called, "should_rate_limit should be called when team rate limits are configured" + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when team rate limits are configured" # Test 5: End user rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1144,7 +1219,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert should_rate_limit_called, "should_rate_limit should be called when end user rate limits are configured" + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when end user rate limits are configured" # Test 6: Max parallel requests configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1160,7 +1237,9 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert should_rate_limit_called, "should_rate_limit should be called when max parallel requests are configured" + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when max parallel requests are configured" @pytest.mark.asyncio @@ -1176,7 +1255,9 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -1192,7 +1273,9 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): should_rate_limit_called = False user_api_key_dict_with_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={"model_tpm_limit": {"gpt-4": 1000}}, # Rate limit for gpt-4, not gpt-3.5-turbo + metadata={ + "model_tpm_limit": {"gpt-4": 1000} + }, # Rate limit for gpt-4, not gpt-3.5-turbo ) await parallel_request_handler.async_pre_call_hook( @@ -1202,15 +1285,17 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert not should_rate_limit_called, ( - "should_rate_limit should not be called when model-specific limits don't match requested model" - ) + assert ( + not should_rate_limit_called + ), "should_rate_limit should not be called when model-specific limits don't match requested model" # Test 2: Model-specific rate limits configured for requested model - SHOULD be called should_rate_limit_called = False user_api_key_dict_with_matching_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={"model_tpm_limit": {"gpt-3.5-turbo": 1000}}, # Rate limit for requested model + metadata={ + "model_tpm_limit": {"gpt-3.5-turbo": 1000} + }, # Rate limit for requested model ) await parallel_request_handler.async_pre_call_hook( @@ -1220,9 +1305,9 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert should_rate_limit_called, ( - "should_rate_limit should be called when model-specific limits match requested model" - ) + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when model-specific limits match requested model" @pytest.mark.asyncio @@ -1249,7 +1334,9 @@ async def test_tpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1307,11 +1394,15 @@ async def test_tpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( - "Api-Key value should combine api_key and model" - ) - assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" - assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" + assert ( + model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" + ), "Api-Key value should combine api_key and model" + assert ( + model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit + ), "Api-Key RPM limit should be set" + assert ( + model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit + ), "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1338,7 +1429,9 @@ async def test_rpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1396,11 +1489,15 @@ async def test_rpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( - "Api-Key value should combine api_key and model" - ) - assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" - assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" + assert ( + model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" + ), "Api-Key value should combine api_key and model" + assert ( + model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit + ), "Api-Key RPM limit should be set" + assert ( + model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit + ), "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1422,7 +1519,9 @@ async def test_team_member_rate_limits_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1454,12 +1553,18 @@ async def test_team_member_rate_limits_v3(): team_member_descriptor = descriptor break - assert team_member_descriptor is not None, "Team member descriptor should be present" - assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", ( - "Team member value should combine team_id and user_id" - ) - assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set" - assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set" + assert ( + team_member_descriptor is not None + ), "Team member descriptor should be present" + assert ( + team_member_descriptor["value"] == f"{_team_id}:{_user_id}" + ), "Team member value should combine team_id and user_id" + assert ( + team_member_descriptor["rate_limit"]["requests_per_unit"] == 10 + ), "Team member RPM limit should be set" + assert ( + team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000 + ), "Team member TPM limit should be set" @pytest.mark.asyncio @@ -1482,7 +1587,9 @@ async def test_team_member_rate_limits_v3_raises_429_when_over_limit(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = None @@ -1558,7 +1665,9 @@ async def test_dynamic_rate_limiting_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to track if limits are enforced captured_descriptors = [] @@ -1591,9 +1700,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert api_key_descriptor["rate_limit"]["requests_per_unit"] is None, ( - "RPM limit should be None when dynamic mode and no failures" - ) + assert ( + api_key_descriptor["rate_limit"]["requests_per_unit"] is None + ), "RPM limit should be None when dynamic mode and no failures" # Test 2: With failures - rate limits SHOULD be enforced (rpm_limit should be set) async def mock_check_with_failures(*args, **kwargs): @@ -1617,9 +1726,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert api_key_descriptor["rate_limit"]["requests_per_unit"] == 2, ( - "RPM limit should be enforced when dynamic mode and failures detected" - ) + assert ( + api_key_descriptor["rate_limit"]["requests_per_unit"] == 2 + ), "RPM limit should be enforced when dynamic mode and failures detected" @pytest.mark.asyncio @@ -1630,30 +1739,44 @@ async def test_async_increment_tokens_fallback_behavior(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock the token_increment_script to None to simulate unavailable script parallel_request_handler.token_increment_script = None # Mock the fallback method fallback_called = False - original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline + original_method = ( + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline + ) async def mock_fallback(*args, **kwargs): nonlocal fallback_called fallback_called = True return await original_method(*args, **kwargs) - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_fallback + ) # Test operations - pipeline_operations = [RedisPipelineIncrementOperation(key="test_fallback_key", increment_value=10.0, ttl=60)] + pipeline_operations = [ + RedisPipelineIncrementOperation( + key="test_fallback_key", increment_value=10.0, ttl=60 + ) + ] # Execute increment - await parallel_request_handler.async_increment_tokens_with_ttl_preservation(pipeline_operations=pipeline_operations) + await parallel_request_handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations + ) # Verify fallback was called - assert fallback_called, "Fallback method should be called when Lua script is not available" + assert ( + fallback_called + ), "Fallback method should be called when Lua script is not available" # Redis Cluster Compatibility Tests @@ -1664,7 +1787,9 @@ def test_group_keys_by_hash_tag_regular_redis(): For regular Redis, all keys should be grouped together under a single group. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Test keys with different hash tags test_keys = [ @@ -1684,7 +1809,9 @@ def test_group_keys_by_hash_tag_regular_redis(): # Verify all keys are in single group for regular Redis assert len(groups) == 1, f"Expected 1 group for regular Redis, got {len(groups)}" assert "all_keys" in groups, "Expected 'all_keys' group for regular Redis" - assert set(groups["all_keys"]) == set(test_keys), "All keys should be in single group" + assert set(groups["all_keys"]) == set( + test_keys + ), "All keys should be in single group" def test_group_keys_by_hash_tag_redis_cluster(): @@ -1696,7 +1823,9 @@ def test_group_keys_by_hash_tag_redis_cluster(): from unittest.mock import patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock _is_redis_cluster to return True with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -1716,13 +1845,17 @@ def test_group_keys_by_hash_tag_redis_cluster(): # All group keys should start with "slot_" for group_key in groups.keys(): - assert group_key.startswith("slot_"), f"Group key {group_key} should start with 'slot_'" + assert group_key.startswith( + "slot_" + ), f"Group key {group_key} should start with 'slot_'" # Verify all original keys are present across groups all_grouped_keys = [] for group_keys in groups.values(): all_grouped_keys.extend(group_keys) - assert set(all_grouped_keys) == set(test_keys), "All keys should be present in groups" + assert set(all_grouped_keys) == set( + test_keys + ), "All keys should be present in groups" def test_keyslot_for_redis_cluster(): @@ -1730,7 +1863,9 @@ def test_keyslot_for_redis_cluster(): Test the keyslot calculation for Redis cluster. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Test basic key slot1 = handler.keyslot_for_redis_cluster("user:1000") @@ -1758,14 +1893,18 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): from unittest.mock import AsyncMock, patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script that simulates Redis cluster slot conflict mock_script = AsyncMock() mock_script.side_effect = [ - Exception("EVALSHA - all keys must map to the same key slot"), # First group fails + Exception( + "EVALSHA - all keys must map to the same key slot" + ), # First group fails [1234, 2], # Second group succeeds ] handler.batch_rate_limiter_script = mock_script @@ -1782,7 +1921,9 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): ] # Execute the method - results = await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=test_keys, now_int=1234) + results = await handler._execute_redis_batch_rate_limiter_script( + keys_to_fetch=test_keys, now_int=1234 + ) assert results == [1234, 1, 1234, 2] @@ -1806,7 +1947,9 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): # Should have processed all keys (some might be duplicated due to fallback) unique_processed_keys = set(all_processed_keys) - assert len(unique_processed_keys) >= 2, "Should have processed at least some keys" + assert ( + len(unique_processed_keys) >= 2 + ), "Should have processed at least some keys" @pytest.mark.asyncio @@ -1832,7 +1975,9 @@ async def test_multiple_rate_limits_per_descriptor(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to return a response with multiple statuses where one hits the limit # This simulates the case where we have more statuses than descriptors due to multiple rate limit types @@ -1910,7 +2055,9 @@ async def test_missing_descriptor_fallback(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock should_rate_limit to return a status with descriptor_key that doesn't match descriptors async def mock_should_rate_limit(descriptors, **kwargs): @@ -1954,7 +2101,9 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): This verifies the change from 'output' to 'total' as the default value. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock general_settings to return empty dict (no token_rate_limit_type set) import litellm.proxy.proxy_server as proxy_server @@ -1964,7 +2113,9 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): try: result = parallel_request_handler.get_rate_limit_type() - assert result == "total", f"Default rate limit type should be 'total', got '{result}'" + assert ( + result == "total" + ), f"Default rate limit type should be 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -1975,17 +2126,23 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): Test that get_rate_limit_type falls back to 'total' when an invalid value is specified. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock general_settings to return an invalid token_rate_limit_type import litellm.proxy.proxy_server as proxy_server original_settings = getattr(proxy_server, "general_settings", {}) - monkeypatch.setattr(proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"}) + monkeypatch.setattr( + proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"} + ) try: result = parallel_request_handler.get_rate_limit_type() - assert result == "total", f"Invalid rate limit type should fall back to 'total', got '{result}'" + assert ( + result == "total" + ), f"Invalid rate limit type should fall back to 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -1999,7 +2156,9 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): ], ) @pytest.mark.asyncio -async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_limit_type, expected_field): +async def test_async_log_success_event_with_dict_usage( + monkeypatch, token_rate_limit_type, expected_field +): """ Test that async_log_success_event correctly handles usage as a dict (Responses API format). @@ -2011,13 +2170,17 @@ async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_l _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) # Create a mock response object with usage as a dict (Responses API format) from litellm.types.utils import BaseLiteLLMOpenAIResponseObject @@ -2080,9 +2243,9 @@ async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_l "total": 60, # total_tokens } - assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( - f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" - ) + assert ( + tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] + ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" @pytest.mark.asyncio @@ -2097,13 +2260,17 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc _api_key = "sk-98765" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return "output" - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) # Create a mock response object with usage as a dict missing some fields mock_response = MagicMock() @@ -2114,7 +2281,9 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc } from litellm.types.utils import BaseLiteLLMOpenAIResponseObject - mock_response.__class__ = type("MockResponse", (BaseLiteLLMOpenAIResponseObject,), {}) + mock_response.__class__ = type( + "MockResponse", (BaseLiteLLMOpenAIResponseObject,), {} + ) # Create mock kwargs for the success event mock_kwargs = { @@ -2170,7 +2339,9 @@ async def test_execute_token_increment_script_cluster_compatibility(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -2208,14 +2379,18 @@ async def test_execute_token_increment_script_cluster_compatibility(): "{api_key:sk-123}:max_parallel_requests", "{user:user-456}:tokens", } - assert set(all_processed_keys) == expected_keys, "All operation keys should be processed" + assert ( + set(all_processed_keys) == expected_keys + ), "All operation keys should be processed" # Verify args structure is correct for each call for call_args in call_args_list: keys = call_args[1]["keys"] args = call_args[1]["args"] # Each key should have 2 args (increment_value, ttl) - assert len(args) == len(keys) * 2, f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" + assert ( + len(args) == len(keys) * 2 + ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" @pytest.mark.asyncio @@ -2238,7 +2413,9 @@ async def test_agent_level_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2303,7 +2480,9 @@ async def test_agent_session_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2370,7 +2549,9 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2403,7 +2584,8 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): # should_rate_limit should not have been called (no agent-level limits, only session limits # but no session_id) assert captured_descriptors is None, ( - "No descriptors should be created when agent has only session limits but no session_id in request" + "No descriptors should be created when agent has only session limits " + "but no session_id in request" ) @@ -2426,7 +2608,9 @@ async def test_agent_rate_limit_from_metadata_agent_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2467,7 +2651,9 @@ async def test_agent_rate_limit_from_metadata_agent_id(): agent_descriptor = d break - assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" + assert ( + agent_descriptor is not None + ), "Agent descriptor should be created from metadata agent_id" assert agent_descriptor["value"] == _agent_id assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 @@ -2493,7 +2679,9 @@ async def test_agent_both_agent_and_session_rate_limits(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2560,12 +2748,16 @@ async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): _session_id = "sess_tpm_test" local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) def mock_get_rate_limit_type(): return "total" - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) mock_response = ModelResponse( @@ -2676,12 +2868,18 @@ async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) results.append(now) results.append(new_counter) return results @@ -2819,7 +3017,9 @@ async def test_project_model_rate_limits_enforced_v3(): """ _api_key = hash_token("sk-project-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = [] @@ -2847,9 +3047,13 @@ async def test_project_model_rate_limits_enforced_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert "model_per_project" in descriptor_keys, f"Expected model_per_project descriptor, got: {descriptor_keys}" + assert ( + "model_per_project" in descriptor_keys + ), f"Expected model_per_project descriptor, got: {descriptor_keys}" - model_per_project = next(d for d in captured_descriptors if d["key"] == "model_per_project") + model_per_project = next( + d for d in captured_descriptors if d["key"] == "model_per_project" + ) assert model_per_project["value"] == "proj-abc123:gpt-4" assert model_per_project["rate_limit"]["requests_per_unit"] == 5 assert model_per_project["rate_limit"]["tokens_per_unit"] == 1000 @@ -2860,7 +3064,9 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): """Project model limits should not trigger for a model not in project_metadata.""" _api_key = hash_token("sk-project-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = [] @@ -2887,9 +3093,9 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert "model_per_project" not in descriptor_keys, ( - f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" - ) + assert ( + "model_per_project" not in descriptor_keys + ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" @pytest.mark.asyncio @@ -2900,7 +3106,9 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): """ _api_key = hash_token("sk-project-io-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = [] @@ -2931,8 +3139,12 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): assert "model_per_project_otpm" in descriptor_keys assert "model_per_project" not in descriptor_keys - itpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_itpm") - otpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_otpm") + itpm_descriptor = next( + d for d in captured_descriptors if d["key"] == "model_per_project_itpm" + ) + otpm_descriptor = next( + d for d in captured_descriptors if d["key"] == "model_per_project_otpm" + ) assert itpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" assert itpm_descriptor["rate_limit"]["tokens_per_unit"] == 20000000 assert otpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" @@ -2944,7 +3156,9 @@ async def test_project_model_itpm_otpm_limits_not_triggered_for_other_model_v3() """Split project limits must not apply to an unrelated model.""" _api_key = hash_token("sk-project-io-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = [] @@ -2979,7 +3193,9 @@ async def test_project_model_itpm_and_tpm_limits_coexist_v3(): """Combined project TPM and split ITPM/OTPM limits are enforced together.""" _api_key = hash_token("sk-project-io-test-3") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) captured_descriptors = [] @@ -3021,7 +3237,9 @@ async def test_enforce_project_io_token_quota_for_frame_blocks_over_limit_otpm() limit and reject once a frame's estimated output tokens exceed it.""" _api_key = hash_token("sk-ws-frame-otpm") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, project_id="proj-mantle-ws", @@ -3053,7 +3271,9 @@ async def test_enforce_project_io_token_quota_for_frame_noop_without_project_lim the per-frame check (no descriptors to reserve against).""" _api_key = hash_token("sk-ws-frame-no-limits") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) await handler.enforce_project_io_token_quota_for_frame( @@ -3188,13 +3408,17 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-refund-lifecycle") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=_api_key, rate_limit_type="tokens" + ) parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3211,18 +3435,28 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): reserved = get_request_stash().reserved_tokens assert reserved > 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 kwargs = {"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}} - await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 assert get_request_stash().reservation_released is True - await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 @pytest.mark.asyncio @@ -3266,7 +3500,9 @@ async def test_pre_call_hook_ignores_caller_supplied_stash_values(): async def spy_increment_pipeline(increment_list, **kwargs): refund_calls.append(increment_list) - handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = spy_increment_pipeline + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + spy_increment_pipeline + ) await handler.async_post_call_failure_hook( request_data=data, @@ -3292,13 +3528,17 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-nested-guard") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=_api_key, rate_limit_type="tokens" + ) parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3323,23 +3563,33 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) "litellm_call_id": "nested-guardrail-call", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_success_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) - await handler.async_log_failure_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_success_event( + kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None + ) + await handler.async_log_failure_event( + kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None + ) assert stash.parallel_slot is not None assert stash.reservation_released is False - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved owner_kwargs = { "litellm_call_id": "owner-call-id", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_failure_event(kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_failure_event( + kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None + ) assert stash.parallel_slot is None assert stash.reservation_released is True - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 @@ -3353,7 +3603,9 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): do not thread ``litellm_call_id`` into their logging kwargs. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) unclaimed = get_or_create_request_stash() unclaimed.reserved_tokens = 42 @@ -3365,7 +3617,9 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): ) assert unclaimed.reservation_released is True - claimed = RequestRateLimiterStash(owner_litellm_call_id="owner-1", reserved_tokens=42) + claimed = RequestRateLimiterStash( + owner_litellm_call_id="owner-1", reserved_tokens=42 + ) _request_stash.set(claimed) await handler.async_log_failure_event( kwargs={"standard_logging_object": {}}, @@ -3548,7 +3802,9 @@ async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usag def _make_mcp_handler() -> tuple[PROXY_MaxParallelRequestsHandler_v3, DualCache]: local_cache: Final = DualCache() - handler: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(local_cache)) + handler: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(local_cache) + ) return handler, local_cache @@ -3572,7 +3828,11 @@ async def test_mcp_per_key_rate_limit_uses_trusted_server_alias_v3() -> None: with pytest.raises(ProxyRateLimitError, match="mcp_per_key"): await handler.enforce_mcp_server_rate_limits(user_api_key_dict, server) - assert all(value == 0 for key, value in local_cache.in_memory_cache.cache_dict.items() if key.endswith(":tokens")) + assert all( + value == 0 + for key, value in local_cache.in_memory_cache.cache_dict.items() + if key.endswith(":tokens") + ) @pytest.mark.asyncio @@ -3887,7 +4147,9 @@ def test_get_key_mcp_rpm_limit_precedence(): _TEST_SLOT_ID = "slot-disconnect-test" -async def _seed_max_parallel_requests_slots(dual_cache: DualCache, counter_key: str, slot_ids: List[str]) -> None: +async def _seed_max_parallel_requests_slots( + dual_cache: DualCache, counter_key: str, slot_ids: List[str] +) -> None: await dual_cache.async_set_cache( key=counter_key, value={slot_id: time.time() for slot_id in slot_ids}, @@ -3899,7 +4161,9 @@ async def _seed_max_parallel_requests_slots(dual_cache: DualCache, counter_key: async def test_completed_responses_post_call_releases_parallel_slot() -> None: api_key = hash_token("sk-responses-post-call") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) data = { "model": "gpt-4o-mini", @@ -3914,7 +4178,9 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: data=data, call_type="aresponses", ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 await handler.async_post_call_success_hook( data=data, @@ -3928,7 +4194,9 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: status="completed", ), ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 await handler.async_log_success_event( kwargs={"litellm_call_id": data["litellm_call_id"]}, @@ -3936,20 +4204,26 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 @pytest.mark.asyncio async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: api_key: Final = hash_token("sk-mcp-auto-execute") local_cache: Final = DualCache() - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + 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" owner_call_id: Final = "mcp-auto-execute-owner" async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) def mcp_tool_call() -> dict[str, object]: return { @@ -3966,7 +4240,9 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() call_type=call_type, ) - async def pre_call_from_another_request(data: dict[str, object], call_type: str) -> None: + async def pre_call_from_another_request( + data: dict[str, object], call_type: str + ) -> None: await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) await pre_call( @@ -3978,7 +4254,9 @@ async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() assert await in_flight() == 1 with pytest.raises(HTTPException): - await pre_call_from_another_request(mcp_tool_call(), CallTypes.call_mcp_tool.value) + await pre_call_from_another_request( + mcp_tool_call(), CallTypes.call_mcp_tool.value + ) await handler.async_log_success_event( kwargs={"litellm_call_id": owner_call_id}, @@ -4056,15 +4334,16 @@ async def test_failed_mcp_tool_call_keeps_its_requests_parallel_slot_and_reserva @pytest.mark.asyncio -async def test_round_logged_before_a_tool_call_fails_unserved_keeps_the_requests_parallel_slot() -> None: +@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("sk-messages-unserved-tool") + 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 = "messages-unserved-tool-owner" + 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( @@ -4074,14 +4353,6 @@ async def test_round_logged_before_a_tool_call_fails_unserved_keeps_the_requests 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": []})) @@ -4092,16 +4363,22 @@ async def test_round_logged_before_a_tool_call_fails_unserved_keeps_the_requests 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)) + with request_slot_held_by_tool_loop(): + 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 + assert await next_request_status() == 429 - await log_round(now + timedelta(milliseconds=1)) + await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert await next_request_status() == 200 @@ -4154,7 +4431,9 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request owner_call_id: Final = f"mcp-late-round-log-{stream}" async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) async def pre_call(data: dict[str, object], call_type: str) -> None: await handler.async_pre_call_hook( @@ -4215,7 +4494,9 @@ async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kep messages: Final = [{"role": "user", "content": "add 2 and 3"}] async def in_flight() -> int: - return handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) async def pre_call(data: dict[str, object], call_type: str) -> None: await handler.async_pre_call_hook( @@ -4285,7 +4566,9 @@ async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kep async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: api_key: Final = hash_token("sk-mcp-reuse-headers") local_cache: Final = DualCache() - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) data: Final[dict[str, object]] = { "model": "gpt-4o-mini", @@ -4315,7 +4598,9 @@ async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_hea ) response._hidden_params = {} - await handler.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) + await handler.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) headers: Final = response._hidden_params["additional_headers"] assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 @@ -4331,14 +4616,18 @@ async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redi api_key = hash_token("sk-concurrent-release") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) call_id = "concurrent-release-owner" parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" release_started = asyncio.Event() allow_redis_failure = asyncio.Event() - async def failing_release(keys: Sequence[str], args: Sequence[object]) -> list[int]: + async def failing_release( + keys: Sequence[str], args: Sequence[object] + ) -> list[int]: release_started.set() await allow_redis_failure.wait() raise ConnectionError("redis unavailable") @@ -4403,7 +4692,9 @@ async def _build_seeded_limiter(): """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") cache = DualCache() - limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + limiter = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache) + ) counter_key = f"{{api_key:{api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(cache, counter_key, [_TEST_SLOT_ID]) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) @@ -4438,12 +4729,16 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 get_or_create_request_stash().parallel_slot = ParallelSlotAcquisition( slot_id=_TEST_SLOT_ID, @@ -4452,7 +4747,9 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert get_request_stash().parallel_slot is None - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 @pytest.mark.asyncio @@ -4465,7 +4762,9 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) @@ -4476,7 +4775,9 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): await handler.async_release_max_parallel_requests_on_disconnect( UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=None) ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 @pytest.mark.asyncio @@ -4492,7 +4793,9 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4503,14 +4806,18 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 await handler.async_post_call_failure_hook( request_data=admitted_data, original_exception=Exception("guardrail rejected the request"), user_api_key_dict=user_api_key_dict, ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 await handler.async_log_failure_event( kwargs={ @@ -4520,7 +4827,9 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4539,7 +4848,9 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4551,17 +4862,23 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 await handler.async_log_success_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, }, - response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), + response_obj=ModelResponse( + usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) + ), start_time=datetime.now(), end_time=datetime.now(), ) - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4578,7 +4895,9 @@ _SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1) @pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"]) @pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) @pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) -async def test_sse_keepalive_releases_parallel_capacity(after_ping: bool, fails: bool, preadmitted: bool) -> None: +async def test_sse_keepalive_releases_parallel_capacity( + after_ping: bool, fails: bool, preadmitted: bool +) -> None: from fastapi.responses import StreamingResponse from litellm.proxy.common_request_processing import open_sse_before_first_byte @@ -4597,7 +4916,9 @@ async def test_sse_keepalive_releases_parallel_capacity(after_ping: bool, fails: upstream_answers.set() async def admit() -> None: - await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion") + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" + ) async def failure(exc: Exception) -> None: await handler.async_post_call_failure_hook( @@ -4711,7 +5032,9 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" descriptors = [ { @@ -4730,7 +5053,9 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): handler.parallel_count_script = fake_count response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) - assert captured_calls == [([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS])] + assert captured_calls == [ + ([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS]) + ] assert response["overall_code"] == "OK" assert response["statuses"] == [ { @@ -4747,7 +5072,9 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): raise ConnectionError("redis unavailable") handler.parallel_count_script = failing_count - await _seed_max_parallel_requests_slots(local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"]) + await _seed_max_parallel_requests_slots( + local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"] + ) response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) assert response["overall_code"] == "OVER_LIMIT" assert response["statuses"][0]["rate_limit_type"] == "max_parallel_requests" @@ -4762,7 +5089,9 @@ async def test_redis_release_script_updates_local_mirror_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" captured_calls = [] @@ -4800,8 +5129,12 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5, tpm_limit=100) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, max_parallel_requests=5, tpm_limit=100 + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None, **_kwargs): @@ -4828,7 +5161,9 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): call_type="", ) assert exc_info.value.status_code == 429 - assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 @pytest.mark.asyncio @@ -4843,7 +5178,9 @@ async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4893,7 +5230,9 @@ async def test_release_max_parallel_requests_on_disconnect_noop_v3(): """ _api_key = hash_token("sk-98765") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_release_max_parallel_requests_on_disconnect( @@ -4924,7 +5263,9 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() - assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 1 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 1 proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter @@ -4955,7 +5296,9 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 @pytest.mark.parametrize("disconnect", ["cancel", "aclose"]) @@ -5002,10 +5345,14 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect else: await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( + saved_hook + ) else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -5023,7 +5370,9 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): from litellm.integrations.custom_logger import CustomLogger class _PassthroughIteratorOverride(CustomLogger): - async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict, response, request_data + ): async for chunk in response: yield chunk @@ -5051,10 +5400,14 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): await gen.__anext__() await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( + saved_hook + ) else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -5062,14 +5415,18 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): def test_tpm_reservation_enabled_by_default(monkeypatch): """Upfront TPM reservation is on unless explicitly disabled via env.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) assert handler.tpm_reservation_enabled is True @pytest.mark.parametrize("value", ["false", "False", "FALSE"]) def test_tpm_reservation_disabled_via_env(monkeypatch, value): monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", value) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) assert handler.tpm_reservation_enabled is False @@ -5081,7 +5438,9 @@ async def test_pre_call_hook_reserves_tpm_when_enabled(monkeypatch): only the reservation path owns it. """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -5121,7 +5480,9 @@ async def test_pre_call_hook_skips_reservation_when_disabled(monkeypatch): pre-v1.82 post-call accounting behavior. """ monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -5169,7 +5530,9 @@ async def test_per_tag_rate_limit_independent_counters_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) async def call(tag: str) -> None: await handler.async_pre_call_hook( @@ -5203,7 +5566,9 @@ async def test_per_tag_descriptor_creation_v3(): api_key=_api_key, metadata={"tag_rpm_limit": {"cell-1": 5}}, ) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -5227,7 +5592,9 @@ async def test_per_tag_descriptor_absent_without_config_v3(): api_key=hash_token("sk-no-tag"), rpm_limit=10, ) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -5258,7 +5625,9 @@ async def test_per_tag_untagged_request_governed_by_key_limit_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) async def call(metadata: dict) -> None: await handler.async_pre_call_hook( @@ -5455,7 +5824,9 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) tpm_limit=10000, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) # Real pre-call: populates data and stashes the response into metadata # so the success callback can find it via litellm_params.metadata. @@ -5512,16 +5883,27 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) call_type="acompletion", ) - additional_headers = mock_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) + additional_headers = ( + mock_kwargs["standard_logging_object"] + .get("hidden_params", {}) + .get("additional_headers", {}) + ) # api_key-scoped remaining/limit values are the baseline every request # emits and must always reach the SLP. - remaining_keys = [k for k in additional_headers if "-remaining-" in k] - assert remaining_keys, f"streaming success must populate remaining values, got {additional_headers!r}" + remaining_keys = [ + k for k in additional_headers if "-remaining-" in k + ] + assert ( + remaining_keys + ), f"streaming success must populate remaining values, got {additional_headers!r}" limit_keys = [k for k in additional_headers if "-limit-" in k] assert limit_keys, "streaming success must also populate limit values" - assert additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99, ( - f"api_key remaining requests should reflect the just-consumed slot; got {additional_headers!r}" + assert ( + additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99 + ), ( + "api_key remaining requests should reflect the just-consumed slot;" + f" got {additional_headers!r}" ) @@ -5541,7 +5923,9 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) async def _noop_increment(increment_list, **_): return True @@ -5591,7 +5975,9 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} - assert additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99, f"got {additional_headers!r}" + assert ( + additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99 + ), f"got {additional_headers!r}" assert additional_headers.get("x-ratelimit-model_per_key-limit-requests") == 100 # response._hidden_params is also updated for late readers. @@ -5608,7 +5994,9 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): """ monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-stream-no-mirror") - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) async def _noop_increment(increment_list, **_): return True @@ -5650,7 +6038,9 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} ratelimit_keys = [k for k in additional_headers if k.startswith("x-ratelimit-")] - assert not ratelimit_keys, f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" + assert ( + not ratelimit_keys + ), f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" @pytest.mark.asyncio @@ -5670,7 +6060,9 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) async def _noop_increment(increment_list, **_): return True @@ -5705,11 +6097,15 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): user_api_key_dict=user_api_key_dict, response=non_stream_response, ) - non_stream_headers = non_stream_response._hidden_params.get("additional_headers", {}) + non_stream_headers = non_stream_response._hidden_params.get( + "additional_headers", {} + ) # Streaming path: async_logging_hook mirrors into standard_logging_object. stream_kwargs: Dict[str, Any] = { - "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + "standard_logging_object": { + "metadata": {"user_api_key_hash": _api_key} + }, "litellm_params": {"metadata": data["metadata"]}, "model": "gpt-4o-mini", } @@ -5726,13 +6122,18 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): result=stream_response, call_type="acompletion", ) - stream_slp_headers = stream_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) + stream_slp_headers = ( + stream_kwargs["standard_logging_object"] + .get("hidden_params", {}) + .get("additional_headers", {}) + ) def _rl_only(headers: Dict[str, Any]) -> Dict[str, Any]: return {k: v for k, v in headers.items() if k.startswith("x-ratelimit-")} assert _rl_only(stream_slp_headers) == _rl_only(non_stream_headers), ( - f"streaming={_rl_only(stream_slp_headers)} non_streaming={_rl_only(non_stream_headers)}" + f"streaming={_rl_only(stream_slp_headers)}" + f" non_streaming={_rl_only(non_stream_headers)}" ) assert "x-ratelimit-model_per_key-remaining-requests" in stream_slp_headers @@ -5748,8 +6149,12 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-passthrough") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", lambda: "total" + ) captured_operations = [] @@ -5785,7 +6190,9 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey @pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"]) @pytest.mark.asyncio -async def test_aggregate_only_usage_charges_tpm_under_every_limit_type(monkeypatch, rate_limit_type): +async def test_aggregate_only_usage_charges_tpm_under_every_limit_type( + monkeypatch, rate_limit_type +): """ A pass-through target reports one total for the whole request and cannot split it into prompt/completion. Reading a split out of it yields 0, which @@ -5795,8 +6202,12 @@ async def test_aggregate_only_usage_charges_tpm_under_every_limit_type(monkeypat monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-aggregate-only") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type + ) captured_operations = [] @@ -5831,7 +6242,9 @@ async def test_split_usage_still_respects_the_configured_limit_type(monkeypatch) monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-split-usage") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "output") captured_operations = [] @@ -5873,7 +6286,9 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): """ from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) descriptor = RateLimitDescriptor( key="model_saturation_check", value="tpm-only-model", @@ -5888,9 +6303,15 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): assert under_limit["overall_code"] == "OK" assert [s["rate_limit_type"] for s in under_limit["statuses"]] == ["tokens"] - counter_key = handler.create_rate_limit_keys("model_saturation_check", "tpm-only-model", "tokens") + counter_key = handler.create_rate_limit_keys( + "model_saturation_check", "tpm-only-model", "tokens" + ) await handler.async_increment_tokens_with_ttl_preservation( - pipeline_operations=[RedisPipelineIncrementOperation(key=counter_key, increment_value=100, ttl=60)], + pipeline_operations=[ + RedisPipelineIncrementOperation( + key=counter_key, increment_value=100, ttl=60 + ) + ], ) at_limit = await handler.atomic_check_and_increment_by_n( @@ -5935,7 +6356,9 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): its descriptors or an exhausted RPM budget would double-enforce here.""" from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) descriptor = RateLimitDescriptor( key="api_key", value="reserve-test-key", @@ -5947,7 +6370,8 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): estimated_tokens=10, ) assert response["overall_code"] == "OK", ( - f"an exhausted requests budget (limit 0) must not block the token reservation pass, got: {response}" + "an exhausted requests budget (limit 0) must not block the token " + f"reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] @@ -6198,7 +6622,9 @@ async def test_estimated_output_tokens_resolution_precedence( """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-{expected_output_estimate}-{tier}"), tpm_limit=1_000_000, @@ -6221,7 +6647,9 @@ async def test_request_max_tokens_outranks_configured_estimate(monkeypatch): """An explicit request-level max_tokens stays the top of the precedence order.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-explicit-max-tokens"), tpm_limit=1_000_000, @@ -6243,7 +6671,9 @@ async def test_configured_estimate_does_not_apply_to_embeddings(monkeypatch): """Embeddings generate no output, so a declared output estimate must not be reserved.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-embeddings"), tpm_limit=1_000_000, @@ -6270,7 +6700,9 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) configured = UserAPIKeyAuth( api_key=hash_token("sk-estimate-contentless-configured"), tpm_limit=1_000_000, @@ -6282,9 +6714,17 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): ) assert ( - await _reserved_tokens_for(handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []}) == 2002 + await _reserved_tokens_for( + handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []} + ) + == 2002 + ) + assert ( + await _reserved_tokens_for( + handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []} + ) + == 1 ) - assert await _reserved_tokens_for(handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []}) == 1 @pytest.mark.asyncio @@ -6301,7 +6741,9 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} raised_reserved = await _reserved_tokens_for( handler, @@ -6353,7 +6795,9 @@ async def test_one_malformed_estimate_field_does_not_discard_the_other(monkeypat """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) broken_map = await _reserved_tokens_for( handler, @@ -6402,7 +6846,9 @@ async def test_declared_estimate_over_the_tpm_budget_is_honored_and_explained(mo """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-over-budget-{declared}"), tpm_limit=5000, @@ -6441,7 +6887,9 @@ async def test_a_key_that_declared_nothing_is_never_blamed_for_a_declaration(mon """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): with pytest.raises(HTTPException) as exc_info: @@ -6469,7 +6917,9 @@ async def test_declared_estimate_inside_the_tpm_budget_is_not_explained(monkeypa """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): await handler.async_pre_call_hook( @@ -6502,7 +6952,9 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo async def admitted(metadata): local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-overrun-{metadata}"), tpm_limit=8000, @@ -6530,7 +6982,9 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo @pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"]) def test_internal_call_origin_success_ops_are_skipped(origin, rate_limit_type): """Foreground compaction charges the same scopes as ordinary caller traffic.""" - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) response = ModelResponse( id="internal-origin-tpm", object="chat.completion", @@ -6582,10 +7036,15 @@ def test_conflicting_token_limits_reserve_the_larger_declared_budget(): A request declaring max_tokens=1 alongside max_completion_tokens=10000 previously reserved one output token while the provider stayed free to emit ten thousand. """ - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) bodies = _conflicting_budget_bodies() - reserved = {label: handler._estimate_tokens_for_request(data=body) for label, body in bodies.items()} + reserved = { + label: handler._estimate_tokens_for_request(data=body) + for label, body in bodies.items() + } assert reserved["both"] == reserved["only_large"] assert reserved["both"] > reserved["only_small"] @@ -6599,7 +7058,9 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): output floor, so dropping it from the estimate under-reserves and reopens the same TPM bypass that reading both spellings was meant to close. """ - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]} reserved = handler._estimate_tokens_for_request(data={**base, "max_tokens": declared}) @@ -6612,8 +7073,12 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): """The pre-call hook must refuse a request whose larger declared budget exceeds the TPM limit.""" local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[]) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[] + ) bodies = _conflicting_budget_bodies() await handler.async_pre_call_hook( @@ -6640,9 +7105,7 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): def _enqueued_test_handler() -> _PROXY_MaxParallelRequestsHandler: - return PROXY_MaxParallelRequestsHandler_v3( - internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)) - ) + return PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60))) def _batch_response(batch_id: str, status: str): @@ -7026,7 +7489,9 @@ async def test_redis_batch_rate_limiter_returns_values_in_caller_order(): "{end_user:u28551}:requests", ] - results: Final = await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=keys_to_fetch, now_int=now) + results: Final = await handler._execute_redis_batch_rate_limiter_script( + keys_to_fetch=keys_to_fetch, now_int=now + ) assert results == [now, 1, now, 2, now, 1] @@ -7137,7 +7602,9 @@ async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired( acquired = get_or_create_request_stash().parallel_slot assert acquired is not None - await handler.async_post_call_failure_hook(request_data=data, original_exception=exc.value, user_api_key_dict=auth) + await handler.async_post_call_failure_hook( + request_data=data, original_exception=exc.value, user_api_key_dict=auth + ) assert redis.released_slots == [(list(acquired["counter_keys"]), [acquired["slot_id"]])] assert get_or_create_request_stash().parallel_slot is None @@ -7322,11 +7789,12 @@ async def test_batch_increment_refunds_pipelined_groups_declared_after_the_one_t ) @pytest.mark.parametrize("request_kind", ["count", "generation"]) @pytest.mark.asyncio -async def test_request_capacity_enforces_shared_rpm_scopes(limits, request_data, counter_scope, request_kind): +async def test_request_capacity_enforces_shared_rpm_scopes( + limits, request_data, counter_scope, request_kind +): cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) auth = UserAPIKeyAuth(api_key=hash_token("sk-count-rpm"), **limits) - async def request(): if request_kind == "generation": await handler.async_pre_call_hook( @@ -7350,15 +7818,11 @@ async def test_request_capacity_enforces_shared_rpm_scopes(limits, request_data, async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch): import litellm.proxy.proxy_server as proxy_server - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - } - ] - ) + router = Router(model_list=[{ + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, + }]) monkeypatch.setattr(proxy_server, "llm_router", router) handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) auth = UserAPIKeyAuth( @@ -7615,9 +8079,7 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li choices=[], ) kwargs = { - "standard_logging_object": { - "metadata": {"user_api_key_hash": hash_token("sk-pool"), "user_api_key_team_id": "t"} - }, + "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-pool"), "user_api_key_team_id": "t"}}, "litellm_params": { "metadata": { "model_group": "test-model", @@ -7628,14 +8090,10 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li "model": "test-model", } - ops = handler._build_success_event_pipeline_operations( - kwargs=kwargs, response_obj=response, rate_limit_type="output" - ) + ops = handler._build_success_event_pipeline_operations(kwargs=kwargs, response_obj=response, rate_limit_type="output") charged_keys = {op["key"] for op in ops} - assert ( - handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys - ) + assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys team_pool_key = handler.create_rate_limit_keys("model_per_team", "t:test-model", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool @@ -7689,7 +8147,9 @@ def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_t def test_per_model_rate_limit_error_carries_descriptor_key() -> None: - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) over_limit: Final[RateLimitResponse] = { "overall_code": "OVER_LIMIT", "statuses": [ @@ -7707,7 +8167,9 @@ def test_per_model_rate_limit_error_carries_descriptor_key() -> None: with pytest.raises(ProxyRateLimitError) as exc_info: handler._handle_rate_limit_error( response=over_limit, - descriptors=[{"key": "model_per_key", "value": "gpt-4o-mini", "rate_limit": None}], + descriptors=[ + {"key": "model_per_key", "value": "gpt-4o-mini", "rate_limit": None} + ], requested_model="gpt-4o-mini", ) @@ -7718,12 +8180,8 @@ def _resolve_alias_to_target(model: str) -> str | None: return "target" if model == "alias" else None -async def _rpm_request( - handler: _PROXY_MaxParallelRequestsHandler, cache: DualCache, auth: UserAPIKeyAuth, model: str -) -> None: - await handler.async_pre_call_hook( - user_api_key_dict=auth, cache=cache, data={"model": model}, call_type="acompletion" - ) +async def _rpm_request(handler: _PROXY_MaxParallelRequestsHandler, cache: DualCache, auth: UserAPIKeyAuth, model: str) -> None: + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data={"model": model}, call_type="acompletion") @pytest.mark.asyncio @@ -7850,10 +8308,7 @@ def test_success_tpm_accounting_charges_the_alias_target_bucket( ) charged_keys: Final = {op["key"] for op in ops} - assert ( - handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-alias-tpm')}:target", "tokens") - in charged_keys - ) + assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-alias-tpm')}:target", "tokens") in charged_keys assert not any(":alias" in key for key in charged_keys) team_pool_key: Final = handler.create_rate_limit_keys("model_per_team", "t:target", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool @@ -8746,17 +9201,11 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( from litellm.types.agents import AgentResponse expensive: Final = AgentResponse( - agent_id="expensive", - agent_name="Expensive", - agent_card_params={}, - rpm_limit=1, + agent_id="expensive", agent_name="Expensive", agent_card_params={}, rpm_limit=1, litellm_params={"cost_per_query": 0.25}, ) cheap: Final = AgentResponse( - agent_id="cheap", - agent_name="Cheap", - agent_card_params={}, - rpm_limit=100, + agent_id="cheap", agent_name="Cheap", agent_card_params={}, rpm_limit=100, litellm_params={"cost_per_query": 0.01}, ) registry: Final = agent_registry.AgentRegistry() @@ -8770,9 +9219,7 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( monkeypatch.setattr(proxy_server, "prisma_client", database) auth: Final = UserAPIKeyAuth(agent_id="caller") auth.managed_agent_policy = AgentResponse( - agent_id="caller", - agent_name="Caller", - agent_card_params={}, + agent_id="caller", agent_name="Caller", agent_card_params={}, object_permission={"object_permission_id": "both-targets", "agents": ["expensive", "cheap"]}, ) body: Final = {"model": "a2a/cheap"} diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..8862c1a1483 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/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." From 9c74725a61fc237ac3726ae532a5c0dba836131b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:08:13 -0700 Subject: [PATCH 5/6] fix(mcp): free the held parallel slot when a streamed MCP client hangs up --- .../responses/mcp/chat_completions_handler.py | 16 ++++++++--- .../messages/native/test_mcp_handler.py | 7 ++++- .../mcp/test_chat_completions_handler.py | 27 ++++++++++++++++--- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a43cda98815..c926c57b548 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/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: 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 0f8bebf3aa1..f3c6ec2f359 100644 --- a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py +++ b/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" diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 8862c1a1483..5cb74535dce 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 @@ -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 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 6/6] 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