Commit graph

156 commits

Author SHA1 Message Date
yuneng-jiang
6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802)
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00
yuneng-jiang
693797420d
test: unwind environment writes in tests/test_litellm with monkeypatch (#37806)
* test: use monkeypatch.setenv for env writes in tests/test_litellm

`os.environ["X"] = v` inside a test leaks the value into every test that runs
after it in the same worker, so ordering decides the result. 262 of those
writes across 40 files now go through pytest's `monkeypatch` fixture, which
restores the previous value at teardown.

The rewrite skips any test that a mock.patch-family decorator wraps, any test
with defaulted positional parameters, any test whose own name is called
directly elsewhere, and rebinds nothing inside nested defs, because in each of
those cases appending a fixture parameter changes what pytest or mock binds.

Ratchets the TQ004 ceiling from 768 to 506.

* fix(test): delete the key through monkeypatch instead of popping it first

Five tests popped a key straight out of `os.environ`, ran, then restored it with
`monkeypatch.setenv`. By the time monkeypatch saw the name it was already gone,
so it recorded "absent" as the value to go back to and deleted the key at
teardown. On a worker that inherited a real `RESEND_API_KEY`, `SENDGRID_API_KEY`,
`UI_PASSWORD`, `LITELLM_SALT_KEY` or `OPENAI_API_KEY`, every test after the first
one ran without it.

`monkeypatch.delenv(..., raising=False)` removes the key and restores whatever
was there, so the try/finally the manual restore needed goes with it.

* chore(test): leave the two cost-calc files to the PR that rewrites them fully

Both files are also in #37815, which converts the module-global writes as well
as the env writes and folds them into one fixture. Two PRs rewriting the same
lines differently is a conflict nobody benefits from resolving, so this one
drops back to staging on those two and keeps the other 39.

TQ004 clears 200 here instead of 275; the rest moves with #37815.
2026-08-21 20:28:37 -07:00
ryan-crabbe-berri
243ed4393d test: reject assertions on a caught error inside except (ruff PT017)
A test that asserts on the error inside its own except block passes when the
call stops raising, because nothing runs the handler. That is the exact case
the test exists to catch, so the regression lands green.

Rewrites all 111 such blocks into pytest.raises, which fails when the call
succeeds, and selects PT017 in ruff-tests.toml so no new one lands.
2026-08-21 13:35:08 -07:00
ryan-crabbe-berri
e9d40a8f73 test: enforce F811 so a duplicate definition cannot silently replace the first
A name bound twice keeps only the second binding. In `tests/` that is nearly
always a repeated import, harmless but misleading, and the same rule is what
catches the cases that are not harmless: a local that shadows an import the
module still calls, and a second `def test_x` that quietly replaces the first.

311 of the 344 sites were repeated imports and came out with ruff's own fix.
The remaining 33 needed a decision. Four modules imported a name they never
used because a local definition below already shadowed it. Two comprehensions
bound `call` over `unittest.mock.call`, which those modules import and use.
One test rebound the two module handles its nested reload closure had captured.
One class attribute shadowed an unused `status` import.

The load-test fixtures move to a conftest, which is how pytest is meant to share
them, so the test module no longer imports three fixture names it never calls.
The nine `prisma_client` parameters keep a narrow `noqa`: pytest resolves that
fixture by name before the body runs, so the parameter never shadows anything.
2026-08-21 12:06:19 -07:00
ryan-crabbe-berri
b76def0e5d
test: require a match= on broad pytest.raises, and drop duplicate parametrize cases (#37769)
`pytest.raises(Exception)` with no `match=` passes on any error that broad. A
TypeError from a refactor, a botched fixture, an import that moved: all of them
read as the rejection the test claims to police, so the test goes green for the
wrong reason and stays green after the behaviour it guards is gone.

PT011 closes that gap for the 317 sites B017 could not reach, because B017 only
fires on a single-statement body with no `as e` binding. Each pattern here is the
message the code actually raised, recorded by running the sites under a plugin
that logged the concrete type and text per call site, so the assertions describe
observed behaviour rather than a guess. Where a site raises more than one message
across its parametrize cases, the pattern is an alternation of what was seen;
where the exception carries an empty `str()` and puts the text on `.message`, the
site keeps a narrow `noqa` with the reason.

PT014 removes four parametrize cases that were listed twice. The duplicate re-runs
an assertion that already passed, and it usually marks a case someone meant to
vary and forgot to edit.
2026-08-20 20:24:49 -07:00
mateo-berri
d5ac49588a Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_batch_enqueued_token_limit 2026-08-19 19:31:46 -07:00
hiraku-miyoshi
2a771caf02 fix(proxy): clamp reservation record TTL so stale records never outlive their counters 2026-08-19 19:12:22 -07:00
devin-ai-integration[bot]
3a04860122
feat(proxy)!: default audit logs on for enterprise licenses (#37518)
* feat(proxy): enable audit logs by premium license

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): support premium audit logging mocks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): disable audit logging for key rotation mocks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-08-19 18:49:21 -07:00
mateo-berri
5ab20c3678 fix(batch_enqueued_tokens): tombstone popped Redis reservation records so local ghosts cannot double-refund 2026-08-19 17:47:06 -07:00
mateo-berri
504112d5ca fix(batch_enqueued_tokens): keep the over-limit verdict on rollback failure, find locally saved records on pop
A Redis over-limit verdict now survives a failing rollback DECRBY instead of
escaping into the in-memory fallback and granting tokens the counter already
rejected; the unrolled increments expire with the TTL. pop_reservation now
falls through to the local record when the Redis pop succeeds but finds
nothing, so a reservation saved in memory after a transient Redis save
failure still refunds on cancel or completion.
2026-08-19 16:52:46 -07:00
mateo-berri
4333d52813 fix(batch_enqueued_tokens): scope in-memory refunds to the granting worker
In-memory grants now record an owner token, and a refund only debits local
counters when the popping worker is the one that granted them, so a terminal
response handled elsewhere can no longer shrink another worker's unrelated
fallback reservations. A Redis-granted refund that fails no longer falls back
to decrementing local counters either: the leaked Redis increments expire
with the TTL and only tighten the allowance.
2026-08-19 16:36:27 -07:00
mateo-berri
50896f21b3 fix(batch_enqueued_tokens): roll back partial reserves, route refunds by backend, lowercase terminal statuses
Reserve-script failures now roll back the scopes already incremented before
re-raising into the in-memory fallback, so a partial redis outage no longer
leaks counter increments that shrink the shared allowance. Reservations
record which backend granted them, so a refund never debits redis counters
an in-memory grant did not charge. Terminal-status matching is now
case-insensitive because the Bedrock async-invoke retrieve path returns raw
AWS-cased statuses like Completed.
2026-08-19 16:16:08 -07:00
mateo-berri
160d3dac42 fix(proxy): issue enqueued-token Lua calls one key at a time for Redis Cluster compatibility 2026-08-19 15:40:15 -07:00
mateo-berri
7a6a677b72 feat(proxy): enqueued-token rate limiting for batches with refund on completion and cancellation 2026-08-19 15:09:05 -07:00
mateo-berri
608d749983 fix(batches): stop one bad output line from zeroing an entire batch's spend 2026-08-18 20:46:59 -07:00
mateo-berri
c435c25da2 Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_pr35110_itpm_otpm
# Conflicts:
#	type-discipline-budget.json
2026-08-18 16:11:12 -07:00
mateo-berri
3c34c34459 fix(proxy): guard candidate-count and batch cap coercion against float overflow 2026-08-18 15:42:02 -07:00
mateo-berri
69ea1c6599 fix(proxy): coerce batch candidate counts like the live limiter path 2026-08-18 15:04:18 -07:00
mateo-berri
72960d10e9 fix(proxy): address review findings on project ITPM/OTPM quotas
- scale batch output-token reservations by the row's n / best_of candidate count
- parse client-supplied output caps defensively instead of 500ing on unparseable values
- exclude project IO descriptors from the first should_rate_limit pass when TPM
  reservation is disabled so their buckets are not double-charged
2026-08-18 14:45:05 -07:00
mateo-berri
be594f5984 feat(guardrails): count bedrock guardrail cost against spend and budgets
Price ApplyGuardrail usage units recorded by PR #37225 with a new
bedrock/guardrails entry in the model cost map (regional override via
bedrock/{region}/guardrails), add the per-request guardrail_cost to the
standard logging payload's response_cost and CostBreakdown, surface it in
the x-litellm-response-cost header, and bill blocked requests through the
failure hook so key and team budgets see what AWS bills
2026-08-18 14:16:07 -07:00
mateo-berri
c1c23bf39a fix(proxy): reserve measured input tokens for multimodal project ITPM
Image, file, video, and previous_response_id requests reserved the whole
project ITPM limit up front, so any window with existing usage rejected
them and one in-flight multimodal request blocked the entire project.
Reserve the token_counter estimate instead, like every other request;
post-call reconciliation already charges actual usage.
2026-08-18 13:51:47 -07:00
mateo-berri
842bb7e917 Merge branch 'litellm_internal_staging' into feature/bedrock-mantle-quota-project-itr1 2026-08-18 13:14:14 -07:00
mateo-berri
55e80849d1 feat(guardrails): track bedrock guardrail usage units per invocation 2026-08-17 15:18:53 -07:00
Shivi Jain
a6dc447470 fix(proxy): stop Responses batch rows from bypassing project OTPM
Embeddings rows were identified by body shape (has `input`, no
`messages`/`prompt`), which also matches a `/v1/responses` batch row
and reserved zero output tokens for it -- letting a project caller run
large Responses generations against a quota-limited model without
consuming OTPM. Classify embeddings by the row's own `url` instead,
and read `max_output_tokens` as a Responses output cap alongside
`max_tokens`/`max_completion_tokens`.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-16 16:26:30 +05:30
Shivi Jain
8a44c14928 Merge upstream/litellm_internal_staging and fix batch quota review comments
Resolves conflicts from the upstream merge and addresses the Veria-AI
review comment on this PR: batch rows could bypass a project's
per-model ITPM/OTPM quota when the batch's file-bound/routing model
had no quota configured. Charges each row's own model against its own
project quota instead of only the routing model's, and fixes rate
limit error messages to attribute the correct model via a new
descriptor_value field on RateLimitStatus/AtomicCounterMeta. Also
re-syncs the ruff-strict, type-discipline, and basedpyright budgets
against the correct (non-stale) merge base.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-16 15:47:46 +05:30
Yassin Kortam
33e9f54dc8
fix(proxy): reserve the larger declared output budget for TPM limits (#37001)
The TPM pre-call reservation read `max_tokens or max_completion_tokens`, so a
request declaring both was charged for whichever field came first. A caller
sending `max_tokens=1` with `max_completion_tokens=10000` reserved 2 tokens and
was then free to consume ten thousand, since the provider honours the modern
field and litellm's own param mapping drops the legacy one for the gpt-5 and
o-series families.

Reserve against the larger of the declared budgets instead. Over-reserving is
the safe direction for a limiter: post-call reconciliation refunds the
difference between the reservation and actual usage, while under-reserving lets
the window be exceeded before anything notices.
2026-08-15 11:54:22 -07:00
Shivi Jain
2b23295f82 fix(proxy): reconcile project quota reservations 2026-08-14 23:31:18 +05:30
Shivi Jain
312d12fe0c fix(proxy): enforce project ITPM/OTPM quota on every Responses WebSocket frame
The connection-level pre-call hook only ran once per WebSocket
connection, so a project caller could send unlimited high-token
response.create frames after a single minimal reservation. Adds
enforce_project_io_token_quota_for_frame to the v3 rate limiter and
wires it into both the native and managed WebSocket handlers via a
duck-typed litellm.callbacks lookup, so the SDK layer stays free of
proxy imports. A rejected frame gets an error event; the connection
stays open for the client to retry.

Also fixes the RET504 and BLE001 strict-lint-budget violations the
litellm_internal_staging merge introduced in
parallel_request_limiter_v3.py, which were failing the lint check.
2026-08-14 21:39:34 +05:30
Shivi Jain
a9227057a1 Resolving merge conflicts and verai comment for batch 2026-08-14 20:52:05 +05:30
tin-berri
d8fda675cc
feat: pre-adoption shadow eval for the auto-router (blind pairwise judge, derived state) (#36587) 2026-08-13 13:15:45 -07:00
mubashir1osmani
b0fac57fe4
fix(email): stop duplicate legacy invitation email and fix its onboarding link (#36455) 2026-08-10 23:26:27 -07:00
Yassin Kortam
ade805ef0c
feat(rate limiting): configurable estimated output tokens per key, team and model (#36143) 2026-08-10 12:51:14 -07:00
yucheng-berri
efc4e6f28c
fix(batches): keep batch state in sync on a poll without claiming attribution (#34456)
A poll of a Vertex passthrough batch wrote nothing to the managed-object row,
so status and file_object stayed frozen at the create-time snapshot and
GET /v1/batches served a stale status and an empty output file id for the life
of the batch. Only the create may claim a batch, but every observation of one
may refresh its state.

store_unified_object_id takes create_if_missing, which the poll clears: it
refreshes status and file_object through update_many, and leaves a row that is
absent absent rather than creating one owned by the observer, since created_by
and team_id are written by whoever reaches the create branch. The update payload
is now shared with the upsert so it cannot drift into writing api_key,
request_tags, created_by or team_id.

The passthrough identity re-assertion that was previously part of this PR ships
separately in #36121, so this PR keeps only the batch attribution work.

The creating key owns user_api_key_alias only when it actually has one. Guarding
the overwrite on the presence of a key rather than on a resolved alias nulled the
field out for every key generated without key_alias, and for any key rotated or
deleted before its batch finished, losing the creating user's alias that the spend
row previously carried. The guard now matches the team-alias line below it.
2026-08-08 16:01:47 -07:00
ryan-crabbe-berri
527dc0a8bb
feat(proxy): add apply_user_budget_to_team_keys opt-in (#36102)
* feat(proxy): add apply_user_budget_to_team_keys opt-in

PR #32005 made a user's personal max_budget apply to their team-scoped keys
too, and PR #35271 reverted the whole thing (behavior plus the
skip_user_budget_on_team_key opt-out) because that flipped the default for
everyone. This brings the behavior back the other way round: default is
unchanged, and general_settings.apply_user_budget_to_team_keys opts a
deployment into charging the key owner's personal budget on team keys.

The flag reaches all three personal-budget gates so an opted-in deployment
enforces consistently: the read-time check in common_checks, the optimistic
reservation counter in _get_budget_counters, and the _PROXY_MaxBudgetLimiter
pre-call hook. It is also in the /config/list allowed args and, unlike the
reverted flag, in the _update_general_settings propagation allowlist, so the
Admin UI General Settings toggle actually takes effect at runtime; an explicit
YAML value still wins over the DB value on reload.

get_config_list's allowed_args moves to a module-level frozen mapping of
field name to type string, dropping 18 LIT002 violations and rebuilding one
less dict per request.

* style(proxy): drop explanatory comments from the budget flag paths
2026-08-07 15:40:13 +00:00
Shivi Jain
9dbe61aa6d feat(proxy): add project-level ITPM and OTPM quotas
Add model_itpm_limit and model_otpm_limit to project create and update requests, storing both quota maps in project metadata without a database migration

Reserve input and output tokens independently before provider dispatch, expose separate project rate-limit headers, and reconcile counters across successful calls, failures, retries, fallbacks, streaming, caching, and cancellation

Harden token estimation for pre-tokenized embeddings, multimodal inputs, Responses API requests, native Gemini requests, multiple candidates, and conflicting output-cap aliases

Reject negative output caps, preserve conservative reservations when usage is missing or zero, bind reconciliation and refunds to the reservation window, prevent double refunds or negative counters, update generated API types, and add regression coverage
2026-08-05 12:53:22 +05:30
elinacse
833670f7db fix(batch): track cost for managed batches with no attributable key/user/team
LiteLLM_ManagedObjectTable only stores created_by (user_id) and team_id,
never the raw API key hash. A batch created with the master key or a
team-less key has both null, so CheckBatchCost's synthetic logging_obj
for the completed batch carried no attributable key/user/team/end-user.
_should_track_cost_callback silently skipped the DB write in that case
(by design, to avoid tracking truly anonymous requests), with no error
or warning: batch_processed still became true, but no LiteLLM_SpendLogs
row was ever written despite real, already-incurred provider cost.

Extend the same allowance already made for unauthenticated pass-through
requests to aretrieve_batch's cost event, and pass job.team_id through
so a batch's team gets real attribution when one exists.
2026-08-02 12:20:46 +05:30
Shivam Rawat
e204e629e0
Merge pull request #35422 from BerriAI/litellm_fix_tpm_only_dynamic_rate_limit
fix(rate-limit): enforce token limits when the pre-call increment is zero
2026-08-01 13:02:11 -07:00
Shivam Rawat
334805990c fix(rate-limit): block check-only counters at the limit and scope tpm reservation to tokens
Aligns the fix with the constraints in LIT-4800. A zero-increment
counter now blocks at current >= limit, matching RPM's semantics; the
previous current > limit let a pool sitting exactly at its reservation
admit one extra request. reserve_tpm_tokens rebuilds its descriptors
with only tokens_per_unit so the requests dimension stays out of the
reservation pass, which deliberately leaves RPM to the separate
should_rate_limit check.
2026-08-01 11:56:18 -07:00
Shivam Rawat
836bd927b6 fix(rate-limit): skip negative increments in the atomic payload builder
Review feedback: the relaxed predicate admitted negative increments,
which both atomic backends would apply as decrements. Restrict the new
behavior to zero-valued pure checks and assert negatives neither check
nor mutate counters.
2026-07-31 18:05:50 -07:00
Shivam Rawat
d640ace6d8 fix(rate-limit): enforce token limits when the pre-call increment is zero
The atomic check-and-increment path skipped any counter whose increment
was <= 0. The dynamic rate limiter always passes a zero token increment
pre-call because usage lands on the counters post-response, so on a model
configured with only tpm the limiter evaluated no counters at all: no
model-wide TPM cap and no priority reservation, in either generous or
strict mode. Regressed in dd57ae6691 when the pre-call flow moved off the
read-only should_rate_limit check, which did evaluate token limits.

Keep zero-increment counters in the payload so they act as a pure check
(current + 0 > limit), matching the pre-regression semantics in both the
Lua and in-memory paths. Adds unit regressions at the primitive and hook
level plus a live e2e covering the priority_generous/priority_strict
registry rows.
2026-07-31 18:05:50 -07:00
mateo-berri
26ace642be
test(proxy): cover the user-created audit hook's database read-back
The hook resolves the newly created user through UserRepository and builds the
audit entry from that row. Pin both halves: the entry carries the persisted
row's fields rather than the /user/new response, and a user id that resolves to
nothing produces no entry at all.
2026-07-31 19:48:30 +00:00
mateo-berri
f507a118af fix(rate-limits): pin the request stash to its owning litellm_call_id so nested calls cannot release it 2026-07-30 14:54:28 -07:00
mateo-berri
631c02fe12 refactor(rate-limits): move the v3 limiter per-request stash off request metadata onto a ContextVar
The v3 parallel-request limiter stashed its per-request bookkeeping (TPM
reservation, descriptors, parallel slot, rate-limit response snapshot,
released flag) in the request body's metadata channels. On routes where
metadata is a provider request parameter (Responses API and the other
LITELLM_METADATA_ROUTES) that leaked internal keys upstream and produced
HTTP 400s, and it required denylist stripping plus dual-channel writes to
contain.

The stash now lives on an asyncio ContextVar holding a single typed
RequestRateLimiterStash per request. The pre-call hook writes it, and the
success/failure callbacks, disconnect release, and post-call hooks read
and clear the same shared instance, which keeps the refund and slot
release idempotent across sibling callbacks. The request body is never
touched, so the stash-key stripping, the metadata mirror writes, and the
all_litellm_params denylist entries are removed
2026-07-30 14:00:20 -07:00
yucheng-berri
6cc136de90
fix(proxy): hash caller-supplied key in key update audit log object_id (#34632)
* fix(proxy): hash caller-supplied key in key update audit log object_id

* test: bound audit-log wait to the captured task instead of gathering the loop
2026-07-25 10:45:51 -07:00
Yuneng Jiang
9c48ad41ac
fix(passthrough): honor the zero fallback and aggregate-only TPM usage
Two follow-ups from review on the upstream-reported usage contract.

An unusable cost header fell through to the endpoint's flat cost_per_request
instead of the zero the contract promises, so a target that contradicted itself
got billed an estimate it had just disowned. A target that speaks this contract
now owns the cost for the request whether or not the value it sent parsed.

The reported total also cannot be split into prompt and completion, so reading
one out of it under token_rate_limit_type input or output yielded zero and left
the TPM window uncharged; pass-through traffic then ran past a limit it is
meant to share with the general API. Usage that carries no split now charges
its total under every limit type, while usage that does carry one is untouched.
2026-07-24 18:49:41 -07:00
Yuneng Jiang
838c7a7ea7
feat(passthrough): record upstream-reported cost and token usage
A pass-through target that fans a single HTTP request out to several models
internally cannot be priced from its response body, so LiteLLM had nothing to
record and every such request landed in the spend logs with zero cost and zero
tokens. The target now reports the totals for the whole request in
x-litellm-response-cost and x-litellm-total-tokens response headers, and
LiteLLM records those values as-is rather than recomputing them.

The headers are read on every upstream response, so a request that burned
tokens before failing still books its spend on the failure row instead of
being dropped for having a 4xx/5xx status. Only what the upstream actually
reported is written, so a target that sends a cost but no token count keeps
the token count LiteLLM derived on its own; a target that sends neither header
is untouched, which is the normal case for Anthropic, Vertex and friends.

Two supporting fixes fall out of this. The rate limiter only pulled token
counts off response shapes it models, so pass-through usage never charged the
TPM window and a team could exceed its shared token limit through pass-through
traffic alone; it now falls back to combined_usage_object. And the streaming
success path reset response_cost unconditionally before the assembled response
recomputed it, which discarded any cost a pass-through handler had already
established (the pass-through branch right below it has always intended to
preserve exactly that).
2026-07-24 18:10:15 -07:00
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
devin-ai-integration[bot]
8936d07be8
fix(proxy): track unauthenticated pass-through requests in spend logs (#32410)
Pass-through endpoints configured with auth=false reach the cost-tracking callback with no key/user/team/end-user, so _should_track_cost_callback returned False and the spend-log write was skipped, leaving the request out of request/usage logs. Track pass-through call types even when unauthenticated so the SpendLog row is still written.

Co-authored-by: Mubashir Osmani <mubashir@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-07-13 13:39:38 -04:00
yucheng-berri
2c1d62ce2b
fix(rate-limit-v3): populate x-ratelimit-* remaining/limit values in standard_logging_object for streaming (LIT-4333) (#32711)
Streaming requests return from common_request_processing before
async_post_call_success_hook runs, so response._hidden_params.additional_headers
never gets the v3 x-ratelimit-{descriptor_key}-{remaining|limit}-{rate_limit_type}
entries. Prometheus / logging callbacks that read those values from
standard_logging_object.hidden_params.additional_headers then see nothing;
combined with the pre-existing gap that Prometheus reads from that same slot
(LIT-2577 / PR #28816), per-key remaining RPM/TPM cannot be monitored for
streaming traffic at all.

Fix in three parts:

- Stash the pre-call RateLimitResponse in the metadata channels the async
  success-logging callback inherits, alongside the existing top-level entry
  the non-streaming path reads.
- Add async_logging_hook to the v3 handler. It fires in a distinct earlier
  loop inside async_success_handler (all callbacks' async_logging_hook
  complete before any async_log_success_event starts), so mirroring the
  pre-call snapshot into standard_logging_object.hidden_params.additional_headers
  and response._hidden_params.additional_headers here guarantees every
  downstream success callback sees the values regardless of registration
  order. Non-streaming keeps the existing async_post_call_success_hook write
  and this hook re-populates the same values idempotently.
- Extract the shared `_merge_ratelimit_statuses_into_additional_headers`
  helper the non-streaming path already had inlined so both callsites emit
  the identical key shape.
2026-07-11 12:28:38 -07:00
Yassin Kortam
bcd52754de
feat(rate_limit): support per-tag rpm limiting on a single key (#31502)
Add a tag_rpm_limit field to virtual keys so each request tag gets its own independent RPM counter on the v3 rate limiter. A key configured with per-tag limits tracks each tag/group separately, and requests whose tag has no configured limit fall back to the key-level limit. Includes the dashboard UI to manage per-tag limits on key create and edit.

Resolves LIT-3147
2026-07-08 09:43:47 +03:00