fix(router): reserve-then-settle spend caps to close concurrency gap

[rohith500]

Fixes #793.

Summary

Org monthly, user monthly, and api-key lifetime spend caps had a check-then-act race: CheckOrgMonthlySpend (a plain SELECT) followed by an unguarded write. Concurrent requests could all pass the same preflight check before any of their debits landed, allowing spend to blow past a configured "hard cap" by a multiple of the limit, not by a bounded amount. Independently reproduced three times against real Postgres and real production code before any implementation began (see #793 for the investigation, including exact reproduction numbers and the corrected reading of PR #769's actual claim).

This PR replaces the check-then-act gate with reserve-then-settle: an atomic conditional reservation at request start, settled (or released) at request end, backed by a TTL sweeper for crash safety.

Design

Core mechanism: a single combined ArmSpendReservations call, made once per request after identity resolves and before dispatch, reserves all applicable caps (org-monthly, api-key-lifetime, user-monthly) in one Postgres transaction. Each scope's reservation is a genuinely atomic conditional UPDATE:

UPDATE ... SET reserved = reserved + @amount
WHERE spent + reserved + @amount <= limit
RETURNING id

Zero rows returned means the cap would be exceeded, mapped to the appropriate Err*SpendLimitReached/Err*SpendCapReached sentinel. This is the crux of the fix: the bound check lives inside the write's own WHERE clause, so Postgres's row-level locking serializes concurrent attempts against the same row rather than two reads-then-writes racing each other.

Wrapping all three scopes in one transaction means a failure on any later scope (e.g. user-monthly fails after org+key succeeded) rolls back everything automatically via normal Postgres rollback, no manual multi-scope release logic needed, which is simpler and more clearly correct than the alternative (three independent reserve points) considered and rejected during design.

Settle/release/sweep share one idempotent primitive: DELETE FROM spend_reservations WHERE id = @id RETURNING ..., only decrementing the denormalized reserved_usd_micros counter when a row was actually returned. A second delete on an already-consumed reservation is a clean no-op, this is what makes settle, a defer-based release-on-failure, and the TTL sweeper all safe to race against each other.

TTL sweeper: every ~1 minute (not the existing session-pin sweeper's 1-hour interval, a stuck reservation has much worse consequences than a stale pin), TTL = request timeout + 5 minutes, following the existing runSessionPinSweep scheduling pattern (every-replica ticker, no leader election, confirmed this router runs multi-replica in managed deployments, so no in-process-only mechanism could be the real backstop).

Reservation amount (R): a fixed, configurable slot (ROUTER_SPEND_RESERVE_USD_MICROS, default $1), not derived from max_tokens (which would false-402 legitimate large-context requests).

Honest worst-case guarantee

With fixed reservation R, monthly limit, N concurrent request starts, and post-response actual cost per request:

admitted <= min(N, floor(limit / R))

worst-case overshoot <= admitted * max(0, actual - R)

Verified against the real implementation (delayed-settle fixture forcing the true worst case, not a favorable fast-settle timing): $10 limit, R=$1, actual=$6.75, N=20, admitted=10 (exactly floor(limit/R)), final spend $67.50, overshoot $57.50, an exact match to the theoretical bound. Against the original bug's unbounded result on the same fixture (N * actual = $135), this is real but modest (~2x) improvement in this specific fixture, because R=$1 is well below actual=$6.75. The structural fix, overshoot no longer scales with N, is what matters; the numerical tightness of the bound depends entirely on R being tuned close to real per-request cost. R=$1 is a v1 default and should be tuned from real p95 turn cost in production, not left as-is.

Shadow-eval interaction (discovered during rebase, not part of original design)

origin/main moved during this PR's development (#787, agent-shadow evaluation routing) and added an explicit skip for shadow-eval traffic to WithBalanceCheck, WithOrgMonthlySpendCap, and WithAPIKeySpendCap, synthetic evaluation requests must not mutate production billing state. Since this PR moves org-monthly and api-key enforcement out of the latter two middlewares into ArmSpendReservations, the rebase required adding the equivalent skip there too (internal/proxy/spend_reserve.go), without it, shadow-eval traffic would have silently started consuming real org spend budget on the real inference paths, regressing what #787 had just protected against. Covered by a new test (TestArmSpendReservations_AgentShadowEvalSkipsReserve) with positive and negative controls. /v1/route's check-only middleware (unaffected by this PR) retains its own independent shadow-eval skip.

Commits

1. test: permanent regression test reproducing the TOCTOU (confirmed failing against pre-fix code with the actual overshoot numbers)
2. feat(db): migration 0040_spend-reservations, new table + denormalized reserved_usd_micros columns
3. fix(router): the reserve-then-settle mechanism + combined arming at all four handler entry points (Anthropic, OpenAI chat, OpenAI responses, Gemini)
4. feat(router): TTL sweeper
5. style: gofmt
6. refactor: removed a dead conditional found during review
7. test: shadow-eval skip (post-rebase)
8. fix(db): single Go-computed month for spend reservations, fixing a Bugbot-flagged clock-skew edge case

Testing

• Permanent regression test: unbounded-before / bounded-after, exact numbers in commit history.
• Sweeper test: seeded expired reservation, confirmed cleanup + correct counter decrement.
• Shadow-eval skip test with controls.
• Full HTTP end-to-end verification against the local Docker stack (not just Postgres-level or unit tests): release path (pre-dispatch failure, reservation correctly released), gate path (limit at cap, 402, no upstream call attempted), settle path (one real completion, spent incremented by real notional cost, not R, confirmed via ledger), and concurrent HTTP path (N=20 real concurrent requests against a $3 limit, peak reservations = 3 = floor(limit/R), matching theory over real HTTP, not just in isolation).
• Both the original Postgres-only concurrency test and the sweeper test re-run and reconfirmed passing after the rebase onto feat: add isolated agent shadow routing #787, to make sure the merge didn't disturb correctness.
• Bugbot flagged a real clock-skew edge case (commit 8a73a4a), fixed and covered by a dedicated regression test (TestReserveSpendCaps_SingleGoComputedMonthAcrossEnsureBumpInsert) that freezes Go's clock to a different UTC month than Postgres's live NOW() to prove the fix.
• go test ./... -count=1, clean, zero regressions anywhere in the repo.
• make check, clean.

Known, deliberately unaddressed

• Compaction/handover summary billing (_compaction_summary) remains unreserved soft overshoot, pre-existing, separate from the caps this PR fixes.
• A request that legitimately runs longer than the TTL: its reservation gets swept while still in-flight; the real settle still records true spend correctly, but the temporary reservation bookkeeping was already cleared for that window, narrows, doesn't eliminate, protection during an edge case bounded by the TTL margin (5 min beyond request timeout).
• Process crash before the defer registers: TTL is the only backstop, irreducible.
• R=$1 default needs production tuning against real p95 request cost to tighten the worst-case bound described above, flagging explicitly rather than presenting the current default as final.

Related

#478 (prepaid balance TOCTOU, same bug class, different table/product surface, not yet fixed there), #564 (unrelated, key rotation), #769 (the PR whose "bounded by in-flight cost" claim this issue showed was false under concurrency), #787 (unrelated feature that moved main during this PR's development; required the shadow-eval fix above during rebase), #794 (unrelated feature that also moved main, no file overlap).

[greptile-apps]

PR author is not in the allowed authors list.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 8a73a4a. Configure here.}

[devin-ai-integration]

Thank you so much for this contribution — the investigation behind it is genuinely impressive. Independently reproducing the TOCTOU three times against real Postgres before writing a line of the fix, the atomic WHERE spent + reserved + R <= limit reservation design, the idempotent DELETE … RETURNING settle/release/sweep primitive, and the honest worst-case-bound analysis are all exactly the kind of engineering this repo aims for. The production code, migrations, and SQLC queries needed no changes at all.

We've re-landed the change as #799 (with full credit to you) to align one thing with our internal conventions — and please know this is not your fault: these conventions are internal, evolving, and genuinely not obvious from the outside.

The one convention applied: root AGENTS.md ("Tests" section) has a hard rule — "No DB-backed integration tests in internal/. If need real Postgres, docker compose stack is runtime fixture; write scripts under scripts/ rather than *_test.go." The three live-Postgres tests under internal/postgres/ (org_monthly_spend_toctou_test.go, spend_reservation_month_internal_test.go, spend_reservation_sweep_test.go) were therefore ported verbatim into a runnable check at scripts/spend_reservation_check/ (mirroring the existing scripts/feedback_integration_check/ pattern), so go test ./... never touches Postgres. All your pure unit tests were kept as-is.

Everything else — the reserve-then-settle mechanism, the single-TX multi-scope reservation, the sweeper cadence rationale, the shadow-eval skip parity with #787 — carried over untouched. Thanks again, and we'd love more contributions like this one.