diff --git a/tests/e2e/coverage_registry/rate_limiting.yaml b/tests/e2e/coverage_registry/rate_limiting.yaml new file mode 100644 index 00000000000..2ed1a74953c --- /dev/null +++ b/tests/e2e/coverage_registry/rate_limiting.yaml @@ -0,0 +1,3 @@ +# Rate limiting behaviors. Grounded in litellm/proxy/hooks/parallel_request_limiter.py. +- {id: rate_limiting.key.tpm.under_limit_allows, module: rate_limiting, tier: P0, scope: key, limit: tpm, behavior: allows_under_limit, assertions: [persists_limit, allows_request], source: "parallel_request_limiter.py:145-260", rationale: "A virtual key below its TPM quota can call the data plane"} +- {id: rate_limiting.key.tpm.over_limit_blocks, module: rate_limiting, tier: P0, scope: key, limit: tpm, behavior: blocks_over_limit, assertions: [persists_limit, records_usage, blocks_next_request], source: "parallel_request_limiter.py:145-260", rationale: "A virtual key above its TPM quota is rejected with 429"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 54902051ca0..b7252f51298 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -112,6 +112,13 @@ class GuardrailCell(_Base): exercised_on: tuple[str, ...] +class RateLimitingCell(_Base): + module: Literal["rate_limiting"] + scope: Literal["key", "team", "user", "customer", "model_per_key"] + limit: Literal["tpm", "rpm", "parallel_requests"] + behavior: Literal["allows_under_limit", "blocks_over_limit"] + + class OtherCell(_Base): module: Literal["other"] area: str @@ -124,6 +131,7 @@ Cell = Annotated[ | ReliabilityCell | LoggingCell | GuardrailCell + | RateLimitingCell | OtherCell, Field(discriminator="module"), ] @@ -144,6 +152,7 @@ PREFIX_ROLLUP: dict[str, str] = { "reliability": "reliability_performance", "logging": "logging_guardrails", "guardrail": "logging_guardrails", + "rate_limiting": "rate_limiting", "other": "other", } @@ -152,6 +161,7 @@ MODULE_ORDER: tuple[str, ...] = ( "non_core_llms", "mcp", "management_ui", + "rate_limiting", "reliability_performance", "logging_guardrails", "other", diff --git a/tests/e2e/rate_limiting/conftest.py b/tests/e2e/rate_limiting/conftest.py new file mode 100644 index 00000000000..1bbd2374ac0 --- /dev/null +++ b/tests/e2e/rate_limiting/conftest.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import pytest + +from rate_limiting_client import RateLimitingClient, build_client + + +@pytest.fixture +def client() -> RateLimitingClient: + return build_client() diff --git a/tests/e2e/rate_limiting/rate_limiting_client.py b/tests/e2e/rate_limiting/rate_limiting_client.py new file mode 100644 index 00000000000..47a505b185f --- /dev/null +++ b/tests/e2e/rate_limiting/rate_limiting_client.py @@ -0,0 +1,31 @@ +"""Client for rate-limiting e2e checks.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway +from e2e_http import StreamingResponse +from models import ChatBody, ChatMessage + + +@dataclass(frozen=True, slots=True) +class RateLimitingClient: + gateway: Gateway + + def chat_status( + self, key: str, model: str, content: str, max_tokens: int = 8 + ) -> StreamingResponse: + return self.gateway.transport.send( + "/chat/completions", + headers=self.gateway.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + +def build_client() -> RateLimitingClient: + return RateLimitingClient(gateway=build_gateway()) diff --git a/tests/e2e/rate_limiting/test_rate_limiting_e2e.py b/tests/e2e/rate_limiting/test_rate_limiting_e2e.py new file mode 100644 index 00000000000..42df65cff5a --- /dev/null +++ b/tests/e2e/rate_limiting/test_rate_limiting_e2e.py @@ -0,0 +1,100 @@ +"""Live e2e coverage for virtual-key TPM enforcement.""" + +from __future__ import annotations + +import time + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody +from rate_limiting_client import RateLimitingClient + +pytestmark = pytest.mark.e2e + + +RATE_LIMIT_MARKERS = ("rate limit", "tpm limit", "current tpm") + + +def _generate_tpm_key( + client: RateLimitingClient, resources: ResourceManager, tpm_limit: int +) -> str: + key = client.gateway.generate_key( + KeyGenerateBody( + models=["gemini-2.5-flash"], + key_alias=f"e2e-rate-limit-{unique_marker()}", + tpm_limit=tpm_limit, + ) + ) + resources.defer(lambda: client.gateway.delete_key(key)) + return key + + +def _assert_tpm_limit_persisted( + client: RateLimitingClient, key: str, tpm_limit: int +) -> None: + info = client.gateway.key_info(key) + assert ( + info.tpm_limit == tpm_limit + ), f"/key/info reports tpm_limit {info.tpm_limit}, configured {tpm_limit}" + + +def _chat(client: RateLimitingClient, key: str, marker: str) -> StreamingResponse: + return client.chat_status( + key, "gemini-2.5-flash", f"reply with one short word {marker}" + ) + + +class TestKeyTpmRateLimiting: + @pytest.mark.covers("rate_limiting.key.tpm.under_limit_allows") + def test_tpm_key_under_limit_allows_request( + self, client: RateLimitingClient, resources: ResourceManager + ) -> None: + key = _generate_tpm_key(client, resources, tpm_limit=100_000) + _assert_tpm_limit_persisted(client, key, 100_000) + + outcome = _chat(client, key, unique_marker()) + + require_successful_call(outcome) + assert ( + outcome.call_id is not None + ), "successful chat response should include x-litellm-call-id for spend-log correlation" + + @pytest.mark.covers("rate_limiting.key.tpm.over_limit_blocks") + def test_tpm_key_over_limit_blocks_next_request( + self, client: RateLimitingClient, resources: ResourceManager + ) -> None: + key = _generate_tpm_key(client, resources, tpm_limit=1) + _assert_tpm_limit_persisted(client, key, 1) + + first = _chat(client, key, unique_marker()) + require_successful_call(first) + assert ( + first.call_id is not None + ), "first request should return a call id so the TPM-consuming success can be observed" + + rows = client.gateway.poll_logs_for_request_id( + first.call_id, + predicate=lambda found: any((row.total_tokens or 0) > 1 for row in found), + ) + assert any( + (row.total_tokens or 0) > 1 for row in rows + ), f"expected first request to record more than the 1 TPM limit, got rows={rows}" + + deadline = time.monotonic() + client.gateway.poll_timeout + last = StreamingResponse(status_code=-1, body="not attempted") + while time.monotonic() < deadline: + last = _chat(client, key, unique_marker()) + if last.status_code == 429: + break + time.sleep(client.gateway.poll_interval) + + assert last.status_code == 429, ( + f"request after exhausting a 1 TPM key should be blocked with 429, got " + f"{last.status_code}: {last.body[:300]}" + ) + assert any( + marker in last.body.lower() for marker in RATE_LIMIT_MARKERS + ), f"429 body should identify TPM/rate limiting, got: {last.body[:300]}"