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

.env.example

@@ -57,6 +57,12 @@ PORT=8080
 # falls back to BYOK-only behavior (safe rollback path).
 ROUTER_DEPLOYMENT_MODE=selfhosted
 
+# Spend-cap reserve-then-settle (#793). Fixed slot R reserved against org /
+# api-key / user caps before inference; settled on debit or released on
+# early exit. TTL should stay ≥ messages/chat request timeout (600s) + grace.
+# ROUTER_SPEND_RESERVE_USD_MICROS=1000000
+# ROUTER_SPEND_RESERVE_TTL=15m
+
 # Admin dashboard password. Defaults to "admin" when unset (logs a warning) —
 # DEV ONLY behavior. REQUIRED in any deployment exposed beyond localhost.
 # ROUTER_ADMIN_PASSWORD=change-me

cmd/router/main.go

@@ -161,6 +161,33 @@ func main() {
 			billingSvc = billing.NewService(billingRepo)
 			logger.Info("Router billing enabled", "min_balance_usd_micros", billing.MinBalanceMicros)
 		}
+		if billingSvc != nil {
+			reserveAmount := billing.DefaultReserveAmountMicros
+			if v := config.GetOr("ROUTER_SPEND_RESERVE_USD_MICROS", ""); v != "" {
+				if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
+					reserveAmount = n
+				} else {
+					logger.Warn("Invalid ROUTER_SPEND_RESERVE_USD_MICROS; using default", "value", v, "default", billing.DefaultReserveAmountMicros)
+				}
+			}
+			reserveTTL := billing.DefaultReserveTTL
+			if v := config.GetOr("ROUTER_SPEND_RESERVE_TTL", ""); v != "" {
+				if d, err := time.ParseDuration(v); err == nil && d > 0 {
+					reserveTTL = d
+				} else {
+					logger.Warn("Invalid ROUTER_SPEND_RESERVE_TTL; using default", "value", v, "default", billing.DefaultReserveTTL.String())
+				}
+			}
+			billingSvc.WithReservationConfig(reserveAmount, reserveTTL)
+			logger.Info("Spend reservation config",
+				"reserve_usd_micros", reserveAmount,
+				"reserve_ttl", reserveTTL.String(),
+			)
+			safeGo(logger, "spend-reservation-sweep", func() {
+				runSpendReservationSweep(context.Background(), billingSvc)
+			})
+			logger.Info("Spend reservation TTL sweeper enabled", "interval", "1m")
+		}
 	}
 
 	// Managed without billing stays BYOK-only (avoids spending platform-key
@@ -1363,6 +1390,31 @@ func runSessionPinSweep(ctx context.Context, store sessionpin.Store) {
 	}
 }
 
