fix(auth): apply temp_budget_increase for cache-hit keys

[devin-ai-integration]

Relevant issues

Fixes #25760

Linear ticket

Resolves LIT-4575

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

Live proof against a local proxy (Postgres + Redis key-auth cache enabled, real anthropic/claude-haiku-4-5 calls costing real $). The same script is run twice against the same proxy config, once on pre-fix code and once with the fix, so the only variable is the one-line change. Unique prompts are used because identical prompts get served from the LLM response cache at $0 and never accumulate spend, which masks the repro. Spend propagation to the enforcement layer is async, so the block shows up after a short stream of requests rather than on the first over-budget one

Commands (same for both runs):

MASTER=sk-1234; BASE=http://localhost:4000

# 1. key with a low max_budget
KEY=$(curl -s -X POST $BASE/key/generate -H "Authorization: Bearer $MASTER" \
  -H "Content-Type: application/json" \
  -d '{"max_budget": 0.0001, "models": ["anthropic-haiku-4-5"]}' \
  | python3 -c "import sys,json;print(json.load(sys.stdin)['key'])")

# 2. grant temp_budget_increase = 100  => effective budget 100.0001
curl -s -X POST $BASE/key/update -H "Authorization: Bearer $MASTER" \
  -H "Content-Type: application/json" \
  -d "{\"key\": \"$KEY\", \"temp_budget_increase\": 100.0, \"temp_budget_expiry\": \"2027-01-01T00:00:00\"}"

# 3. 12 UNIQUE chat/completions; spend crosses the original 0.0001 but stays under 100.0001
for i in $(seq 1 12); do
  curl -s -X POST $BASE/v1/chat/completions -H "Authorization: Bearer $KEY" \
    -H "Content-Type: application/json" \
    -d "{\"model\":\"anthropic-haiku-4-5\",\"messages\":[{\"role\":\"user\",\"content\":\"unique fun fact $i (run $RANDOM)\"}],\"max_tokens\":60}"
done

Before the fix (commit fdf380d0e3, the fix reverted): once spend crosses the original budget, cache-hit requests are wrongly blocked at the original Max budget: 0.0001 even though current cost is far under the effective 100.0001

Request 1: OK (allowed)
Request 2: BLOCKED -> Budget has been exceeded! Current cost: 0.000329, Max budget: 0.0001
Request 3: BLOCKED -> Budget has been exceeded! Current cost: 0.000329, Max budget: 0.0001
...
Request 12: BLOCKED -> Budget has been exceeded! Current cost: 0.000329, Max budget: 0.0001
key/info: spend = 0.000329 | max_budget = 0.0001
RESULT: BUG PRESENT -> 11 cache-hit request(s) wrongly blocked

After the fix (commit af1ca72f1c): every request is allowed, spend reaches 0.002238 (over 22x the original 0.0001) while staying under the effective 100.0001

Request 1: OK (allowed)
Request 2: OK (allowed)
...
Request 12: OK (allowed)
key/info: spend = 0.002238 | max_budget = 0.0001
RESULT: FIX WORKING -> no cache-hit request was wrongly blocked

Type

🐛 Bug Fix

Changes

temp_budget_increase (a temporary budget bump stored in a key's metadata) was only applied on the DB-fetch branch of _user_api_key_auth_builder. When the same key was served from the auth cache, _update_key_budget_with_temp_budget_increase never ran, so the effective max_budget reverted to the original low value and requests were wrongly blocked with BudgetExceededError once spend crossed the original budget while staying far under the effective budget

The fix moves the call out of the if valid_token is None: (DB-only) branch so it runs for every resolved token regardless of whether it came from cache or DB

        if valid_token is not None:
            valid_token = _update_key_budget_with_temp_budget_increase(valid_token)

This never double-applies: the cache stores the original budget (the increase is applied to the returned object, not the cached one) and each cache hit hands back a fresh model_copy(), so the increase is recomputed per request from the original max_budget

Regression test test_temp_budget_increase_applied_for_cached_key seeds the auth cache with a key whose spend (5.0) exceeds its original max_budget (2.0) but stays under the effective budget (2.0 + 100.0), then drives the cache-hit request through _user_api_key_auth_builder and asserts it is not blocked and the resolved token carries max_budget == 102.0. It fails on the pre-fix code (raises BudgetExceededError) and passes with the fix

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/aa9a3156545a4d39b7b3ea9c2d204c6a
Requested by: @shivamrawat1

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[greptile-apps]

Greptile Summary

This PR fixes a bug where temp_budget_increase was only applied when a key was fetched from the database, but not when it was served from the auth cache. A key with a temporary budget bump could be wrongly blocked once its spend crossed the original max_budget, even if it was well under the effective budget.

• The fix is a one-line move: _update_key_budget_with_temp_budget_increase is relocated from inside the if valid_token is None: (DB-only) block to a shared if valid_token is not None: guard that runs for every resolved token, whether it came from cache or DB.
• Double-application is not a risk: the cache layer always returns a shallow model_copy() via _copy_user_api_key_auth_for_cache, so the increase is recomputed against the original max_budget on each request without mutating the cached object.
• A new mock-only regression test (test_temp_budget_increase_applied_for_cached_key) seeds the auth cache with a key where spend (5.0) > max_budget (2.0) but spend < max_budget + temp_budget_increase (102.0), drives it through the builder, and asserts no exception is raised and result.max_budget == 102.0.

Confidence Score: 5/5

Safe to merge — the change is a minimal, targeted relocation of a single function call from a DB-only branch to a shared branch, with no new logic introduced.

The one-line move is correct: _update_key_budget_with_temp_budget_increase now runs for every resolved token regardless of source. The cache correctly returns a model_copy() each time, so repeated cache hits cannot accumulate the increase. The regression test directly exercises the exact failure scenario and passes only with the fix in place. No existing tests were weakened or removed.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/auth/user_api_key_auth.py	Moves _update_key_budget_with_temp_budget_increase outside the DB-only branch so it executes for both cache-hit and DB-fetch paths; the change is minimal and correct.
tests/test_litellm/proxy/auth/test_user_api_key_auth.py	Adds a new regression test that seeds the auth cache with a key whose spend exceeds its original budget but falls under the effective budget after the temp increase, then drives it through the builder and asserts no exception and the correct resolved max_budget; all mocks, no real network calls.

_{Reviews (1): Last reviewed commit: "fix(auth): apply temp_budget_increase fo..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_fix_temp_budget_increase_cached_key (af1ca72) with litellm_internal_staging (66dea7d)}