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] 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."