litellm/tests/test_litellm/proxy/hooks
Yassin Kortam 561b6796bc
fix(proxy): enforce max_parallel_requests as a per-slot concurrency gauge (#32441)
* fix(proxy): enforce max_parallel_requests as a per-slot concurrency gauge

The v3 rate limiter tracked max_parallel_requests with the same
sliding-window machinery as RPM/TPM. A concurrency gauge cannot live on a
windowed counter: every window roll reset the counter to 1 while requests
were still in flight, the completion decrements for those forgotten
requests then drove the counter negative, and rejected requests left
stranded increments that nothing released. Under sustained load a key with
max_parallel_requests=5 let backend concurrency climb to the full client
concurrency (observed 60 on a live proxy) while the proxy kept returning
429s for everyone else

Replace the windowed counter with a per-slot registry (Redis sorted set of
slot ids scored by acquire time, with an asyncio-locked in-memory fallback):
admission atomically prunes expired slots and registers a new slot id only
when in_flight + 1 <= limit, so rejected requests never occupy a slot;
success, failure, and client-disconnect paths release exactly the slot id
this request acquired (stashed in the request metadata channels), so a
release without a matching acquire or a double-fired callback can never
free another request's slot; and a slot leaked by a crashed worker is
pruned individually after its TTL even under continuous traffic

Resolves LIT-4259
Fixes #16011

* fix(proxy): release every acquired gauge and respect mirrored counts in the in-memory fallback

Address review findings on the slot-registry gauge: the acquisition stash
now carries the gauge counter keys alongside the slot id, so the release
paths free the slot from every gauge it was registered under instead of
hardcoding the api_key scope, and the disconnect release keys off the
stashed acquisition instead of the key object's current
max_parallel_requests configuration (which can change mid-request). The
in-memory fallback now treats a cached integer (the count mirrored from
the last successful Redis script call) as real occupancy, carrying it
forward as a floored counter during a Redis outage instead of restarting
from an empty registry

* fix(proxy): release the parallel slot on proxy-level rejections

async_post_call_failure_hook is the only callback that fires when a
downstream hook (guardrail, budget check) rejects a request after the rate
limiter's pre-call hook acquired a slot; async_log_failure_event is a
completion-level callback and never runs for proxy-side rejections.
Release the stashed acquisition at the top of the hook, before the TPM
reservation guard, so those slots do not linger for the full slot TTL and
wedge the key at its limit under moderate rejection rates. Clearing the
acquisition marker keeps the release idempotent when a later failure
callback runs in the same flow

* test(proxy): cover success release, read-only count, Redis release mirror, and TPM rejection release

Four behaviors of the slot-registry gauge had no direct test: a successful
completion releasing exactly its acquired slot, read_only callers counting
in-flight slots through the count script (and degrading to the local
mirror when the script fails) without acquiring, the Redis release script
mirroring returned counts into the local cache, and the TPM reservation
rejection releasing the already-acquired slot before raising

* style(proxy): use builtin generics and union syntax in new rate limiter annotations

The slot-gauge code added Tuple/List/Dict and Optional[...] annotations, pushing
the UP006 and UP045 strict-rule totals past their ceilings in ruff-strict-budget.json.
Convert only the annotations this branch introduces to builtin generics and PEP 604
unions, leaving the rest of the module untouched.
2026-07-17 09:29:08 -07:00
..
litellm_skills Litellm oss 090626 (#30021) 2026-06-10 10:34:07 -07:00
test_async_post_call_streaming_iterator_hook.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_batch_file_validation.py revert: undo teamless all-team-models denial from #32022 and #29746 (#32032) 2026-07-03 18:56:07 -07:00
test_dynamic_rate_limiter.py fix(dynamic_rate_limiter): inject clock so active-project window is stable within a request (#32299) 2026-07-06 18:12:47 -07:00
test_dynamic_rate_limiter_v3.py Include model name + configured TPM/RPM in priority rate-limit 429 er… (#27216) 2026-05-05 19:05:22 -07:00
test_image_generation_guardrails.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_key_management_event_hooks.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_max_budget_limiter.py [Fix] Proxy: Skip Personal Budget Hook When Reservation Covers Counter 2026-05-01 15:57:42 -07:00
test_max_budget_per_session_limiter.py Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) 2026-03-07 19:12:42 -08:00
test_max_iterations_limiter.py Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) 2026-03-07 19:12:42 -08:00
test_parallel_request_limiter.py Litellm oss 090626 (#30021) 2026-06-10 10:34:07 -07:00
test_parallel_request_limiter_v3.py fix(proxy): enforce max_parallel_requests as a per-slot concurrency gauge (#32441) 2026-07-17 09:29:08 -07:00
test_post_call_failure_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_post_call_response_headers_hook.py fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232) 2026-06-09 22:10:23 +02:00
test_post_call_streaming_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_post_call_success_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_prompt_injection_detection.py fix: prompt injection not working (#16701) 2025-11-17 20:04:57 -08:00
test_proxy_hooks_init.py [Fix] Proxy: Address Greptile feedback on hook-cycle PR 2026-05-04 20:13:31 -07:00
test_proxy_rate_limit_provider_field.py feat: standardize rate limit errors with category, rate_limit_type, model, and llm_provider fields (#27687) 2026-06-06 17:50:29 -07:00
test_proxy_track_cost_callback.py fix(proxy): track unauthenticated pass-through requests in spend logs (#32410) 2026-07-13 13:39:38 -04:00
test_rate_limiter_toctou.py fix(rate-limit): fail closed on unrecognized OVER_LIMIT descriptor 2026-05-01 12:19:43 -07:00
test_send_invite_email.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_sensitive_data_routing.py [internal copy of #29511] feat(guardrails): add sensitive data routing to on-premise models (#29531) 2026-06-04 22:22:28 -07:00
test_tpm_concurrent.py fix(v3 limiter): cap no-max_tokens TPM floor at smallest configured limit (#28805) 2026-05-30 19:36:04 -07:00