+// runSpendReservationSweep releases expired spend-cap reservations every
+// minute. Stuck reserved incorrectly blocks spend, so this is much more
+// aggressive than the session-pin GC. Every replica runs its own ticker;
+// DELETE … RETURNING is idempotent across replicas.
+func runSpendReservationSweep(ctx context.Context, svc *billing.Service) {
+	logger := observability.Get()
+	ticker := time.NewTicker(1 * time.Minute)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			sweepCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+			n, err := svc.SweepExpiredReservations(sweepCtx, time.Now().UTC())
+			if err != nil {
+				logger.Error("Spend reservation sweep failed", "err", err)
+			} else if n > 0 {
+				logger.Info("Spend reservation sweep released expired holds", "released", n)
+			}
+			cancel()
+		}
+	}
+}
+
 // defaultHardPinProvider and defaultHardPinModel are the fallback (provider,
 // model) used by resolveHardPinModel when the cluster bundle can't be loaded.
 const (

db/migrations/0040_spend-reservations.down.sql

@@ -0,0 +1,14 @@
+BEGIN;
+
+ALTER TABLE router.model_router_api_keys
+    DROP COLUMN reserved_usd_micros;
+
+ALTER TABLE router.model_router_user_monthly_spend
+    DROP COLUMN reserved_usd_micros;
+
+ALTER TABLE router.organization_monthly_spend
+    DROP COLUMN reserved_usd_micros;
+
+DROP TABLE router.spend_reservations;
+
+COMMIT;

db/migrations/0040_spend-reservations.up.sql

@@ -0,0 +1,41 @@
+BEGIN;
+
+-- Open spend-cap reservations for reserve-then-settle (#793). One row per
+-- in-flight request per applicable scope. Sweeper deletes expired rows and
+-- decrements the denormalized reserved_usd_micros counters; settle/release
+-- delete by id and only decrement when DELETE … RETURNING yields a row.
+CREATE TABLE router.spend_reservations (
+    id                  UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    scope_kind          VARCHAR(16) NOT NULL
+        CHECK (scope_kind IN ('org_month', 'user_month', 'api_key')),
+    scope_id            VARCHAR(64) NOT NULL,
+    -- First day of the UTC month for *_month scopes; NULL for lifetime api_key.
+    month               DATE,
+    amount_usd_micros   BIGINT NOT NULL CHECK (amount_usd_micros > 0),
+    expires_at          TIMESTAMPTZ NOT NULL,
+    created_at          TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    router_request_id   VARCHAR(64),
+    CONSTRAINT spend_reservations_month_null_for_api_key CHECK (
+        (scope_kind = 'api_key' AND month IS NULL)
+        OR (scope_kind IN ('org_month', 'user_month') AND month IS NOT NULL)
+    )
+);
+
+CREATE INDEX spend_reservations_expires_at_idx
+    ON router.spend_reservations (expires_at);
+
+CREATE INDEX spend_reservations_scope_idx
+    ON router.spend_reservations (scope_kind, scope_id, month);
+
+-- Denormalized in-flight reserved totals so the gate is a single-row
+-- spent + reserved + R <= limit check without summing open reservations.
+ALTER TABLE router.organization_monthly_spend
+    ADD COLUMN reserved_usd_micros BIGINT NOT NULL DEFAULT 0;
+
+ALTER TABLE router.model_router_user_monthly_spend
+    ADD COLUMN reserved_usd_micros BIGINT NOT NULL DEFAULT 0;
+
+ALTER TABLE router.model_router_api_keys
+    ADD COLUMN reserved_usd_micros BIGINT NOT NULL DEFAULT 0;
+
+COMMIT;

db/queries/model_router_api_keys.sql

@@ -54,7 +54,7 @@ WHERE id = @id::uuid
 -- the per-request balance read in WithBalanceCheck. Returns sql.ErrNoRows when
 -- the key is missing/soft-deleted.
 -- name: GetModelRouterAPIKeySpend :one
-SELECT spend_cap_usd_micros, spent_usd_micros
+SELECT spend_cap_usd_micros, spent_usd_micros, reserved_usd_micros
 FROM router.model_router_api_keys
 WHERE id = @id::uuid
   AND deleted_at IS NULL;

db/queries/spend_limits.sql

@@ -27,7 +27,13 @@ SELECT
         FROM router.model_router_user_monthly_spend sp
         WHERE sp.router_user_id = @router_user_id::uuid
           AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date
-    ), 0)::bigint AS spent_usd_micros;
+    ), 0)::bigint AS spent_usd_micros,
+    COALESCE((
+        SELECT sp.reserved_usd_micros
+        FROM router.model_router_user_monthly_spend sp
+        WHERE sp.router_user_id = @router_user_id::uuid
+          AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date
+    ), 0)::bigint AS reserved_usd_micros;
 
 -- Reads the org's month-to-date spend alongside its monthly cap for the
 -- org-wide gate. Zero-spend months have no row; COALESCE keeps the read
@@ -42,4 +48,10 @@ SELECT
         FROM router.organization_monthly_spend sp
         WHERE sp.organization_id = @organization_id::varchar
           AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date
