From bbe828c8faefa31750aa6447fdd941960c8cc253 Mon Sep 17 00:00:00 2001 From: Steven Tohme Date: Mon, 20 Jul 2026 13:06:57 +0000 Subject: [PATCH 1/5] feat: move spend reservation checks to script Co-authored-by: rohith500 --- .env.example | 6 + cmd/router/main.go | 52 +++ .../0040_spend-reservations.down.sql | 14 + db/migrations/0040_spend-reservations.up.sql | 41 ++ db/queries/model_router_api_keys.sql | 2 +- db/queries/spend_limits.sql | 16 +- db/queries/spend_reservations.sql | 119 +++++ 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 | 358 +++++++++++++-- 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 | 39 ++ internal/proxy/spend_reserve_internal_test.go | 76 ++++ .../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 | 428 ++++++++++++++++++ scripts/spend_reservation_check/README.md | 25 + scripts/spend_reservation_check/main.go | 281 ++++++++++++ 33 files changed, 1842 insertions(+), 131 deletions(-) create mode 100644 db/migrations/0040_spend-reservations.down.sql create mode 100644 db/migrations/0040_spend-reservations.up.sql 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/proxy/spend_reserve_internal_test.go create mode 100644 internal/sqlc/spend_reservations.sql.go create mode 100644 scripts/spend_reservation_check/README.md create mode 100644 scripts/spend_reservation_check/main.go diff --git a/.env.example b/.env.example index f8a275168..973c6f39a 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 8e95dbed6..22381e687 100644 --- a/cmd/router/main.go +++ b/cmd/router/main.go @@ -161,6 +161,33 @@ func main() { billingSvc = billing.NewService(billingRepo) logger.Info("Router billing enabled", "min_balance_usd_micros", billing.MinBalanceMicros) } + if billingSvc != nil { + reserveAmount := billing.DefaultReserveAmountMicros + if v := config.GetOr("ROUTER_SPEND_RESERVE_USD_MICROS", ""); v != "" { + if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { + reserveAmount = n + } else { + logger.Warn("Invalid ROUTER_SPEND_RESERVE_USD_MICROS; using default", "value", v, "default", billing.DefaultReserveAmountMicros) + } + } + reserveTTL := billing.DefaultReserveTTL + if v := config.GetOr("ROUTER_SPEND_RESERVE_TTL", ""); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + reserveTTL = d + } else { + logger.Warn("Invalid ROUTER_SPEND_RESERVE_TTL; using default", "value", v, "default", billing.DefaultReserveTTL.String()) + } + } + billingSvc.WithReservationConfig(reserveAmount, reserveTTL) + logger.Info("Spend reservation config", + "reserve_usd_micros", reserveAmount, + "reserve_ttl", reserveTTL.String(), + ) + safeGo(logger, "spend-reservation-sweep", func() { + runSpendReservationSweep(context.Background(), billingSvc) + }) + logger.Info("Spend reservation TTL sweeper enabled", "interval", "1m") + } } // Managed without billing stays BYOK-only (avoids spending platform-key @@ -1363,6 +1390,31 @@ func runSessionPinSweep(ctx context.Context, store sessionpin.Store) { } } +// runSpendReservationSweep releases expired spend-cap reservations every +// minute. Stuck reserved incorrectly blocks spend, so this is much more +// aggressive than the session-pin GC. Every replica runs its own ticker; +// DELETE … RETURNING is idempotent across replicas. +func runSpendReservationSweep(ctx context.Context, svc *billing.Service) { + logger := observability.Get() + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + sweepCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + n, err := svc.SweepExpiredReservations(sweepCtx, time.Now().UTC()) + if err != nil { + logger.Error("Spend reservation sweep failed", "err", err) + } else if n > 0 { + logger.Info("Spend reservation sweep released expired holds", "released", n) + } + cancel() + } + } +} + // defaultHardPinProvider and defaultHardPinModel are the fallback (provider, // model) used by resolveHardPinModel when the cluster bundle can't be loaded. const ( diff --git a/db/migrations/0040_spend-reservations.down.sql b/db/migrations/0040_spend-reservations.down.sql new file mode 100644 index 000000000..7b165438e --- /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 000000000..8ba1a9d4b --- /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; diff --git a/db/queries/model_router_api_keys.sql b/db/queries/model_router_api_keys.sql index 5d3821aa6..ada494c0c 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 3898f4411..e12bc4c73 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 000000000..1ea839cc4 --- /dev/null +++ b/db/queries/spend_reservations.sql @@ -0,0 +1,119 @@ +-- Ensures the org spend row for @month exists so a reserve UPDATE can target +-- it. Month is supplied by Go (utcMonthDate) so Ensure/Bump/Insert share one +-- clock reading for the whole ReserveSpendCaps transaction. +-- name: EnsureOrgMonthlySpendRow :exec +INSERT INTO router.organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) +VALUES ( + @organization_id::varchar, + @month::date, + 0, + 0 +) +ON CONFLICT (organization_id, month) DO NOTHING; + +-- Atomically bumps org-month reserved when spent+reserved+amount still fits +-- under the configured org monthly limit. Returns the organization_id on +-- success; zero rows means limit reached (or no limit configured — caller +-- must skip reserve when limit is NULL before calling this). Month must be +-- the same Go-computed value passed to Ensure and InsertSpendReservation. +-- name: TryBumpOrgMonthReserved :one +UPDATE router.organization_monthly_spend sp +SET reserved_usd_micros = sp.reserved_usd_micros + @amount_usd_micros::bigint, + updated_at = NOW() +FROM router.organization_spend_limits lim +WHERE sp.organization_id = @organization_id::varchar + AND sp.month = @month::date + AND lim.organization_id = sp.organization_id + AND lim.org_monthly_limit_usd_micros IS NOT NULL + AND sp.spent_usd_micros + sp.reserved_usd_micros + @amount_usd_micros::bigint + <= lim.org_monthly_limit_usd_micros +RETURNING sp.organization_id; + +-- name: EnsureUserMonthlySpendRow :exec +INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, reserved_usd_micros) +VALUES ( + @router_user_id::uuid, + @month::date, + 0, + 0 +) +ON CONFLICT (router_user_id, month) DO NOTHING; + +-- Bumps user-month reserved under the effective limit (per-user override when +-- present, else org default). Caller supplies the already-resolved effective +-- limit; NULL limit means the caller should skip this scope. Month must match +-- the Go-computed value used for Ensure and InsertSpendReservation. +-- name: TryBumpUserMonthReserved :one +UPDATE router.model_router_user_monthly_spend sp +SET reserved_usd_micros = sp.reserved_usd_micros + @amount_usd_micros::bigint, + updated_at = NOW() +WHERE sp.router_user_id = @router_user_id::uuid + AND sp.month = @month::date + AND sp.spent_usd_micros + sp.reserved_usd_micros + @amount_usd_micros::bigint + <= @limit_usd_micros::bigint +RETURNING sp.router_user_id; + +-- Bumps api-key lifetime reserved under the key's spend_cap. Zero rows when +-- the key is missing, uncapped, or the bump would exceed the cap. +-- name: TryBumpAPIKeyReserved :one +UPDATE router.model_router_api_keys k +SET reserved_usd_micros = k.reserved_usd_micros + @amount_usd_micros::bigint +WHERE k.id = @api_key_id::uuid + AND k.deleted_at IS NULL + AND k.spend_cap_usd_micros IS NOT NULL + AND k.spent_usd_micros + k.reserved_usd_micros + @amount_usd_micros::bigint + <= k.spend_cap_usd_micros +RETURNING k.id; + +-- name: InsertSpendReservation :one +INSERT INTO router.spend_reservations ( + scope_kind, + scope_id, + month, + amount_usd_micros, + expires_at, + router_request_id +) VALUES ( + @scope_kind::varchar, + @scope_id::varchar, + sqlc.narg('month')::date, + @amount_usd_micros::bigint, + @expires_at::timestamptz, + sqlc.narg('router_request_id')::varchar +) +RETURNING id, scope_kind, scope_id, month, amount_usd_micros, expires_at; + +-- Atomic consume: DELETE … RETURNING is the sole settle/release/sweep +-- primitive. Zero rows = already released/settled/swept (idempotent no-op). +-- name: DeleteSpendReservation :one +DELETE FROM router.spend_reservations +WHERE id = @id::uuid +RETURNING id, scope_kind, scope_id, month, amount_usd_micros; + +-- name: DecrementOrgMonthReserved :exec +UPDATE router.organization_monthly_spend +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint), + updated_at = NOW() +WHERE organization_id = @organization_id::varchar + AND month = @month::date; + +-- name: DecrementUserMonthReserved :exec +UPDATE router.model_router_user_monthly_spend +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint), + updated_at = NOW() +WHERE router_user_id = @router_user_id::uuid + AND month = @month::date; + +-- name: DecrementAPIKeyReserved :exec +UPDATE router.model_router_api_keys +SET reserved_usd_micros = GREATEST(0, reserved_usd_micros - @amount_usd_micros::bigint) +WHERE id = @api_key_id::uuid; + +-- Deletes every expired reservation and returns the doomed rows so the +-- adapter can decrement denormalized reserved counters. Prefer calling +-- DeleteSpendReservation per id from Go when settling a known hold; this +-- batch path is for the TTL sweeper only. +-- name: DeleteExpiredSpendReservations :many +DELETE FROM router.spend_reservations +WHERE expires_at < @now::timestamptz +RETURNING id, scope_kind, scope_id, month, amount_usd_micros; diff --git a/internal/api/anthropic/messages.go b/internal/api/anthropic/messages.go index a9ef2839d..763f5cb5e 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 6f8ed26f5..de3f19bdf 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 43408d9e1..395f7e811 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 206def474..75229c55a 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 bbddafe5d..f1d38ac57 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 2851fa90a..6a3b0f6e3 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 000000000..a2eae6efa --- /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 476b29e8d..d34403b0b 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 e4ae3cfbd..db10f3771 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 f6d13a7a4..dcbb1e5cf 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,270 @@ 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 != "" { + _, _, limit, err := r.getOrgMonthlyInTx(ctx, q, p.OrganizationID) + if err != nil { + return nil, err + } + if limit != 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) { + 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, 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 { + 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) + } +} + +// 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 := 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/proxy/dispatch_error.go b/internal/proxy/dispatch_error.go index 644b051aa..3faaaefba 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 8d2c1e256..a70f23f43 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 4f654ffbc..2fda771bc 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 646b474c5..062f7f081 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 c3a6ce138..fdcaff348 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 000000000..a71d22490 --- /dev/null +++ b/internal/proxy/spend_reserve.go @@ -0,0 +1,39 @@ +package proxy + +import ( + "context" + + "workweave/router/internal/auth" + "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, + ) + return ctx, nil, err + } + return ctx, release, nil +} diff --git a/internal/proxy/spend_reserve_internal_test.go b/internal/proxy/spend_reserve_internal_test.go new file mode 100644 index 000000000..20873f13d --- /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()) +} diff --git a/internal/server/middleware/api_key_spend_cap.go b/internal/server/middleware/api_key_spend_cap.go index 8199b43a5..dfc8faa4e 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 44fcf8aca..23c4dd412 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 bb9c45bf5..1742ecdf7 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 60df1f17a..36bea158e 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 831a73328..d9f9d59ef 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 f8a46cf64..34cb84276 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 000000000..943eb69b9 --- /dev/null +++ b/internal/sqlc/spend_reservations.sql.go @@ -0,0 +1,428 @@ +// 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, + $2::date, + 0, + 0 +) +ON CONFLICT (organization_id, month) DO NOTHING +` + +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, +// $2::date, +// 0, +// 0 +// ) +// ON CONFLICT (organization_id, month) DO NOTHING +func (q *Queries) EnsureOrgMonthlySpendRow(ctx context.Context, arg EnsureOrgMonthlySpendRowParams) error { + _, err := q.db.Exec(ctx, ensureOrgMonthlySpendRow, arg.OrganizationID, arg.Month) + 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, + $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, +// $2::date, +// 0, +// 0 +// ) +// ON CONFLICT (router_user_id, month) DO NOTHING +func (q *Queries) EnsureUserMonthlySpendRow(ctx context.Context, arg EnsureUserMonthlySpendRowParams) error { + _, err := q.db.Exec(ctx, ensureUserMonthlySpendRow, arg.RouterUserID, arg.Month) + 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 = $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 +` + +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). 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 = $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, arg.Month) + 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 = $3::date + AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::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. 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 = $3::date +// AND sp.spent_usd_micros + sp.reserved_usd_micros + $1::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.Month, + arg.LimitUsdMicros, + ) + var router_user_id uuid.UUID + err := row.Scan(&router_user_id) + return router_user_id, err +} diff --git a/scripts/spend_reservation_check/README.md b/scripts/spend_reservation_check/README.md new file mode 100644 index 000000000..a524afed8 --- /dev/null +++ b/scripts/spend_reservation_check/README.md @@ -0,0 +1,25 @@ +# spend_reservation_check + +Exercises spend-cap reservations against the real Postgres database used by +the local `docker compose` stack. It covers the #793 concurrent +reserve-then-settle regression, verifies that reservation and monthly spend +rows use the same Go-computed UTC month, and checks that the TTL sweeper +releases expired reservations. + +Per [`../../AGENTS.md`](../../AGENTS.md) ("No DB-backed integration tests in +`internal/`"), this lives here as a runnable script rather than +`internal/postgres/*_test.go`. It is **not** a test (`package main`), so +`go test ./...` never touches Postgres. It is gated on +`ROUTER_TEST_DATABASE_URL` and is a no-op without it. + +## Usage (from the repo root) + +Bring up the local `docker compose` Postgres stack with router migrations +applied, then: + +```bash +ROUTER_TEST_DATABASE_URL="postgres://router:router@localhost:5432/router?search_path=router" \ + go run ./scripts/spend_reservation_check +``` + +Exits non-zero and logs which check failed on any mismatch. diff --git a/scripts/spend_reservation_check/main.go b/scripts/spend_reservation_check/main.go new file mode 100644 index 000000000..fd932277d --- /dev/null +++ b/scripts/spend_reservation_check/main.go @@ -0,0 +1,281 @@ +// Command spend_reservation_check exercises spend reservations against the +// docker-compose Postgres fixture without making database access part of +// `go test ./...`. +// +// It is gated on ROUTER_TEST_DATABASE_URL and is a no-op without it. +// +// Usage (from the repo root): +// +// ROUTER_TEST_DATABASE_URL="postgres://router:router@localhost:5432/router?search_path=router" \ +// go run ./scripts/spend_reservation_check +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + "os" + "sync" + "sync/atomic" + "time" + + "workweave/router/internal/billing" + "workweave/router/internal/postgres" + "workweave/router/internal/router/catalog" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +const ( + reservationLimitMicros int64 = 1_000_000 + reservationAmountMicros int64 = 1_000_000 +) + +func main() { + dsn := os.Getenv("ROUTER_TEST_DATABASE_URL") + if dsn == "" { + slog.Info("ROUTER_TEST_DATABASE_URL not set; skipping live-DB spend reservation check") + return + } + + ctx := context.Background() + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + fail("parse database URL", 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) + if err != nil { + fail("connect to database", err) + } + defer pool.Close() + if err := pool.Ping(ctx); err != nil { + fail("ping database", err) + } + + checks := []struct { + name string + run func(context.Context, *pgxpool.Pool) error + }{ + {"concurrent arm and settle", checkConcurrentArmAndSettle}, + {"Go-computed month consistency", checkGoComputedMonth}, + {"expired reservation sweep", checkExpiredSweep}, + } + for _, check := range checks { + if err := check.run(ctx, pool); err != nil { + fail(check.name, err) + } + slog.Info("spend reservation check passed", "check", check.name) + } +} + +func fail(step string, err error) { + slog.Error("spend reservation check failed", "step", step, "err", err) + os.Exit(1) +} + +func checkConcurrentArmAndSettle(ctx context.Context, pool *pgxpool.Pool) error { + orgID := fmt.Sprintf("spend-check-toctou-%d", time.Now().UnixNano()) + defer cleanupOrg(ctx, pool, orgID) + + const ( + startSpent int64 = 900_000 + 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), + ) + if perCallMicros != 6_750_000 { + return fmt.Errorf("fixture cost changed: got %d, want 6750000", perCallMicros) + } + if _, err := pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, reservationLimitMicros); err != nil { + return fmt.Errorf("insert spend limit: %w", err) + } + if _, 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); err != nil { + return fmt.Errorf("insert monthly spend: %w", err) + } + if _, err := pool.Exec(ctx, ` + INSERT INTO organization_credit_balance (organization_id, balance_usd_micros) + VALUES ($1, $2)`, orgID, int64(1_000_000_000)); err != nil { + return fmt.Errorf("insert credit balance: %w", err) + } + + svc := billing.NewService(postgres.NewBillingRepo(pool)). + WithReservationConfig(billing.DefaultReserveAmountMicros, billing.DefaultReserveTTL) + var wg sync.WaitGroup + var debited, rejected atomic.Int32 + errs := make(chan error, concurrency) + wg.Add(concurrency) + for i := 0; i < concurrency; i++ { + i := i + go func() { + defer wg.Done() + reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + reqID := fmt.Sprintf("%s-%d", orgID, i) + reqCtx, release, err := svc.ArmSpendReservations(reqCtx, orgID, "", "", reqID) + if err != nil { + rejected.Add(1) + return + } + defer release() + _, err = svc.DebitForInference(reqCtx, billing.DebitInferenceParams{ + OrganizationID: orgID, + RouterRequestID: reqID, + Model: "spend-reservation-fixture", + InputTokens: inputTokens, + OutputTokens: outputTokens, + Pricing: pricing, + ReservationIDs: billing.SpendHoldFrom(reqCtx).IDs, + }) + if err != nil { + errs <- fmt.Errorf("debit: %w", err) + return + } + billing.SpendHoldFrom(reqCtx).MarkSettled() + debited.Add(1) + }() + } + wg.Wait() + close(errs) + for err := range errs { + return err + } + + spent, reserved, _, err := postgres.NewBillingRepo(pool).GetOrgMonthlySpendAndLimit(ctx, orgID) + if err != nil { + return fmt.Errorf("read final monthly spend: %w", err) + } + slog.Info("concurrent reservation fixture complete", "organization_id", orgID, + "debited", debited.Load(), "rejected", rejected.Load(), "spent_usd_micros", spent, + "reserved_usd_micros", reserved) + if spent > reservationLimitMicros+perCallMicros { + return fmt.Errorf("monthly spend overshot cap by more than one turn: spent=%d max=%d", + spent, reservationLimitMicros+perCallMicros) + } + if reserved != 0 { + return fmt.Errorf("reservations remain after settle: %d", reserved) + } + return nil +} + +func checkGoComputedMonth(ctx context.Context, pool *pgxpool.Pool) error { + orgID := fmt.Sprintf("spend-check-month-%d", time.Now().UnixNano()) + defer cleanupOrg(ctx, pool, orgID) + now := time.Now().UTC() + month := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + if _, err := pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, int64(10_000_000)); err != nil { + return fmt.Errorf("insert spend limit: %w", err) + } + if _, err := pool.Exec(ctx, ` + INSERT INTO organization_monthly_spend (organization_id, month, spent_usd_micros, reserved_usd_micros) + VALUES ($1, $2, 0, 0)`, orgID, month); err != nil { + return fmt.Errorf("insert monthly spend: %w", err) + } + ids, err := postgres.NewBillingRepo(pool).ReserveSpendCaps(ctx, billing.ReserveSpendCapsParams{ + OrganizationID: orgID, RouterRequestID: orgID, AmountUsdMicros: reservationAmountMicros, + TTL: 15 * time.Minute, SkipKey: true, SkipUser: true, + }) + if err != nil { + return fmt.Errorf("reserve spend caps: %w", err) + } + if len(ids) != 1 { + return fmt.Errorf("want one reservation, got %d", len(ids)) + } + var reservationMonth time.Time + if err := pool.QueryRow(ctx, `SELECT month FROM spend_reservations WHERE id = $1`, ids[0]).Scan(&reservationMonth); err != nil { + return fmt.Errorf("read reservation month: %w", err) + } + reservationMonth = reservationMonth.UTC() + reservationMonth = time.Date(reservationMonth.Year(), reservationMonth.Month(), 1, 0, 0, 0, 0, time.UTC) + var reserved int64 + if err := pool.QueryRow(ctx, ` + SELECT reserved_usd_micros FROM organization_monthly_spend + WHERE organization_id = $1 AND month = $2`, orgID, month).Scan(&reserved); err != nil { + return fmt.Errorf("read reserved month amount: %w", err) + } + if !reservationMonth.Equal(month) { + return fmt.Errorf("reservation month=%s, Go-computed month=%s", reservationMonth, month) + } + if reserved != reservationAmountMicros { + return fmt.Errorf("reserved amount=%d, want %d", reserved, reservationAmountMicros) + } + return nil +} + +func checkExpiredSweep(ctx context.Context, pool *pgxpool.Pool) error { + orgID := fmt.Sprintf("spend-check-sweep-%d", time.Now().UnixNano()) + defer cleanupOrg(ctx, pool, orgID) + const amount int64 = 1_000_000 + reservationID := uuid.New() + if _, err := pool.Exec(ctx, ` + INSERT INTO organization_spend_limits (organization_id, org_monthly_limit_usd_micros) + VALUES ($1, $2)`, orgID, int64(10_000_000)); err != nil { + return fmt.Errorf("insert spend limit: %w", err) + } + if _, 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); err != nil { + return fmt.Errorf("insert monthly spend: %w", err) + } + if _, 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')`, + reservationID, orgID, amount); err != nil { + return fmt.Errorf("insert expired reservation: %w", err) + } + released, err := billing.NewService(postgres.NewBillingRepo(pool)).SweepExpiredReservations(ctx, time.Now().UTC()) + if err != nil { + return fmt.Errorf("sweep expired reservations: %w", err) + } + if released < 1 { + return fmt.Errorf("sweeper released %d reservations", released) + } + var count int + if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM spend_reservations WHERE id = $1`, reservationID).Scan(&count); err != nil { + return fmt.Errorf("check swept reservation: %w", err) + } + if count != 0 { + return errors.New("expired reservation row remains") + } + var reserved int64 + if err := pool.QueryRow(ctx, ` + SELECT reserved_usd_micros FROM organization_monthly_spend + WHERE organization_id = $1 AND month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date`, + orgID).Scan(&reserved); err != nil { + return fmt.Errorf("read swept monthly spend: %w", err) + } + if reserved != 0 { + return fmt.Errorf("reserved amount after sweep=%d", reserved) + } + return nil +} + +func cleanupOrg(ctx context.Context, pool *pgxpool.Pool, orgID string) { + cleanupCtx, cancel := context.WithTimeout(ctx, 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) +} From 1df65817420914e137abf3950d9038f44bd66808 Mon Sep 17 00:00:00 2001 From: Steven Tohme Date: Mon, 20 Jul 2026 13:12:08 +0000 Subject: [PATCH 2/5] docs: apply spend reservation review nits --- cmd/router/main.go | 6 ++---- internal/billing/service.go | 9 +++------ internal/postgres/billing_repo.go | 5 ++--- internal/proxy/spend_limit.go | 5 ++--- internal/proxy/spend_reserve.go | 7 ++----- 5 files changed, 11 insertions(+), 21 deletions(-) diff --git a/cmd/router/main.go b/cmd/router/main.go index 22381e687..0f3a29efd 100644 --- a/cmd/router/main.go +++ b/cmd/router/main.go @@ -1390,10 +1390,8 @@ 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. +// runSpendReservationSweep releases expired spend reservations every minute. +// Stuck reserved headroom incorrectly blocks spend; DELETE … RETURNING is idempotent across replicas. func runSpendReservationSweep(ctx context.Context, svc *billing.Service) { logger := observability.Get() ticker := time.NewTicker(1 * time.Minute) diff --git a/internal/billing/service.go b/internal/billing/service.go index d34403b0b..54ad0e0e2 100644 --- a/internal/billing/service.go +++ b/internal/billing/service.go @@ -275,9 +275,7 @@ 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. +// ReserveApplicableCaps reserves all applicable caps in one TX. 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 @@ -315,9 +313,8 @@ func (s *Service) SweepExpiredReservations(ctx context.Context, now time.Time) ( 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). +// ArmSpendReservations reserves applicable caps and returns a defer-safe release func +// that no-ops once MarkSettled is called by DebitForInference. func (s *Service) ArmSpendReservations(ctx context.Context, organizationID, apiKeyID, routerUserID, requestID string) (context.Context, func(), error) { if s == nil { return ctx, func() {}, nil diff --git a/internal/postgres/billing_repo.go b/internal/postgres/billing_repo.go index dcbb1e5cf..5e2ce623e 100644 --- a/internal/postgres/billing_repo.go +++ b/internal/postgres/billing_repo.go @@ -404,9 +404,8 @@ 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. +// utcNow is the clock for reservation month bucketing; tests override it so +// Ensure/Bump/Insert share one Go-computed month even across a UTC midnight. var utcNow = time.Now func utcMonthDate() pgtype.Date { diff --git a/internal/proxy/spend_limit.go b/internal/proxy/spend_limit.go index 062f7f081..dee5b475c 100644 --- a/internal/proxy/spend_limit.go +++ b/internal/proxy/spend_limit.go @@ -9,9 +9,8 @@ import ( "workweave/router/internal/observability" ) -// 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*. +// checkUserMonthlySpendLimit is retained for unit tests only; request-path +// enforcement moved to ArmSpendReservations in API handlers. func (s *Service) checkUserMonthlySpendLimit(ctx context.Context) error { if s.billing == nil { return nil diff --git a/internal/proxy/spend_reserve.go b/internal/proxy/spend_reserve.go index a71d22490..ce5c1c7eb 100644 --- a/internal/proxy/spend_reserve.go +++ b/internal/proxy/spend_reserve.go @@ -10,11 +10,8 @@ import ( ) // 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. +// before Proxy*. Shadow eval skips entirely — synthetic traffic must not consume +// production spend budget (parity with billing middleware, #787). func (s *Service) ArmSpendReservations(ctx context.Context) (context.Context, func(), error) { if s == nil || s.billing == nil { return ctx, func() {}, nil From cb876f60811ded3b77a3e696aa8b7b040b8a5940 Mon Sep 17 00:00:00 2001 From: Steven Tohme Date: Mon, 20 Jul 2026 13:16:45 +0000 Subject: [PATCH 3/5] docs: tighten reservation interface comments --- internal/billing/repo.go | 10 ++++------ internal/billing/reservation.go | 5 ++--- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/internal/billing/repo.go b/internal/billing/repo.go index 6a3b0f6e3..57dde4b1d 100644 --- a/internal/billing/repo.go +++ b/internal/billing/repo.go @@ -41,14 +41,12 @@ type Repo interface { // 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 atomically reserves every applicable scope in one TX; + // any scope failure rolls back the whole TX (no partial holds). 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 deletes reservation ids and decrements denormalized + // reserved. Idempotent: already-gone ids are no-ops. ReleaseSpendReservations(ctx context.Context, ids []uuid.UUID) error // SweepExpiredSpendReservations deletes expired reservation rows and diff --git a/internal/billing/reservation.go b/internal/billing/reservation.go index a2eae6efa..be044d529 100644 --- a/internal/billing/reservation.go +++ b/internal/billing/reservation.go @@ -15,9 +15,8 @@ const ( 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. +// DefaultReserveAmountMicros is the v1 fixed reservation slot ($1) — p95-ish turn cost, +// not exact. Override via ROUTER_SPEND_RESERVE_USD_MICROS. const DefaultReserveAmountMicros int64 = 1_000_000 // DefaultReserveTTL is request timeout (600s) + 5m grace. Override via From 47012296356c5e659c3f1c25d512d189d1a3ae3d Mon Sep 17 00:00:00 2001 From: Steven Tohme Date: Mon, 20 Jul 2026 13:20:52 +0000 Subject: [PATCH 4/5] docs: clarify monthly spend error --- internal/billing/errors.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/internal/billing/errors.go b/internal/billing/errors.go index f1d38ac57..5e79b7bd0 100644 --- a/internal/billing/errors.go +++ b/internal/billing/errors.go @@ -22,8 +22,7 @@ var ErrBalanceRowMissing = errors.New("billing: balance row missing") 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. +// 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 From 073aa053bd4643b3fa3c3d2a16d9ac1b708f7d5b Mon Sep 17 00:00:00 2001 From: Steven Tohme Date: Mon, 20 Jul 2026 13:25:58 +0000 Subject: [PATCH 5/5] docs: shorten reservation godoc --- internal/proxy/spend_reserve.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/internal/proxy/spend_reserve.go b/internal/proxy/spend_reserve.go index ce5c1c7eb..1507a8174 100644 --- a/internal/proxy/spend_reserve.go +++ b/internal/proxy/spend_reserve.go @@ -9,9 +9,8 @@ import ( "github.com/google/uuid" ) -// ArmSpendReservations reserves all applicable spend caps in one transaction -// before Proxy*. Shadow eval skips entirely — synthetic traffic must not consume -// production spend budget (parity with billing middleware, #787). +// ArmSpendReservations reserves applicable spend caps before Proxy*. +// Shadow eval skips — synthetic traffic must not consume production budget. func (s *Service) ArmSpendReservations(ctx context.Context) (context.Context, func(), error) { if s == nil || s.billing == nil { return ctx, func() {}, nil