From cc3eae8699db31979519b501fa563c882a0b98cf Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Sun, 19 Jul 2026 23:41:00 -0400 Subject: [PATCH 1/8] test: reproduce org monthly spend-limit TOCTOU under concurrency (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- .../postgres/org_monthly_spend_toctou_test.go | 147 ++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 internal/postgres/org_monthly_spend_toctou_test.go diff --git a/internal/postgres/org_monthly_spend_toctou_test.go b/internal/postgres/org_monthly_spend_toctou_test.go new file mode 100644 index 00000000..aa67fd90 --- /dev/null +++ b/internal/postgres/org_monthly_spend_toctou_test.go @@ -0,0 +1,147 @@ +package postgres_test + +import ( + "context" + "fmt" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "workweave/router/internal/billing" + "workweave/router/internal/postgres" + "workweave/router/internal/router/catalog" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" +) + +// TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot is the +// permanent #793 regression: N concurrent CheckOrgMonthlySpend + DebitForInference +// must not multiply past the hard monthly cap unboundedly. +// +// Fixture matches the issue reproduction: $1 limit, $0.90 starting spend, +// $6.75 notional per debit (1M in @ $3/MTok + 250K out @ $15/MTok), N=20. +// Soft-overshoot bound after the reserve-then-settle fix is roughly one +// under-reserved turn (or zero when remaining headroom < R); before the fix +// this assertion fails with ~$100+ overshoot. +// +// Gated on ROUTER_TEST_DATABASE_URL (falls back to DATABASE_URL). Skips when +// neither is set so unit CI without Postgres stays green. Requires a migrated +// router schema (docker compose postgres on :5433 is the local fixture). +func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + dsn = os.Getenv("DATABASE_URL") + } + if dsn == "" { + t.Skip("ROUTER_TEST_DATABASE_URL / DATABASE_URL not set; need live Postgres for #793 concurrency regression") + } + + ctx := context.Background() + cfg, err := pgxpool.ParseConfig(dsn) + require.NoError(t, err) + cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + _, err := conn.Exec(ctx, "SET search_path TO router, public") + return err + } + cfg.MaxConns = 20 + pool, err := pgxpool.NewWithConfig(ctx, cfg) + require.NoError(t, err) + t.Cleanup(pool.Close) + + require.NoError(t, pool.Ping(ctx)) + + orgID := fmt.Sprintf("toctou-793-%d", time.Now().UnixNano()) + const ( + limitMicros int64 = 1_000_000 // $1.00 + startSpent int64 = 900_000 // $0.90 + inputTokens = 1_000_000 + outputTokens = 250_000 + concurrency = 20 + ) + pricing := catalog.Pricing{InputUSDPer1M: 3, OutputUSDPer1M: 15} + perCallMicros := catalog.USDToMicros( + catalog.EffectiveInputCost(inputTokens, 0, 0, pricing.InputUSDPer1M, pricing, "") + + catalog.EffectiveOutputCost(outputTokens, pricing.OutputUSDPer1M), + ) + require.Equal(t, int64(6_750_000), perCallMicros, "fixture must stay $6.75 for continuity with #793") + + _, err = pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, limitMicros) + require.NoError(t, err) + _, err = pool.Exec(ctx, ` + INSERT INTO organization_monthly_spend (organization_id, month, spent_usd_micros) + VALUES ($1, DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, $2)`, orgID, startSpent) + require.NoError(t, err) + _, err = pool.Exec(ctx, ` + INSERT INTO organization_credit_balance (organization_id, balance_usd_micros) + VALUES ($1, $2)`, orgID, int64(1_000_000_000)) // $1000 — not under test + require.NoError(t, err) + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_credit_ledger WHERE organization_id = $1`, orgID) + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_credit_balance WHERE organization_id = $1`, orgID) + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_monthly_spend WHERE organization_id = $1`, orgID) + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_spend_limits WHERE organization_id = $1`, orgID) + }) + + svc := billing.NewService(postgres.NewBillingRepo(pool)) + + var ( + wg sync.WaitGroup + debited atomic.Int32 + rejected atomic.Int32 + ) + wg.Add(concurrency) + for i := 0; i < concurrency; i++ { + i := i + go func() { + defer wg.Done() + reqCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + res, err := svc.CheckOrgMonthlySpend(reqCtx, orgID) + if err != nil { + t.Errorf("CheckOrgMonthlySpend: %v", err) + return + } + if res.LimitReached() { + rejected.Add(1) + return + } + _, err = svc.DebitForInference(reqCtx, billing.DebitInferenceParams{ + OrganizationID: orgID, + RouterRequestID: fmt.Sprintf("%s-%d", orgID, i), + Model: "toctou-fixture", + InputTokens: inputTokens, + OutputTokens: outputTokens, + Pricing: pricing, + }) + if err != nil { + t.Errorf("DebitForInference: %v", err) + return + } + debited.Add(1) + }() + } + wg.Wait() + + final, _, err := postgres.NewBillingRepo(pool).GetOrgMonthlySpendAndLimit(ctx, orgID) + require.NoError(t, err) + + // Soft overshoot bound: at most one in-flight turn past the limit (the + // product promise before concurrency multiplied it). Equivalent to + // start + N*cost when every check races — that path must fail this assert. + maxAllowed := limitMicros + perCallMicros + t.Logf("org=%s debited=%d rejected=%d final_spent_usd_micros=%d overshoot_usd_micros=%d max_allowed=%d", + orgID, debited.Load(), rejected.Load(), final, final-limitMicros, maxAllowed) + + require.LessOrEqual(t, final, maxAllowed, + "org monthly spend TOCTOU (#793): final=%d ($%.2f) exceeds limit+one_turn=%d ($%.2f); overshoot=$%.2f from %d concurrent debits", + final, float64(final)/1e6, maxAllowed, float64(maxAllowed)/1e6, float64(final-limitMicros)/1e6, debited.Load()) +} From 5d4d9b2241c372527c03e2a17614dc86c5145f61 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Sun, 19 Jul 2026 23:41:40 -0400 Subject: [PATCH 2/8] feat(db): add spend_reservations table for #793 (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- .../0040_spend-reservations.down.sql | 14 +++++++ db/migrations/0040_spend-reservations.up.sql | 41 +++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 db/migrations/0040_spend-reservations.down.sql create mode 100644 db/migrations/0040_spend-reservations.up.sql diff --git a/db/migrations/0040_spend-reservations.down.sql b/db/migrations/0040_spend-reservations.down.sql new file mode 100644 index 00000000..7b165438 --- /dev/null +++ b/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; diff --git a/db/migrations/0040_spend-reservations.up.sql b/db/migrations/0040_spend-reservations.up.sql new file mode 100644 index 00000000..8ba1a9d4 --- /dev/null +++ b/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; From 6f9e0b38de14023583266f53484800873b7bd05d Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Sun, 19 Jul 2026 23:48:43 -0400 Subject: [PATCH 3/8] fix(router): reserve-then-settle spend caps to close concurrency gap (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- .env.example | 6 + cmd/router/main.go | 23 + db/queries/model_router_api_keys.sql | 2 +- db/queries/spend_limits.sql | 16 +- db/queries/spend_reservations.sql | 116 +++++ internal/api/anthropic/messages.go | 13 + internal/api/gemini/generate_content.go | 13 + internal/api/openai/completions.go | 13 + internal/api/openai/responses.go | 13 + internal/billing/errors.go | 9 + internal/billing/repo.go | 38 +- internal/billing/reservation.go | 73 ++++ internal/billing/service.go | 132 +++++- internal/billing/service_test.go | 30 +- internal/postgres/billing_repo.go | 346 ++++++++++++--- .../postgres/org_monthly_spend_toctou_test.go | 44 +- internal/proxy/dispatch_error.go | 16 + internal/proxy/gemini.go | 3 - internal/proxy/service.go | 15 +- internal/proxy/spend_limit.go | 6 +- internal/proxy/spend_limit_internal_test.go | 21 +- internal/proxy/spend_reserve.go | 47 ++ .../server/middleware/api_key_spend_cap.go | 3 +- .../server/middleware/balance_check_test.go | 23 +- internal/server/server.go | 6 +- internal/sqlc/model_router_api_keys.sql.go | 22 +- internal/sqlc/models.go | 30 +- internal/sqlc/spend_limits.sql.go | 37 +- internal/sqlc/spend_reservations.sql.go | 408 ++++++++++++++++++ 29 files changed, 1372 insertions(+), 152 deletions(-) create mode 100644 db/queries/spend_reservations.sql create mode 100644 internal/billing/reservation.go create mode 100644 internal/proxy/spend_reserve.go create mode 100644 internal/sqlc/spend_reservations.sql.go diff --git a/.env.example b/.env.example index f8a27516..973c6f39 100644 --- a/.env.example +++ b/.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 diff --git a/cmd/router/main.go b/cmd/router/main.go index 8e95dbed..e64e815a 100644 --- a/cmd/router/main.go +++ b/cmd/router/main.go @@ -161,6 +161,29 @@ 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(), + ) + } } // Managed without billing stays BYOK-only (avoids spending platform-key diff --git a/db/queries/model_router_api_keys.sql b/db/queries/model_router_api_keys.sql index 5d3821aa..ada494c0 100644 --- a/db/queries/model_router_api_keys.sql +++ b/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; diff --git a/db/queries/spend_limits.sql b/db/queries/spend_limits.sql index 3898f441..e12bc4c7 100644 --- a/db/queries/spend_limits.sql +++ b/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; diff --git a/db/queries/spend_reservations.sql b/db/queries/spend_reservations.sql new file mode 100644 index 00000000..2dc98e72 --- /dev/null +++ b/db/queries/spend_reservations.sql @@ -0,0 +1,116 @@ +-- Ensures the current UTC-month org spend row exists so a reserve UPDATE can +-- target it. No-op when the row is already present. +-- name: EnsureOrgMonthlySpendRow :exec +INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) +VALUES ( + @organization_id::varchar, + DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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). +-- 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 = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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, + DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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. +-- 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 = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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; diff --git a/internal/api/anthropic/messages.go b/internal/api/anthropic/messages.go index a9ef2839..763f5cb5 100644 --- a/internal/api/anthropic/messages.go +++ b/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 { diff --git a/internal/api/gemini/generate_content.go b/internal/api/gemini/generate_content.go index 6f8ed26f..de3f19bd 100644 --- a/internal/api/gemini/generate_content.go +++ b/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 { diff --git a/internal/api/openai/completions.go b/internal/api/openai/completions.go index 43408d9e..395f7e81 100644 --- a/internal/api/openai/completions.go +++ b/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 { diff --git a/internal/api/openai/responses.go b/internal/api/openai/responses.go index 206def47..75229c55 100644 --- a/internal/api/openai/responses.go +++ b/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 { diff --git a/internal/billing/errors.go b/internal/billing/errors.go index bbddafe5..f1d38ac5 100644 --- a/internal/billing/errors.go +++ b/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") diff --git a/internal/billing/repo.go b/internal/billing/repo.go index 2851fa90..6a3b0f6e 100644 --- a/internal/billing/repo.go +++ b/internal/billing/repo.go @@ -1,6 +1,11 @@ package billing -import "context" +import ( + "context" + "time" + + "github.com/google/uuid" +) // Repo is the adapter-boundary contract for credit billing reads and writes. // Implementations live in internal/postgres/billing_repo.go. @@ -17,20 +22,38 @@ type Repo interface { // DebitInference atomically decrements the balance and appends the // ledger row. Under an active override, delta is 0 but the ledger row // still records notional_cost_micros for the shadow billing trail. + // When ReservationIDs is non-empty, those reservations are settled + // (DELETE … RETURNING + reserved decrement) in the same transaction. DebitInference(ctx context.Context, p DebitParams) (balanceAfterMicros int64, err error) // GetAPIKeySpend reads a key's spend cap and spend-to-date fresh from // Postgres, bypassing the auth cache. found is false if the key was // deleted mid-request, treated as "no cap to enforce". - GetAPIKeySpend(ctx context.Context, apiKeyID string) (spentMicros int64, capMicros *int64, found bool, err error) + GetAPIKeySpend(ctx context.Context, apiKeyID string) (spentMicros, reservedMicros int64, capMicros *int64, found bool, err error) // GetUserMonthlySpendAndLimit returns the engineer's current UTC-month spend - // and effective monthly limit: per-user override when set (NULL = explicitly uncapped), else org default. - GetUserMonthlySpendAndLimit(ctx context.Context, organizationID, routerUserID string) (spentMicros int64, limitMicros *int64, err error) + // (and in-flight reserved) and effective monthly limit: per-user override + // when set (NULL = explicitly uncapped), else org default. + GetUserMonthlySpendAndLimit(ctx context.Context, organizationID, routerUserID string) (spentMicros, reservedMicros int64, limitMicros *int64, err error) // GetOrgMonthlySpendAndLimit returns the org's current UTC-month spend - // and its configured monthly cap. nil limitMicros means no cap is set. - GetOrgMonthlySpendAndLimit(ctx context.Context, organizationID string) (spentMicros int64, limitMicros *int64, err error) + // (and in-flight reserved) and its configured monthly cap. nil limitMicros + // means no cap is set. + GetOrgMonthlySpendAndLimit(ctx context.Context, organizationID string) (spentMicros, reservedMicros int64, limitMicros *int64, err error) + + // ReserveSpendCaps atomically reserves every applicable scope in one + // transaction. On any scope failure the whole TX rolls back (no partial + // holds). Returns reservation ids (possibly empty when no caps apply). + ReserveSpendCaps(ctx context.Context, p ReserveSpendCapsParams) ([]uuid.UUID, error) + + // ReleaseSpendReservations consumes reservation ids via DELETE … RETURNING + // and decrements denormalized reserved only when a row was returned. + // Idempotent: already-gone ids are no-ops. + ReleaseSpendReservations(ctx context.Context, ids []uuid.UUID) error + + // SweepExpiredSpendReservations deletes expired reservation rows and + // decrements reserved counters for each DELETE … RETURNING hit. + SweepExpiredSpendReservations(ctx context.Context, now time.Time) (released int, err error) // GetAutopayConfig reports the org's autopay state and recharge // threshold. A missing config row returns enabled=false, nil error @@ -56,4 +79,7 @@ type DebitParams struct { // RouterUserID, if non-empty, attributes the debit to that engineer and bumps // the org's monthly counter; org counter is always bumped regardless. RouterUserID string + // ReservationIDs, when non-empty, are settled in the same transaction as + // the debit (DELETE … RETURNING + reserved decrement). + ReservationIDs []uuid.UUID } diff --git a/internal/billing/reservation.go b/internal/billing/reservation.go new file mode 100644 index 00000000..a2eae6ef --- /dev/null +++ b/internal/billing/reservation.go @@ -0,0 +1,73 @@ +package billing + +import ( + "context" + "sync/atomic" + "time" + + "github.com/google/uuid" +) + +// Spend reservation scope_kind values (match spend_reservations CHECK). +const ( + ScopeOrgMonth = "org_month" + ScopeUserMonth = "user_month" + ScopeAPIKey = "api_key" +) + +// DefaultReserveAmountMicros is the v1 fixed reservation slot ($1). Tuned as a +// p95-ish turn cost approximation — not exact. Override via +// ROUTER_SPEND_RESERVE_USD_MICROS. +const DefaultReserveAmountMicros int64 = 1_000_000 + +// DefaultReserveTTL is request timeout (600s) + 5m grace. Override via +// ROUTER_SPEND_RESERVE_TTL. Keep ≥ messagesTimeout / chatCompletionTimeout. +const DefaultReserveTTL = 15 * time.Minute + +// ReserveSpendCapsParams is the input to Repo.ReserveSpendCaps. +type ReserveSpendCapsParams struct { + OrganizationID string + APIKeyID string // empty skips api-key scope + RouterUserID string // empty skips user-month scope + RouterRequestID string + AmountUsdMicros int64 + TTL time.Duration + // SkipOrg / SkipKey / SkipUser force-skip a scope (e.g. billing override). + SkipOrg bool + SkipKey bool + SkipUser bool +} + +// SpendHold tracks reservation ids for one request. MarkSettled after a +// successful DebitForInference so the handler defer's ReleaseAll no-ops. +type SpendHold struct { + IDs []uuid.UUID + settled atomic.Bool +} + +// MarkSettled records that reservations were consumed by settle (or need not +// be released). Safe to call more than once. +func (h *SpendHold) MarkSettled() { + if h == nil { + return + } + h.settled.Store(true) +} + +// Settled reports whether MarkSettled has been called. +func (h *SpendHold) Settled() bool { + return h != nil && h.settled.Load() +} + +type spendHoldContextKeyT struct{} + +// WithSpendHold stashes the request's reservation hold on ctx. +func WithSpendHold(ctx context.Context, h *SpendHold) context.Context { + return context.WithValue(ctx, spendHoldContextKeyT{}, h) +} + +// SpendHoldFrom returns the hold stashed by WithSpendHold, or nil. +func SpendHoldFrom(ctx context.Context) *SpendHold { + h, _ := ctx.Value(spendHoldContextKeyT{}).(*SpendHold) + return h +} diff --git a/internal/billing/service.go b/internal/billing/service.go index 476b29e8..d34403b0 100644 --- a/internal/billing/service.go +++ b/internal/billing/service.go @@ -2,9 +2,12 @@ package billing import ( "context" + "time" "workweave/router/internal/observability" "workweave/router/internal/router/catalog" + + "github.com/google/uuid" ) // hasOverrideContextKeyT lives in billing (not middleware/proxy) so both @@ -81,14 +84,32 @@ const SubscriptionOverdraftFloorMicros int64 = -5_000_000 // Service orchestrates balance reads and debits. No I/O of its own — all // persistence flows through the Repo interface. type Service struct { - repo Repo - autopay AutopayNotifier + repo Repo + autopay AutopayNotifier + reserveAmountMicros int64 + reserveTTL time.Duration } // NewService constructs a billing service. The Repo is required; nil panics // at request time, so the composition root must guard against it. func NewService(repo Repo) *Service { - return &Service{repo: repo} + return &Service{ + repo: repo, + reserveAmountMicros: DefaultReserveAmountMicros, + reserveTTL: DefaultReserveTTL, + } +} + +// WithReservationConfig sets the fixed reserve slot R and TTL. Zero values +// keep the defaults. Wired from env in the composition root. +func (s *Service) WithReservationConfig(amountMicros int64, ttl time.Duration) *Service { + if amountMicros > 0 { + s.reserveAmountMicros = amountMicros + } + if ttl > 0 { + s.reserveTTL = ttl + } + return s } // AutopayNotifier signals the control plane that an org's balance just @@ -117,21 +138,22 @@ type CheckResult struct { // APIKeySpendCapResult is the outcome of a preflight per-key spend-cap // check. Found is false if the key was deleted mid-request; CapMicros is // nil for an uncapped key. Middleware blocks when Found && CapMicros != -// nil && SpentMicros >= *CapMicros. +// nil && SpentMicros+ReservedMicros >= *CapMicros. type APIKeySpendCapResult struct { - Found bool - SpentMicros int64 - CapMicros *int64 + Found bool + SpentMicros int64 + ReservedMicros int64 + CapMicros *int64 } // CheckAPIKeySpendCap reads fresh from the repo (not the auth cache) so a // hot cached key can't overrun its cap within the cache TTL. func (s *Service) CheckAPIKeySpendCap(ctx context.Context, apiKeyID string) (APIKeySpendCapResult, error) { - spent, cap, found, err := s.repo.GetAPIKeySpend(ctx, apiKeyID) + spent, reserved, cap, found, err := s.repo.GetAPIKeySpend(ctx, apiKeyID) if err != nil { return APIKeySpendCapResult{}, err } - return APIKeySpendCapResult{Found: found, SpentMicros: spent, CapMicros: cap}, nil + return APIKeySpendCapResult{Found: found, SpentMicros: spent, ReservedMicros: reserved, CapMicros: cap}, nil } // CheckBalance short-circuits the balance read when an override is active @@ -157,34 +179,37 @@ func (s *Service) CheckBalance(ctx context.Context, orgID string) (CheckResult, // MonthlySpendResult carries a current-UTC-month spend counter alongside the // limit that applies to it. LimitMicros nil means no limit is configured. +// ReservedMicros is in-flight reserve-then-settle hold. type MonthlySpendResult struct { - SpentMicros int64 - LimitMicros *int64 + SpentMicros int64 + ReservedMicros int64 + LimitMicros *int64 } -// LimitReached reports whether a configured limit has been met or exceeded. +// LimitReached reports whether spent + reserved has met or exceeded a +// configured limit (in-flight reservations count against headroom). func (r MonthlySpendResult) LimitReached() bool { - return r.LimitMicros != nil && r.SpentMicros >= *r.LimitMicros + return r.LimitMicros != nil && r.SpentMicros+r.ReservedMicros >= *r.LimitMicros } // CheckUserMonthlySpend reads the engineer's current-month spend and // effective monthly limit fresh from Postgres. func (s *Service) CheckUserMonthlySpend(ctx context.Context, organizationID, routerUserID string) (MonthlySpendResult, error) { - spent, limit, err := s.repo.GetUserMonthlySpendAndLimit(ctx, organizationID, routerUserID) + spent, reserved, limit, err := s.repo.GetUserMonthlySpendAndLimit(ctx, organizationID, routerUserID) if err != nil { return MonthlySpendResult{}, err } - return MonthlySpendResult{SpentMicros: spent, LimitMicros: limit}, nil + return MonthlySpendResult{SpentMicros: spent, ReservedMicros: reserved, LimitMicros: limit}, nil } // CheckOrgMonthlySpend reads the org's current-month spend and configured // monthly cap fresh from Postgres. func (s *Service) CheckOrgMonthlySpend(ctx context.Context, organizationID string) (MonthlySpendResult, error) { - spent, limit, err := s.repo.GetOrgMonthlySpendAndLimit(ctx, organizationID) + spent, reserved, limit, err := s.repo.GetOrgMonthlySpendAndLimit(ctx, organizationID) if err != nil { return MonthlySpendResult{}, err } - return MonthlySpendResult{SpentMicros: spent, LimitMicros: limit}, nil + return MonthlySpendResult{SpentMicros: spent, ReservedMicros: reserved, LimitMicros: limit}, nil } // DebitInferenceParams is the input to DebitForInference. Token counts @@ -211,6 +236,9 @@ type DebitInferenceParams struct { // RouterUserID attributes the debit to the resolved engineer identity for // monthly spend-limit tracking; empty leaves per-user spend untouched. RouterUserID string + // ReservationIDs are settled in the same TX as the debit when non-empty. + // Prefer setting via fireBilling from the request SpendHold. + ReservationIDs []uuid.UUID } // DebitForInference writes one ledger row at cost — no markup math here; @@ -238,6 +266,7 @@ func (s *Service) DebitForInference(ctx context.Context, p DebitInferenceParams) RouterModel: p.Model, APIKeyID: p.APIKeyID, RouterUserID: p.RouterUserID, + ReservationIDs: p.ReservationIDs, }) if err != nil { return balanceAfter, err @@ -246,6 +275,75 @@ func (s *Service) DebitForInference(ctx context.Context, p DebitInferenceParams) return balanceAfter, nil } +// ReserveApplicableCaps reserves every applicable spend cap (org monthly, +// api-key lifetime, user monthly) in one transaction. Returns a SpendHold +// (possibly with empty IDs when no caps apply). Billing override skips all. +func (s *Service) ReserveApplicableCaps(ctx context.Context, organizationID, apiKeyID, routerUserID, requestID string) (*SpendHold, error) { + if s == nil || s.repo == nil { + return &SpendHold{}, nil + } + if HasOverrideFromContext(ctx) { + return &SpendHold{}, nil + } + ids, err := s.repo.ReserveSpendCaps(ctx, ReserveSpendCapsParams{ + OrganizationID: organizationID, + APIKeyID: apiKeyID, + RouterUserID: routerUserID, + RouterRequestID: requestID, + AmountUsdMicros: s.reserveAmountMicros, + TTL: s.reserveTTL, + }) + if err != nil { + return nil, err + } + return &SpendHold{IDs: ids}, nil +} + +// ReleaseHold releases unsettled reservations. Idempotent against settle/sweep. +func (s *Service) ReleaseHold(ctx context.Context, hold *SpendHold) error { + if s == nil || s.repo == nil || hold == nil || hold.Settled() || len(hold.IDs) == 0 { + return nil + } + return s.repo.ReleaseSpendReservations(ctx, hold.IDs) +} + +// SweepExpiredReservations runs one TTL sweep cycle. +func (s *Service) SweepExpiredReservations(ctx context.Context, now time.Time) (int, error) { + if s == nil || s.repo == nil { + return 0, nil + } + return s.repo.SweepExpiredSpendReservations(ctx, now) +} + +// ArmSpendReservations reserves applicable caps and returns a release func +// for defer. Call after identity is resolved and before Proxy*. The release +// func no-ops after a successful DebitForInference (MarkSettled). +func (s *Service) ArmSpendReservations(ctx context.Context, organizationID, apiKeyID, routerUserID, requestID string) (context.Context, func(), error) { + if s == nil { + return ctx, func() {}, nil + } + hold, err := s.ReserveApplicableCaps(ctx, organizationID, apiKeyID, routerUserID, requestID) + if err != nil { + return ctx, nil, err + } + ctx = WithSpendHold(ctx, hold) + release := func() { + if hold.Settled() { + return + } + relCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if relErr := s.ReleaseHold(relCtx, hold); relErr != nil { + observability.Get().Error("spend reservation release failed", + "err", relErr, + "organization_id", organizationID, + "reservation_count", len(hold.IDs), + ) + } + } + return ctx, release, nil +} + // maybeSignalRecharge fires once, on the debit that crosses the org's // balance from at-or-above its autopay threshold to below it. No-ops if // autopay isn't wired, the debit moved nothing, or autopay is disabled. diff --git a/internal/billing/service_test.go b/internal/billing/service_test.go index e4ae3cfb..db10f377 100644 --- a/internal/billing/service_test.go +++ b/internal/billing/service_test.go @@ -8,12 +8,14 @@ import ( "sync" "sync/atomic" "testing" + "time" "workweave/router/internal/billing" "workweave/router/internal/observability" "workweave/router/internal/providers" "workweave/router/internal/router/catalog" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -92,22 +94,34 @@ func (r *fakeRepo) DebitInference(_ context.Context, p billing.DebitParams) (int func (r *fakeRepo) BillingTablesExist(_ context.Context) (bool, error) { return true, nil } -func (r *fakeRepo) GetAPIKeySpend(_ context.Context, _ string) (int64, *int64, bool, error) { - return 0, nil, false, nil +func (r *fakeRepo) GetAPIKeySpend(_ context.Context, _ string) (int64, int64, *int64, bool, error) { + return 0, 0, nil, false, nil } -func (r *fakeRepo) GetUserMonthlySpendAndLimit(_ context.Context, _, _ string) (int64, *int64, error) { +func (r *fakeRepo) GetUserMonthlySpendAndLimit(_ context.Context, _, _ string) (int64, int64, *int64, error) { if r.userMonthErr != nil { - return 0, nil, r.userMonthErr + return 0, 0, nil, r.userMonthErr } - return r.userMonthSpent, r.userMonthLimit, nil + return r.userMonthSpent, 0, r.userMonthLimit, nil } -func (r *fakeRepo) GetOrgMonthlySpendAndLimit(_ context.Context, _ string) (int64, *int64, error) { +func (r *fakeRepo) GetOrgMonthlySpendAndLimit(_ context.Context, _ string) (int64, int64, *int64, error) { if r.orgMonthErr != nil { - return 0, nil, r.orgMonthErr + return 0, 0, nil, r.orgMonthErr } - return r.orgMonthSpent, r.orgMonthLimit, nil + return r.orgMonthSpent, 0, r.orgMonthLimit, nil +} + +func (r *fakeRepo) ReserveSpendCaps(_ context.Context, _ billing.ReserveSpendCapsParams) ([]uuid.UUID, error) { + return nil, nil +} + +func (r *fakeRepo) ReleaseSpendReservations(_ context.Context, _ []uuid.UUID) error { + return nil +} + +func (r *fakeRepo) SweepExpiredSpendReservations(_ context.Context, _ time.Time) (int, error) { + return 0, nil } func (r *fakeRepo) GetAutopayConfig(_ context.Context, _ string) (bool, int64, error) { diff --git a/internal/postgres/billing_repo.go b/internal/postgres/billing_repo.go index f6d13a7a..e6bb2c42 100644 --- a/internal/postgres/billing_repo.go +++ b/internal/postgres/billing_repo.go @@ -3,33 +3,38 @@ package postgres import ( "context" "errors" + "fmt" + "time" "workweave/router/internal/billing" "workweave/router/internal/sqlc" "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" ) // BillingRepo implements billing.Repo against the router-schema credit // tables via SQLC. type BillingRepo struct { - tx sqlc.DBTX + pool *pgxpool.Pool } -// NewBillingRepo constructs a BillingRepo backed by the given connection. -func NewBillingRepo(tx sqlc.DBTX) *BillingRepo { - return &BillingRepo{tx: tx} +// NewBillingRepo constructs a BillingRepo backed by the given pool. +func NewBillingRepo(pool *pgxpool.Pool) *BillingRepo { + return &BillingRepo{pool: pool} } var _ billing.Repo = (*BillingRepo)(nil) +func (r *BillingRepo) q() *sqlc.Queries { return sqlc.New(r.pool) } + // GetBalance returns the org's current credit balance in USD micros. // Maps pgx.ErrNoRows to billing.ErrBalanceRowMissing so middleware can // distinguish "row missing" from "balance == 0". func (r *BillingRepo) GetBalance(ctx context.Context, orgID string) (int64, error) { - q := sqlc.New(r.tx) - balance, err := q.GetOrgCreditBalance(ctx, orgID) + balance, err := r.q().GetOrgCreditBalance(ctx, orgID) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return 0, billing.ErrBalanceRowMissing @@ -42,19 +47,19 @@ func (r *BillingRepo) GetBalance(ctx context.Context, orgID string) (int64, erro // HasActiveOverride reports whether the org has an unexpired billing // override row. EXISTS-based query — true means the org bypasses billing. func (r *BillingRepo) HasActiveOverride(ctx context.Context, orgID string) (bool, error) { - q := sqlc.New(r.tx) - override, err := q.GetActiveBillingOverride(ctx, orgID) - if err != nil { - return false, err - } - return override, nil + return r.q().GetActiveBillingOverride(ctx, orgID) } -// DebitInference performs the atomic UPDATE + INSERT CTE. Returns the -// post-debit balance, or billing.ErrBalanceRowMissing if no balance row -// existed (the CTE returns zero rows in that case). +// DebitInference performs the atomic UPDATE + INSERT CTE, then settles any +// reservation ids in the same transaction. func (r *BillingRepo) DebitInference(ctx context.Context, p billing.DebitParams) (int64, error) { - q := sqlc.New(r.tx) + tx, err := r.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer func() { _ = tx.Rollback(ctx) }() + + q := sqlc.New(tx) balanceAfter, err := q.DebitOrgCredits(ctx, sqlc.DebitOrgCreditsParams{ OrganizationID: p.OrganizationID, DeltaUsdMicros: p.DeltaUsdMicros, @@ -71,71 +76,67 @@ func (r *BillingRepo) DebitInference(ctx context.Context, p billing.DebitParams) } return 0, err } + for _, id := range p.ReservationIDs { + if err := consumeSpendReservation(ctx, q, id); err != nil { + return 0, err + } + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } return balanceAfter, nil } // GetAPIKeySpend reads a key's cap and spend-to-date fresh from Postgres. // Returns found=false (nil error) when no active key matches the id. -func (r *BillingRepo) GetAPIKeySpend(ctx context.Context, apiKeyID string) (int64, *int64, bool, error) { +func (r *BillingRepo) GetAPIKeySpend(ctx context.Context, apiKeyID string) (int64, int64, *int64, bool, error) { parsed, err := uuid.Parse(apiKeyID) if err != nil { - // A malformed id can't match any row; treat as "no cap to enforce" - // rather than failing the request closed on a client-shaped value. - return 0, nil, false, nil + return 0, 0, nil, false, nil } - q := sqlc.New(r.tx) - row, err := q.GetModelRouterAPIKeySpend(ctx, parsed) + row, err := r.q().GetModelRouterAPIKeySpend(ctx, parsed) if err != nil { if errors.Is(err, pgx.ErrNoRows) { - return 0, nil, false, nil + return 0, 0, nil, false, nil } - return 0, nil, false, err + return 0, 0, nil, false, err } - return row.SpentUsdMicros, row.SpendCapUsdMicros, true, nil + return row.SpentUsdMicros, row.ReservedUsdMicros, row.SpendCapUsdMicros, true, nil } // GetUserMonthlySpendAndLimit resolves the effective monthly limit // (per-user override > org default; NULL override = explicitly uncapped). -func (r *BillingRepo) GetUserMonthlySpendAndLimit(ctx context.Context, organizationID, routerUserID string) (int64, *int64, error) { +func (r *BillingRepo) GetUserMonthlySpendAndLimit(ctx context.Context, organizationID, routerUserID string) (int64, int64, *int64, error) { parsed, err := uuid.Parse(routerUserID) if err != nil { - // A malformed id can't match any row; treat as "no limit to enforce" - // rather than failing the request closed on a client-shaped value. - return 0, nil, nil + return 0, 0, nil, nil } - q := sqlc.New(r.tx) - row, err := q.GetUserMonthlySpendAndLimit(ctx, sqlc.GetUserMonthlySpendAndLimitParams{ + row, err := r.q().GetUserMonthlySpendAndLimit(ctx, sqlc.GetUserMonthlySpendAndLimitParams{ RouterUserID: parsed, OrganizationID: organizationID, }) if err != nil { - return 0, nil, err + return 0, 0, nil, err } limit := row.OrgDefaultLimitUsdMicros if row.HasOverride { limit = row.OverrideLimitUsdMicros } - return row.SpentUsdMicros, limit, nil + return row.SpentUsdMicros, row.ReservedUsdMicros, limit, nil } // GetOrgMonthlySpendAndLimit reads the org's current UTC-month spend and cap. -// Scalar subqueries guarantee a row; missing config/spend -> nil limit / zero spend, not an error. -func (r *BillingRepo) GetOrgMonthlySpendAndLimit(ctx context.Context, organizationID string) (int64, *int64, error) { - q := sqlc.New(r.tx) - row, err := q.GetOrgMonthlySpendAndLimit(ctx, organizationID) +func (r *BillingRepo) GetOrgMonthlySpendAndLimit(ctx context.Context, organizationID string) (int64, int64, *int64, error) { + row, err := r.q().GetOrgMonthlySpendAndLimit(ctx, organizationID) if err != nil { - return 0, nil, err + return 0, 0, nil, err } - return row.SpentUsdMicros, row.OrgLimitUsdMicros, nil + return row.SpentUsdMicros, row.ReservedUsdMicros, row.OrgLimitUsdMicros, nil } // GetAutopayConfig reads the org's autopay enabled flag and recharge threshold. -// Maps pgx.ErrNoRows (org never configured autopay) to enabled=false with a nil -// error so the debit hook skips the crossing check rather than treating a -// missing row as a failure. func (r *BillingRepo) GetAutopayConfig(ctx context.Context, orgID string) (bool, int64, error) { - q := sqlc.New(r.tx) - row, err := q.GetAutopayConfig(ctx, orgID) + row, err := r.q().GetAutopayConfig(ctx, orgID) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return false, 0, nil @@ -145,9 +146,258 @@ func (r *BillingRepo) GetAutopayConfig(ctx context.Context, orgID string) (bool, return row.Enabled, row.ThresholdUsdMicros, nil } -// BillingTablesExist runs the boot-time health check. Returns true when -// all three billing tables exist in the router schema. +// BillingTablesExist runs the boot-time health check. func (r *BillingRepo) BillingTablesExist(ctx context.Context) (bool, error) { - q := sqlc.New(r.tx) - return q.CheckBillingTablesExist(ctx) + return r.q().CheckBillingTablesExist(ctx) +} + +// ReserveSpendCaps reserves all applicable scopes in one transaction. +func (r *BillingRepo) ReserveSpendCaps(ctx context.Context, p billing.ReserveSpendCapsParams) ([]uuid.UUID, error) { + if p.AmountUsdMicros <= 0 { + return nil, fmt.Errorf("billing: reserve amount must be positive") + } + if p.TTL <= 0 { + return nil, fmt.Errorf("billing: reserve TTL must be positive") + } + + tx, err := r.pool.Begin(ctx) + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback(ctx) }() + q := sqlc.New(tx) + + expires := time.Now().UTC().Add(p.TTL) + expiresTS := pgtype.Timestamptz{Time: expires, Valid: true} + month := utcMonthDate() + var ids []uuid.UUID + + if !p.SkipOrg && p.OrganizationID != "" { + spent, _, limit, err := r.getOrgMonthlyInTx(ctx, q, p.OrganizationID) + if err != nil { + return nil, err + } + _ = spent + if limit != nil { + if err := q.EnsureOrgMonthlySpendRow(ctx, p.OrganizationID); err != nil { + return nil, err + } + _, err := q.TryBumpOrgMonthReserved(ctx, sqlc.TryBumpOrgMonthReservedParams{ + AmountUsdMicros: p.AmountUsdMicros, + OrganizationID: p.OrganizationID, + }) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, billing.ErrOrgMonthlySpendLimitReached + } + return nil, err + } + row, err := q.InsertSpendReservation(ctx, sqlc.InsertSpendReservationParams{ + ScopeKind: billing.ScopeOrgMonth, + ScopeID: p.OrganizationID, + Month: month, + AmountUsdMicros: p.AmountUsdMicros, + ExpiresAt: expiresTS, + RouterRequestID: stringPtrOrNil(p.RouterRequestID), + }) + if err != nil { + return nil, err + } + ids = append(ids, row.ID) + } + } + + if !p.SkipKey && p.APIKeyID != "" { + keyUUID, err := uuid.Parse(p.APIKeyID) + if err != nil { + return nil, fmt.Errorf("billing: invalid api_key_id: %w", err) + } + spentRow, err := q.GetModelRouterAPIKeySpend(ctx, keyUUID) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return nil, err + } + if err == nil && spentRow.SpendCapUsdMicros != nil { + _, err := q.TryBumpAPIKeyReserved(ctx, sqlc.TryBumpAPIKeyReservedParams{ + AmountUsdMicros: p.AmountUsdMicros, + APIKeyID: keyUUID, + }) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, billing.ErrAPIKeySpendCapReached + } + return nil, err + } + row, err := q.InsertSpendReservation(ctx, sqlc.InsertSpendReservationParams{ + ScopeKind: billing.ScopeAPIKey, + ScopeID: p.APIKeyID, + Month: pgtype.Date{}, // NULL for lifetime + AmountUsdMicros: p.AmountUsdMicros, + ExpiresAt: expiresTS, + RouterRequestID: stringPtrOrNil(p.RouterRequestID), + }) + if err != nil { + return nil, err + } + ids = append(ids, row.ID) + } + } + + if !p.SkipUser && p.RouterUserID != "" && p.OrganizationID != "" { + userUUID, err := uuid.Parse(p.RouterUserID) + if err != nil { + return nil, fmt.Errorf("billing: invalid router_user_id: %w", err) + } + _, _, limit, err := r.getUserMonthlyInTx(ctx, q, p.OrganizationID, userUUID) + if err != nil { + return nil, err + } + if limit != nil { + if err := q.EnsureUserMonthlySpendRow(ctx, userUUID); err != nil { + return nil, err + } + _, err := q.TryBumpUserMonthReserved(ctx, sqlc.TryBumpUserMonthReservedParams{ + AmountUsdMicros: p.AmountUsdMicros, + RouterUserID: userUUID, + LimitUsdMicros: *limit, + }) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, billing.ErrUserMonthlySpendLimitReached + } + return nil, err + } + row, err := q.InsertSpendReservation(ctx, sqlc.InsertSpendReservationParams{ + ScopeKind: billing.ScopeUserMonth, + ScopeID: p.RouterUserID, + Month: month, + AmountUsdMicros: p.AmountUsdMicros, + ExpiresAt: expiresTS, + RouterRequestID: stringPtrOrNil(p.RouterRequestID), + }) + if err != nil { + return nil, err + } + ids = append(ids, row.ID) + } + } + + if err := tx.Commit(ctx); err != nil { + return nil, err + } + return ids, nil +} + +func (r *BillingRepo) getOrgMonthlyInTx(ctx context.Context, q *sqlc.Queries, orgID string) (int64, int64, *int64, error) { + row, err := q.GetOrgMonthlySpendAndLimit(ctx, orgID) + if err != nil { + return 0, 0, nil, err + } + return row.SpentUsdMicros, row.ReservedUsdMicros, row.OrgLimitUsdMicros, nil +} + +func (r *BillingRepo) getUserMonthlyInTx(ctx context.Context, q *sqlc.Queries, orgID string, userID uuid.UUID) (int64, int64, *int64, error) { + row, err := q.GetUserMonthlySpendAndLimit(ctx, sqlc.GetUserMonthlySpendAndLimitParams{ + RouterUserID: userID, + OrganizationID: orgID, + }) + if err != nil { + return 0, 0, nil, err + } + limit := row.OrgDefaultLimitUsdMicros + if row.HasOverride { + limit = row.OverrideLimitUsdMicros + } + return row.SpentUsdMicros, row.ReservedUsdMicros, limit, nil +} + +// ReleaseSpendReservations consumes ids via DELETE … RETURNING. +func (r *BillingRepo) ReleaseSpendReservations(ctx context.Context, ids []uuid.UUID) error { + if len(ids) == 0 { + return nil + } + tx, err := r.pool.Begin(ctx) + if err != nil { + return err + } + defer func() { _ = tx.Rollback(ctx) }() + q := sqlc.New(tx) + for _, id := range ids { + if err := consumeSpendReservation(ctx, q, id); err != nil { + return err + } + } + return tx.Commit(ctx) +} + +// SweepExpiredSpendReservations deletes expired rows and decrements reserved. +func (r *BillingRepo) SweepExpiredSpendReservations(ctx context.Context, now time.Time) (int, error) { + tx, err := r.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer func() { _ = tx.Rollback(ctx) }() + q := sqlc.New(tx) + rows, err := q.DeleteExpiredSpendReservations(ctx, pgtype.Timestamptz{Time: now.UTC(), Valid: true}) + if err != nil { + return 0, err + } + for _, row := range rows { + if err := decrementReservedForScope(ctx, q, row.ScopeKind, row.ScopeID, row.Month, row.AmountUsdMicros); err != nil { + return 0, err + } + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return len(rows), nil +} + +// consumeSpendReservation is the shared DELETE … RETURNING + conditional +// reserved decrement used by settle, release, and (per-id) callers. +func consumeSpendReservation(ctx context.Context, q *sqlc.Queries, id uuid.UUID) error { + row, err := q.DeleteSpendReservation(ctx, id) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil + } + return err + } + return decrementReservedForScope(ctx, q, row.ScopeKind, row.ScopeID, row.Month, row.AmountUsdMicros) +} + +func decrementReservedForScope(ctx context.Context, q *sqlc.Queries, scopeKind, scopeID string, month pgtype.Date, amount int64) error { + switch scopeKind { + case billing.ScopeOrgMonth: + return q.DecrementOrgMonthReserved(ctx, sqlc.DecrementOrgMonthReservedParams{ + AmountUsdMicros: amount, + OrganizationID: scopeID, + Month: month, + }) + case billing.ScopeUserMonth: + uid, err := uuid.Parse(scopeID) + if err != nil { + return fmt.Errorf("billing: sweep user scope_id: %w", err) + } + return q.DecrementUserMonthReserved(ctx, sqlc.DecrementUserMonthReservedParams{ + AmountUsdMicros: amount, + RouterUserID: uid, + Month: month, + }) + case billing.ScopeAPIKey: + uid, err := uuid.Parse(scopeID) + if err != nil { + return fmt.Errorf("billing: sweep api_key scope_id: %w", err) + } + return q.DecrementAPIKeyReserved(ctx, sqlc.DecrementAPIKeyReservedParams{ + AmountUsdMicros: amount, + APIKeyID: uid, + }) + default: + return fmt.Errorf("billing: unknown spend reservation scope_kind %q", scopeKind) + } +} + +func utcMonthDate() pgtype.Date { + now := time.Now().UTC() + return pgtype.Date{Time: time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC), Valid: true} } diff --git a/internal/postgres/org_monthly_spend_toctou_test.go b/internal/postgres/org_monthly_spend_toctou_test.go index aa67fd90..b77c7c17 100644 --- a/internal/postgres/org_monthly_spend_toctou_test.go +++ b/internal/postgres/org_monthly_spend_toctou_test.go @@ -19,18 +19,15 @@ import ( ) // TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot is the -// permanent #793 regression: N concurrent CheckOrgMonthlySpend + DebitForInference -// must not multiply past the hard monthly cap unboundedly. +// permanent #793 regression: N concurrent reserve-then-settle turns must not +// multiply past the hard monthly cap unboundedly. // // Fixture matches the issue reproduction: $1 limit, $0.90 starting spend, // $6.75 notional per debit (1M in @ $3/MTok + 250K out @ $15/MTok), N=20. -// Soft-overshoot bound after the reserve-then-settle fix is roughly one -// under-reserved turn (or zero when remaining headroom < R); before the fix -// this assertion fails with ~$100+ overshoot. +// With fixed R=$1 and only $0.10 headroom, zero reserves succeed and spend +// stays at $0.90. Soft overshoot bound remains limit + one turn. // -// Gated on ROUTER_TEST_DATABASE_URL (falls back to DATABASE_URL). Skips when -// neither is set so unit CI without Postgres stays green. Requires a migrated -// router schema (docker compose postgres on :5433 is the local fixture). +// Gated on ROUTER_TEST_DATABASE_URL (falls back to DATABASE_URL). func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) { dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") if dsn == "" { @@ -79,18 +76,20 @@ func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) require.NoError(t, err) _, err = pool.Exec(ctx, ` INSERT INTO organization_credit_balance (organization_id, balance_usd_micros) - VALUES ($1, $2)`, orgID, int64(1_000_000_000)) // $1000 — not under test + VALUES ($1, $2)`, orgID, int64(1_000_000_000)) require.NoError(t, err) t.Cleanup(func() { cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() + _, _ = pool.Exec(cleanupCtx, `DELETE FROM spend_reservations WHERE scope_id = $1`, orgID) _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_credit_ledger WHERE organization_id = $1`, orgID) _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_credit_balance WHERE organization_id = $1`, orgID) _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_monthly_spend WHERE organization_id = $1`, orgID) _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_spend_limits WHERE organization_id = $1`, orgID) }) - svc := billing.NewService(postgres.NewBillingRepo(pool)) + svc := billing.NewService(postgres.NewBillingRepo(pool)). + WithReservationConfig(billing.DefaultReserveAmountMicros, billing.DefaultReserveTTL) var ( wg sync.WaitGroup @@ -105,15 +104,13 @@ func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) reqCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() - res, err := svc.CheckOrgMonthlySpend(reqCtx, orgID) + reqCtx, release, err := svc.ArmSpendReservations(reqCtx, orgID, "", "", fmt.Sprintf("%s-%d", orgID, i)) if err != nil { - t.Errorf("CheckOrgMonthlySpend: %v", err) - return - } - if res.LimitReached() { rejected.Add(1) return } + defer release() + _, err = svc.DebitForInference(reqCtx, billing.DebitInferenceParams{ OrganizationID: orgID, RouterRequestID: fmt.Sprintf("%s-%d", orgID, i), @@ -121,27 +118,32 @@ func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) InputTokens: inputTokens, OutputTokens: outputTokens, Pricing: pricing, + ReservationIDs: billing.SpendHoldFrom(reqCtx).IDs, }) if err != nil { t.Errorf("DebitForInference: %v", err) return } + if hold := billing.SpendHoldFrom(reqCtx); hold != nil { + hold.MarkSettled() + } debited.Add(1) }() } wg.Wait() - final, _, err := postgres.NewBillingRepo(pool).GetOrgMonthlySpendAndLimit(ctx, orgID) + final, reserved, err := func() (int64, int64, error) { + s, r, _, e := postgres.NewBillingRepo(pool).GetOrgMonthlySpendAndLimit(ctx, orgID) + return s, r, e + }() require.NoError(t, err) - // Soft overshoot bound: at most one in-flight turn past the limit (the - // product promise before concurrency multiplied it). Equivalent to - // start + N*cost when every check races — that path must fail this assert. maxAllowed := limitMicros + perCallMicros - t.Logf("org=%s debited=%d rejected=%d final_spent_usd_micros=%d overshoot_usd_micros=%d max_allowed=%d", - orgID, debited.Load(), rejected.Load(), final, final-limitMicros, maxAllowed) + t.Logf("org=%s debited=%d rejected=%d final_spent_usd_micros=%d reserved=%d overshoot_usd_micros=%d max_allowed=%d", + orgID, debited.Load(), rejected.Load(), final, reserved, final-limitMicros, maxAllowed) require.LessOrEqual(t, final, maxAllowed, "org monthly spend TOCTOU (#793): final=%d ($%.2f) exceeds limit+one_turn=%d ($%.2f); overshoot=$%.2f from %d concurrent debits", final, float64(final)/1e6, maxAllowed, float64(maxAllowed)/1e6, float64(final-limitMicros)/1e6, debited.Load()) + require.Equal(t, int64(0), reserved, "all reservations must be settled or released") } diff --git a/internal/proxy/dispatch_error.go b/internal/proxy/dispatch_error.go index 644b051a..3faaaefb 100644 --- a/internal/proxy/dispatch_error.go +++ b/internal/proxy/dispatch_error.go @@ -131,6 +131,22 @@ func ClassifyDispatchError(err error) (DispatchErrorClass, bool) { LogLevel: "warn", LogMessage: "Request refused: engineer monthly spend limit reached", }, true + case errors.Is(err, billing.ErrOrgMonthlySpendLimitReached): + return DispatchErrorClass{ + Kind: DispatchErrorUserSpendLimitReached, // same 402 kind bucket + Status: http.StatusPaymentRequired, + Message: "Your organization has reached its monthly Weave Router spend limit. An org admin can raise the limit, or it resets next month.", + LogLevel: "warn", + LogMessage: "Request refused: org monthly spend limit reached", + }, true + case errors.Is(err, billing.ErrAPIKeySpendCapReached): + return DispatchErrorClass{ + Kind: DispatchErrorUserSpendLimitReached, + Status: http.StatusPaymentRequired, + Message: "This router key has reached its spend cap. Mint a new key or raise the cap to continue.", + LogLevel: "warn", + LogMessage: "Request refused: api key spend cap reached", + }, true case errors.Is(err, billing.ErrSpendLimitCheckUnavailable): return DispatchErrorClass{ Kind: DispatchErrorSpendLimitUnavailable, diff --git a/internal/proxy/gemini.go b/internal/proxy/gemini.go index 8d2c1e25..a70f23f4 100644 --- a/internal/proxy/gemini.go +++ b/internal/proxy/gemini.go @@ -31,9 +31,6 @@ var ErrGeminiCrossFormatUnsupported = errors.New("gemini cross-format emit not i // and "stream" (true for :streamGenerateContent) fields into body before // calling; both are stripped before forwarding upstream. func (s *Service) ProxyGeminiGenerateContent(ctx context.Context, body []byte, w http.ResponseWriter, r *http.Request) error { - if err := s.checkUserMonthlySpendLimit(ctx); err != nil { - return err - } log := observability.FromContext(ctx) requestStart := time.Now() requestID := uuid.New().String() diff --git a/internal/proxy/service.go b/internal/proxy/service.go index 4f654ffb..2fda771b 100644 --- a/internal/proxy/service.go +++ b/internal/proxy/service.go @@ -1907,9 +1907,6 @@ func (s *Service) maybeRepinOnRefusal(ctx context.Context, obs *refusalObserver, } func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.ResponseWriter, r *http.Request) error { - if err := s.checkUserMonthlySpendLimit(ctx); err != nil { - return err - } ctx = s.withUsageObserver(ctx, r.Header) log := observability.FromContext(ctx) requestStart := time.Now() @@ -3931,10 +3928,19 @@ func (s *Service) fireBilling(ctx context.Context, p billing.DebitInferenceParam observability.Get().Debug("Billing debit skipped: no organization_id on request") return } + hold := billing.SpendHoldFrom(ctx) + // Settle the main-request hold only once (first debit). Compaction / + // handover summary debits stay unreserved soft overshoot (#793 residual). + if hold != nil && !hold.Settled() && len(p.ReservationIDs) == 0 { + p.ReservationIDs = hold.IDs + } dbCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() balance, err := s.billing.DebitForInference(dbCtx, p) if err == nil { + if hold != nil && len(p.ReservationIDs) > 0 { + hold.MarkSettled() + } observability.Get().Debug("Billing debit complete", "organization_id", p.OrganizationID, "router_request_id", p.RouterRequestID, @@ -4003,9 +4009,6 @@ func finalizeAfterProxy(proxyErr error, fn func() error) error { // ProxyOpenAIChatCompletion routes an OpenAI Chat Completion request, // translating cross-format when the decision picks a non-OpenAI provider. func (s *Service) ProxyOpenAIChatCompletion(ctx context.Context, body []byte, w http.ResponseWriter, r *http.Request) error { - if err := s.checkUserMonthlySpendLimit(ctx); err != nil { - return err - } ctx = s.withUsageObserver(ctx, r.Header) log := observability.FromContext(ctx) requestStart := time.Now() diff --git a/internal/proxy/spend_limit.go b/internal/proxy/spend_limit.go index 646b474c..062f7f08 100644 --- a/internal/proxy/spend_limit.go +++ b/internal/proxy/spend_limit.go @@ -9,9 +9,9 @@ import ( "workweave/router/internal/observability" ) -// checkUserMonthlySpendLimit gates a turn on the resolved engineer's monthly -// spend limit. Runs inside the proxy (not middleware) because user identity is -// resolved by the handler after the middleware chain; no identity passes through. +// checkUserMonthlySpendLimit is retained for unit tests and as a read-only +// helper. Request-path enforcement uses billing.ArmSpendReservations (combined +// reserve of org/key/user caps) in API handlers before Proxy*. func (s *Service) checkUserMonthlySpendLimit(ctx context.Context) error { if s.billing == nil { return nil diff --git a/internal/proxy/spend_limit_internal_test.go b/internal/proxy/spend_limit_internal_test.go index c3a6ce13..fdcaff34 100644 --- a/internal/proxy/spend_limit_internal_test.go +++ b/internal/proxy/spend_limit_internal_test.go @@ -4,10 +4,12 @@ import ( "context" "errors" "testing" + "time" "workweave/router/internal/auth" "workweave/router/internal/billing" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -27,14 +29,21 @@ func (r *spendLimitRepo) HasActiveOverride(context.Context, string) (bool, error func (r *spendLimitRepo) DebitInference(context.Context, billing.DebitParams) (int64, error) { return 0, nil } -func (r *spendLimitRepo) GetAPIKeySpend(context.Context, string) (int64, *int64, bool, error) { - return 0, nil, false, nil +func (r *spendLimitRepo) GetAPIKeySpend(context.Context, string) (int64, int64, *int64, bool, error) { + return 0, 0, nil, false, nil } -func (r *spendLimitRepo) GetUserMonthlySpendAndLimit(context.Context, string, string) (int64, *int64, error) { - return r.spent, r.limit, r.err +func (r *spendLimitRepo) GetUserMonthlySpendAndLimit(context.Context, string, string) (int64, int64, *int64, error) { + return r.spent, 0, r.limit, r.err } -func (r *spendLimitRepo) GetOrgMonthlySpendAndLimit(context.Context, string) (int64, *int64, error) { - return 0, nil, nil +func (r *spendLimitRepo) GetOrgMonthlySpendAndLimit(context.Context, string) (int64, int64, *int64, error) { + return 0, 0, nil, nil +} +func (r *spendLimitRepo) ReserveSpendCaps(context.Context, billing.ReserveSpendCapsParams) ([]uuid.UUID, error) { + return nil, nil +} +func (r *spendLimitRepo) ReleaseSpendReservations(context.Context, []uuid.UUID) error { return nil } +func (r *spendLimitRepo) SweepExpiredSpendReservations(context.Context, time.Time) (int, error) { + return 0, nil } func (r *spendLimitRepo) GetAutopayConfig(context.Context, string) (bool, int64, error) { return false, 0, nil diff --git a/internal/proxy/spend_reserve.go b/internal/proxy/spend_reserve.go new file mode 100644 index 00000000..5db08679 --- /dev/null +++ b/internal/proxy/spend_reserve.go @@ -0,0 +1,47 @@ +package proxy + +import ( + "context" + "errors" + + "workweave/router/internal/auth" + "workweave/router/internal/billing" + "workweave/router/internal/observability" + + "github.com/google/uuid" +) + +// ArmSpendReservations reserves all applicable spend caps in one transaction +// after identity is on ctx and before Proxy*. Returns a release func for +// defer (no-ops after DebitForInference settles). Nil billing is a no-op. +// Agent-shadow evaluation traffic skips entirely (no reservation, no gate), +// matching the billing middleware skips from #787 — synthetic eval must not +// consume production spend budget. +func (s *Service) ArmSpendReservations(ctx context.Context) (context.Context, func(), error) { + if s == nil || s.billing == nil { + return ctx, func() {}, nil + } + if _, ok := AgentShadowEvalFromContext(ctx); ok { + return ctx, func() {}, nil + } + orgID, _ := ctx.Value(ExternalIDContextKey{}).(string) + apiKeyID, _ := ctx.Value(APIKeyIDContextKey{}).(string) + userID := auth.UserIDFrom(ctx) + requestID := uuid.New().String() + ctx, release, err := s.billing.ArmSpendReservations(ctx, orgID, apiKeyID, userID, requestID) + if err != nil { + observability.FromContext(ctx).Info("Spend reservation refused", + "err", err, + "organization_id", orgID, + "router_user_id", userID, + ) + // Preserve sentinel identity for ClassifyDispatchError. + if errors.Is(err, billing.ErrOrgMonthlySpendLimitReached) || + errors.Is(err, billing.ErrAPIKeySpendCapReached) || + errors.Is(err, billing.ErrUserMonthlySpendLimitReached) { + return ctx, nil, err + } + return ctx, nil, err + } + return ctx, release, nil +} diff --git a/internal/server/middleware/api_key_spend_cap.go b/internal/server/middleware/api_key_spend_cap.go index 8199b43a..dfc8faa4 100644 --- a/internal/server/middleware/api_key_spend_cap.go +++ b/internal/server/middleware/api_key_spend_cap.go @@ -53,10 +53,11 @@ func WithAPIKeySpendCap(svc *billing.Service) gin.HandlerFunc { return } - if result.SpentMicros >= *result.CapMicros { + if result.SpentMicros+result.ReservedMicros >= *result.CapMicros { log.Info("Request rejected: api key spend cap reached", "api_key_id", apiKey.ID, "spent_usd_micros", result.SpentMicros, + "reserved_usd_micros", result.ReservedMicros, "spend_cap_usd_micros", *result.CapMicros, ) c.AbortWithStatusJSON(http.StatusPaymentRequired, gin.H{ diff --git a/internal/server/middleware/balance_check_test.go b/internal/server/middleware/balance_check_test.go index 44fcf8ac..23c4dd41 100644 --- a/internal/server/middleware/balance_check_test.go +++ b/internal/server/middleware/balance_check_test.go @@ -7,6 +7,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "workweave/router/internal/auth" "workweave/router/internal/billing" @@ -14,6 +15,7 @@ import ( "workweave/router/internal/server/middleware" "github.com/gin-gonic/gin" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -46,14 +48,23 @@ func (r *stubBillingRepo) HasActiveOverride(_ context.Context, _ string) (bool, func (r *stubBillingRepo) DebitInference(_ context.Context, _ billing.DebitParams) (int64, error) { return 0, nil } -func (r *stubBillingRepo) GetAPIKeySpend(_ context.Context, _ string) (int64, *int64, bool, error) { - return r.spendMicros, r.capMicros, r.spendFound, r.spendErr +func (r *stubBillingRepo) GetAPIKeySpend(_ context.Context, _ string) (int64, int64, *int64, bool, error) { + return r.spendMicros, 0, r.capMicros, r.spendFound, r.spendErr } -func (r *stubBillingRepo) GetUserMonthlySpendAndLimit(_ context.Context, _, _ string) (int64, *int64, error) { - return 0, nil, nil +func (r *stubBillingRepo) GetUserMonthlySpendAndLimit(_ context.Context, _, _ string) (int64, int64, *int64, error) { + return 0, 0, nil, nil } -func (r *stubBillingRepo) GetOrgMonthlySpendAndLimit(_ context.Context, _ string) (int64, *int64, error) { - return r.orgMonthSpent, r.orgMonthLimit, r.orgMonthErr +func (r *stubBillingRepo) GetOrgMonthlySpendAndLimit(_ context.Context, _ string) (int64, int64, *int64, error) { + return r.orgMonthSpent, 0, r.orgMonthLimit, r.orgMonthErr +} +func (r *stubBillingRepo) ReserveSpendCaps(_ context.Context, _ billing.ReserveSpendCapsParams) ([]uuid.UUID, error) { + return nil, nil +} +func (r *stubBillingRepo) ReleaseSpendReservations(_ context.Context, _ []uuid.UUID) error { + return nil +} +func (r *stubBillingRepo) SweepExpiredSpendReservations(_ context.Context, _ time.Time) (int, error) { + return 0, nil } func (r *stubBillingRepo) BillingTablesExist(_ context.Context) (bool, error) { return true, nil diff --git a/internal/server/server.go b/internal/server/server.go index bb9c45bf..1742ecdf 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -150,7 +150,9 @@ func Register(engine *gin.Engine, authSvc *auth.Service, proxySvc *proxy.Service middleware.WithAgentShadowEvaluation(), } if billingSvc != nil { - messagesMiddleware = append(messagesMiddleware, middleware.WithBalanceCheck(billingSvc, billing.MinBalanceMicros), middleware.WithAPIKeySpendCap(billingSvc), middleware.WithOrgMonthlySpendCap(billingSvc)) + // Spend caps (org / key / user) are reserved in one TX after identity + // resolves in the handler (#793). Balance check stays here. + messagesMiddleware = append(messagesMiddleware, middleware.WithBalanceCheck(billingSvc, billing.MinBalanceMicros)) } messagesMiddleware = append(messagesMiddleware, middleware.WithEmbedOnlyUserMessageOverride(), @@ -169,7 +171,7 @@ func Register(engine *gin.Engine, authSvc *auth.Service, proxySvc *proxy.Service middleware.WithAuth(authSvc, byokDisabled), } if billingSvc != nil { - chatCompletionMiddleware = append(chatCompletionMiddleware, middleware.WithBalanceCheck(billingSvc, billing.MinBalanceMicros), middleware.WithAPIKeySpendCap(billingSvc), middleware.WithOrgMonthlySpendCap(billingSvc)) + chatCompletionMiddleware = append(chatCompletionMiddleware, middleware.WithBalanceCheck(billingSvc, billing.MinBalanceMicros)) } chatCompletionMiddleware = append(chatCompletionMiddleware, middleware.WithEmbedOnlyUserMessageOverride(), diff --git a/internal/sqlc/model_router_api_keys.sql.go b/internal/sqlc/model_router_api_keys.sql.go index 60df1f17..36bea158 100644 --- a/internal/sqlc/model_router_api_keys.sql.go +++ b/internal/sqlc/model_router_api_keys.sql.go @@ -30,7 +30,7 @@ VALUES ( $6::varchar, $7 ) -RETURNING id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros +RETURNING id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros, reserved_usd_micros ` type CreateModelRouterAPIKeyParams struct { @@ -63,7 +63,7 @@ type CreateModelRouterAPIKeyParams struct { // $6::varchar, // $7 // ) -// RETURNING id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros +// RETURNING id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros, reserved_usd_micros func (q *Queries) CreateModelRouterAPIKey(ctx context.Context, arg CreateModelRouterAPIKeyParams) (RouterModelRouterAPIKey, error) { row := q.db.QueryRow(ctx, createModelRouterAPIKey, arg.InstallationID, @@ -89,12 +89,13 @@ func (q *Queries) CreateModelRouterAPIKey(ctx context.Context, arg CreateModelRo &i.CreatedBy, &i.SpendCapUsdMicros, &i.SpentUsdMicros, + &i.ReservedUsdMicros, ) return i, err } const getActiveModelRouterAPIKeyWithInstallationByHash = `-- name: GetActiveModelRouterAPIKeyWithInstallationByHash :one -SELECT k.id, k.installation_id, k.external_id, k.name, k.key_prefix, k.key_hash, k.key_suffix, k.last_used_at, k.created_at, k.deleted_at, k.created_by, k.spend_cap_usd_micros, k.spent_usd_micros, i.id, i.external_id, i.name, i.created_at, i.updated_at, i.deleted_at, i.created_by, i.excluded_models, i.excluded_providers, i.routing_quality_weight, i.usage_bypass_enabled, i.usage_bypass_threshold, i.preferred_models, i.subscription_routing_disabled, i.routing_strategy, i.routing_rollout_id, i.policy_shadow_strategy, i.policy_debug_enabled, i.policy_header_overrides_enabled, i.policy_routing_intent, i.ai_training_allowed +SELECT k.id, k.installation_id, k.external_id, k.name, k.key_prefix, k.key_hash, k.key_suffix, k.last_used_at, k.created_at, k.deleted_at, k.created_by, k.spend_cap_usd_micros, k.spent_usd_micros, k.reserved_usd_micros, i.id, i.external_id, i.name, i.created_at, i.updated_at, i.deleted_at, i.created_by, i.excluded_models, i.excluded_providers, i.routing_quality_weight, i.usage_bypass_enabled, i.usage_bypass_threshold, i.preferred_models, i.subscription_routing_disabled, i.routing_strategy, i.routing_rollout_id, i.policy_shadow_strategy, i.policy_debug_enabled, i.policy_header_overrides_enabled, i.policy_routing_intent, i.ai_training_allowed FROM router.model_router_api_keys k INNER JOIN router.model_router_installations i ON i.id = k.installation_id WHERE k.key_hash = $1::varchar @@ -113,7 +114,7 @@ type GetActiveModelRouterAPIKeyWithInstallationByHashRow struct { // both sides so a soft-deleted installation invalidates all its keys without per-key // updates. // -// SELECT k.id, k.installation_id, k.external_id, k.name, k.key_prefix, k.key_hash, k.key_suffix, k.last_used_at, k.created_at, k.deleted_at, k.created_by, k.spend_cap_usd_micros, k.spent_usd_micros, i.id, i.external_id, i.name, i.created_at, i.updated_at, i.deleted_at, i.created_by, i.excluded_models, i.excluded_providers, i.routing_quality_weight, i.usage_bypass_enabled, i.usage_bypass_threshold, i.preferred_models, i.subscription_routing_disabled, i.routing_strategy, i.routing_rollout_id, i.policy_shadow_strategy, i.policy_debug_enabled, i.policy_header_overrides_enabled, i.policy_routing_intent, i.ai_training_allowed +// SELECT k.id, k.installation_id, k.external_id, k.name, k.key_prefix, k.key_hash, k.key_suffix, k.last_used_at, k.created_at, k.deleted_at, k.created_by, k.spend_cap_usd_micros, k.spent_usd_micros, k.reserved_usd_micros, i.id, i.external_id, i.name, i.created_at, i.updated_at, i.deleted_at, i.created_by, i.excluded_models, i.excluded_providers, i.routing_quality_weight, i.usage_bypass_enabled, i.usage_bypass_threshold, i.preferred_models, i.subscription_routing_disabled, i.routing_strategy, i.routing_rollout_id, i.policy_shadow_strategy, i.policy_debug_enabled, i.policy_header_overrides_enabled, i.policy_routing_intent, i.ai_training_allowed // FROM router.model_router_api_keys k // INNER JOIN router.model_router_installations i ON i.id = k.installation_id // WHERE k.key_hash = $1::varchar @@ -136,6 +137,7 @@ func (q *Queries) GetActiveModelRouterAPIKeyWithInstallationByHash(ctx context.C &i.RouterModelRouterAPIKey.CreatedBy, &i.RouterModelRouterAPIKey.SpendCapUsdMicros, &i.RouterModelRouterAPIKey.SpentUsdMicros, + &i.RouterModelRouterAPIKey.ReservedUsdMicros, &i.RouterModelRouterInstallation.ID, &i.RouterModelRouterInstallation.ExternalID, &i.RouterModelRouterInstallation.Name, @@ -162,7 +164,7 @@ func (q *Queries) GetActiveModelRouterAPIKeyWithInstallationByHash(ctx context.C } const getModelRouterAPIKeySpend = `-- 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 = $1::uuid AND deleted_at IS NULL @@ -171,6 +173,7 @@ WHERE id = $1::uuid type GetModelRouterAPIKeySpendRow struct { SpendCapUsdMicros *int64 SpentUsdMicros int64 + ReservedUsdMicros int64 } // Fresh per-request read of a key's lifetime spend cap and spend-to-date for @@ -179,19 +182,19 @@ type GetModelRouterAPIKeySpendRow struct { // the per-request balance read in WithBalanceCheck. Returns sql.ErrNoRows when // the key is missing/soft-deleted. // -// 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 = $1::uuid // AND deleted_at IS NULL func (q *Queries) GetModelRouterAPIKeySpend(ctx context.Context, id uuid.UUID) (GetModelRouterAPIKeySpendRow, error) { row := q.db.QueryRow(ctx, getModelRouterAPIKeySpend, id) var i GetModelRouterAPIKeySpendRow - err := row.Scan(&i.SpendCapUsdMicros, &i.SpentUsdMicros) + err := row.Scan(&i.SpendCapUsdMicros, &i.SpentUsdMicros, &i.ReservedUsdMicros) return i, err } const listModelRouterAPIKeysForInstallation = `-- name: ListModelRouterAPIKeysForInstallation :many -SELECT id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros +SELECT id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros, reserved_usd_micros FROM router.model_router_api_keys WHERE installation_id = $1::uuid AND deleted_at IS NULL @@ -200,7 +203,7 @@ ORDER BY created_at DESC // Lists active keys for an installation (dashboard / CRUD; not on the request path). // -// SELECT id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros +// SELECT id, installation_id, external_id, name, key_prefix, key_hash, key_suffix, last_used_at, created_at, deleted_at, created_by, spend_cap_usd_micros, spent_usd_micros, reserved_usd_micros // FROM router.model_router_api_keys // WHERE installation_id = $1::uuid // AND deleted_at IS NULL @@ -228,6 +231,7 @@ func (q *Queries) ListModelRouterAPIKeysForInstallation(ctx context.Context, ins &i.CreatedBy, &i.SpendCapUsdMicros, &i.SpentUsdMicros, + &i.ReservedUsdMicros, ); err != nil { return nil, err } diff --git a/internal/sqlc/models.go b/internal/sqlc/models.go index 831a7332..d9f9d59e 100644 --- a/internal/sqlc/models.go +++ b/internal/sqlc/models.go @@ -44,6 +44,7 @@ type RouterModelRouterAPIKey struct { CreatedBy *string SpendCapUsdMicros *int64 SpentUsdMicros int64 + ReservedUsdMicros int64 } // Customer-owned provider API keys for BYOK routing @@ -189,10 +190,11 @@ type RouterModelRouterUser struct { } type RouterModelRouterUserMonthlySpend struct { - RouterUserID uuid.UUID - Month pgtype.Date - SpentUsdMicros int64 - UpdatedAt pgtype.Timestamptz + RouterUserID uuid.UUID + Month pgtype.Date + SpentUsdMicros int64 + UpdatedAt pgtype.Timestamptz + ReservedUsdMicros int64 } type RouterModelRouterUserSpendLimit struct { @@ -253,10 +255,11 @@ type RouterOrganizationCreditLedger struct { } type RouterOrganizationMonthlySpend struct { - OrganizationID string - Month pgtype.Date - SpentUsdMicros int64 - UpdatedAt pgtype.Timestamptz + OrganizationID string + Month pgtype.Date + SpentUsdMicros int64 + UpdatedAt pgtype.Timestamptz + ReservedUsdMicros int64 } type RouterOrganizationSpendLimit struct { @@ -411,6 +414,17 @@ type RouterSessionPin struct { PairedModel string } +type RouterSpendReservation struct { + ID uuid.UUID + ScopeKind string + ScopeID string + Month pgtype.Date + AmountUsdMicros int64 + ExpiresAt pgtype.Timestamptz + CreatedAt pgtype.Timestamptz + RouterRequestID *string +} + // Shadow-mode spiral (death-march) detections: log-only fire-rate corpus measured on live traffic before escalation is armed type RouterSpiralShadowEvent struct { ID uuid.UUID diff --git a/internal/sqlc/spend_limits.sql.go b/internal/sqlc/spend_limits.sql.go index f8a46cf6..34cb8427 100644 --- a/internal/sqlc/spend_limits.sql.go +++ b/internal/sqlc/spend_limits.sql.go @@ -21,12 +21,19 @@ SELECT FROM router.organization_monthly_spend sp WHERE sp.organization_id = $1::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 = $1::varchar + AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + ), 0)::bigint AS reserved_usd_micros ` type GetOrgMonthlySpendAndLimitRow struct { OrgLimitUsdMicros *int64 SpentUsdMicros int64 + ReservedUsdMicros int64 } // Reads the org's month-to-date spend alongside its monthly cap for the @@ -42,11 +49,17 @@ type GetOrgMonthlySpendAndLimitRow struct { // FROM router.organization_monthly_spend sp // WHERE sp.organization_id = $1::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 = $1::varchar +// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date +// ), 0)::bigint AS reserved_usd_micros func (q *Queries) GetOrgMonthlySpendAndLimit(ctx context.Context, organizationID string) (GetOrgMonthlySpendAndLimitRow, error) { row := q.db.QueryRow(ctx, getOrgMonthlySpendAndLimit, organizationID) var i GetOrgMonthlySpendAndLimitRow - err := row.Scan(&i.OrgLimitUsdMicros, &i.SpentUsdMicros) + err := row.Scan(&i.OrgLimitUsdMicros, &i.SpentUsdMicros, &i.ReservedUsdMicros) return i, err } @@ -91,7 +104,13 @@ SELECT FROM router.model_router_user_monthly_spend sp WHERE sp.router_user_id = $1::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 = $1::uuid + AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + ), 0)::bigint AS reserved_usd_micros ` type GetUserMonthlySpendAndLimitParams struct { @@ -104,6 +123,7 @@ type GetUserMonthlySpendAndLimitRow struct { OverrideLimitUsdMicros *int64 OrgDefaultLimitUsdMicros *int64 SpentUsdMicros int64 + ReservedUsdMicros int64 } // Resolves a user's effective monthly limit and current-month spend in one @@ -128,7 +148,13 @@ type GetUserMonthlySpendAndLimitRow struct { // FROM router.model_router_user_monthly_spend sp // WHERE sp.router_user_id = $1::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 = $1::uuid +// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date +// ), 0)::bigint AS reserved_usd_micros func (q *Queries) GetUserMonthlySpendAndLimit(ctx context.Context, arg GetUserMonthlySpendAndLimitParams) (GetUserMonthlySpendAndLimitRow, error) { row := q.db.QueryRow(ctx, getUserMonthlySpendAndLimit, arg.RouterUserID, arg.OrganizationID) var i GetUserMonthlySpendAndLimitRow @@ -137,6 +163,7 @@ func (q *Queries) GetUserMonthlySpendAndLimit(ctx context.Context, arg GetUserMo &i.OverrideLimitUsdMicros, &i.OrgDefaultLimitUsdMicros, &i.SpentUsdMicros, + &i.ReservedUsdMicros, ) return i, err } diff --git a/internal/sqlc/spend_reservations.sql.go b/internal/sqlc/spend_reservations.sql.go new file mode 100644 index 00000000..2488f0df --- /dev/null +++ b/internal/sqlc/spend_reservations.sql.go @@ -0,0 +1,408 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: spend_reservations.sql + +package sqlc + +import ( + "context" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" +) + +const decrementAPIKeyReserved = `-- name: DecrementAPIKeyReserved :exec +UPDATE router.model_router_api_keys +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint) +WHERE id = $2::uuid +` + +type DecrementAPIKeyReservedParams struct { + AmountUsdMicros int64 + APIKeyID uuid.UUID +} + +// DecrementAPIKeyReserved +// +// UPDATE router.model_router_api_keys +// SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint) +// WHERE id = $2::uuid +func (q *Queries) DecrementAPIKeyReserved(ctx context.Context, arg DecrementAPIKeyReservedParams) error { + _, err := q.db.Exec(ctx, decrementAPIKeyReserved, arg.AmountUsdMicros, arg.APIKeyID) + return err +} + +const decrementOrgMonthReserved = `-- name: DecrementOrgMonthReserved :exec +UPDATE router.organization_monthly_spend +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint), + updated_at = NOW() +WHERE organization_id = $2::varchar + AND month = $3::date +` + +type DecrementOrgMonthReservedParams struct { + AmountUsdMicros int64 + OrganizationID string + Month pgtype.Date +} + +// DecrementOrgMonthReserved +// +// UPDATE router.organization_monthly_spend +// SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint), +// updated_at = NOW() +// WHERE organization_id = $2::varchar +// AND month = $3::date +func (q *Queries) DecrementOrgMonthReserved(ctx context.Context, arg DecrementOrgMonthReservedParams) error { + _, err := q.db.Exec(ctx, decrementOrgMonthReserved, arg.AmountUsdMicros, arg.OrganizationID, arg.Month) + return err +} + +const decrementUserMonthReserved = `-- name: DecrementUserMonthReserved :exec +UPDATE router.model_router_user_monthly_spend +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint), + updated_at = NOW() +WHERE router_user_id = $2::uuid + AND month = $3::date +` + +type DecrementUserMonthReservedParams struct { + AmountUsdMicros int64 + RouterUserID uuid.UUID + Month pgtype.Date +} + +// DecrementUserMonthReserved +// +// UPDATE router.model_router_user_monthly_spend +// SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - $1::bigint), +// updated_at = NOW() +// WHERE router_user_id = $2::uuid +// AND month = $3::date +func (q *Queries) DecrementUserMonthReserved(ctx context.Context, arg DecrementUserMonthReservedParams) error { + _, err := q.db.Exec(ctx, decrementUserMonthReserved, arg.AmountUsdMicros, arg.RouterUserID, arg.Month) + return err +} + +const deleteExpiredSpendReservations = `-- name: DeleteExpiredSpendReservations :many +DELETE FROM router.spend_reservations +WHERE expires_at < $1::timestamptz +RETURNING id, scope_kind, scope_id, month, amount_usd_micros +` + +type DeleteExpiredSpendReservationsRow struct { + ID uuid.UUID + ScopeKind string + ScopeID string + Month pgtype.Date + AmountUsdMicros int64 +} + +// 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. +// +// DELETE FROM router.spend_reservations +// WHERE expires_at < $1::timestamptz +// RETURNING id, scope_kind, scope_id, month, amount_usd_micros +func (q *Queries) DeleteExpiredSpendReservations(ctx context.Context, now pgtype.Timestamptz) ([]DeleteExpiredSpendReservationsRow, error) { + rows, err := q.db.Query(ctx, deleteExpiredSpendReservations, now) + if err != nil { + return nil, err + } + defer rows.Close() + var items []DeleteExpiredSpendReservationsRow + for rows.Next() { + var i DeleteExpiredSpendReservationsRow + if err := rows.Scan( + &i.ID, + &i.ScopeKind, + &i.ScopeID, + &i.Month, + &i.AmountUsdMicros, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const deleteSpendReservation = `-- name: DeleteSpendReservation :one +DELETE FROM router.spend_reservations +WHERE id = $1::uuid +RETURNING id, scope_kind, scope_id, month, amount_usd_micros +` + +type DeleteSpendReservationRow struct { + ID uuid.UUID + ScopeKind string + ScopeID string + Month pgtype.Date + AmountUsdMicros int64 +} + +// Atomic consume: DELETE … RETURNING is the sole settle/release/sweep +// primitive. Zero rows = already released/settled/swept (idempotent no-op). +// +// DELETE FROM router.spend_reservations +// WHERE id = $1::uuid +// RETURNING id, scope_kind, scope_id, month, amount_usd_micros +func (q *Queries) DeleteSpendReservation(ctx context.Context, id uuid.UUID) (DeleteSpendReservationRow, error) { + row := q.db.QueryRow(ctx, deleteSpendReservation, id) + var i DeleteSpendReservationRow + err := row.Scan( + &i.ID, + &i.ScopeKind, + &i.ScopeID, + &i.Month, + &i.AmountUsdMicros, + ) + return i, err +} + +const ensureOrgMonthlySpendRow = `-- name: EnsureOrgMonthlySpendRow :exec +INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) +VALUES ( + $1::varchar, + DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + 0, + 0 +) +ON CONFLICT (organization_id, month) DO NOTHING +` + +// Ensures the current UTC-month org spend row exists so a reserve UPDATE can +// target it. No-op when the row is already present. +// +// INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) +// VALUES ( +// $1::varchar, +// DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, +// 0, +// 0 +// ) +// ON CONFLICT (organization_id, month) DO NOTHING +func (q *Queries) EnsureOrgMonthlySpendRow(ctx context.Context, organizationID string) error { + _, err := q.db.Exec(ctx, ensureOrgMonthlySpendRow, organizationID) + return err +} + +const ensureUserMonthlySpendRow = `-- name: EnsureUserMonthlySpendRow :exec +INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) +VALUES ( + $1::uuid, + DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + 0, + 0 +) +ON CONFLICT (router_user_id, month) DO NOTHING +` + +// EnsureUserMonthlySpendRow +// +// INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) +// VALUES ( +// $1::uuid, +// DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, +// 0, +// 0 +// ) +// ON CONFLICT (router_user_id, month) DO NOTHING +func (q *Queries) EnsureUserMonthlySpendRow(ctx context.Context, routerUserID uuid.UUID) error { + _, err := q.db.Exec(ctx, ensureUserMonthlySpendRow, routerUserID) + return err +} + +const insertSpendReservation = `-- name: InsertSpendReservation :one +INSERT INTO router.spend_reservations ( + scope_kind, + scope_id, + month, + amount_usd_micros, + expires_at, + router_request_id +) VALUES ( + $1::varchar, + $2::varchar, + $3::date, + $4::bigint, + $5::timestamptz, + $6::varchar +) +RETURNING id, scope_kind, scope_id, month, amount_usd_micros, expires_at +` + +type InsertSpendReservationParams struct { + ScopeKind string + ScopeID string + Month pgtype.Date + AmountUsdMicros int64 + ExpiresAt pgtype.Timestamptz + RouterRequestID *string +} + +type InsertSpendReservationRow struct { + ID uuid.UUID + ScopeKind string + ScopeID string + Month pgtype.Date + AmountUsdMicros int64 + ExpiresAt pgtype.Timestamptz +} + +// InsertSpendReservation +// +// INSERT INTO router.spend_reservations ( +// scope_kind, +// scope_id, +// month, +// amount_usd_micros, +// expires_at, +// router_request_id +// ) VALUES ( +// $1::varchar, +// $2::varchar, +// $3::date, +// $4::bigint, +// $5::timestamptz, +// $6::varchar +// ) +// RETURNING id, scope_kind, scope_id, month, amount_usd_micros, expires_at +func (q *Queries) InsertSpendReservation(ctx context.Context, arg InsertSpendReservationParams) (InsertSpendReservationRow, error) { + row := q.db.QueryRow(ctx, insertSpendReservation, + arg.ScopeKind, + arg.ScopeID, + arg.Month, + arg.AmountUsdMicros, + arg.ExpiresAt, + arg.RouterRequestID, + ) + var i InsertSpendReservationRow + err := row.Scan( + &i.ID, + &i.ScopeKind, + &i.ScopeID, + &i.Month, + &i.AmountUsdMicros, + &i.ExpiresAt, + ) + return i, err +} + +const tryBumpAPIKeyReserved = `-- name: TryBumpAPIKeyReserved :one +UPDATE router.model_router_api_keys k +SET reserved_usd_micros = k.reserved_usd_micros + $1::bigint +WHERE k.id = $2::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 + $1::bigint + <= k.spend_cap_usd_micros +RETURNING k.id +` + +type TryBumpAPIKeyReservedParams struct { + AmountUsdMicros int64 + APIKeyID uuid.UUID +} + +// 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. +// +// UPDATE router.model_router_api_keys k +// SET reserved_usd_micros = k.reserved_usd_micros + $1::bigint +// WHERE k.id = $2::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 + $1::bigint +// <= k.spend_cap_usd_micros +// RETURNING k.id +func (q *Queries) TryBumpAPIKeyReserved(ctx context.Context, arg TryBumpAPIKeyReservedParams) (uuid.UUID, error) { + row := q.db.QueryRow(ctx, tryBumpAPIKeyReserved, arg.AmountUsdMicros, arg.APIKeyID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const tryBumpOrgMonthReserved = `-- name: TryBumpOrgMonthReserved :one +UPDATE router.organization_monthly_spend sp +SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, + updated_at = NOW() +FROM router.organization_spend_limits lim +WHERE sp.organization_id = $2::varchar + AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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 + $1::bigint + <= lim.org_monthly_limit_usd_micros +RETURNING sp.organization_id +` + +type TryBumpOrgMonthReservedParams struct { + AmountUsdMicros int64 + OrganizationID string +} + +// 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). +// +// UPDATE router.organization_monthly_spend sp +// SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, +// updated_at = NOW() +// FROM router.organization_spend_limits lim +// WHERE sp.organization_id = $2::varchar +// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::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 + $1::bigint +// <= lim.org_monthly_limit_usd_micros +// RETURNING sp.organization_id +func (q *Queries) TryBumpOrgMonthReserved(ctx context.Context, arg TryBumpOrgMonthReservedParams) (string, error) { + row := q.db.QueryRow(ctx, tryBumpOrgMonthReserved, arg.AmountUsdMicros, arg.OrganizationID) + var organization_id string + err := row.Scan(&organization_id) + return organization_id, err +} + +const tryBumpUserMonthReserved = `-- name: TryBumpUserMonthReserved :one +UPDATE router.model_router_user_monthly_spend sp +SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, + updated_at = NOW() +WHERE sp.router_user_id = $2::uuid + AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::bigint + <= $3::bigint +RETURNING sp.router_user_id +` + +type TryBumpUserMonthReservedParams struct { + AmountUsdMicros int64 + RouterUserID uuid.UUID + LimitUsdMicros int64 +} + +// 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. +// +// UPDATE router.model_router_user_monthly_spend sp +// SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, +// updated_at = NOW() +// WHERE sp.router_user_id = $2::uuid +// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date +// AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::bigint +// <= $3::bigint +// RETURNING sp.router_user_id +func (q *Queries) TryBumpUserMonthReserved(ctx context.Context, arg TryBumpUserMonthReservedParams) (uuid.UUID, error) { + row := q.db.QueryRow(ctx, tryBumpUserMonthReserved, arg.AmountUsdMicros, arg.RouterUserID, arg.LimitUsdMicros) + var router_user_id uuid.UUID + err := row.Scan(&router_user_id) + return router_user_id, err +} From a1cbbd676ccd351fdef2b47a593f5a0216ce5259 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Sun, 19 Jul 2026 23:49:52 -0400 Subject: [PATCH 4/8] feat(router): TTL sweeper for orphaned spend reservations (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- cmd/router/main.go | 29 +++++++ .../postgres/spend_reservation_sweep_test.go | 81 +++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 internal/postgres/spend_reservation_sweep_test.go diff --git a/cmd/router/main.go b/cmd/router/main.go index e64e815a..22381e68 100644 --- a/cmd/router/main.go +++ b/cmd/router/main.go @@ -183,6 +183,10 @@ func main() { "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") } } @@ -1386,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 ( diff --git a/internal/postgres/spend_reservation_sweep_test.go b/internal/postgres/spend_reservation_sweep_test.go new file mode 100644 index 00000000..088b110e --- /dev/null +++ b/internal/postgres/spend_reservation_sweep_test.go @@ -0,0 +1,81 @@ +package postgres_test + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "workweave/router/internal/billing" + "workweave/router/internal/postgres" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" +) + +// TestSweepExpiredSpendReservations_DecrementsReserved seeds an already- +// expired org-month reservation, runs one sweep cycle, and asserts the row +// is gone and denormalized reserved_usd_micros is decremented. +func TestSweepExpiredSpendReservations_DecrementsReserved(t *testing.T) { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + dsn = os.Getenv("DATABASE_URL") + } + if dsn == "" { + t.Skip("ROUTER_TEST_DATABASE_URL / DATABASE_URL not set") + } + + ctx := context.Background() + cfg, err := pgxpool.ParseConfig(dsn) + require.NoError(t, err) + cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + _, err := conn.Exec(ctx, "SET search_path TO router, public") + return err + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + require.NoError(t, err) + t.Cleanup(pool.Close) + + orgID := fmt.Sprintf("sweep-793-%d", time.Now().UnixNano()) + const amount int64 = 1_000_000 + resID := uuid.New() + + _, err = pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, int64(10_000_000)) + require.NoError(t, err) + _, err = pool.Exec(ctx, ` + INSERT INTO organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) + VALUES ($1, DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, 0, $2)`, orgID, amount) + require.NoError(t, err) + _, err = pool.Exec(ctx, ` + INSERT INTO spend_reservations (id, scope_kind, scope_id, month, amount_usd_micros, expires_at) + VALUES ($1, 'org_month', $2, DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, $3, NOW() - INTERVAL '1 minute')`, + resID, orgID, amount) + require.NoError(t, err) + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, _ = pool.Exec(cleanupCtx, `DELETE FROM spend_reservations WHERE scope_id = $1`, orgID) + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_monthly_spend WHERE organization_id = $1`, orgID) + _, _ = pool.Exec(cleanupCtx, `DELETE FROM organization_spend_limits WHERE organization_id = $1`, orgID) + }) + + svc := billing.NewService(postgres.NewBillingRepo(pool)) + released, err := svc.SweepExpiredReservations(ctx, time.Now().UTC()) + require.NoError(t, err) + require.GreaterOrEqual(t, released, 1, "sweeper must consume at least the seeded expired row") + + var stillThere int + err = pool.QueryRow(ctx, `SELECT COUNT(*) FROM spend_reservations WHERE id = $1`, resID).Scan(&stillThere) + require.NoError(t, err) + require.Equal(t, 0, stillThere, "expired reservation row must be deleted") + + spent, reserved, _, err := postgres.NewBillingRepo(pool).GetOrgMonthlySpendAndLimit(ctx, orgID) + require.NoError(t, err) + require.Equal(t, int64(0), spent) + require.Equal(t, int64(0), reserved, "denormalized reserved_usd_micros must drop by the swept amount") +} From 4184bde8dd9015ba1b5ce3596fd39f004ce5f4d1 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Sun, 19 Jul 2026 23:50:51 -0400 Subject: [PATCH 5/8] style: gofmt org monthly spend TOCTOU test (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- internal/postgres/org_monthly_spend_toctou_test.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/internal/postgres/org_monthly_spend_toctou_test.go b/internal/postgres/org_monthly_spend_toctou_test.go index b77c7c17..725ef63f 100644 --- a/internal/postgres/org_monthly_spend_toctou_test.go +++ b/internal/postgres/org_monthly_spend_toctou_test.go @@ -53,11 +53,11 @@ func TestOrgMonthlySpend_ConcurrentCheckThenDebit_BoundedOvershoot(t *testing.T) orgID := fmt.Sprintf("toctou-793-%d", time.Now().UnixNano()) const ( - limitMicros int64 = 1_000_000 // $1.00 - startSpent int64 = 900_000 // $0.90 - inputTokens = 1_000_000 - outputTokens = 250_000 - concurrency = 20 + limitMicros int64 = 1_000_000 // $1.00 + startSpent int64 = 900_000 // $0.90 + inputTokens = 1_000_000 + outputTokens = 250_000 + concurrency = 20 ) pricing := catalog.Pricing{InputUSDPer1M: 3, OutputUSDPer1M: 15} perCallMicros := catalog.USDToMicros( From 3bd48d40bd96c7b36b2e02c9f24f441531e79590 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Mon, 20 Jul 2026 00:38:19 -0400 Subject: [PATCH 6/8] refactor: remove dead conditional in ArmSpendReservations (#793) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- internal/proxy/spend_reserve.go | 8 -------- 1 file changed, 8 deletions(-) diff --git a/internal/proxy/spend_reserve.go b/internal/proxy/spend_reserve.go index 5db08679..a71d2249 100644 --- a/internal/proxy/spend_reserve.go +++ b/internal/proxy/spend_reserve.go @@ -2,10 +2,8 @@ package proxy import ( "context" - "errors" "workweave/router/internal/auth" - "workweave/router/internal/billing" "workweave/router/internal/observability" "github.com/google/uuid" @@ -35,12 +33,6 @@ func (s *Service) ArmSpendReservations(ctx context.Context) (context.Context, fu "organization_id", orgID, "router_user_id", userID, ) - // Preserve sentinel identity for ClassifyDispatchError. - if errors.Is(err, billing.ErrOrgMonthlySpendLimitReached) || - errors.Is(err, billing.ErrAPIKeySpendCapReached) || - errors.Is(err, billing.ErrUserMonthlySpendLimitReached) { - return ctx, nil, err - } return ctx, nil, err } return ctx, release, nil From 8a73a4a572d135888d31712658fd44ee97d25fca Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Mon, 20 Jul 2026 00:59:45 -0400 Subject: [PATCH 7/8] test: agent-shadow eval skips ArmSpendReservations (#793/#787) Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- internal/proxy/spend_reserve_internal_test.go | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 internal/proxy/spend_reserve_internal_test.go diff --git a/internal/proxy/spend_reserve_internal_test.go b/internal/proxy/spend_reserve_internal_test.go new file mode 100644 index 00000000..20873f13 --- /dev/null +++ b/internal/proxy/spend_reserve_internal_test.go @@ -0,0 +1,76 @@ +package proxy + +import ( + "context" + "sync/atomic" + "testing" + + "workweave/router/internal/billing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// armCountingRepo records ReserveSpendCaps calls so tests can assert the +// agent-shadow path never reaches billing reservation. +type armCountingRepo struct { + spendLimitRepo + reserveCalls atomic.Int64 + reserveErr error +} + +func (r *armCountingRepo) ReserveSpendCaps(context.Context, billing.ReserveSpendCapsParams) ([]uuid.UUID, error) { + r.reserveCalls.Add(1) + if r.reserveErr != nil { + return nil, r.reserveErr + } + return []uuid.UUID{uuid.New()}, nil +} + +func agentShadowCtx(orgID string) context.Context { + ctx := spendLimitCtx("u1", orgID) + ctx = context.WithValue(ctx, APIKeyIDContextKey{}, "key-1") + return context.WithValue(ctx, AgentShadowEvalContextKey{}, AgentShadowEvaluation{ + Model: "claude-opus-4-8", + RolloutID: "pilot-1", + StateID: "state-1", + }) +} + +func TestArmSpendReservations_AgentShadowEvalSkipsReserve(t *testing.T) { + // Repo would refuse a real reserve (tight/exceeded limit). Shadow eval must + // still pass with zero reservations and never call ReserveSpendCaps. + repo := &armCountingRepo{reserveErr: billing.ErrOrgMonthlySpendLimitReached} + s := &Service{billing: billing.NewService(repo)} + + ctx, release, err := s.ArmSpendReservations(agentShadowCtx("org-1")) + require.NoError(t, err) + require.NotNil(t, release) + assert.Equal(t, int64(0), repo.reserveCalls.Load(), "shadow eval must not call ReserveSpendCaps") + assert.Nil(t, billing.SpendHoldFrom(ctx), "shadow eval must leave ctx unarmed") + assert.NotPanics(t, release) +} + +func TestArmSpendReservations_NonShadowStillReserves(t *testing.T) { + repo := &armCountingRepo{} + s := &Service{billing: billing.NewService(repo)} + + ctx, release, err := s.ArmSpendReservations(spendLimitCtx("u1", "org-1")) + require.NoError(t, err) + require.NotNil(t, release) + assert.Equal(t, int64(1), repo.reserveCalls.Load()) + assert.NotNil(t, billing.SpendHoldFrom(ctx)) + release() +} + +func TestArmSpendReservations_NonShadowPropagatesCapError(t *testing.T) { + repo := &armCountingRepo{reserveErr: billing.ErrOrgMonthlySpendLimitReached} + s := &Service{billing: billing.NewService(repo)} + + _, release, err := s.ArmSpendReservations(spendLimitCtx("u1", "org-1")) + require.Error(t, err) + assert.ErrorIs(t, err, billing.ErrOrgMonthlySpendLimitReached) + assert.Nil(t, release) + assert.Equal(t, int64(1), repo.reserveCalls.Load()) +} From 00f012f6abe11ee50fcc95e76a028ab3f34c3622 Mon Sep 17 00:00:00 2001 From: N Rohith Reddy Date: Mon, 20 Jul 2026 01:10:51 -0400 Subject: [PATCH 8/8] fix(db): use a single Go-computed month for spend reservations, not independent SQL NOW() (#793) Bugbot flagged (commit 8a73a4a): TryBumpOrgMonthReserved/ TryBumpUserMonthReserved computed their target month via SQL's own NOW(), while InsertSpendReservation recorded month from Go's utcMonthDate(). A clock disagreement across a UTC month boundary could bump one month's row while recording a different month on the reservation, permanently stranding reserved_usd_micros that settle/release/sweep could never find. Now a single Go-computed month value is passed explicitly into every query in one ReserveSpendCaps invocation, eliminating the possibility of divergence. Signed-off-by: N Rohith Reddy Co-authored-by: Cursor --- db/queries/spend_reservations.sql | 19 +-- internal/postgres/billing_repo.go | 19 ++- .../spend_reservation_month_internal_test.go | 118 ++++++++++++++++++ internal/sqlc/spend_reservations.sql.go | 60 ++++++--- 4 files changed, 185 insertions(+), 31 deletions(-) create mode 100644 internal/postgres/spend_reservation_month_internal_test.go diff --git a/db/queries/spend_reservations.sql b/db/queries/spend_reservations.sql index 2dc98e72..1ea839cc 100644 --- a/db/queries/spend_reservations.sql +++ b/db/queries/spend_reservations.sql @@ -1,10 +1,11 @@ --- Ensures the current UTC-month org spend row exists so a reserve UPDATE can --- target it. No-op when the row is already present. +-- 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, - DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + @month::date, 0, 0 ) @@ -13,14 +14,15 @@ 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). +-- 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 = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + 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 @@ -31,7 +33,7 @@ RETURNING sp.organization_id; INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) VALUES ( @router_user_id::uuid, - DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + @month::date, 0, 0 ) @@ -39,13 +41,14 @@ 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. +-- 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 = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + 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; diff --git a/internal/postgres/billing_repo.go b/internal/postgres/billing_repo.go index e6bb2c42..e1568ea6 100644 --- a/internal/postgres/billing_repo.go +++ b/internal/postgres/billing_repo.go @@ -179,12 +179,16 @@ func (r *BillingRepo) ReserveSpendCaps(ctx context.Context, p billing.ReserveSpe } _ = spent if limit != nil { - if err := q.EnsureOrgMonthlySpendRow(ctx, p.OrganizationID); err != nil { + if err := q.EnsureOrgMonthlySpendRow(ctx, sqlc.EnsureOrgMonthlySpendRowParams{ + OrganizationID: p.OrganizationID, + Month: month, + }); err != nil { return nil, err } _, err := q.TryBumpOrgMonthReserved(ctx, sqlc.TryBumpOrgMonthReservedParams{ AmountUsdMicros: p.AmountUsdMicros, OrganizationID: p.OrganizationID, + Month: month, }) if err != nil { if errors.Is(err, pgx.ErrNoRows) { @@ -252,12 +256,16 @@ func (r *BillingRepo) ReserveSpendCaps(ctx context.Context, p billing.ReserveSpe return nil, err } if limit != nil { - if err := q.EnsureUserMonthlySpendRow(ctx, userUUID); err != nil { + if err := q.EnsureUserMonthlySpendRow(ctx, sqlc.EnsureUserMonthlySpendRowParams{ + RouterUserID: userUUID, + Month: month, + }); err != nil { return nil, err } _, err := q.TryBumpUserMonthReserved(ctx, sqlc.TryBumpUserMonthReservedParams{ AmountUsdMicros: p.AmountUsdMicros, RouterUserID: userUUID, + Month: month, LimitUsdMicros: *limit, }) if err != nil { @@ -397,7 +405,12 @@ func decrementReservedForScope(ctx context.Context, q *sqlc.Queries, scopeKind, } } +// utcNow is the clock for reservation month bucketing. Tests override it to +// prove Ensure/Bump/Insert share one Go-computed month even when Postgres +// NOW() would land in a different UTC month. +var utcNow = time.Now + func utcMonthDate() pgtype.Date { - now := time.Now().UTC() + now := utcNow().UTC() return pgtype.Date{Time: time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC), Valid: true} } diff --git a/internal/postgres/spend_reservation_month_internal_test.go b/internal/postgres/spend_reservation_month_internal_test.go new file mode 100644 index 00000000..c529f645 --- /dev/null +++ b/internal/postgres/spend_reservation_month_internal_test.go @@ -0,0 +1,118 @@ +package postgres + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "workweave/router/internal/billing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestReserveSpendCaps_SingleGoComputedMonthAcrossEnsureBumpInsert freezes the +// Go reservation clock on a UTC month that differs from Postgres NOW()'s +// month. The PG-NOW month row is already at its cap (no headroom); the +// Go-computed month row has headroom. After the fix, Ensure/Bump/Insert all +// target the Go month, so the reserve succeeds and the reservation's recorded +// month matches the bumped spend row. Pre-fix (SQL NOW() in Ensure/Bump), +// the bump would hit the capped PG-NOW month and either refuse or strand +// reserved_usd_micros on a month the reservation record does not name. +func TestReserveSpendCaps_SingleGoComputedMonthAcrossEnsureBumpInsert(t *testing.T) { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + dsn = os.Getenv("DATABASE_URL") + } + if dsn == "" { + t.Skip("ROUTER_TEST_DATABASE_URL / DATABASE_URL not set") + } + + ctx := context.Background() + cfg, err := pgxpool.ParseConfig(dsn) + require.NoError(t, err) + cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + _, err := conn.Exec(ctx, "SET search_path TO router, public") + return err + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + require.NoError(t, err) + t.Cleanup(pool.Close) + require.NoError(t, pool.Ping(ctx)) + + var pgMonth time.Time + require.NoError(t, pool.QueryRow(ctx, + `SELECT DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date`).Scan(&pgMonth)) + pgMonth = time.Date(pgMonth.Year(), pgMonth.Month(), 1, 0, 0, 0, 0, time.UTC) + goMonth := pgMonth.AddDate(0, -1, 0) + require.NotEqual(t, pgMonth, goMonth) + + prev := utcNow + utcNow = func() time.Time { return goMonth.Add(12 * time.Hour) } + t.Cleanup(func() { utcNow = prev }) + + orgID := fmt.Sprintf("month-skew-793-%d", time.Now().UnixNano()) + const ( + limitMicros int64 = 10_000_000 // $10 + reserveR int64 = 1_000_000 // $1 + ) + t.Cleanup(func() { + _, _ = pool.Exec(ctx, `DELETE FROM spend_reservations WHERE scope_id = $1`, orgID) + _, _ = pool.Exec(ctx, `DELETE FROM organization_monthly_spend WHERE organization_id = $1`, orgID) + _, _ = pool.Exec(ctx, `DELETE FROM organization_spend_limits WHERE organization_id = $1`, orgID) + }) + + _, err = pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, limitMicros) + require.NoError(t, err) + // PG-NOW month: fully spent — an independent SQL NOW() bump would fail here. + _, err = pool.Exec(ctx, ` + INSERT INTO organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) + VALUES ($1, $2, $3, 0)`, orgID, pgMonth, limitMicros) + require.NoError(t, err) + // Go-computed month: empty headroom for R. + _, err = pool.Exec(ctx, ` + INSERT INTO organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) + VALUES ($1, $2, 0, 0)`, orgID, goMonth) + require.NoError(t, err) + + repo := NewBillingRepo(pool) + ids, err := repo.ReserveSpendCaps(ctx, billing.ReserveSpendCapsParams{ + OrganizationID: orgID, + RouterRequestID: "month-skew-req", + AmountUsdMicros: reserveR, + TTL: 15 * time.Minute, + SkipKey: true, + SkipUser: true, + }) + require.NoError(t, err) + require.Len(t, ids, 1) + + var ( + resMonth time.Time + resAmount int64 + goReserved int64 + pgReserved int64 + ) + require.NoError(t, pool.QueryRow(ctx, ` + SELECT month, amount_usd_micros FROM spend_reservations WHERE id = $1`, ids[0]). + Scan(&resMonth, &resAmount)) + resMonth = time.Date(resMonth.Year(), resMonth.Month(), 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, pool.QueryRow(ctx, ` + SELECT reserved_usd_micros FROM organization_monthly_spend + WHERE organization_id = $1 AND month = $2`, orgID, goMonth).Scan(&goReserved)) + require.NoError(t, pool.QueryRow(ctx, ` + SELECT reserved_usd_micros FROM organization_monthly_spend + WHERE organization_id = $1 AND month = $2`, orgID, pgMonth).Scan(&pgReserved)) + + assert.Equal(t, goMonth, resMonth, "reservation must record the Go-computed month") + assert.Equal(t, reserveR, resAmount) + assert.Equal(t, reserveR, goReserved, "bump must land on the same Go-computed month as the reservation") + assert.Equal(t, int64(0), pgReserved, "PG-NOW month must stay untouched (would be bumped pre-fix)") + assert.Equal(t, goMonth, utcMonthDate().Time.UTC(), "utcMonthDate must match the frozen Go clock") +} diff --git a/internal/sqlc/spend_reservations.sql.go b/internal/sqlc/spend_reservations.sql.go index 2488f0df..943eb69b 100644 --- a/internal/sqlc/spend_reservations.sql.go +++ b/internal/sqlc/spend_reservations.sql.go @@ -170,26 +170,32 @@ const ensureOrgMonthlySpendRow = `-- name: EnsureOrgMonthlySpendRow :exec INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) VALUES ( $1::varchar, - DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + $2::date, 0, 0 ) ON CONFLICT (organization_id, month) DO NOTHING ` -// Ensures the current UTC-month org spend row exists so a reserve UPDATE can -// target it. No-op when the row is already present. +type EnsureOrgMonthlySpendRowParams struct { + OrganizationID string + Month pgtype.Date +} + +// 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. // // INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) // VALUES ( // $1::varchar, -// DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, +// $2::date, // 0, // 0 // ) // ON CONFLICT (organization_id, month) DO NOTHING -func (q *Queries) EnsureOrgMonthlySpendRow(ctx context.Context, organizationID string) error { - _, err := q.db.Exec(ctx, ensureOrgMonthlySpendRow, organizationID) +func (q *Queries) EnsureOrgMonthlySpendRow(ctx context.Context, arg EnsureOrgMonthlySpendRowParams) error { + _, err := q.db.Exec(ctx, ensureOrgMonthlySpendRow, arg.OrganizationID, arg.Month) return err } @@ -197,25 +203,30 @@ const ensureUserMonthlySpendRow = `-- name: EnsureUserMonthlySpendRow :exec INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) VALUES ( $1::uuid, - DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, + $2::date, 0, 0 ) ON CONFLICT (router_user_id, month) DO NOTHING ` +type EnsureUserMonthlySpendRowParams struct { + RouterUserID uuid.UUID + Month pgtype.Date +} + // EnsureUserMonthlySpendRow // // INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) // VALUES ( // $1::uuid, -// DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date, +// $2::date, // 0, // 0 // ) // ON CONFLICT (router_user_id, month) DO NOTHING -func (q *Queries) EnsureUserMonthlySpendRow(ctx context.Context, routerUserID uuid.UUID) error { - _, err := q.db.Exec(ctx, ensureUserMonthlySpendRow, routerUserID) +func (q *Queries) EnsureUserMonthlySpendRow(ctx context.Context, arg EnsureUserMonthlySpendRowParams) error { + _, err := q.db.Exec(ctx, ensureUserMonthlySpendRow, arg.RouterUserID, arg.Month) return err } @@ -335,7 +346,7 @@ SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, updated_at = NOW() FROM router.organization_spend_limits lim WHERE sp.organization_id = $2::varchar - AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + AND sp.month = $3::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 + $1::bigint @@ -346,26 +357,28 @@ RETURNING sp.organization_id type TryBumpOrgMonthReservedParams struct { AmountUsdMicros int64 OrganizationID string + Month pgtype.Date } // 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). +// must skip reserve when limit is NULL before calling this). Month must be +// the same Go-computed value passed to Ensure and InsertSpendReservation. // // UPDATE router.organization_monthly_spend sp // SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, // updated_at = NOW() // FROM router.organization_spend_limits lim // WHERE sp.organization_id = $2::varchar -// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date +// AND sp.month = $3::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 + $1::bigint // <= lim.org_monthly_limit_usd_micros // RETURNING sp.organization_id func (q *Queries) TryBumpOrgMonthReserved(ctx context.Context, arg TryBumpOrgMonthReservedParams) (string, error) { - row := q.db.QueryRow(ctx, tryBumpOrgMonthReserved, arg.AmountUsdMicros, arg.OrganizationID) + row := q.db.QueryRow(ctx, tryBumpOrgMonthReserved, arg.AmountUsdMicros, arg.OrganizationID, arg.Month) var organization_id string err := row.Scan(&organization_id) return organization_id, err @@ -376,32 +389,39 @@ UPDATE router.model_router_user_monthly_spend sp SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, updated_at = NOW() WHERE sp.router_user_id = $2::uuid - AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date + AND sp.month = $3::date AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::bigint - <= $3::bigint + <= $4::bigint RETURNING sp.router_user_id ` type TryBumpUserMonthReservedParams struct { AmountUsdMicros int64 RouterUserID uuid.UUID + Month pgtype.Date LimitUsdMicros int64 } // 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. +// limit; NULL limit means the caller should skip this scope. Month must match +// the Go-computed value used for Ensure and InsertSpendReservation. // // UPDATE router.model_router_user_monthly_spend sp // SET reserved_usd_micros = sp.reserved_usd_micros + $1::bigint, // updated_at = NOW() // WHERE sp.router_user_id = $2::uuid -// AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date +// AND sp.month = $3::date // AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::bigint -// <= $3::bigint +// <= $4::bigint // RETURNING sp.router_user_id func (q *Queries) TryBumpUserMonthReserved(ctx context.Context, arg TryBumpUserMonthReservedParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, tryBumpUserMonthReserved, arg.AmountUsdMicros, arg.RouterUserID, arg.LimitUsdMicros) + row := q.db.QueryRow(ctx, tryBumpUserMonthReserved, + arg.AmountUsdMicros, + arg.RouterUserID, + arg.Month, + arg.LimitUsdMicros, + ) var router_user_id uuid.UUID err := row.Scan(&router_user_id) return router_user_id, err