litellm/tests
Sameer Kankute 5b93ba0ada
Some checks are pending
GitHub Actions Security Analysis / zizmor (push) Waiting to run
feat(router): add separate ITPM/OTPM deployment rate limits (#31952)
* feat(router): add separate ITPM/OTPM deployment rate limits

Support input/output tokens per minute on deployments via enforce_model_rate_limits, with reservation, reconciliation, refund on failure, and rate-limit headers.

Co-authored-by: Cursor <cursoragent@cursor.com>

* chore(router): keep ITPM/OTPM diff minimal in router.py

Drop unrelated Black reformatting from router.py and types/router.py so the PR only contains functional ITPM/OTPM changes.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(router): make ITPM/OTPM limits separate and atomic

Address Greptile review on separate ITPM/OTPM deployment rate limits.

- OTPM is now reserved atomically pre-call with rollback, matching the ITPM
  path, so concurrent requests can no longer overshoot the configured output
  limit before reconciliation
- ITPM counts input tokens only; it no longer accumulates completion tokens,
  so the input-token limit and x-ratelimit-limit-input-tokens header describe
  input usage as their names imply
- _read_reservation_from_kwargs only falls back to litellm_params.metadata when
  the top-level metadata channel is absent, so production requests carrying a
  litellm_params.metadata dict still reconcile and refund their reservation

Adds regression tests for OTPM atomicity under concurrency, input-only ITPM
enforcement, and reservation lookup when litellm_params.metadata is present.

* fix(router): subtract input tokens only from remaining-input-tokens header

The in-flight replay for x-ratelimit-remaining-input-tokens subtracted total
tokens (input + output) instead of input tokens only, so clients saw remaining
input quota understated by the completion token count on every response. Now
consistent with the input-only ITPM counter.

* fix(router): make itpm/otpm vs tpm/rpm precedence explicit

When a deployment configures itpm/otpm alongside tpm/rpm, the io-token path
takes over and the tpm/rpm limits are not enforced. Log a warning the first
time such a conflicting deployment is seen so the supersession is not silent,
and document the mutual exclusivity.

Post-call reconciliation now only trues up a counter that was actually
reserved against, so the itpm/otpm keys are no longer incremented for
deployments that never configured that limit.

* fix(router): track actual io-token usage on the reservation-minute key

Post-call reconciliation now keys off the exact cache key stashed at pre-call
time rather than one recomputed from the response-time minute. This fixes two
issues: a request whose pre-call estimate was 0 now still writes its actual
billable input to the ITPM counter (previously it was skipped, leaving the
limit unenforceable for that request), and a call that finishes in a later
minute reconciles against the minute it reserved against instead of pushing a
negative delta into the next minute. Counters are only touched when their
limit is configured.

* fix(router): run io-token reconciliation before the model_id guard

async_log_success_event gated IO reconciliation behind the model_id guard that
only the TPM tracking path needs. Since reconciliation works entirely from the
cache keys stashed in kwargs, a success event whose standard_logging_object
lacks model_id would skip reconciliation and leave the reservation on the
counter until the TTL expired, wasting quota. Route the IO path first.

* fix(router): don't replay in-flight delta for itpm/otpm headers

For ITPM/OTPM model groups the counter is incremented at reservation time
(pre-call), so the remaining values returned by get_remaining_model_group_usage
already account for the current request. Replaying the in-flight delta on top
double-counted it and understated x-ratelimit-remaining-input/output-tokens by
up to max_tokens on every response. Skip the delta for io-token groups; the
legacy TPM/RPM replay path is unchanged.

* fix(router): clear io-token reservation after reconcile/refund

async_io_token_refund_failure and async_io_token_reconcile_success now clear
the stashed reservation keys from the request metadata once done. Otherwise, on
a model group mixing IO-limited and non-IO deployments, a failed IO call that
retries on a non-IO fallback left the stale sentinel in the shared request
metadata; the fallback's success handler would divert into IO reconciliation
against the already-refunded key, driving the ITPM counter negative and
skipping the non-IO deployment's TPM tracking.

* fix(router): tidy reservation channel lookup and header guard

Consolidate the reservation channel lookup into a single ordered helper shared
by read and clear, so top-level metadata always wins over litellm_params
metadata without the tangled per-iteration fallback.

Also stop gating the router rate-limit header block on the presence of
x-ratelimit-remaining-input/output-tokens. That block only emits those headers
for ITPM/OTPM groups; for a non-IO group backed by a provider that natively
returns input/output token headers, the extra conditions suppressed the
router's own remaining-tokens/requests headers.

* fix(router): strip client-supplied io-token reservation keys

The reservation sentinels (_litellm_itpm_reserved, _litellm_itpm_cache_key,
and the otpm equivalents) are server-only, but metadata is caller-controlled on
proxy requests. An authenticated caller could forge these fields with an
arbitrary cache key so the post-call reconcile/refund path would decrement any
deployment's ITPM/OTPM counter and let it exceed the configured limit. Strip
the reserved keys from the request metadata in set_io_token_rate_limit_request_kwargs,
which runs before the router stashes its own reservation, so only a genuine
server-side reservation is ever read post-call.

* fix(router): track TPM routing load for io-limited deployments

deployment_callback_on_success early-returned for any deployment with itpm/otpm
set, so its total-token usage never landed in the router's TPM routing counter.
TPM-aware routing strategies then saw 0 load for IO deployments and over-routed
to them in mixed model groups. Only skip tracking when neither tpm/rpm nor
itpm/otpm are configured; itpm/otpm enforcement still runs separately in
ModelRateLimitingCheck, so the routing counter and the enforcement counters
stay independent.

* fix(router): expose standard tpm/rpm headers for io-limited groups

get_remaining_model_group_usage returned early for ITPM/OTPM groups, so a group
that also set tpm/rpm never emitted x-ratelimit-remaining-tokens / -requests;
clients and prometheus gauges reading those saw no data. Build both header sets
instead of returning early.

Also simplify the in-flight header replay: only the tpm/rpm counters are
incremented post-response, so the delta now adjusts just those. The itpm/otpm
counters are incremented at reservation time (pre-call), so the input/output
token headers already reflect the request and are left untouched - which
removes the need for the separate io-group special case.

* fix(router): roll back ITPM on any OTPM reservation error; dedup warning per instance

Two follow-ups from review. The pre-call OTPM reservation only rolled back the
ITPM reservation on a RateLimitError, so a transient cache error while reserving
OTPM left the ITPM counter inflated until the TTL expired; catch any exception,
release the ITPM reservation, then re-raise.

Replace the module-level lru_cache warn-once (caching a logging side effect,
which never re-warns in a long-lived process) with an instance-scoped set of
already-warned deployment ids on ModelRateLimitingCheck.

* fix(router): always clear reservation stash on reconcile; don't collapse id-less warning dedup

Clear the reservation in a finally block so a mid-reconciliation cache error
still removes the stash and a duplicate success event can't re-process it.

Dedup the itpm/otpm-vs-tpm/rpm conflict warning per real deployment id; a
deployment with no id no longer collapses every id-less deployment onto the
str(None) key (which would suppress all but the first warning).

* fix(router): skip io reservation when deployment can't be keyed

_get_cache_keys returned a shared 'global_router:None:None:...' key when a
deployment was missing model_info.id or litellm_params.model, so misconfigured
deployments could share one rate-limit bucket. Return None in that case and
skip io reservation for the request.

* fix(router): honor explicit max_tokens=0 in io reservation

_resolve_max_tokens used 'max_tokens or max_completion_tokens', so an explicit
max_tokens=0 fell through to the model default. Only fall back to
max_completion_tokens when max_tokens is absent.

* fix(ci): satisfy lint budget, router coverage, and dashboard schema sync

- Modernize the new itpm/otpm module's type hints to PEP 585 lowercase
  generics (Dict/Tuple/List -> dict/tuple/list) to clear the added UP006
  violations; ratchet ruff-strict-budget.json's UP006 ceiling down to match.
- Replace three try/except Exception blocks that must stay broad by design
  (token_counter and litellm.get_model_info raise untyped exceptions, and an
  io-token refund failure must never break the logging pipeline) with
  contextlib.suppress(Exception), matching the codebase's existing resolution
  for this exact BLE001 pattern.
- Add direct unit tests for get_model_group_io_token_usage (multi-deployment
  aggregation and the empty-model-list case) in test_router_helper_utils.py,
  satisfying the router function-coverage check.
- Regenerate the dashboard's schema.d.ts so the new itpm/otpm fields on
  GenericLiteLLMParams and ModelGroupInfo are reflected in the OpenAPI types.

* fix: enforce io token rate limits consistently

* fix: honor zero max tokens in otpm reservation

* fix(lint): fix UP007 violation and resync ruff-strict-budget.json to base

Convert Union[_Span, Any] to _Span | Any (safe on this repo's Python >=3.10
floor) to clear the new UP007 violation from the TYPE_CHECKING-gated Span
alias.

The previously committed ruff-strict-budget.json ratcheted UP006 down from a
stale base; litellm_internal_staging has since tightened that same ceiling
further on its own. Reset the file to the current base's committed values and
re-ratchet from there so the budget only ever moves down relative to the
actual merge-base, never against a stale snapshot.

* fix(router): attach ITPM/OTPM headers on dict responses and harden reservation

Strip itpm/otpm from provider kwargs, ensure messages are available for ITPM
estimation, honor max_output_tokens on /v1/responses, and propagate rate-limit
headers through /v1/messages dict responses via _hidden_params.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(router): attach ITPM/OTPM headers to streaming /v1/messages responses

Wrap bare async iterators in HiddenParamsAsyncIteratorWrapper so
set_response_headers can attach rate-limit headers to streaming Anthropic
messages responses that lack a _hidden_params slot.

Co-authored-by: Cursor <cursoragent@cursor.com>

* style: ruff format add_retry_fallback_headers.py

Fix CI ruff format check failure on get_hidden_params_dict call site.

Co-authored-by: Cursor <cursoragent@cursor.com>

* refactor(router): extract set_response_headers helpers to fix C901 budget

Move header-attachment logic into add_retry_fallback_headers helpers so
set_response_headers stays under the strict complexity ceiling.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: keep IO token reservation when response usage is missing

Missing usage was reconciled as zero and fully refunded the pre-call
reservation, allowing limit bypass on repeated successful calls. Only
adjust counters when usage is resolved from the response or standard
logging fields; otherwise keep the reservation until TTL expires.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: enforce RPM/TPM alongside IO-token limits on mixed deployments

Deployments with both itpm/otpm and tpm/rpm previously returned after the
IO reservation and skipped RPM/TPM checks. Run both paths and refund the
IO reservation only when RPM/TPM rejects after a successful reservation.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: track TPM usage on success for mixed IO+TPM deployments

The early return after IO-token reconciliation in log_success_event and
async_log_success_event skipped the TPM counter increment, so the tpm_key
the pre-call check reads was never written and tpm_limit was never
actually enforced on deployments that also configure itpm/otpm.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: treat total-only usage as unresolved in IO-token reconcile

usage/standard_logging_object entries carrying only total_tokens (no
prompt/completion or input/output breakdown) were treated as resolved
usage, resolving to (0, 0) and refunding the full reservation. Both
_usage_is_present and the standard_logging_object fallback now require an
actual input/output breakdown before reconciling, keeping the reservation
otherwise.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: reserve minimal token when input/output estimation fails

_reservation_value(0, limit) reserved the entire limit whenever token
estimation failed (empty/unsupported input, tokenizer error), letting one
such request claim the whole bucket and 429 every concurrent request to
the deployment until it completed. Reserve 1 token instead so estimation
failures no longer serialize traffic.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: refund IO reservation synchronously before retry deployment pick

On retry, set_io_token_rate_limit_request_kwargs clears reservation
sentinels from the shared kwargs dict before a background failure handler
can refund them, stranding the counter until TTL. Refund and clear any
stale reservation in _update_kwargs_with_deployment before stripping
sentinels for the next attempt.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(io_token_rate_limit_check): use model-specific tokenizer for ITPM estimate; document sync-refund Redis ceiling

Pass the deployment litellm_params.model to token_counter so it uses the
model's native tokenizer instead of the generic fallback, narrowing the
reservation over/under-estimate window between pre-call and post-call
reconcile.

Add a ponytail: comment to refund_stale_reservation_before_retry explaining
the known ceiling: the synchronous DualCache.increment_cache issues a
blocking Redis INCR when a Redis backend is configured. This only fires on
streaming mid-stream retries (non-streaming failures await their failure
handler before the retry picks a new deployment, leaving no sentinels to
refund). Upgrade path: make _update_kwargs_with_deployment async.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-05 21:58:35 +05:30
..
agent_tests feat(a2a): support a2a-sdk 1.x proxy routing for 0.3 and 1.0 agents (#30950) 2026-06-29 09:32:39 +05:30
audio_tests fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) 2026-05-18 09:15:39 -07:00
basic_proxy_startup_tests build: migrate packaging, CI, and Docker from Poetry to uv (#25007) 2026-04-09 11:46:23 -07:00
batches_tests test(e2e): add live batches suite across providers and routing scenarios (#30958) 2026-07-02 08:05:23 -07:00
benchmarks test(benchmarks): add CodSpeed benchmarks for inference, MCP and A2A hot paths (#31716) 2026-06-30 10:27:12 -07:00
code_coverage_tests fix(otel): hashable scope for _emit_once when guardrail_mode is list (#31262) 2026-06-25 11:59:35 -07:00
documentation_tests feat: add Rust OCR providers (#31272) 2026-06-25 15:12:30 -07:00
e2e Merge pull request #32166 from BerriAI/litellm_e2e_batches_ocr_model_registration 2026-07-05 02:15:10 +00:00
enterprise feat(prometheus): add api_provider label to token, latency, request and cache metrics (#32126) 2026-07-04 15:16:08 -07:00
guardrails_tests feat: litellm oss 110626 (#30202) 2026-06-11 22:30:26 -07:00
image_gen_tests test(bedrock): switch image gen live test off EOL Titan to Nova Canvas (#31937) 2026-07-01 22:57:29 -07:00
integration feat(sandbox): e2b code execution primitive (#30898) 2026-06-20 16:30:01 -07:00
litellm chore: litellm oss staging (#31185) 2026-06-26 09:17:44 -07:00
litellm-proxy-extras style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
litellm_core_utils Merge branch 'litellm_internal_staging' into litellm_staging_03_22_2026 2026-04-20 19:56:00 +05:30
litellm_utils_tests feat(fireworks_ai): sync chat completions endpoint with full API surface (#30885) 2026-06-20 19:49:07 -07:00
llm_responses_api_testing test(responses): bump deprecated gemini-3-pro-preview to gemini-3.1-pro-preview (#29433) 2026-06-01 09:54:30 -07:00
llm_translation feat(guardrails/headroom): add CCR (compress-cache-retrieve) via agentic loop (#31681) 2026-06-30 19:19:34 -07:00
load_tests style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
local_testing fix(bedrock): drop strict/additionalProperties from toolSpec for Claude Sonnet 4 (#31943) 2026-07-01 23:56:25 -07:00
logging_callback_tests fix(proxy): restore admin key/team callback_vars.turn_off_message_logging override (LIT-3587) (#31905) 2026-07-02 15:01:58 -07:00
mcp_tests feat(mcp): make token_exchange (OBO) production-ready - discovery threading + audit hardening + RFC 9728 challenge (#31622) 2026-07-03 17:12:25 -07:00
multi_instance_e2e_tests fix: use fastuuid helper (#14903) 2025-09-25 15:47:01 -07:00
ocr_tests fix(azure_ai): preserve content, tables, and keyValuePairs in doc-intelligence /v1/ocr (#32018) 2026-07-02 22:38:36 -07:00
old_proxy_tests/tests fix: cleanup tests 2026-03-30 16:24:35 -07:00
openai_endpoints_tests chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
otel_tests fix(cli): mint per-session agent credential on lite login (#31072) 2026-06-26 09:05:15 -07:00
pass_through_tests test(pass-through): de-flake vertex spend-log test by routing through the proxy (#31689) 2026-06-30 15:27:48 -07:00
pass_through_unit_tests test(proxy): set status_code on chunk_processor response mocks 2026-07-04 12:49:54 -07:00
proxy_admin_ui_tests Revert "chore: remove _experimental/out (#31546)" 2026-07-01 13:25:47 -07:00
proxy_behavior feat(proxy): add AES-256-GCM at-rest credential encryption with versioned format and re-encryption migration (#31215) 2026-06-29 20:14:22 +02:00
proxy_e2e_anthropic_messages_tests Extend the record/replay proxy to chat, embeddings, moderations, rerank, and Anthropic (#29847) 2026-06-06 14:33:42 -07:00
proxy_migration_tests test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700) 2026-06-04 14:56:02 -07:00
proxy_security_tests test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700) 2026-06-04 14:56:02 -07:00
proxy_unit_tests feat(proxy): add key-level budget_fallbacks to reroute requests when a per-model budget is exceeded (#31783) 2026-07-03 12:20:12 -07:00
router_unit_tests feat(router): add separate ITPM/OTPM deployment rate limits (#31952) 2026-07-05 21:58:35 +05:30
scim_tests
search_tests feat(tinyfish): make search provider permissive, attribute errors (#31997) 2026-07-03 10:17:11 -07:00
spend_tracking_tests chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
store_model_in_db_tests test: de-flake langfuse callbacks-in-db e2e test 2026-07-04 19:14:53 -07:00
test_litellm feat(router): add separate ITPM/OTPM deployment rate limits (#31952) 2026-07-05 21:58:35 +05:30
unified_google_tests fix(google_genai): preserve complete SSE events in Vertex/Gemini image streaming (#30270) 2026-06-12 07:49:30 -07:00
vector_store_tests Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) 2026-05-30 11:26:24 -07:00
windows_tests ci: reproduce default-Windows wheel install to guard MAX_PATH (#29597) 2026-06-03 11:28:08 -07:00
__init__.py
_fake_openai_endpoint_server.py test: point router/completion/triton tests at the local fake OpenAI endpoint (#30900) 2026-06-20 16:20:35 -07:00
_flush_vcr_cache.py tests(vcr): isolate cassette redis to CASSETTE_REDIS_URL 2026-05-01 12:32:59 -07:00
_live_test_helpers.py test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
_openai_record_replay_proxy.py Extend the record/replay proxy to chat, embeddings, moderations, rerank, and Anthropic (#29847) 2026-06-06 14:33:42 -07:00
_vcr_conftest_common.py test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
_vcr_redis_persister.py test(vcr): stop refreshing cassette TTL on read so cassettes lapse after 24h (#29784) 2026-06-05 10:22:41 -07:00
eval_swe_bench.py Prompt Compression - add it to the proxy (#25729) 2026-04-20 15:08:00 -07:00
fake_openai_endpoint.py test: point router/completion/triton tests at the local fake OpenAI endpoint (#30900) 2026-06-20 16:20:35 -07:00
gettysburg.wav
large_text.py
openai_batch_completions.jsonl
pyrightconfig.json test: add e2e tests for spend, budgets and llms (#30869) 2026-06-24 15:01:57 -07:00
README.MD [Feat] MCP Gateway Fine-grained Tools Addition (#15153) 2025-10-03 10:16:29 -07:00
test_anthropic_compaction_usage.py chore(oss): litellm oss staging 150626 (#30463) 2026-06-16 12:06:41 -07:00
test_budget_management.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_callbacks_on_proxy.py test(callbacks): harden flaky proxy callback-leak detector (#28195) 2026-05-18 16:39:02 -07:00
test_config.py fix: use fastuuid helper (#14903) 2025-09-25 15:47:01 -07:00
test_debug_warning.py fix(utils.py): fix togetherai streaming cost calculation 2024-08-01 15:03:08 -07:00
test_default_encoding_non_root.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_end_users.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_entrypoint.py
test_fallbacks.py Revert "fix: prevent error when max_fallbacks exceeds available models (#20071)" 2026-02-03 15:15:30 +05:30
test_gpt5_azure_temperature_support.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_health.py fix(tests): swap dall-e to gpt-image-1 after openai deprecation 2026-05-12 16:55:18 -07:00
test_keys.py test: harden remaining pass-through CI flakes (image-gen spend poll, ruby assistants timeout) (#30685) 2026-06-17 14:35:47 -07:00
test_litellm_proxy_responses_config.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_logging.conf
test_models.py test: replace test_add_and_delete_models integration test with mock 2026-03-30 21:30:57 -07:00
test_new_vector_store_endpoints.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_openai_endpoints.py Extend the record/replay proxy to chat, embeddings, moderations, rerank, and Anthropic (#29847) 2026-06-06 14:33:42 -07:00
test_organizations.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_otel_thread_leak.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_passthrough_endpoints.py fix: update authorization header to use 'Bearer' instead of 'bearer' 2025-09-21 10:44:47 +00:00
test_presidio_latency.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_proxy_server_non_root.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_ratelimit.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_resource_cleanup.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_service_logger_otel.py fix(langfuse_otel): prevent empty proxy request spans from being sent to Langfuse 2026-01-28 15:35:35 +01:00
test_spend_logs.py Litellm oss staging 04 21 2026 2 (#26569) 2026-05-20 21:25:19 -07:00
test_team.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_team_logging.py test: cleanup dead tests 2026-03-28 20:49:02 -07:00
test_team_members.py Litellm oss staging 04 21 2026 2 (#26569) 2026-05-20 21:25:19 -07:00
test_users.py Fix: tag budget reset must drop stale management-cache entry (#27568) 2026-05-10 00:18:55 +00:00

In total litellm runs 1000+ tests

[02/20/2025] Update:

To make it easier to contribute and map what behavior is tested,

we've started mapping the litellm directory in tests/test_litellm

This folder can only run mock tests.