-    ), 0)::bigint AS spent_usd_micros;
+    ), 0)::bigint AS spent_usd_micros,
+    COALESCE((
+        SELECT sp.reserved_usd_micros
+        FROM router.organization_monthly_spend sp
+        WHERE sp.organization_id = @organization_id::varchar
+          AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date
+    ), 0)::bigint AS reserved_usd_micros;

db/queries/spend_reservations.sql

@@ -0,0 +1,119 @@
+-- Ensures the org spend row for @month exists so a reserve UPDATE can target
+-- it. Month is supplied by Go (utcMonthDate) so Ensure/Bump/Insert share one
+-- clock reading for the whole ReserveSpendCaps transaction.
+-- name: EnsureOrgMonthlySpendRow :exec
+INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros)
+VALUES (
+    @organization_id::varchar,
+    @month::date,
+    0,
+    0
+)
+ON CONFLICT (organization_id, month) DO NOTHING;
+
+-- Atomically bumps org-month reserved when spent+reserved+amount still fits
+-- under the configured org monthly limit. Returns the organization_id on
+-- success; zero rows means limit reached (or no limit configured — caller
+-- must skip reserve when limit is NULL before calling this). Month must be
+-- the same Go-computed value passed to Ensure and InsertSpendReservation.
+-- name: TryBumpOrgMonthReserved :one
+UPDATE router.organization_monthly_spend sp
+SET reserved_usd_micros = sp.reserved_usd_micros + @amount_usd_micros::bigint,
+    updated_at = NOW()
+FROM router.organization_spend_limits lim
+WHERE sp.organization_id = @organization_id::varchar
+  AND sp.month = @month::date
+  AND lim.organization_id = sp.organization_id
+  AND lim.org_monthly_limit_usd_micros IS NOT NULL
+  AND sp.spent_usd_micros + sp.reserved_usd_micros + @amount_usd_micros::bigint
+      <= lim.org_monthly_limit_usd_micros
+RETURNING sp.organization_id;
+
+-- name: EnsureUserMonthlySpendRow :exec
+INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros)
+VALUES (
+    @router_user_id::uuid,
+    @month::date,
+    0,
+    0
+)
+ON CONFLICT (router_user_id, month) DO NOTHING;
+
+-- Bumps user-month reserved under the effective limit (per-user override when
+-- present, else org default). Caller supplies the already-resolved effective
+-- limit; NULL limit means the caller should skip this scope. Month must match
+-- the Go-computed value used for Ensure and InsertSpendReservation.
+-- name: TryBumpUserMonthReserved :one
+UPDATE router.model_router_user_monthly_spend sp
+SET reserved_usd_micros = sp.reserved_usd_micros + @amount_usd_micros::bigint,
+    updated_at = NOW()
+WHERE sp.router_user_id = @router_user_id::uuid
+  AND sp.month = @month::date
+  AND sp.spent_usd_micros + sp.reserved_usd_micros + @amount_usd_micros::bigint
+      <= @limit_usd_micros::bigint
+RETURNING sp.router_user_id;
+
+-- Bumps api-key lifetime reserved under the key's spend_cap. Zero rows when
+-- the key is missing, uncapped, or the bump would exceed the cap.
+-- name: TryBumpAPIKeyReserved :one
+UPDATE router.model_router_api_keys k
+SET reserved_usd_micros = k.reserved_usd_micros + @amount_usd_micros::bigint
+WHERE k.id = @api_key_id::uuid
+  AND k.deleted_at IS NULL
+  AND k.spend_cap_usd_micros IS NOT NULL
+  AND k.spent_usd_micros + k.reserved_usd_micros + @amount_usd_micros::bigint
+      <= k.spend_cap_usd_micros
+RETURNING k.id;
+
+-- name: InsertSpendReservation :one
+INSERT INTO router.spend_reservations (
+    scope_kind,
+    scope_id,
+    month,
+    amount_usd_micros,
+    expires_at,
+    router_request_id
+) VALUES (
+    @scope_kind::varchar,
+    @scope_id::varchar,
+    sqlc.narg('month')::date,
+    @amount_usd_micros::bigint,
+    @expires_at::timestamptz,
+    sqlc.narg('router_request_id')::varchar
+)
+RETURNING id, scope_kind, scope_id, month, amount_usd_micros, expires_at;
+
+-- Atomic consume: DELETE … RETURNING is the sole settle/release/sweep
+-- primitive. Zero rows = already released/settled/swept (idempotent no-op).
+-- name: DeleteSpendReservation :one
+DELETE FROM router.spend_reservations
+WHERE id = @id::uuid
+RETURNING id, scope_kind, scope_id, month, amount_usd_micros;
+
+-- name: DecrementOrgMonthReserved :exec
+UPDATE router.organization_monthly_spend
+SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint),
+    updated_at = NOW()
+WHERE organization_id = @organization_id::varchar
+  AND month = @month::date;
+
+-- name: DecrementUserMonthReserved :exec
+UPDATE router.model_router_user_monthly_spend
+SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint),
+    updated_at = NOW()
+WHERE router_user_id = @router_user_id::uuid
+  AND month = @month::date;
+
+-- name: DecrementAPIKeyReserved :exec
+UPDATE router.model_router_api_keys
+SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint)
+WHERE id = @api_key_id::uuid;
+
+-- Deletes every expired reservation and returns the doomed rows so the
+-- adapter can decrement denormalized reserved counters. Prefer calling
+-- DeleteSpendReservation per id from Go when settling a known hold; this
+-- batch path is for the TTL sweeper only.
+-- name: DeleteExpiredSpendReservations :many
+DELETE FROM router.spend_reservations
+WHERE expires_at < @now::timestamptz
+RETURNING id, scope_kind, scope_id, month, amount_usd_micros;

