""" Unit tests for TPM rate limit for concurrent requests Verifies token-reservation pattern: - Concurrent requests cannot all observe "under limit" before any of them has incremented the counter (atomic reservation via ``atomic_check_and_increment_by_n``). - After a successful request, the counter is reconciled to actual usage. - After a failed request, the full reservation is released. The reservation path delegates atomicity to ``atomic_check_and_increment_by_n``, which uses Redis Lua when available and an asyncio-locked in-memory check otherwise. These tests exercise the in-memory fallback so they run without Redis. """ import asyncio from datetime import datetime, timedelta from typing import Any, Dict import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, _AUDIO_BYTES_PER_TOKEN, _PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) from litellm.proxy.utils import InternalUsageCache, hash_token from litellm.types.llms.openai import ( InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse, ) from litellm.types.rerank import RerankResponse from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture def rate_limiter(): cache = DualCache() handler = RateLimitHandler(internal_usage_cache=InternalUsageCache(cache)) return handler, cache @pytest.fixture(autouse=True) def _isolated_request_stash(): token = _request_stash.set(None) yield _request_stash.reset(token) @pytest.mark.asyncio async def test_token_reservation_prevents_concurrent_bypass(rate_limiter): """ With a 100 TPM limit and 5 concurrent requests each estimated at ~50+ tokens, upfront reservation must reject the late arrivals — not let all 5 through. Exercises the in-memory fallback in ``atomic_check_and_increment_by_n``. """ handler, cache = rate_limiter api_key = hash_token("sk-test-key") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, tpm_limit=100, ) request_data = { "model": "gpt-3.5-turbo", "messages": [ { "role": "user", "content": "Hello, this is a test message for concurrent bypass testing.", } ], "max_tokens": 50, } async def make_request(request_id: int) -> Dict[str, Any]: data = request_data.copy() try: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) return { "request_id": request_id, "success": True, "reserved_tokens": get_request_stash().reserved_tokens, } except Exception as e: return { "request_id": request_id, "success": False, "error": str(e), "status_code": getattr(e, "status_code", None), } tasks = [make_request(i) for i in range(5)] results = await asyncio.gather(*tasks) successful = [r for r in results if r["success"]] failed = [r for r in results if not r["success"]] rate_limited = [r for r in failed if r.get("status_code") == 429] assert len(rate_limited) > 0, ( f"Expected some rate-limited requests but all {len(successful)} succeeded — " f"the concurrent bypass bug is still present." ) @pytest.mark.asyncio async def test_no_leak_on_over_limit_rejection(rate_limiter): """ When a reservation would exceed the TPM limit, the counter must NOT be bumped. Otherwise rejected requests would silently consume quota with no path to refund (the failure callback only fires after the reservation was successfully stashed). """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-no-leak"), tpm_limit=10, # tiny limit, easy to blow past ) counter_key = handler.create_rate_limit_keys( key="api_key", value=user_api_key_dict.api_key, rate_limit_type="tokens" ) # Reservation will estimate >> 10 tokens, so this should be rejected. data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "x" * 200}], "max_tokens": 200, } estimated = handler._estimate_tokens_for_request(data=data) assert estimated > user_api_key_dict.tpm_limit, ( "Test assumes the reservation amount blows past the limit; " f"estimated={estimated}, limit={user_api_key_dict.tpm_limit}" ) with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429 # The reservation bump (estimated_tokens) must NOT have committed. The # counter may carry a tiny pre-existing bump from should_rate_limit's # per-request +1 sliding-window logic, but it must be far below the # reservation amount — proving the all-or-nothing primitive rolled back # cleanly on rejection. cached_value = await cache.async_get_cache(key=counter_key, local_only=True) cached_int = int(cached_value or 0) assert cached_int < estimated, ( f"Reservation leaked: counter={cached_int} after rejection of an " f"estimated_tokens={estimated} reservation." ) @pytest.mark.asyncio async def test_token_adjustment_on_success(rate_limiter): """ On success a reserved scope's counter is reconciled to actual via `actual - reserved`. With actual=50 and reserved=100, the api_key counter should see a -50 delta — and only because api_key was reserved against. Unreserved scopes get the full +actual instead. """ handler, _cache = rate_limiter api_key = hash_token("sk-test-adjust") stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("api_key", api_key)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, } }, "model": "gpt-3.5-turbo", } mock_response = ModelResponse( id="test", object="chat.completion", created=int(datetime.now().timestamp()), model="gpt-3.5-turbo", usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), choices=[], ) increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append( { "key": op["key"], "increment": op["increment_value"], } ) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_success_event( kwargs=mock_kwargs, response_obj=mock_response, start_time=datetime.now(), end_time=datetime.now(), ) token_adjustments = [i for i in increments if "tokens" in i["key"]] assert any(i["increment"] == -50 for i in token_adjustments), ( f"Expected a -50 token adjustment (50 actual - 100 reserved) but got: " f"{token_adjustments}" ) @pytest.mark.asyncio async def test_token_release_on_failure(rate_limiter): """On failure the entire reservation must be refunded — but only against scopes that were actually charged at pre-call. Unreserved scopes were never incremented and must not receive a -reserved op (would drift negative).""" handler, _cache = rate_limiter api_key = hash_token("sk-test-fail") stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("api_key", api_key)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, } }, } increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append( { "key": op["key"], "increment": op["increment_value"], } ) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_failure_event( kwargs=mock_kwargs, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) token_releases = [i for i in increments if "tokens" in i["key"]] assert any( i["increment"] == -100 for i in token_releases ), f"Expected the full reservation (-100) to be released, got: {token_releases}" @pytest.mark.asyncio async def test_model_scope_refund_targets_reserved_model(rate_limiter): """ The pre-call reservation is charged against ``data["model"]`` but the router later writes ``model_group`` into ``litellm_params.metadata``, which can be ``None`` or a different value. Reconciliation MUST refund the same model-scoped counter that was incremented; otherwise model-level counters (model_per_team / model_per_key / etc.) drift up forever. This test makes ``model_group`` absent from kwargs (the failure mode in the Greptile P1) and asserts the refund still targets the model the reservation used. """ handler, _cache = rate_limiter api_key = hash_token("sk-test-model-mismatch") team_id = "team-abc" reserved_model = "gpt-4o-mini" stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_model = reserved_model stash.reserved_scopes = frozenset({("model_per_team", f"{team_id}:{reserved_model}")}) mock_kwargs = { # NOTE: no litellm_params.metadata.model_group — get_model_group_from_litellm_kwargs # returns None on this kwargs dict. "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, "user_api_key_team_id": team_id, } }, } increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_failure_event( kwargs=mock_kwargs, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) expected_model_per_team_key = handler.create_rate_limit_keys( key="model_per_team", value=f"{team_id}:{reserved_model}", rate_limit_type="tokens", ) matching = [i for i in increments if i["key"] == expected_model_per_team_key] assert matching, ( f"Expected a refund on the reserved model_per_team counter " f"({expected_model_per_team_key}) but got: " f"{[i['key'] for i in increments]}" ) assert matching[0]["increment"] == -100, ( f"Expected full -100 refund on model_per_team counter, got " f"{matching[0]['increment']}" ) @pytest.mark.asyncio async def test_should_rate_limit_does_not_inflate_tokens_counter(rate_limiter): """ The pre-call sliding-window check (`should_rate_limit`) must not bump the `:tokens` counter. That counter is owned exclusively by the atomic `reserve_tpm_tokens` path; double-handling shrinks the effective TPM budget by 1 per concurrent in-flight request. """ handler, cache = rate_limiter api_key = hash_token("sk-no-tokens-inflation") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, rpm_limit=100, tpm_limit=10_000, ) tokens_counter_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 10, } estimated = handler._estimate_tokens_for_request(data=data) await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) cached = int( await cache.async_get_cache(key=tokens_counter_key, local_only=True) or 0 ) # The :tokens counter should reflect ONLY the reservation amount — not # an additional +1 from the should_rate_limit pre-pass. assert cached == estimated, ( f"Expected :tokens counter to equal the reservation ({estimated}) " f"with no +1 inflation from should_rate_limit, got {cached}" ) @pytest.mark.asyncio async def test_concurrent_burst_within_tpm_budget_all_succeed(rate_limiter): """ With a TPM limit comfortably above (N concurrent × per-request reservation), all N requests must succeed. Pre-fix the should_rate_limit +1-per-key inflation could 429 late arrivals on tight budgets. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-burst-budget"), tpm_limit=1000, rpm_limit=100, ) request_data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "x" * 40}], # ~10 input tokens "max_tokens": 100, } estimated_per_request = handler._estimate_tokens_for_request(data=request_data) n_concurrent = 3 # Sanity: total reservation must fit within tpm_limit and we want enough # headroom that any +1 inflation would NOT push us over. assert estimated_per_request * n_concurrent < user_api_key_dict.tpm_limit async def make_request(request_id: int): data = request_data.copy() try: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) return True except Exception: return False results = await asyncio.gather(*[make_request(i) for i in range(n_concurrent)]) assert all(results), ( f"All {n_concurrent} requests should fit within tpm_limit=" f"{user_api_key_dict.tpm_limit} (estimated_per_request=" f"{estimated_per_request}), but only {sum(results)} succeeded — " f"the should_rate_limit :tokens-counter inflation bug is back." ) @pytest.mark.asyncio async def test_org_scope_refund_on_failure(rate_limiter): """ The plain `organization` scope is reserved upfront (it carries tokens_per_unit) — so on failure, the full reservation must be released against {organization:org_id}:tokens. Pre-fix this scope was missing from `_build_tpm_scope_pipeline_operations`, leaking forever. """ handler, _cache = rate_limiter api_key = hash_token("sk-org-refund") org_id = "org-acme" stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("organization", org_id)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, "user_api_key_org_id": org_id, } }, } increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_failure_event( kwargs=mock_kwargs, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) expected_org_key = handler.create_rate_limit_keys( key="organization", value=org_id, rate_limit_type="tokens" ) matching = [i for i in increments if i["key"] == expected_org_key] assert matching, ( f"Expected a refund on the org tokens counter ({expected_org_key}) " f"but got keys: {[i['key'] for i in increments]}" ) assert ( matching[0]["increment"] == -100 ), f"Expected full -100 refund on org counter, got {matching[0]['increment']}" @pytest.mark.asyncio async def test_org_scope_reconciled_on_success(rate_limiter): """ On success the org tokens counter must be reconciled to actual usage. With reserved=100 and actual=50, the org scope should see a -50 delta. """ handler, _cache = rate_limiter api_key = hash_token("sk-org-success") org_id = "org-acme" stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("organization", org_id)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, "user_api_key_org_id": org_id, } }, "model": "gpt-3.5-turbo", } mock_response = ModelResponse( id="test", object="chat.completion", created=int(datetime.now().timestamp()), model="gpt-3.5-turbo", usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), choices=[], ) increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_success_event( kwargs=mock_kwargs, response_obj=mock_response, start_time=datetime.now(), end_time=datetime.now(), ) expected_org_key = handler.create_rate_limit_keys( key="organization", value=org_id, rate_limit_type="tokens" ) matching = [i for i in increments if i["key"] == expected_org_key] assert matching, ( f"Expected a reconciliation op on the org tokens counter " f"({expected_org_key}), got keys: {[i['key'] for i in increments]}" ) assert matching[0]["increment"] == -50, ( f"Expected -50 delta on org counter (50 actual - 100 reserved), got " f"{matching[0]['increment']}" ) @pytest.mark.asyncio async def test_estimate_tokens_uses_max_tokens_when_explicit(rate_limiter): """When max_tokens is set explicitly, reservation should equal input + max_tokens.""" handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={ "messages": [ {"role": "user", "content": "abcd" * 4} ], # 16 chars ~ 4 tokens "max_tokens": 25, } ) # input ~= 16/4 = 4 tokens; max_tokens = 25; total ~= 29 assert estimate == 4 + 25 @pytest.mark.asyncio async def test_estimate_tokens_honors_explicit_zero_max_tokens(rate_limiter): """ Regression for a Greptile finding: explicit_max_tokens was resolved via `data.get("max_tokens") or data.get("max_completion_tokens") or data.get("max_output_tokens")`, so an explicit 0 in the first field was falsy and fell through to the next field (or the no-max_tokens floor), silently discarding a caller's explicit zero-output request. """ handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={ "messages": [ {"role": "user", "content": "abcd" * 4} ], # 16 chars ~ 4 tokens "max_tokens": 0, } ) assert estimate == 4, ( f"expected input-only reservation (4) for an explicit max_tokens=0, got {estimate}" ) @pytest.mark.asyncio async def test_estimate_tokens_honors_explicit_zero_max_output_tokens_for_responses( rate_limiter, ): handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={ "input": "describe this image in detail", # 29 chars ~ 7 tokens "max_output_tokens": 0, }, min_configured_tpm_limit=40, call_type="aresponses", ) assert estimate == 23 @pytest.mark.asyncio async def test_estimate_tokens_zero_for_empty_embeddings(rate_limiter): """Embeddings have no output budget — reservation should equal input only.""" handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={"input": "hello world"} # 11 chars ) # input ~= 11/4 = 2 tokens (max(1, 11//4)); max_tokens = 0 assert estimate == 2 @pytest.mark.asyncio async def test_contentless_request_reserves_minimum(rate_limiter): """ A contentless request (no messages/prompt/input — e.g. /responses, tool-call continuations) must still hit the atomic counter so concurrent contentless requests don't all observe "under limit". Pre-fix the `has_estimable_content` short-circuit skipped the reservation entirely and post-call reconciliation provided no backpressure. """ handler, cache = rate_limiter api_key = hash_token("sk-contentless") user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpm_limit=2) counter_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) # Two contentless requests should consume two slots of the 2-token # budget. The third must 429. for _ in range(2): data = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert ( get_request_stash().reserved_tokens == 1 ), "Contentless request should reserve the floor of 1 token" counter_after_two = int( await cache.async_get_cache(key=counter_key, local_only=True) or 0 ) assert counter_after_two == 2, ( f"After two contentless requests at the floor, the api_key tokens " f"counter should be 2, got {counter_after_two}" ) with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={"model": "gpt-3.5-turbo"}, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429, ( "Third contentless request must be rate-limited; pre-fix it would " "have bypassed the TPM check entirely." ) @pytest.mark.asyncio async def test_atomic_keys_share_hash_tag_per_descriptor(rate_limiter): """ Cluster safety: every key in a single descriptor's Lua payload must share a `{key:value}` hash tag so the call lands on a single Redis Cluster slot. Otherwise the Lua script raises CROSSSLOT in cluster mode. """ handler, _cache = rate_limiter descriptors = [ { "key": "api_key", "value": "abc", "rate_limit": { "requests_per_unit": 10, "tokens_per_unit": 100, "window_size": 60, }, }, { "key": "user", "value": "xyz", "rate_limit": {"tokens_per_unit": 200, "window_size": 60}, }, ] increments = [{"requests": 1, "tokens": 10}, {"tokens": 10}] for descriptor, inc in zip(descriptors, increments): keys, _args, _meta = handler._build_descriptor_atomic_payload( descriptor=descriptor, increment_amounts=inc, ) # All keys in a descriptor's payload must share the same {tag} # — that's the prefix between the first '{' and '}'. tags = {k[: k.index("}") + 1] for k in keys} assert len(tags) == 1, ( f"Descriptor {descriptor['key']}:{descriptor['value']} produced " f"keys spanning multiple hash tags: {tags}. Redis Cluster would " f"reject this Lua call with CROSSSLOT." ) expected_tag = f"{{{descriptor['key']}:{descriptor['value']}}}" assert tags == {expected_tag}, f"Expected hash tag {expected_tag}, got {tags}" @pytest.mark.asyncio async def test_reservation_released_on_proxy_rejection(rate_limiter): """ If the request is rejected after the pre-call reservation succeeds but before the LLM call (e.g. a downstream guardrail/auth hook raises), `async_post_call_failure_hook` must release the reservation. Otherwise the tokens leak — `async_log_failure_event` is a litellm completion callback and never fires for proxy-side rejections. """ handler, cache = rate_limiter api_key = hash_token("sk-leak-fix") user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpm_limit=1000) data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) reserved = get_request_stash().reserved_tokens assert reserved > 0 counter_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) counter_after_reserve = int( await cache.async_get_cache(key=counter_key, local_only=True) or 0 ) assert counter_after_reserve == reserved # Simulate a downstream guardrail rejecting the request. await handler.async_post_call_failure_hook( request_data=data, original_exception=Exception("guardrail rejected"), user_api_key_dict=user_api_key_dict, ) counter_after_release = int( await cache.async_get_cache(key=counter_key, local_only=True) or 0 ) assert counter_after_release == 0, ( f"Reservation leaked: counter={counter_after_release} after " f"proxy-level rejection refund (expected 0)." ) assert get_request_stash().reservation_released is True, ( "Released flag must be set to prevent " "async_log_failure_event from double-refunding." ) @pytest.mark.asyncio async def test_reservation_release_idempotent(rate_limiter): """ If both `async_post_call_failure_hook` and `async_log_failure_event` end up firing for the same request, only the first refund applies — the second sees the released marker and no-ops. """ handler, _cache = rate_limiter api_key = hash_token("sk-idempotent") increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) # Both hooks read the same per-request ContextVar stash: the # post-call-failure-hook flips reservation_released on it, and the # log-failure-event observes the flip. stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("api_key", api_key)}) await handler.async_post_call_failure_hook( request_data={}, original_exception=Exception("rejected"), user_api_key_dict=UserAPIKeyAuth(api_key=api_key), ) first_refund_count = len([i for i in increments if "tokens" in i["key"]]) assert first_refund_count > 0, "First refund should have applied" # Now simulate async_log_failure_event firing afterwards. It must see # the released flag on the stash and not double-refund. await handler.async_log_failure_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, }, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) second_refund_count = len([i for i in increments if "tokens" in i["key"]]) assert second_refund_count == first_refund_count, ( f"Idempotency violated: refund count went from {first_refund_count} " f"to {second_refund_count} after second hook fired." ) @pytest.mark.asyncio async def test_unreserved_scopes_charged_actual_not_delta_on_success(rate_limiter): """ Counter-drift fix: a scope present in metadata but NOT reserved at pre-call (no configured TPM limit for it) must be charged the full `actual_tokens` on success — never the `delta = actual - reserved`. Otherwise that scope's counter goes negative whenever `actual < reserved` (the common case, since the reservation includes a conservative output pad). """ handler, _cache = rate_limiter api_key = hash_token("sk-mixed-scopes") team_id = "team-no-tpm-limit" # Reservation ONLY hit api_key — team had no TPM limit configured. stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("api_key", api_key)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, "user_api_key_team_id": team_id, } }, "model": "gpt-3.5-turbo", } mock_response = ModelResponse( id="t", object="chat.completion", created=int(datetime.now().timestamp()), model="gpt-3.5-turbo", usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), choices=[], ) increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_success_event( kwargs=mock_kwargs, response_obj=mock_response, start_time=datetime.now(), end_time=datetime.now(), ) api_key_token_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) team_token_key = handler.create_rate_limit_keys( key="team", value=team_id, rate_limit_type="tokens" ) api_key_ops = [i for i in increments if i["key"] == api_key_token_key] team_ops = [i for i in increments if i["key"] == team_token_key] assert api_key_ops and api_key_ops[0]["increment"] == -50, ( f"Reserved api_key scope must reconcile via delta (50-100=-50), " f"got {api_key_ops}" ) assert team_ops and team_ops[0]["increment"] == 50, ( f"Unreserved team scope must be charged full actual (+50), not the " f"-50 delta (which would drift its counter negative). Got {team_ops}" ) @pytest.mark.asyncio async def test_unreserved_scopes_not_refunded_on_failure(rate_limiter): """ Failure refund must only emit ops against scopes the reservation actually charged. Refunding an unreserved scope (which was never incremented at pre-call) would drive its counter to -reserved. """ handler, _cache = rate_limiter api_key = hash_token("sk-mixed-fail") team_id = "team-no-tpm" stash = get_or_create_request_stash() stash.reserved_tokens = 100 stash.reserved_scopes = frozenset({("api_key", api_key)}) mock_kwargs = { "standard_logging_object": { "metadata": { "user_api_key_hash": api_key, "user_api_key_team_id": team_id, } }, } increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_failure_event( kwargs=mock_kwargs, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) team_token_key = handler.create_rate_limit_keys( key="team", value=team_id, rate_limit_type="tokens" ) api_key_token_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) team_ops = [i for i in increments if i["key"] == team_token_key] api_key_ops = [i for i in increments if i["key"] == api_key_token_key] assert not team_ops, ( f"Unreserved team scope must NOT be refunded (would drift negative), " f"got {team_ops}" ) assert ( api_key_ops and api_key_ops[0]["increment"] == -100 ), f"Reserved api_key scope must be refunded -100, got {api_key_ops}" @pytest.mark.asyncio async def test_token_rate_limit_headers_present_in_stored_response(rate_limiter): """ With `skip_tpm_check=True` on the RPM sliding-window pass, token statuses only come from `reserve_tpm_tokens`. They must be merged into the stashed rate-limit response so the post-call hook can emit `x-ratelimit-{key}-remaining-tokens` / `-limit-tokens` headers to the client. """ handler, cache = rate_limiter api_key = hash_token("sk-headers") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, rpm_limit=100, tpm_limit=10_000, ) data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) response = get_request_stash().rate_limit_response assert isinstance( response, dict ), "Expected the stashed rate-limit response to be set after pre-call" statuses = response.get("statuses") or [] token_statuses = [s for s in statuses if s.get("rate_limit_type") == "tokens"] request_statuses = [s for s in statuses if s.get("rate_limit_type") == "requests"] assert token_statuses, ( f"Token rate-limit status missing from stored response. Without it, " f"x-ratelimit-*-tokens headers never reach the client. Got " f"statuses: {[(s.get('descriptor_key'), s.get('rate_limit_type')) for s in statuses]}" ) assert request_statuses, ( "RPM rate-limit status was clobbered by the TPM merge — both must " "coexist in the stored response." ) # The token status carries the limit and a positive remaining budget. api_key_tokens = next( (s for s in token_statuses if s.get("descriptor_key") == "api_key"), None, ) assert api_key_tokens is not None, f"api_key token status absent: {token_statuses}" assert api_key_tokens["current_limit"] == 10_000 assert api_key_tokens["limit_remaining"] >= 0 @pytest.mark.asyncio async def test_estimate_tokens_floor_caps_at_smallest_configured_tpm(rate_limiter): """ Regression: with a small configured TPM cap and no max_tokens, the output-budget floor must be capped at a fraction of that limit so the reservation alone can't trip the limit. """ handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={"messages": [{"role": "user", "content": "hello"}]}, min_configured_tpm_limit=1000, ) # input ~= 5//4 = 1 token; output floor capped at 1000//4 = 250; # total ~= 251 (well under 1000). assert ( estimate <= 1000 // 2 ), f"With TPM=1000, reservation must stay well under the limit; got {estimate}" assert estimate >= 1, "Estimate must be at least the call-site floor of 1" @pytest.mark.asyncio async def test_estimate_tokens_floor_unchanged_for_large_tpm(rate_limiter): """ Large TPM budgets must keep the 1024-token floor so a stream of small concurrent requests can't collectively bypass the limit. """ handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={"messages": [{"role": "user", "content": "hello"}]}, min_configured_tpm_limit=100_000, ) # input ~= 1; output floor = min(1024, 100_000//4=25_000) = 1024; # total ~= 1025. assert estimate == 1 + 1024 @pytest.mark.asyncio async def test_estimate_tokens_floor_unchanged_when_kwarg_omitted(rate_limiter): """ Callers that don't pass min_configured_tpm_limit (legacy path, tests that stub the estimator) must observe the pre-fix floor. """ handler, _cache = rate_limiter estimate = handler._estimate_tokens_for_request( data={"messages": [{"role": "user", "content": "hello"}]}, ) assert estimate == 1 + 1024 @pytest.mark.asyncio async def test_small_tpm_cap_admits_no_max_tokens_request(rate_limiter): """ Regression (end-to-end at the hook level): a project-level model_tpm_limit of 1000 with a tiny no-max_tokens request must not 429 on the first call. Pre-fix the 1024-token floor tripped OVER_LIMIT against the 1000-token cap on every request. """ handler, cache = rate_limiter api_key = hash_token("sk-small-tpm") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-small-tpm", project_metadata={ "model_tpm_limit": {"gpt-3.5-turbo": 1000}, "model_rpm_limit": {"gpt-3.5-turbo": 60}, }, ) data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], } # Must not raise — pre-fix this was a 429. await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) reserved = get_request_stash().reserved_tokens assert reserved > 0, "Reservation should have been stashed" assert reserved <= 1000 // 2, ( f"Capped floor must keep the reservation well under the 1000 TPM " f"cap; got {reserved}" ) @pytest.mark.asyncio async def test_small_tpm_cap_injects_matching_max_tokens(rate_limiter): """ When a small TPM cap forces the no-max_tokens floor below the baseline, the hook must also write data['max_tokens'] = capped_floor so the actual model output is bounded by the reservation. Without this cap, concurrent no-max_tokens generations can spend past the TPM limit before post-call reconciliation runs. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-small-tpm-cap"), project_id="proj-small-tpm-cap", project_metadata={ "model_tpm_limit": {"gpt-3.5-turbo": 1000}, }, ) data: dict = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert data.get("max_tokens") == 1000 // 4, ( f"Capped floor must be written to max_tokens to bound the actual " f"model output; got {data.get('max_tokens')}" ) @pytest.mark.asyncio async def test_large_tpm_cap_does_not_inject_max_tokens(rate_limiter): """ A TPM cap that doesn't constrain the floor must not silently inject max_tokens — that would change behaviour for tenants who already have plenty of budget. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-large-tpm-cap"), project_id="proj-large-tpm-cap", project_metadata={ "model_tpm_limit": {"gpt-3.5-turbo": 100_000}, }, ) data: dict = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert "max_tokens" not in data, ( f"Large TPM caps should leave max_tokens alone; got " f"{data.get('max_tokens')}" ) @pytest.mark.asyncio async def test_small_tpm_cap_preserves_explicit_max_tokens(rate_limiter): """ Explicit max_tokens from the caller must never be overwritten by the bypass mitigation — the user already declared their budget. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-explicit-max-tokens"), project_id="proj-explicit-max-tokens", project_metadata={ "model_tpm_limit": {"gpt-3.5-turbo": 1000}, }, ) data: dict = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 500, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert data["max_tokens"] == 500 @pytest.mark.asyncio async def test_project_otpm_reservation_prevents_concurrent_bypass(rate_limiter): """ Bedrock Mantle-style OTPM: with a 100 OTPM limit and 5 concurrent requests each reserving 50+ output tokens, upfront reservation must reject the late arrivals -- not let all 5 through. Exercises the in-memory fallback in ``atomic_check_and_increment_by_n`` for the project-scoped ITPM/OTPM descriptors specifically. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-otpm-bypass"), project_id="proj-mantle-bypass", project_metadata={ "model_otpm_limit": {"bedrock_mantle/claude-opus": 100}, }, ) request_data = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50, } async def make_request(request_id: int) -> Dict[str, Any]: data = request_data.copy() try: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) return {"request_id": request_id, "success": True} except Exception as e: return { "request_id": request_id, "success": False, "status_code": getattr(e, "status_code", None), } results = await asyncio.gather(*[make_request(i) for i in range(5)]) successful = [r for r in results if r["success"]] rate_limited = [ r for r in results if not r["success"] and r.get("status_code") == 429 ] assert len(rate_limited) > 0, ( f"Expected some OTPM-rate-limited requests but all {len(successful)} succeeded." ) @pytest.mark.asyncio async def test_project_otpm_rejects_multiple_completion_candidates(rate_limiter): handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-otpm-multiple-candidates"), project_id="proj-multiple-candidates", project_metadata={ "model_otpm_limit": {"bedrock_mantle/claude-opus": 500}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "n": 10, } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="acompletion", ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio async def test_project_otpm_reserves_largest_conflicting_output_cap(rate_limiter): handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-otpm-conflicting-caps"), project_id="proj-conflicting-caps", project_metadata={ "model_otpm_limit": {"bedrock_mantle/claude-opus": 50}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 1, "max_completion_tokens": 100, } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="acompletion", ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio @pytest.mark.parametrize( "call_type", ["agenerate_content", "agenerate_content_stream"], ) @pytest.mark.parametrize("config_field", ["config", "generationConfig"]) async def test_project_otpm_rejects_google_genai_native_output_cap( rate_limiter, call_type, config_field, ): handler, cache = rate_limiter model = "gemini/gemini-3-flash-preview" user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-google-genai-native-otpm"), project_id="project-google-genai-native-otpm", project_metadata={"model_otpm_limit": {model: 50}}, ) with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={ "model": model, "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], config_field: {"maxOutputTokens": 100}, }, call_type=call_type, ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio @pytest.mark.parametrize( "call_type", ["agenerate_content", "agenerate_content_stream"], ) @pytest.mark.parametrize("candidate_count_field", ["candidateCount", "candidate_count"]) async def test_project_otpm_rejects_google_genai_native_candidate_count( rate_limiter, call_type, candidate_count_field, ): handler, cache = rate_limiter model = "gemini/gemini-3-flash-preview" user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-google-genai-native-candidate-count"), project_id="project-google-genai-native-candidate-count", project_metadata={"model_otpm_limit": {model: 150}}, ) with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={ "model": model, "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "config": { "maxOutputTokens": 50, candidate_count_field: 4, }, }, call_type=call_type, ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio @pytest.mark.parametrize( "call_type", ["agenerate_content", "agenerate_content_stream"], ) @pytest.mark.parametrize("config_field", [None, "config", "generationConfig"]) async def test_project_otpm_injects_google_genai_native_output_cap( rate_limiter, call_type, config_field, ): handler, cache = rate_limiter model = "gemini/gemini-3-flash-preview" user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-google-genai-native-implicit-otpm"), project_id="project-google-genai-native-implicit-otpm", project_metadata={"model_otpm_limit": {model: 40}}, ) data = { "model": model, "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], } if config_field is not None: data[config_field] = {"temperature": 0} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type, ) stash = get_request_stash() assert stash is not None assert stash.otpm_reserved_tokens == 10 expected_config_field = config_field or "config" assert data[expected_config_field]["maxOutputTokens"] == 10 assert "max_tokens" not in data @pytest.mark.asyncio async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter): """ When ITPM reserves fine but OTPM is then over limit, the ITPM reservation this same pre-call already made must be rolled back -- otherwise it leaks until the window's TTL, silently shrinking the ITPM budget for every other request in that minute. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-otpm-rollback"), project_id="proj-mantle-rollback", project_metadata={ "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 10}, }, ) itpm_counter_key = handler.create_rate_limit_keys( key="model_per_project_itpm", value="proj-mantle-rollback:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) data = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 500, # blows past the 10-token OTPM limit } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429 cached_value = await cache.async_get_cache(key=itpm_counter_key, local_only=True) assert int(cached_value or 0) == 0, ( f"ITPM reservation leaked after OTPM rejection: counter={cached_value}" ) @pytest.mark.asyncio async def test_project_itpm_reconciled_on_success_excludes_cached_tokens(rate_limiter): """ On success, ITPM reconciles to billable input tokens (prompt_tokens minus cached_tokens) -- not raw prompt_tokens. Cached prompt-read tokens are free under Bedrock Mantle and must not count against the ITPM quota, even though they still appear in usage/cost logging elsewhere. """ handler, _cache = rate_limiter itpm_scope = ("model_per_project_itpm", "proj-mantle:model") otpm_scope = ("model_per_project_otpm", "proj-mantle:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) mock_kwargs = {} mock_response = ModelResponse( id="test", object="chat.completion", created=int(datetime.now().timestamp()), model="bedrock_mantle/claude-opus", usage=Usage( prompt_tokens=80, completion_tokens=40, total_tokens=120, prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=30), ), choices=[], ) increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_success_event( kwargs=mock_kwargs, response_obj=mock_response, start_time=datetime.now(), end_time=datetime.now(), ) itpm_adjustments = [i for i in increments if "model_per_project_itpm" in i["key"]] otpm_adjustments = [i for i in increments if "model_per_project_otpm" in i["key"]] # billable_input = 80 - 30 cached = 50; delta = 50 - 100 reserved = -50 assert any(i["increment"] == -50 for i in itpm_adjustments), ( f"Expected a -50 ITPM adjustment (50 billable - 100 reserved), got: {itpm_adjustments}" ) # delta = 40 actual completion - 60 reserved = -20 assert any(i["increment"] == -20 for i in otpm_adjustments), ( f"Expected a -20 OTPM adjustment (40 actual - 60 reserved), got: {otpm_adjustments}" ) @pytest.mark.asyncio async def test_project_reconciliation_does_not_decrement_later_window(): current_time = datetime(2026, 8, 5, 12, 0, 0) cache = DualCache() handler = RateLimitHandler( internal_usage_cache=InternalUsageCache(cache), time_provider=lambda: current_time, ) handler.window_size = 60 scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") descriptor = { "key": scope[0], "value": scope[1], "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, } reservation = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], increments=[{"tokens": 100}], ) counter_key = handler.create_rate_limit_keys(*scope, rate_limit_type="tokens") window_identity = next( identity for identity in reservation["reservation_windows"] if identity[0] == counter_key ) stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({scope}) stash.itpm_reserved_window_identities = frozenset( {window_identity} ) current_time += timedelta(seconds=61) later_reservation = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], increments=[{"tokens": 20}], ) assert window_identity not in later_reservation["reservation_windows"] await handler.async_log_success_event( kwargs={}, response_obj=ModelResponse( usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) ), start_time=current_time, end_time=current_time, ) assert float(await cache.async_get_cache(key=counter_key, local_only=True) or 0) == 20 @pytest.mark.asyncio async def test_project_reconciliation_decrements_its_active_window(rate_limiter): handler, cache = rate_limiter scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") descriptor = { "key": scope[0], "value": scope[1], "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, } reservation = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], increments=[{"tokens": 100}], ) counter_key = handler.create_rate_limit_keys(*scope, rate_limit_type="tokens") window_identity = next( identity for identity in reservation["reservation_windows"] if identity[0] == counter_key ) stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({scope}) stash.itpm_reserved_window_identities = frozenset( {window_identity} ) await handler.async_log_success_event( kwargs={}, response_obj=ModelResponse( usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) ), start_time=datetime.now(), end_time=datetime.now(), ) assert float(await cache.async_get_cache(key=counter_key, local_only=True) or 0) == 10 @pytest.mark.asyncio async def test_redis_window_guard_uses_reservation_identity_and_never_falls_back_negative( rate_limiter, ): handler, _cache = rate_limiter calls = [] async def failing_guard(*, keys, args): calls.append((keys, args)) raise RuntimeError("redis unavailable") unguarded_calls = [] async def capture_unguarded(pipeline_operations, **_kwargs): unguarded_calls.extend(pipeline_operations) handler.window_guarded_token_increment_script = failing_guard handler.async_increment_tokens_with_ttl_preservation = capture_unguarded await handler.async_increment_reservation_aware_tokens( pipeline_operations=[ { "key": "{model_per_project_itpm:project:model}:tokens", "increment_value": -90, "ttl": 60, "window_key": "{model_per_project_itpm:project:model}:window", "expected_window_start": "1234", "reservation_backend": "redis", } ] ) assert calls == [ ( [ "{model_per_project_itpm:project:model}:window", "{model_per_project_itpm:project:model}:tokens", ], ["1234", -90, 60], ) ] assert unguarded_calls == [] @pytest.mark.asyncio async def test_atomic_lua_response_carries_redis_window_identity(rate_limiter): handler, _cache = rate_limiter counter_key = "{model_per_project_itpm:project:model}:tokens" meta = [ { "descriptor_key": PROJECT_ITPM_DESCRIPTOR_KEY, "descriptor_value": "project:model", "current_limit": 100, "rate_limit_type": "tokens", "counter_key": counter_key, } ] async def successful_reservation(*, keys, args): return [0, 25, 1234] handler.check_and_increment_by_n_script = successful_reservation assert await handler._atomic_lua_per_descriptor([]) == { "overall_code": "OK", "statuses": [], } response = await handler._atomic_lua_per_descriptor( descriptor_groups=[ ( [ "{model_per_project_itpm:project:model}:window", counter_key, ], [100, 25, 60, 60], meta, ) ] ) assert response["statuses"][0]["limit_remaining"] == 75 assert response["reservation_windows"] == frozenset( {(counter_key, "1234", "redis")} ) @pytest.mark.asyncio async def test_project_itpm_otpm_released_on_failure(rate_limiter): """On failure, the full ITPM and OTPM reservations must be refunded.""" handler, _cache = rate_limiter itpm_scope = ("model_per_project_itpm", "proj-mantle:model") otpm_scope = ("model_per_project_otpm", "proj-mantle:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) mock_kwargs = {} increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_failure_event( kwargs=mock_kwargs, response_obj=None, start_time=datetime.now(), end_time=datetime.now(), ) itpm_releases = [i for i in increments if "model_per_project_itpm" in i["key"]] otpm_releases = [i for i in increments if "model_per_project_otpm" in i["key"]] assert any(i["increment"] == -100 for i in itpm_releases), itpm_releases assert any(i["increment"] == -60 for i in otpm_releases), otpm_releases @pytest.mark.asyncio async def test_proxy_rejection_refunds_itpm_otpm_by_their_own_amount_not_combined( rate_limiter, ): """ Regression for a Greptile-flagged bug: when a project configures both a combined model_tpm_limit and split model_itpm_limit/model_otpm_limit for the same model, async_post_call_failure_hook's proxy-side refund path used to decrement every token descriptor -- including the ITPM/OTPM ones -- by the flat combined reservation amount, instead of each bucket's own reserved amount. That drives the split counters negative (or under-refunds them) instead of returning them to exactly zero. """ handler, cache = rate_limiter api_key = hash_token("sk-mixed-tpm-io") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-mixed", project_metadata={ "model_tpm_limit": {"bedrock_mantle/claude-opus": 100000}, "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 100000}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [ {"role": "user", "content": "hello there, this is a test message"} ], "max_tokens": 60, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) tpm_counter_key = handler.create_rate_limit_keys( key="model_per_project", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) itpm_counter_key = handler.create_rate_limit_keys( key="model_per_project_itpm", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) otpm_counter_key = handler.create_rate_limit_keys( key="model_per_project_otpm", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) tpm_reserved = int( await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 ) itpm_reserved = int( await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 ) otpm_reserved = int( await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 ) assert tpm_reserved > 0 and itpm_reserved > 0 and otpm_reserved > 0 await handler.async_post_call_failure_hook( request_data=data, original_exception=Exception("guardrail rejected"), user_api_key_dict=user_api_key_dict, ) tpm_after = int( await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 ) itpm_after = int( await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 ) otpm_after = int( await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 ) assert tpm_after == 0, f"combined TPM counter leaked: {tpm_after}" assert itpm_after == 0, ( f"ITPM counter corrupted by combined-amount refund: {itpm_after}" ) assert otpm_after == 0, ( f"OTPM counter corrupted by combined-amount refund: {otpm_after}" ) @pytest.mark.asyncio async def test_proxy_rejection_refunds_itpm_otpm_only_reservation_with_no_combined_tpm( rate_limiter, ): """ Regression for the second half of the same bug: with only model_itpm_limit/model_otpm_limit configured (no model_tpm_limit), the combined reserved_tokens is 0, and the proxy-side refund path used to return immediately on that -- leaking the ITPM/OTPM reservations until the rate-limit window's TTL expired. """ handler, cache = rate_limiter api_key = hash_token("sk-io-only") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-io-only", project_metadata={ "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 100000}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [ {"role": "user", "content": "hello there, this is a test message"} ], "max_tokens": 60, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) itpm_counter_key = handler.create_rate_limit_keys( key="model_per_project_itpm", value="proj-io-only:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) otpm_counter_key = handler.create_rate_limit_keys( key="model_per_project_otpm", value="proj-io-only:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) assert ( int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0) > 0 ) assert ( int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0) > 0 ) await handler.async_post_call_failure_hook( request_data=data, original_exception=Exception("guardrail rejected"), user_api_key_dict=user_api_key_dict, ) itpm_after = int( await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 ) otpm_after = int( await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 ) assert itpm_after == 0, ( f"ITPM-only reservation leaked on proxy rejection: {itpm_after}" ) assert otpm_after == 0, ( f"OTPM-only reservation leaked on proxy rejection: {otpm_after}" ) @pytest.mark.asyncio async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter): """ Regression for a High-severity review finding: when the project ITPM reservation succeeds but OTPM is then over limit, _reserve_project_io_tokens_or_raise rolls back the combined-TPM reservation that already succeeded earlier in the same pre-call, then raises. If it doesn't also mark that reservation released, async_post_call_failure_hook -- which fires next in the real request lifecycle, since raising from async_pre_call_hook triggers it -- sees the same still-stashed reservation and refunds it a second time, driving the combined TPM counter negative and letting a caller push past the project's real TPM budget by repeatedly triggering OTPM rejections. """ handler, cache = rate_limiter api_key = hash_token("sk-double-refund") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-double-refund", project_metadata={ "model_tpm_limit": {"bedrock_mantle/claude-opus": 100000}, "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 5}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [ {"role": "user", "content": "hello there, this is a test message"} ], "max_tokens": 60, # blows past the 5-token OTPM limit } tpm_counter_key = handler.create_rate_limit_keys( key="model_per_project", value="proj-double-refund:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429 tpm_after_pre_call = int( await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 ) assert tpm_after_pre_call == 0, ( f"combined TPM reservation not rolled back: {tpm_after_pre_call}" ) # In the real request lifecycle, async_post_call_failure_hook fires next # for a pre-call rejection. It must not refund the same reservation again. await handler.async_post_call_failure_hook( request_data=data, original_exception=exc_info.value, user_api_key_dict=user_api_key_dict, ) tpm_after_failure_hook = int( await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 ) assert tpm_after_failure_hook == 0, ( f"combined TPM counter went negative from a double refund: {tpm_after_failure_hook}" ) @pytest.mark.parametrize( "embedding_input", [ list(range(51)), [list(range(25)), list(range(26))], ], ) @pytest.mark.asyncio async def test_project_itpm_rejects_pretokenized_embedding_input( rate_limiter, embedding_input, ): handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-pretokenized-embedding-itpm"), project_id="proj-pretokenized-embedding", project_metadata={ "model_itpm_limit": {"text-embedding-3-small": 50}, }, ) data = { "model": "text-embedding-3-small", "input": embedding_input, } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="aembedding", ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio async def test_responses_api_not_misclassified_as_embedding_for_output_estimate( rate_limiter, ): """ Regression for a High-severity review finding: the Responses API also puts its prompt in data["input"], the same field embeddings use, so the output-token estimate treated every Responses call as an embedding and reserved zero output tokens. call_type now disambiguates the two: the same input-only payload gets zero output tokens for an embedding call but a real floor for a Responses API call. """ handler, _cache = rate_limiter data = {"input": "describe this image in detail"} _, embedding_output_estimate = handler._estimate_input_and_output_tokens( data=data, call_type="aembedding" ) assert embedding_output_estimate == 0 _, responses_output_estimate = handler._estimate_input_and_output_tokens( data=data, call_type="aresponses" ) assert responses_output_estimate > 0, ( "Responses API call was misclassified as an embedding and reserved zero output tokens" ) @pytest.mark.parametrize( ("data", "call_type", "expected_output_tokens"), [ ( { "messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "n": 10, }, "acompletion", 1000, ), ( { "prompt": "hello", "max_tokens": 100, "n": 2, "best_of": 5, }, "text_completion", 500, ), ( { "messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "n": 0, "best_of": "invalid", }, "acompletion", 100, ), ], ) def test_output_estimate_accounts_for_completion_candidates( rate_limiter, data, call_type, expected_output_tokens, ): handler, _cache = rate_limiter _, estimated_output_tokens = handler._estimate_input_and_output_tokens( data=data, call_type=call_type, ) assert estimated_output_tokens == expected_output_tokens @pytest.mark.asyncio async def test_responses_api_usage_reconciles_using_input_output_tokens_fields( rate_limiter, ): """ Regression for the other half of the same finding: ResponseAPIUsage exposes input_tokens/output_tokens, not prompt_tokens/completion_tokens. Before this fix, _resolve_io_token_reconcile_usage couldn't resolve Responses API usage at all, so the reservation was silently kept as-is instead of being trued up to the much larger actual usage. """ handler, _cache = rate_limiter itpm_scope = ("model_per_project_itpm", "proj-responses:model") otpm_scope = ("model_per_project_otpm", "proj-responses:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 10 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 10 stash.otpm_reserved_scopes = frozenset({otpm_scope}) mock_kwargs = {} mock_response = ResponsesAPIResponse( id="resp_test", created_at=int(datetime.now().timestamp()), output=[], usage=ResponseAPIUsage(input_tokens=80, output_tokens=400, total_tokens=480), ) increments = [] async def mock_increment(increment_list, **kwargs): for op in increment_list: increments.append({"key": op["key"], "increment": op["increment_value"]}) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_log_success_event( kwargs=mock_kwargs, response_obj=mock_response, start_time=datetime.now(), end_time=datetime.now(), ) itpm_adjustments = [i for i in increments if "model_per_project_itpm" in i["key"]] otpm_adjustments = [i for i in increments if "model_per_project_otpm" in i["key"]] # delta = 80 actual input - 10 reserved = +70 assert any(i["increment"] == 70 for i in itpm_adjustments), ( f"ITPM reservation was never trued up to actual Responses API usage: {itpm_adjustments}" ) # delta = 400 actual output - 10 reserved = +390 assert any(i["increment"] == 390 for i in otpm_adjustments), ( f"OTPM reservation was never trued up to actual Responses API usage: {otpm_adjustments}" ) @pytest.mark.asyncio async def test_itpm_reservation_accounts_for_audio_content_not_just_text(rate_limiter): """ Regression for the audio half of a Medium-severity review finding: litellm.token_counter has no per-type handling for `input_audio` content blocks (unlike images, which it does count via use_default_image_token_count), so it silently contributes 0 tokens for them. Without DEFAULT_AUDIO_TOKEN_ESTIMATE, a burst of audio-heavy requests with minimal text would each reserve only the one-token floor and blow past the project ITPM limit before post-call reconciliation. """ handler, cache = rate_limiter api_key = hash_token("sk-audio-itpm") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-audio", project_metadata={ # Tighter than DEFAULT_AUDIO_TOKEN_ESTIMATE (300), but far bigger # than the handful of tokens the bare text "hi" would cost. "model_itpm_limit": {"bedrock_mantle/claude-opus": 50}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "hi"}, { "type": "input_audio", "input_audio": {"data": "base64-audio-bytes", "format": "wav"}, }, ], } ], } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429, ( "Expected the audio content to push the ITPM reservation over the " "50-token limit; if this doesn't raise, audio content isn't being " "counted again." ) def test_audio_token_estimate_scales_with_payload_size(): """ Regression for veria-ai Low finding: audio token reservation was flat 300 per block regardless of duration. A short clip and a long clip both reserved the same amount, letting a caller hide long audio in one block to exhaust ITPM quota while reserving almost nothing. The estimate must now grow proportionally with the base64 payload size (len(b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN), floored at DEFAULT_AUDIO_TOKEN_ESTIMATE so reference-only blocks and genuinely short clips still get a non-trivial reservation. To exceed the floor the decoded payload must be > 300 * 1600 = 480 000 bytes. We synthesise a fake b64-length string of 650 000 chars (decoded ≈ 487 500 bytes → 304 tokens) to avoid actually allocating and encoding ~480 kB of audio in every test run. """ large_b64 = "A" * 650_000 very_large_b64 = "A" * 12_900_000 small_b64 = "A" * 1_000 large_block = { "type": "input_audio", "input_audio": {"data": large_b64, "format": "wav"}, } small_block = { "type": "input_audio", "input_audio": {"data": small_b64, "format": "wav"}, } very_large_block = { "type": "input_audio", "input_audio": {"data": very_large_b64, "format": "wav"}, } no_data_block = {"type": "input_audio", "input_audio": {"format": "wav"}} large_estimate = RateLimitHandler._estimate_audio_block_tokens(large_block) very_large_estimate = RateLimitHandler._estimate_audio_block_tokens( very_large_block ) small_estimate = RateLimitHandler._estimate_audio_block_tokens(small_block) no_data_estimate = RateLimitHandler._estimate_audio_block_tokens(no_data_block) assert large_estimate > small_estimate, ( f"Large payload ({large_estimate}) must reserve more than small payload " f"({small_estimate}); flat-rate bug is back" ) assert very_large_estimate == len(very_large_b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN assert very_large_estimate > 6_000 assert no_data_estimate >= 300, ( f"Reference-only block (no data) must use the DEFAULT_AUDIO_TOKEN_ESTIMATE floor; got {no_data_estimate}" ) assert small_estimate >= 300, ( f"Small payload must be floored at DEFAULT_AUDIO_TOKEN_ESTIMATE=300; got {small_estimate}" ) @pytest.mark.asyncio async def test_itpm_rejects_large_audio_payload_that_would_pass_flat_estimate( rate_limiter, ): """ Regression: a caller placing a long audio clip in one block previously reserved only 300 tokens (the flat estimate). With the size-proportional estimate, the same clip now reserves proportionally more and must trip the ITPM limit when the limit is tuned to exactly expose the difference. 1 100 000 b64 chars → decoded ≈ 825 000 bytes → 825 000 // 1600 ≈ 515 tokens > the 400-token limit. The flat estimate (300) would have passed. """ handler, cache = rate_limiter large_b64 = "A" * 1_100_000 user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-large-audio"), project_id="proj-large-audio", project_metadata={ "model_itpm_limit": {"bedrock_mantle/claude-opus": 400}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "transcribe this"}, { "type": "input_audio", "input_audio": {"data": large_b64, "format": "wav"}, }, ], } ], } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) assert getattr(exc_info.value, "status_code", None) == 429, ( "Large audio payload must exceed the 400-token ITPM limit under the " "size-proportional estimate; the old flat-rate estimate (300 tokens) " "would have passed this limit silently" ) @pytest.mark.asyncio @pytest.mark.parametrize( ("call_type", "request_data"), [ ( "acompletion", { "messages": [ { "role": "user", "content": [ {"type": "text", "text": "describe this"}, { "type": "image_url", "image_url": { "url": "https://example.com/high-resolution.png", "detail": "high", }, }, ], } ] }, ), ( "acompletion", { "messages": [ { "role": "user", "content": [ {"type": "text", "text": "summarize this"}, { "type": "file", "file": { "filename": "document.pdf", "file_data": "data:application/pdf;base64,dGVzdA==", }, }, ], } ] }, ), ( "aresponses", { "input": [ { "role": "user", "content": [ {"type": "input_text", "text": "describe this"}, { "type": "input_image", "image_url": "https://example.com/high-resolution.png", "detail": "high", }, ], } ] }, ), ( "aresponses", {"input": "continue", "previous_response_id": "resp-123"}, ), ], ) async def test_multimodal_requests_reserve_measured_project_itpm_not_full_limit( rate_limiter, call_type, request_data, ): """ Regression: image, file, and previous_response_id requests used to reserve the project's whole ITPM limit up front. Because the atomic check is ``current + increment > limit``, that made every such request 429 as soon as the window carried any usage at all and, while in flight, blocked every other request for the same project + model. They now reserve the token_counter estimate like everything else, so two multimodal requests fit in the same window. """ handler, cache = rate_limiter model = "bedrock_mantle/claude-opus" project_itpm_limit = 10_000 user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-multimodal-measured"), project_id="project-multimodal-measured", project_metadata={"model_itpm_limit": {model: project_itpm_limit}}, ) await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={"model": model, **request_data}, call_type=call_type, ) first_stash = get_request_stash() assert first_stash is not None first_reservation = first_stash.itpm_reserved_tokens assert 0 < first_reservation < project_itpm_limit // 2 _request_stash.set(None) await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={"model": model, **request_data}, call_type=call_type, ) second_stash = get_request_stash() assert second_stash is not None assert second_stash is not first_stash assert second_stash.itpm_reserved_tokens == first_reservation @pytest.mark.asyncio async def test_itpm_otpm_reservation_is_kept_on_stream_disconnect(rate_limiter): handler, cache = rate_limiter api_key = hash_token("sk-disconnect-test") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-disconnect", project_metadata={ "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 500}, }, ) data: Dict[str, Any] = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) stash = get_request_stash() assert stash is not None assert stash.itpm_reserved_tokens > 0, ( "pre-call hook must stash an ITPM reservation" ) assert stash.otpm_reserved_tokens > 0, ( "pre-call hook must stash an OTPM reservation" ) increment_calls: list[dict] = [] async def mock_increment(increment_list, litellm_parent_otel_span=None): for op in increment_list: increment_calls.append( {"key": op["key"], "increment": op["increment_value"]} ) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( mock_increment ) await handler.async_release_max_parallel_requests_on_disconnect( user_api_key_dict=user_api_key_dict ) itpm_refunds = [ c for c in increment_calls if "model_per_project_itpm" in c["key"] and c["increment"] < 0 ] otpm_refunds = [ c for c in increment_calls if "model_per_project_otpm" in c["key"] and c["increment"] < 0 ] assert not itpm_refunds assert not otpm_refunds assert stash.reservation_released is False @pytest.mark.asyncio async def test_responses_api_otpm_output_cap_applied_not_skipped_as_embedding( rate_limiter, ): """ Regression for a Greptile P1 finding: _reserve_project_io_tokens_or_raise classified any request with data["input"] set as an embedding (no output tokens), which also misclassifies the Responses API -- it puts its prompt in "input" too, but does generate output. That skipped the output cap applied whenever the configured OTPM limit is small enough to need it, letting an unbounded Responses generation blow past OTPM before post-call reconciliation catches up. The cap must land on data["max_output_tokens"], not data["max_tokens"]: the Responses-to-chat-completion transformation only reads max_output_tokens, so a max_tokens cap is silently dropped before provider dispatch (a second Greptile finding on the same code path). """ handler, cache = rate_limiter api_key = hash_token("sk-responses-otpm-cap") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-responses-otpm", project_metadata={ "model_otpm_limit": {"bedrock_mantle/claude-opus": 40}, }, ) data: Dict[str, Any] = { "model": "bedrock_mantle/claude-opus", "input": "describe this image in detail", } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="aresponses", ) assert data.get("max_output_tokens") is not None, ( "Responses call was misclassified as an embedding and skipped the OTPM output cap" ) assert data["max_output_tokens"] == 16 assert data.get("max_tokens") is None, ( "OTPM output cap was written to max_tokens, which the Responses transformation ignores" ) @pytest.mark.asyncio async def test_explicit_zero_output_responses_call_reserves_effective_provider_minimum( rate_limiter, ): handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-responses-zero-output"), project_id="proj-responses-zero-output", project_metadata={ "model_otpm_limit": {"bedrock_mantle/claude-opus": 5}, }, ) with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={ "model": "bedrock_mantle/claude-opus", "input": "describe this image in detail", "max_output_tokens": 0, }, call_type="aresponses", ) assert getattr(exc_info.value, "status_code", None) == 429 @pytest.mark.asyncio async def test_responses_api_combined_tpm_output_cap_applied_not_skipped_as_embedding( rate_limiter, ): """ Regression for the same misclassification bug in the combined-TPM output-cap block of async_pre_call_hook (a second, independent `is_embedding = data.get("input") is not None` check). A project with only a combined model_tpm_limit (no split itpm/otpm) configured small enough to need the output cap must still apply it to a Responses call, and must write it to max_output_tokens for the same reason as the OTPM case above. """ handler, cache = rate_limiter api_key = hash_token("sk-responses-tpm-cap") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, project_id="proj-responses-tpm", project_metadata={ "model_tpm_limit": {"bedrock_mantle/claude-opus": 40}, }, ) data: Dict[str, Any] = { "model": "bedrock_mantle/claude-opus", "input": "describe this image in detail", } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="aresponses", ) assert data.get("max_output_tokens") is not None, ( "Responses call was misclassified as an embedding and skipped the combined-TPM output cap" ) assert data["max_output_tokens"] == 16 assert data.get("max_tokens") is None, ( "combined-TPM output cap was written to max_tokens, which the Responses transformation ignores" ) @pytest.mark.asyncio async def test_responses_api_multimodal_input_counts_image_content(rate_limiter): """ Regression for a Low-severity veria-ai finding: the Responses API's `input` is commonly a list of message/content-block dicts, but litellm.token_counter's `text` argument only joins plain string entries in a list and silently drops everything else -- so an `input_image` block contributed ~0 tokens to the ITPM estimate instead of the real image token count. _estimate_precise_input_tokens now converts Responses `input` to chat messages first (via the standard transform_responses_api_input_to_messages helper) so image content is counted the same way a chat completion's image content already is. """ handler, _cache = rate_limiter text_only_estimate = handler._estimate_precise_input_tokens( data={"input": "hi"}, model="bedrock_mantle/claude-opus", call_type="aresponses", ) multimodal_estimate = handler._estimate_precise_input_tokens( data={ "input": [ { "role": "user", "content": [ {"type": "input_text", "text": "hi"}, { "type": "input_image", "image_url": "https://example.com/some-image.png", }, ], } ], }, model="bedrock_mantle/claude-opus", call_type="aresponses", ) assert multimodal_estimate > text_only_estimate + 100, ( "Responses API input_image content block was not counted; got " f"text_only={text_only_estimate}, multimodal={multimodal_estimate}" ) @pytest.mark.asyncio async def test_refund_reserved_tokens_noop_when_amount_zero(rate_limiter): """_refund_reserved_tokens returns immediately without calling Redis when amount=0.""" handler, _cache = rate_limiter calls = [] async def mock_increment(pipeline_operations, **kwargs): calls.extend(pipeline_operations) handler.async_increment_tokens_with_ttl_preservation = mock_increment await handler._refund_reserved_tokens( scopes=[("api_key", "sk-test")], amount=0, ) assert not calls, "No Redis ops expected when amount is zero" @pytest.mark.asyncio async def test_reserve_io_tokens_noop_when_no_itpm_otpm_descriptors(rate_limiter): """reserve_io_tokens returns OK immediately when no ITPM/OTPM descriptors present.""" handler, _cache = rate_limiter non_io_descriptor = { "key": "api_key", "value": "sk-test", "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, } response, itpm_reserved, otpm_reserved = await handler.reserve_io_tokens( descriptors=[non_io_descriptor], estimated_input_tokens=50, estimated_output_tokens=50, ) assert response["overall_code"] == "OK" assert itpm_reserved == 0 assert otpm_reserved == 0 @pytest.mark.asyncio async def test_reserve_io_tokens_itpm_only_no_otpm(rate_limiter): """When only ITPM descriptors are present (no OTPM), returns itpm_reserved with otpm=0.""" handler, cache = rate_limiter itpm_descriptor = { "key": PROJECT_ITPM_DESCRIPTOR_KEY, "value": "proj-a:model", "rate_limit": {"tokens_per_unit": 10000, "window_size": 60}, } response, itpm_reserved, otpm_reserved = await handler.reserve_io_tokens( descriptors=[itpm_descriptor], estimated_input_tokens=100, estimated_output_tokens=50, ) assert response["overall_code"] == "OK" assert itpm_reserved == 100 assert otpm_reserved == 0 def test_strip_audio_content_blocks_passthrough_non_list_messages(): """Non-list input is returned unchanged (early return on line 2605).""" result = RateLimitHandler._strip_audio_content_blocks("not a list") assert result == "not a list" def test_strip_audio_content_blocks_passthrough_non_dict_message(): """Non-dict entries in the message list are appended unchanged.""" messages = ["plain string message"] result = RateLimitHandler._strip_audio_content_blocks(messages) assert result == ["plain string message"] def test_strip_audio_content_blocks_passthrough_non_list_content(): """Messages with non-list content (e.g. plain string) pass through unchanged.""" messages = [{"role": "user", "content": "hello"}] result = RateLimitHandler._strip_audio_content_blocks(messages) assert result == [{"role": "user", "content": "hello"}] @pytest.mark.asyncio async def test_otpm_rejection_releases_stashed_parallel_slot(rate_limiter): """ When OTPM is over limit and a parallel slot was already acquired, the disconnect cleanup path in _reserve_project_io_tokens_or_raise must release that slot. Exercises lines 2773-2777. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-otpm-slot"), project_id="proj-slot", project_metadata={"model_otpm_limit": {"m": 5}}, ) data: Dict[str, Any] = { "model": "m", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50, } slot_released = [] async def mock_release(acquisition, parent_otel_span=None): slot_released.append(acquisition) handler._release_parallel_request_slots = mock_release stash = get_or_create_request_stash() stash.parallel_slot = { "slot_id": "test-slot-id", "counter_keys": ["some-key"], } otpm_descriptor = { "key": PROJECT_OTPM_DESCRIPTOR_KEY, "value": "proj-slot:m", "rate_limit": {"tokens_per_unit": 5, "window_size": 60}, } with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler._reserve_project_io_tokens_or_raise( descriptors=[otpm_descriptor], data=data, requested_model="m", user_api_key_dict=user_api_key_dict, tpm_reservation_scopes=[], tpm_reservation_amount=0, ) assert getattr(exc_info.value, "status_code", None) == 429 assert slot_released, "Parallel slot must be released when OTPM rejects" assert stash.parallel_slot is None @pytest.mark.asyncio async def test_itpm_only_status_stored_when_no_prior_rate_limit_response(rate_limiter): """ When only ITPM is configured (no combined TPM/RPM to pre-populate the request stash), a successful ITPM reservation must store its status there so post-call headers can read it. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-itpm-only-store"), project_id="proj-store", ) data: Dict[str, Any] = {"model": "m", "messages": []} itpm_descriptor = { "key": PROJECT_ITPM_DESCRIPTOR_KEY, "value": "proj-store:m", "rate_limit": {"tokens_per_unit": 100000, "window_size": 60}, } await handler._reserve_project_io_tokens_or_raise( descriptors=[itpm_descriptor], data=data, requested_model="m", user_api_key_dict=user_api_key_dict, tpm_reservation_scopes=[], tpm_reservation_amount=0, ) stash = get_request_stash() assert stash is not None stored = stash.rate_limit_response assert stored is not None, ( "ITPM status must be stored in litellm_proxy_rate_limit_response" ) assert stored.get("statuses"), "Stored response must contain statuses" def test_resolve_io_token_usage_responses_api_with_cached_tokens(rate_limiter): """ ResponsesAPIResponse whose usage.input_tokens_details.cached_tokens is set subtracts the cached portion from billable input. Covers line 3501. """ handler, _cache = rate_limiter response_obj = ResponsesAPIResponse( id="resp_cached", created_at=int(datetime.now().timestamp()), output=[], usage=ResponseAPIUsage( input_tokens=100, output_tokens=50, total_tokens=150, input_tokens_details=InputTokensDetails(cached_tokens=25), ), ) billable_input, completion_tokens, resolved = ( handler._resolve_io_token_reconcile_usage(response_obj) ) assert resolved is True assert billable_input == 75, f"Expected 100 - 25 cached = 75, got {billable_input}" assert completion_tokens == 50 def test_resolve_io_token_usage_dict_format(rate_limiter): """ Dict-shaped usage on a ModelResponse (older SDK versions or raw dicts in the usage field) is parsed correctly. Covers lines 3502-3506. """ handler, _cache = rate_limiter response_obj = ModelResponse.model_construct( usage={ "prompt_tokens": 80, "completion_tokens": 40, "prompt_tokens_details": {"cached_tokens": 20}, } ) billable_input, completion_tokens, resolved = ( handler._resolve_io_token_reconcile_usage(response_obj) ) assert resolved is True assert billable_input == 60, f"Expected 80 - 20 cached = 60, got {billable_input}" assert completion_tokens == 40 def test_resolve_io_token_usage_unknown_type_returns_unresolved(rate_limiter): """ A ModelResponse whose usage attribute is not a Usage, ResponseAPIUsage, or dict (e.g. a plain int) returns (0, 0, False) so the reservation is kept rather than guessed. Covers lines 3507-3508. """ handler, _cache = rate_limiter response_obj = ModelResponse.model_construct(usage=42) billable_input, completion_tokens, resolved = ( handler._resolve_io_token_reconcile_usage(response_obj) ) assert resolved is False assert billable_input == 0 assert completion_tokens == 0 @pytest.mark.parametrize( ("combined_usage", "expected_increments"), [ (None, ()), ( Usage(prompt_tokens=40, completion_tokens=15, total_tokens=55), (-60, -45), ), ], ) def test_zero_usage_keeps_reservations_unless_measured_fallback_exists( rate_limiter, combined_usage, expected_increments, ): handler, _cache = rate_limiter itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) kwargs = {} if combined_usage is None else {"combined_usage_object": combined_usage} response_obj = ModelResponse( usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) ) operations = handler._build_io_token_reservation_ops(kwargs, response_obj) assert tuple(operation["increment_value"] for operation in operations) == expected_increments @pytest.mark.parametrize( ("usage", "expected_increments"), [ ( Usage(prompt_tokens=40, completion_tokens=15, total_tokens=55), (40, 15), ), ( Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), (100, 60), ), ], ) def test_retry_success_charges_released_project_io_reservations( rate_limiter, usage, expected_increments, ): handler, _cache = rate_limiter itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) stash.reservation_released = True operations = handler._build_io_token_reservation_ops( {}, ModelResponse(usage=usage), ) assert tuple(operation["increment_value"] for operation in operations) == expected_increments @pytest.mark.asyncio async def test_build_io_token_reservation_ops_skips_unresolvable_usage(rate_limiter): """ When response_obj has no parseable usage, _build_io_token_reservation_ops returns [] to keep the reservation as-is rather than zeroing it out on a bad guess. Covers line 3538. """ handler, _cache = rate_limiter itpm_scope = ("model_per_project_itpm", "proj-b:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 50 stash.itpm_reserved_scopes = frozenset({itpm_scope}) mock_kwargs = {} ops = handler._build_io_token_reservation_ops( kwargs=mock_kwargs, response_obj=object(), ) assert not ops, f"Expected empty ops for unresolvable usage, got {ops}" @pytest.mark.asyncio async def test_post_call_failure_skips_rpm_only_descriptor_in_tpm_refund(rate_limiter): """ async_post_call_failure_hook skips descriptors without tokens_per_unit (e.g. an RPM-only api_key scope) when building the combined-TPM refund ops, so a key with rpm_limit but no tpm_limit doesn't receive a spurious refund that would drive its counter negative. Covers the continue guard at line 4250. """ handler, cache = rate_limiter api_key = hash_token("sk-rpm-only-desc") user_api_key_dict = UserAPIKeyAuth( api_key=api_key, rpm_limit=100, project_id="proj-rpm-only-desc", project_metadata={"model_tpm_limit": {"gpt-3.5-turbo": 100000}}, ) data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) rpm_tokens_key = handler.create_rate_limit_keys( key="api_key", value=api_key, rate_limit_type="tokens" ) await handler.async_post_call_failure_hook( request_data=data, original_exception=Exception("rejected"), user_api_key_dict=user_api_key_dict, ) api_key_tokens_after = int( await cache.async_get_cache(key=rpm_tokens_key, local_only=True) or 0 ) assert api_key_tokens_after >= 0, ( f"RPM-only api_key scope must not receive a negative TPM refund; got {api_key_tokens_after}" ) @pytest.mark.asyncio async def test_max_output_tokens_prevents_cap_injection(rate_limiter): """ Regression for veria-ai comment: when a Responses API request supplies max_output_tokens (the canonical Responses output bound) but not max_tokens or max_completion_tokens, the has_explicit_max_tokens check was False, so the code injected data["max_tokens"] = capped_floor and silently truncated the response. With the fix, max_output_tokens is included in the explicit-cap check and data["max_tokens"] must NOT be injected when max_output_tokens is already set. """ handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-max-output-tokens"), project_id="proj-responses-max-output", project_metadata={ "model_otpm_limit": {"mock-model": 100}, }, ) data: dict = { "model": "mock-model", "input": "Summarise the document", "max_output_tokens": 80, "litellm_call_id": "test-max-output-tokens", "metadata": {}, } try: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="responses", ) except Exception: pass assert "max_tokens" not in data, ( "data['max_tokens'] must not be injected when max_output_tokens is already " "set; the cap injection was overriding the caller's explicit output bound" ) @pytest.mark.asyncio @pytest.mark.parametrize( ("call_type", "request_data", "cap_field", "reserved_tokens"), [ ("aresponses", {"input": "hello", "max_tokens": 1}, "max_output_tokens", 16), ( "acompletion", { "messages": [{"role": "user", "content": "hello"}], "max_output_tokens": 1, }, "max_tokens", 10, ), ], ) async def test_output_reservation_ignores_cap_fields_from_other_endpoints( rate_limiter, call_type, request_data, cap_field, reserved_tokens, ): handler, cache = rate_limiter user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-{call_type}"), project_id=f"project-{call_type}", project_metadata={"model_otpm_limit": {"model": 40}}, ) data = {"model": "model", **request_data} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type, ) assert data[cap_field] == reserved_tokens stash = get_request_stash() assert stash is not None assert stash.otpm_reserved_tokens == reserved_tokens def test_responses_input_is_counted_even_when_messages_is_present(rate_limiter): handler, _cache = rate_limiter small_estimate = handler._estimate_precise_input_tokens( data={"input": "short", "messages": [{"role": "user", "content": "ignored"}]}, model="", call_type="aresponses", ) large_estimate = handler._estimate_precise_input_tokens( data={"input": "large input " * 500, "messages": []}, model="", call_type="aresponses", ) assert large_estimate > small_estimate def test_anthropic_messages_usage_reconciles_split_project_quota(rate_limiter): handler, _cache = rate_limiter billable_input, output_tokens, resolved = handler._resolve_io_token_reconcile_usage( { "usage": { "input_tokens": 100, "output_tokens": 25, "cache_read_input_tokens": 30, } } ) assert resolved is True assert billable_input == 70 assert output_tokens == 25 @pytest.mark.asyncio @pytest.mark.parametrize( "call_type", ["agenerate_content", "agenerate_content_stream"], ) async def test_google_genai_native_contents_reserve_project_itpm( rate_limiter, call_type, ): handler, cache = rate_limiter model = "gemini/gemini-2.5-flash" user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-google-genai-native-itpm"), project_id="project-google-genai-native-itpm", project_metadata={"model_itpm_limit": {model: 10_000}}, ) await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={ "model": model, "contents": [ { "role": "user", "parts": [{"text": "Gemini quota input " * 200}], } ], }, call_type=call_type, ) stash = get_request_stash() assert stash is not None assert stash.itpm_reserved_tokens > 100 @pytest.mark.asyncio @pytest.mark.parametrize("call_type", ["rerank", "arerank"]) async def test_rerank_query_and_documents_enforce_project_itpm( rate_limiter, monkeypatch, call_type, ): handler, cache = rate_limiter captured = {} def token_counter(**kwargs): captured.update(kwargs) return 101 monkeypatch.setattr("litellm.token_counter", token_counter) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-{call_type}-itpm"), project_id=f"project-{call_type}-itpm", project_metadata={"model_itpm_limit": {"rerank-model": 100}}, ) with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data={ "model": "rerank-model", "query": "Which document is most relevant?", "documents": ["first document", {"text": "second document"}], }, call_type=call_type, ) assert getattr(exc_info.value, "status_code", None) == 429 assert captured["text"] == ( "Which document is most relevant?\n" "first document\n" "{'text': 'second document'}" ) def test_rerank_input_estimate_falls_back_to_character_count( rate_limiter, monkeypatch, ): handler, _cache = rate_limiter data = { "query": "query text", "documents": ["first document", "second document"], } def token_counter(**_kwargs): raise ValueError("tokenizer unavailable") monkeypatch.setattr("litellm.token_counter", token_counter) rerank_text = handler._rerank_input_to_text(data) assert handler._estimate_precise_input_tokens( data, model="custom-rerank-model", call_type="rerank", ) == len(rerank_text) // 4 @pytest.mark.parametrize( ("response_obj", "expected"), [ ( RerankResponse( meta={"tokens": {"input_tokens": 42, "output_tokens": 3}} ), (42, 3, True), ), ( RerankResponse( meta={ "tokens": {"input_tokens": 0, "output_tokens": 0}, "billed_units": {"total_tokens": 57}, } ), (57, 0, True), ), ( RerankResponse( meta={ "tokens": {"input_tokens": 0, "output_tokens": 0}, "billed_units": {"total_tokens": 0}, } ), (0, 0, False), ), ], ) def test_rerank_usage_reconciles_project_split_token_quota( rate_limiter, response_obj, expected, ): handler, _cache = rate_limiter assert handler._resolve_io_token_reconcile_usage(response_obj) == expected def test_split_quota_helpers_handle_non_mapping_inputs(rate_limiter): handler, _cache = rate_limiter assert _call_id_from_callback_kwargs(object()) is None assert handler._is_embedding_request(object(), None) is False assert handler._get_explicit_output_cap(object(), None) is None assert handler.get_output_candidate_count(object()) == 1 assert handler.get_output_candidate_count({"n": 1e309}) == 1 assert ( handler._get_explicit_output_cap({"max_output_tokens": []}, "responses") is None ) assert handler._apply_implicit_output_cap(object(), 100, "responses") is None assert handler._estimate_input_and_output_tokens(object()) == (0, 0) assert handler._build_io_token_reservation_ops(object(), object()) == () @pytest.mark.parametrize( ("data", "call_type", "expected"), [ ({"max_tokens": "30.0"}, "", 30), ({"max_tokens": "not-a-number"}, "", None), ({"max_tokens": True}, "", None), ({"max_output_tokens": "30.0"}, "responses", 30), ({"max_output_tokens": "nan"}, "responses", None), ({"generationConfig": {"maxOutputTokens": "12.5"}}, "agenerate_content", 12), ({"generationConfig": {"maxOutputTokens": "oops"}}, "agenerate_content", None), ], ) def test_get_explicit_output_cap_tolerates_unparseable_values( rate_limiter, data, call_type, expected ): """A client-supplied cap the proxy cannot parse must fall back to the no-cap output estimate instead of raising ValueError and 500ing the request before it ever reaches the provider.""" handler, _cache = rate_limiter assert handler._get_explicit_output_cap(data, call_type) == expected @pytest.mark.asyncio async def test_project_io_counters_not_double_charged_when_reservation_disabled( monkeypatch, ): """With LITELLM_TPM_TOKEN_RESERVATION_ENABLED=false the first should_rate_limit pass used to +1 every ITPM/OTPM counter on top of the full reservation _reserve_project_io_tokens_or_raise always makes, permanently inflating each bucket by one token per request.""" monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") cache = DualCache() handler = RateLimitHandler(internal_usage_cache=InternalUsageCache(cache)) assert handler.tpm_reservation_enabled is False user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-io-no-reservation"), project_id="proj-io-no-reservation", project_metadata={ "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000000}, "model_otpm_limit": {"bedrock_mantle/claude-opus": 1000000}, }, ) data = { "model": "bedrock_mantle/claude-opus", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50, } await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="", ) stash = get_request_stash() assert stash is not None assert stash.itpm_reserved_tokens > 0 assert stash.otpm_reserved_tokens > 0 for descriptor_key, reserved in ( ("model_per_project_itpm", stash.itpm_reserved_tokens), ("model_per_project_otpm", stash.otpm_reserved_tokens), ): counter_key = handler.create_rate_limit_keys( key=descriptor_key, value="proj-io-no-reservation:bedrock_mantle/claude-opus", rate_limit_type="tokens", ) cached = await cache.async_get_cache(key=counter_key, local_only=True) assert int(cached or 0) == reserved, ( f"{descriptor_key} counter {cached} != reserved {reserved}: " "first-pass should_rate_limit double-charged the bucket" ) @pytest.mark.parametrize( ("call_type", "data"), [ ( "text_completion", { "messages": [{"role": "user", "content": "ignored"}], "prompt": "abcd", "input": "ignored", "max_tokens": 1, }, ), (None, {"prompt": "abcd", "max_tokens": 1}), (None, {"prompt": ["abcd", "efgh"], "max_tokens": 1}), ], ) def test_split_token_estimate_selects_endpoint_input(rate_limiter, call_type, data): handler, _cache = rate_limiter estimated_input, estimated_output = handler._estimate_input_and_output_tokens( data=data, call_type=call_type, ) assert estimated_input > 0 assert estimated_output == 1 def test_split_quota_multimodal_guards_handle_non_mapping_inputs(rate_limiter): handler, _cache = rate_limiter assert handler._estimate_audio_block_tokens( object() ) == handler._estimate_audio_block_tokens({}) assert handler._responses_input_to_chat_messages(object()) == () assert handler._estimate_precise_input_tokens(object(), model=None) == 0 @pytest.mark.parametrize( ("call_type", "data", "expected_text"), [ ("embedding", {"input": "embedding input"}, "embedding input"), ( "embedding", {"input": ["first embedding", "second embedding"]}, ["first embedding", "second embedding"], ), ("text_completion", {"prompt": "completion prompt"}, "completion prompt"), ], ) def test_precise_input_estimate_selects_endpoint_text( rate_limiter, monkeypatch, call_type, data, expected_text, ): handler, _cache = rate_limiter captured = {} def token_counter(**kwargs): captured.update(kwargs) return 7 monkeypatch.setattr("litellm.token_counter", token_counter) assert ( handler._estimate_precise_input_tokens(data, model="test", call_type=call_type) == 7 ) assert captured["messages"] is None assert captured["text"] == expected_text @pytest.mark.asyncio async def test_project_io_reservation_ignores_non_mapping_request_data(rate_limiter): handler, _cache = rate_limiter await handler._reserve_project_io_tokens_or_raise( descriptors=[], data=object(), requested_model=None, user_api_key_dict=UserAPIKeyAuth(), tpm_reservation_scopes=(), tpm_reservation_amount=0, ) @pytest.mark.asyncio async def test_streaming_combined_usage_reconciles_project_io_reservations( rate_limiter, ): handler, _cache = rate_limiter itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset({itpm_scope}) stash.otpm_reserved_tokens = 60 stash.otpm_reserved_scopes = frozenset({otpm_scope}) kwargs = { "combined_usage_object": Usage( prompt_tokens=40, completion_tokens=15, total_tokens=55, ), } increments = [] async def capture_increments(increment_list, **_kwargs): increments.extend(increment_list) handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( capture_increments ) await handler.async_log_success_event( kwargs=kwargs, response_obj={"response": "stream body"}, start_time=datetime.now(), end_time=datetime.now(), ) itpm_adjustments = [ operation for operation in increments if PROJECT_ITPM_DESCRIPTOR_KEY in operation["key"] ] otpm_adjustments = [ operation for operation in increments if PROJECT_OTPM_DESCRIPTOR_KEY in operation["key"] ] assert [operation["increment_value"] for operation in itpm_adjustments] == [-60] assert [operation["increment_value"] for operation in otpm_adjustments] == [-45] def test_aggregate_only_combined_usage_reconciles_project_io_reservations(rate_limiter): handler, _cache = rate_limiter stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset( {(PROJECT_ITPM_DESCRIPTOR_KEY, "project:model")} ) stash.otpm_reserved_tokens = 80 stash.otpm_reserved_scopes = frozenset( {(PROJECT_OTPM_DESCRIPTOR_KEY, "project:model")} ) kwargs = { "combined_usage_object": Usage(total_tokens=55), } operations = handler._build_io_token_reservation_ops(kwargs, object()) assert [operation["increment_value"] for operation in operations] == [-45, -25] def test_raw_split_usage_dict_reconciles_project_io_tokens(rate_limiter): handler, _cache = rate_limiter assert handler._resolve_io_token_reconcile_usage( { "input_tokens": 30, "output_tokens": 12, "input_tokens_details": {"cached_tokens": 5}, } ) == (25, 12, True) @pytest.mark.asyncio async def test_post_call_success_hook_contains_header_merge_failures( rate_limiter, monkeypatch ): handler, _cache = rate_limiter response = ModelResponse() response._hidden_params = {} def raise_on_merge(**_kwargs): raise RuntimeError("header merge failed") monkeypatch.setattr( handler, "_merge_ratelimit_statuses_into_additional_headers", raise_on_merge, ) await handler.async_post_call_success_hook( data={ "litellm_proxy_rate_limit_response": { "overall_code": "OK", "statuses": (), } }, user_api_key_dict=UserAPIKeyAuth(), response=response, ) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])