How to enforce LLM budgets with atomic Redis Lua scripts

There are two bugs almost everyone ships in their first LLM budget enforcer. One blows through the limit under concurrent traffic. The other permanently bricks a key after the first billing period. Wardin's gateway avoids both, and the fix for each turns out to be small — but only once you stop treating "check the budget" as application logic and push it into Redis as a single atomic operation.

This post walks through the actual implementation: a ~15-line Lua script, a two-phase reserve-then-reconcile flow, and a key-naming trick that replaces the monthly reset cron job entirely.

Bug 1: the read-check-write race

The naive version looks perfectly reasonable:

1. GET budget:key123:spent from Redis
2. Compare against the limit in application code
3. If under, INCRBYFLOAT the spend

Three steps, two round trips, one race. Between the read and the write, nothing stops another request from doing its own read. Two requests for the same key both see spent = $49.50 against a $50 limit, both decide they fit, both increment. The limit is now decorative.

For a web app with occasional API calls this is a theoretical concern. For an LLM gateway it's the normal case: agentic clients like Claude Code fire many requests per second against the same virtual key — a single "fix this test" task can fan out into dozens of near-simultaneous tool-call requests. If your budget check has any window between read and write, agentic traffic will find it in the first hour.

Mutexes in the gateway process don't help either, because a gateway that matters runs as multiple replicas. Instance-local state can't enforce a shared limit — the check has to happen where the state lives.

Redis executes Lua scripts atomically: nothing else touches the keyspace while a script runs. So collapse the read, the comparison, and the increment into one script and the race is gone by construction:

// checkScript atomically reads spent+limit and increments by estimated cost if allowed.
// Sets a 32-day TTL on the spent key on first write so it auto-expires after the period.
var checkScript = redis.NewScript(`
local spent = tonumber(redis.call('GET', KEYS[1]) or '0')
local limit = tonumber(redis.call('GET', KEYS[2]) or '0')
local cost  = tonumber(ARGV[1])
if limit > 0 and (spent + cost) > limit then
  return {0, tostring(limit - spent)}
end
redis.call('INCRBYFLOAT', KEYS[1], cost)
if redis.call('TTL', KEYS[1]) == -1 then
  redis.call('EXPIRE', KEYS[1], 2764800)
end
return {1, tostring(limit - spent - cost)}
`)

There is no moment where two requests both observe stale state, no matter how many gateway replicas are running. The script returns an allow/deny flag plus the remaining budget, so the caller gets an answer and a number in one round trip.

(Why tostring on the remaining value? Redis Lua truncates numbers to integers on the return path, and fractional dollars are the whole point here. Strings survive the trip intact.)

The problem Lua can't solve: you don't know the cost yet

Here's the part that makes LLM budgets genuinely different from rate limits: you don't know what a request costs until after it completes. The bill depends on completion tokens, and — for anything using provider prompt caching — on cache-read and cache-write tokens priced at different rates. None of that exists until the provider's response comes back.

But enforcement has to happen before the request goes out. Checking the budget after the response is just expensive bookkeeping.

The answer is the same one databases and payment systems use: reserve, then reconcile.

1. Check — before forwarding, atomically reserve an estimated cost (derived from things you know up front, like max_tokens and prompt length). If the reservation would exceed the limit, the request is blocked and never reaches the provider.
2. Commit — after the response, compute the real metered cost from the provider's usage data and apply only the delta between estimate and actual.

func (e *Enforcer) Check(ctx context.Context, keyID string, estimatedCost float64) (Result, error) {
	res, err := checkScript.Run(ctx, e.rdb, []string{spentKey(keyID), limitKey(keyID)},
		fmt.Sprintf("%.10f", estimatedCost)).Slice()
	if err != nil {
		return Result{}, err
	}
	allowed := res[0].(int64) == 1
	remaining, _ := res[1].(string)
	var remainingVal float64
	if remaining != "" {
		_, _ = fmt.Sscanf(remaining, "%f", &remainingVal)
	}
	return Result{Allowed: allowed, RemainingUSD: remainingVal}, nil
}
 
// Commit adjusts the reservation from Check to the real cost (actual - estimated delta).
func (e *Enforcer) Commit(ctx context.Context, keyID string, estimated, actual float64) error {
	delta := actual - estimated
	if delta == 0 {
		return nil
	}
	return e.rdb.IncrByFloat(ctx, spentKey(keyID), delta).Err()
}

If the request never reaches the provider — a downstream check blocks it, or the provider call fails — Commit runs with an actual cost of zero, which subtracts the full estimate and releases the reservation cleanly. This is why the enforcer has two methods instead of one: a single "check and charge exact cost" call can't exist when the exact cost is a future fact.

Note that Commit doesn't need the Lua script. The reservation already claimed the budget slot atomically; the commit is just a correction to a number, and INCRBYFLOAT is itself atomic.

Bug 2: the counter that never resets

The first version of a spend counter is usually one key per API key: budget:key123:spent, incremented forever. It works flawlessly through the first billing period — and then every key that hit its limit stays blocked permanently, because nothing ever resets the counter.

The reflex fix is a scheduled job that zeroes counters at month boundary. Now you have a cron job that can fail silently, and when it does, every key in the system is enforcing against last month's spend. The failure mode of the reset job is indistinguishable from correct enforcement until someone complains.

The fix that removes the reset step entirely: put the period in the key name.

// spentKey includes the UTC month so spend resets automatically each billing period.
// Format: budget:<keyID>:<YYYY-MM>:spent
func spentKey(keyID string) string {
	return "budget:" + keyID + ":" + time.Now().UTC().Format("2006-01") + ":spent"
}

On the first write of a new month, the key simply doesn't exist yet — the script's GET ... or '0' treats it as zero spend, and the fresh period starts automatically. The 32-day TTL the script sets on first write (EXPIRE ... 2764800) means last month's key garbage-collects itself. No job to run, no job to fail, no state to migrate at midnight UTC.

Where this sits, and what a denial looks like

This check runs in the request path, inside the gateway, in a fixed order: authenticate the virtual key → rate limits (RPM/TPM, a separate Redis-atomic limiter) → budget check → policy check (model allowlist) → prompt-injection guardrail → forward to the provider → parse actual token usage → commit real spend.

A denial returns HTTP 429 with a structured body:

{ "code": "BUDGET_EXCEEDED", "message": "monthly budget exhausted" }

Machine-readable on purpose. A meaningful share of gateway traffic is unattended — CI pipelines, Claude Code, Cursor — and there's no human there to read an HTML error page. A structured code lets the client script branch on it.

The ordering is the actual point. Most stacks bolt cost tracking on after the fact: an observability dashboard that alerts once the money is gone. In-path enforcement means "over budget" is a request that never went out — not a Monday-morning alert about the $4,000 an agent spent over the weekend.

Honest limitations

Two things worth knowing before you copy this design:

• Reservations can strand. Commit is a separate call after the response; if the gateway crashes between Check and Commit, the estimated reservation sticks until the period key expires. The error is in the conservative direction — you under-spend, never over — but it's real.
• Estimates can undershoot at the boundary. If actual cost exceeds the estimate on the request that lands right at the limit, the commit delta can push spend slightly past it. The overshoot is bounded by one request's estimation error, but it's not zero.

This is the budget enforcer running in Wardin's gateway today. The design question we keep turning over: when Redis itself is unreachable, should the budget gate fail open (availability) or fail closed (cost safety)? For interactive traffic the answer feels different than for unattended CI keys. If you've shipped in-path spend enforcement, I'd genuinely like to hear where you landed.