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