internal/api/anthropic/messages.go

@@ -47,6 +47,19 @@ func MessagesHandler(svc *proxy.Service, authSvc *auth.Service) gin.HandlerFunc
 		if _, agentShadow := proxy.AgentShadowEvalFromContext(ctx); !agentShadow {
 			ctx = proxy.ResolveUserFromContext(ctx, authSvc, middleware.InstallationFrom(c))
 		}
+		ctx, release, armErr := svc.ArmSpendReservations(ctx)
+		if armErr != nil {
+			cls, ok := proxy.ClassifyDispatchError(armErr)
+			if ok {
+				proxy.LogDispatchErrorClass(log, cls, armErr)
+				writeAnthropicError(c, cls.Status, anthropicErrorType(cls.Kind), cls.Message)
+				return
+			}
+			log.Error("Spend reservation failed", "err", armErr)
+			writeAnthropicError(c, http.StatusServiceUnavailable, "api_error", "Billing system is temporarily unavailable. Retry in a few moments.")
+			return
+		}
+		defer release()
 		c.Request = c.Request.WithContext(ctx)
 
 		if err := svc.ProxyMessages(c.Request.Context(), body, c.Writer, c.Request); err != nil {

internal/api/gemini/generate_content.go

@@ -57,6 +57,19 @@ func GenerateContentHandler(svc *proxy.Service, authSvc *auth.Service) gin.Handl
 
 		ctx := context.WithValue(c.Request.Context(), proxy.ClientIdentityContextKey{}, proxy.ClientIdentityFromHeaders(c.Request.Header))
 		ctx = proxy.ResolveUserFromContext(ctx, authSvc, middleware.InstallationFrom(c))
+		ctx, release, armErr := svc.ArmSpendReservations(ctx)
+		if armErr != nil {
+			cls, ok := proxy.ClassifyDispatchError(armErr)
+			if ok {
+				proxy.LogDispatchErrorClass(log, cls, armErr)
+				writeGeminiError(c, cls.Status, "RESOURCE_EXHAUSTED", cls.Message)
+				return
+			}
+			log.Error("Spend reservation failed", "err", armErr)
+			writeGeminiError(c, http.StatusServiceUnavailable, "UNAVAILABLE", "Billing system is temporarily unavailable. Retry in a few moments.")
+			return
+		}
+		defer release()
 		c.Request = c.Request.WithContext(ctx)
 
 		if err := svc.ProxyGeminiGenerateContent(c.Request.Context(), body, c.Writer, c.Request); err != nil {

internal/api/openai/completions.go

@@ -31,6 +31,19 @@ func ChatCompletionHandler(svc *proxy.Service, authSvc *auth.Service) gin.Handle
 
 		ctx := context.WithValue(c.Request.Context(), proxy.ClientIdentityContextKey{}, proxy.ClientIdentityFromHeaders(c.Request.Header))
 		ctx = proxy.ResolveUserFromContext(ctx, authSvc, middleware.InstallationFrom(c))
+		ctx, release, armErr := svc.ArmSpendReservations(ctx)
+		if armErr != nil {
+			cls, ok := proxy.ClassifyDispatchError(armErr)
+			if ok {
+				proxy.LogDispatchErrorClass(log, cls, armErr)
+				writeOpenAIError(c, cls.Status, openAIErrorType(cls.Kind), cls.Message)
+				return
+			}
+			log.Error("Spend reservation failed", "err", armErr)
+			writeOpenAIError(c, http.StatusServiceUnavailable, "api_error", "Billing system is temporarily unavailable. Retry in a few moments.")
+			return
+		}
+		defer release()
 		c.Request = c.Request.WithContext(ctx)
 
 		if err := svc.ProxyOpenAIChatCompletion(c.Request.Context(), body, c.Writer, c.Request); err != nil {

internal/api/openai/responses.go

@@ -35,6 +35,19 @@ func ResponsesHandler(svc *proxy.Service, authSvc *auth.Service) gin.HandlerFunc
 
 		ctx := context.WithValue(c.Request.Context(), proxy.ClientIdentityContextKey{}, proxy.ClientIdentityFromHeaders(c.Request.Header))
 		ctx = proxy.ResolveUserFromContext(ctx, authSvc, middleware.InstallationFrom(c))
+		ctx, release, armErr := svc.ArmSpendReservations(ctx)
+		if armErr != nil {
+			cls, ok := proxy.ClassifyDispatchError(armErr)
+			if ok {
+				proxy.LogDispatchErrorClass(log, cls, armErr)
+				writeOpenAIError(c, cls.Status, openAIErrorType(cls.Kind), cls.Message)
+				return
+			}
+			log.Error("Spend reservation failed", "err", armErr)
+			writeOpenAIError(c, http.StatusServiceUnavailable, "api_error", "Billing system is temporarily unavailable. Retry in a few moments.")
+			return
+		}
+		defer release()
 		c.Request = c.Request.WithContext(ctx)
 
 		if err := svc.ProxyOpenAIResponses(c.Request.Context(), body, c.Writer, c.Request); err != nil {

internal/billing/errors.go

@@ -21,6 +21,15 @@ var ErrBalanceRowMissing = errors.New("billing: balance row missing")
 // UTC-month spend meets their effective monthly limit. Mapped to HTTP 402.
 var ErrUserMonthlySpendLimitReached = errors.New("billing: engineer monthly spend limit reached")
 
+// ErrOrgMonthlySpendLimitReached is returned when the org's current UTC-month
+// spend (including in-flight reservations) cannot fit another reserve slot.
+// Mapped to HTTP 402.
+var ErrOrgMonthlySpendLimitReached = errors.New("billing: organization monthly spend limit reached")
+
+// ErrAPIKeySpendCapReached is returned when a key's lifetime spend (including
+// in-flight reservations) cannot fit another reserve slot. Mapped to HTTP 402.
+var ErrAPIKeySpendCapReached = errors.New("billing: api key spend cap reached")
+
 // ErrSpendLimitCheckUnavailable is returned on a repo failure reading spend
 // limits; gates fail closed (HTTP 503) to prevent unbounded-spend windows.
 var ErrSpendLimitCheckUnavailable = errors.New("billing: spend limit check unavailable")