Enforcing an LLM budget under concurrency: 10 failure modes and their fixes

The obvious way to cap LLM spend is to read how much you've spent, compare it with the limit, and make the call if there's room. It works in a demo. In production, with parallel workers, streaming and retries, it breaks in specific, repeatable ways. This guide walks through each one and the fix that holds, all of which you can build yourself.

The shape every fix builds on: reserve → call → settle. You don't know the real cost before the call, so you reserve a worst-case estimate atomically, make the call only if the reservation fits, then swap the estimate for the actual cost when the bill is known.

1. The check-then-act race

Fifty workers all read spent = $9.80 against a $10 limit, all pass, and you overshoot by fifty calls. The check and the add have to be one atomic step. In Redis, a Lua script:

-- KEYS[1] = budget counter, ARGV[1] = estimate, ARGV[2] = limit
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
local est     = tonumber(ARGV[1])
if current + est > tonumber(ARGV[2]) then
  return redis.error_reply('budget_exceeded')
end
redis.call('INCRBYFLOAT', KEYS[1], est)
return 1

In Postgres, a conditional update, checking the affected row count:

UPDATE budgets
SET reserved = reserved + $1
WHERE id = $2 AND spent + reserved + $1 <= limit_amount
RETURNING id;   -- no row back = refuse the call

2. Floats drift

spent += 0.000015 a million times drifts in a 64-bit float. If you build this on a SQL ledger or in application code, store integers (micro-dollars, or tenths of a cent). Redis's own INCRBYFLOAT uses extended precision, so its drift is negligible at realistic call volumes; what matters more is that the durable ledger you reconcile against is exact.

3. Reservations that never come back

A worker reserves $0.40, then crashes or is OOM-killed before it settles. That money stays reserved, and over weeks the budget locks itself with no real spend behind it. Give every reservation an ID and an expiry (a Redis sorted set scored by expiry works), and run a sweeper that releases expired ones.

The subtle part: a reservation can expire while its call is still alive, and then settle late. Settle and sweep both “release” the same reservation, so they need a single winner. Make both start with the same conditional step:

-- Settle and sweeper both start with the same ZREM: exactly one wins.
local won = redis.call('ZREM', KEYS[2], ARGV[1])   -- the lease
local delta = tonumber(ARGV[3])                    -- actual cost
if won == 1 then delta = delta - tonumber(ARGV[2]) end  -- minus reserved
-- won == 0: the sweeper already released it; just add the real cost

If settle loses, the call still happened, so its actual cost is added without releasing the reservation a second time. Set the expiry well above your longest provider timeout so it only ever catches workers that are really gone.

4. Streams that get cut off

If a client disconnects halfway through a stream, the final usage message often never arrives. OpenAI only sends usage in the last chunk, and only when you ask for it with stream_options.include_usage. Settle at $0 and your ledger under-counts exactly during the incidents you care about. Charge the prompt estimate plus the output you actually streamed (roughly one token per four characters) instead. The test that catches this aborts the client after a few chunks and checks what the ledger recorded; the happy-path test always reads to the end.

5. Nested budgets that partly commit

A call usually has to fit a key, a team and an org limit at once. Reserve them one at a time and a failure at the second level leaves the first one reserved. Either check and reserve every level in one script (on a Redis cluster, put the keys on one slot with hash tags like {org:123}:team:7), or reserve in order and refund everything already reserved the moment a level refuses. In SQL, lock the rows in a fixed order (org → team → key) inside one transaction so two requests can't deadlock.

6. Period rollover

A call reserves at 23:59:58 against today's budget and settles at 00:00:03. If settle works out “today” from the current time, it charges tomorrow and leaves a reservation stuck in yesterday. Store the period key on the reservation and settle against that key, not the clock.

7. Timeouts with an unknown outcome

The naive version releases the reservation on any exception. But a read timeout after the request went out can mean the provider completed the call and billed it, and you simply never got the response. Release only when the request provably never left (a connection error, no pool slot). For anything after that, keep the worst-case charge and correct it later against the provider's usage or billing data. You're pessimistic for a while, which is the right way to be wrong for a spend cap.

8. Retries that open a second reservation

A client gets a 429 or a timeout and retries while the original call is still streaming. Now two live reservations exist for one user request. A replay cache doesn't help, because the original hasn't finished, so there's nothing to replay yet. Take an in-flight claim on the request's idempotency key (SET NX with a TTL longer than your longest stream) before reserving. Better still, hold a non-streaming retry until the original finishes and hand it the original's result, so the retry lands on the original's reservation without the client retrying again. Fall back to a 409 with Retry-After if the original is still running after the wait, or for streams, which can't be replayed.

9. Reserving the worst case locks the budget

Reserving full max_tokens on every call is safe in isolation. With 200 workers and a 4,096-token limit, most of a small budget sits in reservations that are never spent, and real calls get refused while actual spend is far under the cap. Two tweaks keep the safety:

• Switch modes near the cap. Reserve the p95 output length for that route while the budget is far from its limit, and the full worst case once the remaining budget drops inside a margin (a few dozen worst-case calls). Settle catches the rare long reply.
• Clamp the output. Lower each call's max_tokens to what the remaining budget can pay for, so no single reply can overshoot.

The reverse problem shows up near the limit: a call that would cost two cents gets refused because its 64,000-token worst case doesn't fit in the 50 cents left. Instead of refusing, you can shrink its max_tokens to fit, as long as enough output still fits to be useful. Keep that opt-in: a truncated answer is worse than a clean refusal for anything that parses JSON. And if callers don't set max_tokens at all, set it for them, because without an output limit there is no worst case to reserve.

10. The budget store goes down

Decide on purpose whether calls are let through or blocked. Blocking everything means a Redis blip takes your product down. Letting everything through means a loop during the outage is uncapped. The middle ground: let calls through, but give each budget a small dollar allowance that only applies while the store is unreachable, then refuse once it's used.

One limit is never enough

Each kind of limit catches a different leak. A daily cap per key catches a leaked key within hours; a monthly cap alone can be drained in an hour and still be “within limits”. A per-run cap catches a looping agent, and it works best when every response tells the agent how much of its run budget is left, so it can stop itself before the gateway has to. A monthly team cap catches slow drift. Watch spending speed too: “$40 in five minutes on a key that averages $2 an hour” is a better alarm than “80% of the monthly budget”.

How to test it

Set a $1.00 budget, fire 500 concurrent requests at a stub provider with randomized latency and output length, and kill 5% of workers mid-request. Then assert that:

• spend never exceeds the limit once everything has settled;
• reservations return to zero after the sweeper runs;
• no reservation is released twice, even when a settle arrives after the sweep.

The naive version fails the first check within seconds.

What Cognocient AI does

Cognocient AI's proxy implements each fix above in the request path, on every plan including Free. Block budgets reserve atomically across every matching level, each reservation carries a 15-minute lease with a sweeper, and calls settle against the period they were reserved in. Once a model has enough history, calls reserve their p95 output length until the budget gets close, and each call's output limit is clamped to what's left. Read timeouts charge the estimate (marked x-cog-billing-outcome: unknown). A retry with the same X-Idempotency-Key while the original is running waits for it and gets its result. Run calls return their remaining budget. Budgets reset daily or monthly, and each one chooses fail open, fail closed, or an outage allowance.

The details, including every response header, are in the budgets documentation. For the agent side, see stopping a runaway agent, and for why this has to live in the request path, see alerts vs. hard stops.