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,