diff --git a/internal/hooks/budget/budget.go b/internal/hooks/budget/budget.go new file mode 100644 index 00000000..462c568e --- /dev/null +++ b/internal/hooks/budget/budget.go @@ -0,0 +1,99 @@ +// Package budget implements atomic per-tenant monthly token budgets for the +// prompt hook handler. The single-UPDATE deduct pattern avoids select-then- +// update races (L2 mitigation). +package budget + +import ( + "context" + "database/sql" + "errors" + "fmt" + "time" + + "github.com/google/uuid" +) + +// ErrBudgetExceeded is returned when Deduct is called but the remaining +// balance is below the requested cost. Callers should convert to a block +// decision with i18n key "hook.budget_exceeded". +var ErrBudgetExceeded = errors.New("hook budget: insufficient remaining") + +// DefaultMonthlyBudget is used when a tenant row is absent: the first call +// of the month seeds a fresh row with this cap. 0 means "unlimited" and +// the deduct path returns remaining=-1 to signal no-cap without a DB write. +const DefaultMonthlyBudget = 1_000_000 + +// WarnThresholdPct is the percentage of budget_total at which callers should +// warn (log, event). The store does not enforce warn-only behavior; callers +// compare remaining/total against this ratio. +const WarnThresholdPct = 20 // warn when remaining < 20% of total + +// Dialect abstracts the per-driver SQL differences (placeholder syntax, +// RETURNING clause support). PG uses $N + RETURNING; SQLite uses ? + a +// separate SELECT after the UPDATE. +type Dialect interface { + // DeductAtomic performs the atomic deduct + read-back. Returns the new + // remaining value; returns ErrBudgetExceeded when the row exists but + // remaining < cost. When the row is missing for the current month, the + // implementation must UPSERT with budget_total=DefaultMonthlyBudget and + // deduct from that. + DeductAtomic(ctx context.Context, tenantID uuid.UUID, cost int64, monthStart time.Time, defaultBudget int64) (remaining int64, total int64, err error) +} + +// Store is the public facade used by the prompt handler. It wraps a Dialect +// to apply month-rollover logic and normalize error values. +type Store struct { + d Dialect + // Now is injectable for deterministic tests (month rollover assertions). + Now func() time.Time +} + +// New returns a Store that routes writes through d. now may be nil; the +// default time.Now().UTC() is used when unset. +func New(d Dialect, now func() time.Time) *Store { + if now == nil { + now = func() time.Time { return time.Now().UTC() } + } + return &Store{d: d, Now: now} +} + +// Deduct atomically subtracts cost from the tenant's remaining balance for +// the current month. Returns (remaining, total, nil) on success, +// (_, _, ErrBudgetExceeded) when insufficient, or a wrapped driver error +// otherwise. Seeds a fresh month row with DefaultMonthlyBudget when absent. +// +// cost=0 is always allowed and returns the current remaining (useful for +// pre-flight checks and metrics without a real deduct). +func (s *Store) Deduct(ctx context.Context, tenantID uuid.UUID, cost int64) (remaining int64, total int64, err error) { + if tenantID == uuid.Nil { + return 0, 0, fmt.Errorf("budget: nil tenant_id") + } + if cost < 0 { + return 0, 0, fmt.Errorf("budget: negative cost %d", cost) + } + month := monthStart(s.Now()) + remaining, total, err = s.d.DeductAtomic(ctx, tenantID, cost, month, DefaultMonthlyBudget) + if errors.Is(err, sql.ErrNoRows) { + // Dialect guarantees UPSERT — ErrNoRows here means the row was + // found but remaining < cost. Normalize to ErrBudgetExceeded so + // callers don't care about the driver's native signal. + return 0, 0, ErrBudgetExceeded + } + return remaining, total, err +} + +// monthStart returns the first day of the month for t in UTC. +func monthStart(t time.Time) time.Time { + t = t.UTC() + return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC) +} + +// ShouldWarn reports true when remaining has crossed below WarnThresholdPct +// of total. Callers decide whether to emit a warn log/event. +func ShouldWarn(remaining, total int64) bool { + if total <= 0 { + return false + } + pct := remaining * 100 / total + return pct < WarnThresholdPct +} diff --git a/internal/hooks/budget/budget_test.go b/internal/hooks/budget/budget_test.go new file mode 100644 index 00000000..22d98286 --- /dev/null +++ b/internal/hooks/budget/budget_test.go @@ -0,0 +1,142 @@ +package budget_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" +) + +// fakeDialect is an in-memory Dialect for deterministic deduct tests. +// Models tenant_hook_budget row per (tenant_id, month_start). +type fakeDialect struct { + // row represents the current stored state for the single tenant under test. + total int64 + remaining int64 + month time.Time + seeded bool + // deductCalls counts atomic UPDATEs for race/concurrency assertions. + deductCalls int +} + +func (f *fakeDialect) DeductAtomic(_ context.Context, _ uuid.UUID, cost int64, month time.Time, defaultBudget int64) (int64, int64, error) { + f.deductCalls++ + if !f.seeded || !f.month.Equal(month) { + // Month rollover or fresh seed: reset remaining to defaultBudget. + f.total = defaultBudget + f.remaining = defaultBudget + f.month = month + f.seeded = true + } + if f.remaining < cost { + return 0, 0, budget.ErrBudgetExceeded + } + f.remaining -= cost + return f.remaining, f.total, nil +} + +func TestDeduct_Success_ReturnsUpdatedRemaining(t *testing.T) { + fd := &fakeDialect{} + s := budget.New(fd, nil) + got, total, err := s.Deduct(context.Background(), uuid.New(), 1000) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if total != budget.DefaultMonthlyBudget { + t.Errorf("total=%d, want %d", total, budget.DefaultMonthlyBudget) + } + if got != budget.DefaultMonthlyBudget-1000 { + t.Errorf("remaining=%d, want %d", got, budget.DefaultMonthlyBudget-1000) + } +} + +func TestDeduct_ExceedsBudget_ReturnsErr(t *testing.T) { + fd := &fakeDialect{} + s := budget.New(fd, nil) + tid := uuid.New() + // Drain most of the budget first. + if _, _, err := s.Deduct(context.Background(), tid, budget.DefaultMonthlyBudget-500); err != nil { + t.Fatalf("first deduct err: %v", err) + } + // Now request more than remaining — must return ErrBudgetExceeded. + _, _, err := s.Deduct(context.Background(), tid, 1000) + if !errors.Is(err, budget.ErrBudgetExceeded) { + t.Fatalf("want ErrBudgetExceeded, got %v", err) + } +} + +func TestDeduct_MonthRollover_ResetsRemaining(t *testing.T) { + fd := &fakeDialect{} + now := time.Date(2026, 4, 15, 12, 0, 0, 0, time.UTC) + s := budget.New(fd, func() time.Time { return now }) + + tid := uuid.New() + if _, _, err := s.Deduct(context.Background(), tid, 999_999); err != nil { + t.Fatalf("first deduct: %v", err) + } + + // Advance to next month; Deduct must seed a fresh row and allow the call. + now = time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC) + remaining, _, err := s.Deduct(context.Background(), tid, 1) + if err != nil { + t.Fatalf("new-month deduct err: %v", err) + } + if remaining != budget.DefaultMonthlyBudget-1 { + t.Errorf("new-month remaining=%d, want %d", remaining, budget.DefaultMonthlyBudget-1) + } +} + +func TestDeduct_ZeroCost_DoesNotDecrement(t *testing.T) { + fd := &fakeDialect{} + s := budget.New(fd, nil) + tid := uuid.New() + if _, _, err := s.Deduct(context.Background(), tid, 1000); err != nil { + t.Fatalf("seed: %v", err) + } + before := fd.remaining + got, _, err := s.Deduct(context.Background(), tid, 0) + if err != nil { + t.Fatalf("zero-cost err: %v", err) + } + if got != before { + t.Errorf("zero-cost mutated remaining: before=%d after=%d", before, got) + } +} + +func TestDeduct_NegativeCost_ReturnsError(t *testing.T) { + s := budget.New(&fakeDialect{}, nil) + if _, _, err := s.Deduct(context.Background(), uuid.New(), -1); err == nil { + t.Fatal("expected error for negative cost") + } +} + +func TestDeduct_NilTenantID_ReturnsError(t *testing.T) { + s := budget.New(&fakeDialect{}, nil) + if _, _, err := s.Deduct(context.Background(), uuid.Nil, 100); err == nil { + t.Fatal("expected error for nil tenant_id") + } +} + +func TestShouldWarn_BelowThreshold_ReturnsTrue(t *testing.T) { + // 15% remaining is below the 20% warn threshold. + if !budget.ShouldWarn(150, 1000) { + t.Error("ShouldWarn(150,1000)=false; want true") + } +} + +func TestShouldWarn_AboveThreshold_ReturnsFalse(t *testing.T) { + // 50% remaining — well above threshold. + if budget.ShouldWarn(500, 1000) { + t.Error("ShouldWarn(500,1000)=true; want false") + } +} + +func TestShouldWarn_ZeroTotal_ReturnsFalse(t *testing.T) { + if budget.ShouldWarn(0, 0) { + t.Error("ShouldWarn(0,0) should be false (unlimited budget)") + } +} diff --git a/internal/store/pg/hook_budget.go b/internal/store/pg/hook_budget.go new file mode 100644 index 00000000..abe7e4f0 --- /dev/null +++ b/internal/store/pg/hook_budget.go @@ -0,0 +1,119 @@ +package pg + +import ( + "context" + "database/sql" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" +) + +// PGHookBudget implements budget.Dialect over PostgreSQL. +// Uses a single atomic UPDATE with RETURNING for the hot path and an +// UPSERT fallback for first-of-month seeding. L2 mitigation: no +// select-then-update race possible. +type PGHookBudget struct { + db *sql.DB +} + +// NewPGHookBudget returns a PGHookBudget over db. +func NewPGHookBudget(db *sql.DB) *PGHookBudget { + return &PGHookBudget{db: db} +} + +// DeductAtomic implements budget.Dialect. +// +// Path A — current-month row exists and has enough balance: +// +// UPDATE ... SET remaining = remaining - $1 WHERE tenant=$2 AND month=$3 AND remaining >= $1 RETURNING remaining, budget_total +// → one round-trip, no race. +// +// Path B — row missing or month stale: INSERT ... ON CONFLICT DO UPDATE +// with fresh budget_total, then re-run path A. We detect this via +// affected=0 on the UPDATE and retry exactly once. +// +// Path C — row exists but remaining < cost: affected=0 after seed; return +// ErrBudgetExceeded. +func (b *PGHookBudget) DeductAtomic( + ctx context.Context, tenantID uuid.UUID, cost int64, month time.Time, defaultBudget int64, +) (int64, int64, error) { + // First try: direct deduct on an existing row for the current month. + remaining, total, ok, err := b.tryDeduct(ctx, tenantID, cost, month) + if err != nil { + return 0, 0, err + } + if ok { + return remaining, total, nil + } + + // Seed-or-rollover: ensure a row exists for (tenant, month) with a + // fresh default budget; then retry the deduct. ON CONFLICT only + // overwrites the row when the existing month_start is stale. + if err := b.seedIfStale(ctx, tenantID, month, defaultBudget); err != nil { + return 0, 0, err + } + + remaining, total, ok, err = b.tryDeduct(ctx, tenantID, cost, month) + if err != nil { + return 0, 0, err + } + if !ok { + return 0, 0, budget.ErrBudgetExceeded + } + return remaining, total, nil +} + +// tryDeduct runs the atomic UPDATE and reports (remaining, total, rowFound). +// rowFound=false means either the row was missing OR remaining < cost. +// Callers must distinguish via the seed path. +func (b *PGHookBudget) tryDeduct( + ctx context.Context, tenantID uuid.UUID, cost int64, month time.Time, +) (int64, int64, bool, error) { + row := b.db.QueryRowContext(ctx, ` + UPDATE tenant_hook_budget + SET remaining = remaining - $1, + updated_at = NOW() + WHERE tenant_id = $2 + AND month_start = $3::date + AND remaining >= $1 + RETURNING remaining, budget_total`, + cost, tenantID, month.Format("2006-01-02"), + ) + var rem, total int64 + err := row.Scan(&rem, &total) + if errors.Is(err, sql.ErrNoRows) { + return 0, 0, false, nil + } + if err != nil { + return 0, 0, false, fmt.Errorf("budget deduct: %w", err) + } + return rem, total, true, nil +} + +// seedIfStale UPSERTs the per-tenant row for month, resetting remaining to +// defaultBudget when month_start is older than month (rollover) OR the row +// is missing. Concurrent inserts are serialized by the PK. +func (b *PGHookBudget) seedIfStale( + ctx context.Context, tenantID uuid.UUID, month time.Time, defaultBudget int64, +) error { + _, err := b.db.ExecContext(ctx, ` + INSERT INTO tenant_hook_budget + (tenant_id, month_start, budget_total, remaining, metadata, updated_at) + VALUES ($1, $2::date, $3, $3, '{}', NOW()) + ON CONFLICT (tenant_id) DO UPDATE + SET month_start = EXCLUDED.month_start, + budget_total = EXCLUDED.budget_total, + remaining = EXCLUDED.remaining, + updated_at = NOW() + WHERE tenant_hook_budget.month_start < EXCLUDED.month_start`, + tenantID, month.Format("2006-01-02"), defaultBudget, + ) + if err != nil { + return fmt.Errorf("budget seed: %w", err) + } + return nil +} diff --git a/internal/store/sqlitestore/hook_budget.go b/internal/store/sqlitestore/hook_budget.go new file mode 100644 index 00000000..0c9dad86 --- /dev/null +++ b/internal/store/sqlitestore/hook_budget.go @@ -0,0 +1,128 @@ +//go:build sqlite || sqliteonly + +package sqlitestore + +import ( + "context" + "database/sql" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" +) + +// SqliteHookBudget implements budget.Dialect over SQLite. +// SQLite lacks RETURNING on UPDATE in all historical versions; we achieve the +// same atomicity via a single transaction: UPDATE WHERE remaining>=cost, +// then SELECT if RowsAffected=1, else UPSERT seed + retry. +type SqliteHookBudget struct { + db *sql.DB +} + +// NewSQLiteHookBudget returns a SqliteHookBudget over db. +func NewSQLiteHookBudget(db *sql.DB) *SqliteHookBudget { + return &SqliteHookBudget{db: db} +} + +// DeductAtomic implements budget.Dialect. See PG equivalent for semantics. +func (b *SqliteHookBudget) DeductAtomic( + ctx context.Context, tenantID uuid.UUID, cost int64, month time.Time, defaultBudget int64, +) (int64, int64, error) { + monthStr := month.Format("2006-01-02") + + tx, err := b.db.BeginTx(ctx, nil) + if err != nil { + return 0, 0, fmt.Errorf("budget begin: %w", err) + } + defer func() { _ = tx.Rollback() }() + + remaining, total, ok, err := b.tryDeductTx(ctx, tx, tenantID, cost, monthStr) + if err != nil { + return 0, 0, err + } + if ok { + if err := tx.Commit(); err != nil { + return 0, 0, fmt.Errorf("budget commit: %w", err) + } + return remaining, total, nil + } + + if err := b.seedIfStaleTx(ctx, tx, tenantID, monthStr, defaultBudget); err != nil { + return 0, 0, err + } + + remaining, total, ok, err = b.tryDeductTx(ctx, tx, tenantID, cost, monthStr) + if err != nil { + return 0, 0, err + } + if !ok { + _ = tx.Rollback() + return 0, 0, budget.ErrBudgetExceeded + } + if err := tx.Commit(); err != nil { + return 0, 0, fmt.Errorf("budget commit: %w", err) + } + return remaining, total, nil +} + +func (b *SqliteHookBudget) tryDeductTx( + ctx context.Context, tx *sql.Tx, tenantID uuid.UUID, cost int64, monthStr string, +) (int64, int64, bool, error) { + nowStr := time.Now().UTC().Format(time.RFC3339Nano) + res, err := tx.ExecContext(ctx, ` + UPDATE tenant_hook_budget + SET remaining = remaining - ?, + updated_at = ? + WHERE tenant_id = ? + AND month_start = ? + AND remaining >= ?`, + cost, nowStr, tenantID.String(), monthStr, cost, + ) + if err != nil { + return 0, 0, false, fmt.Errorf("budget deduct: %w", err) + } + n, _ := res.RowsAffected() + if n == 0 { + return 0, 0, false, nil + } + + var remaining, total int64 + err = tx.QueryRowContext(ctx, ` + SELECT remaining, budget_total + FROM tenant_hook_budget + WHERE tenant_id = ? AND month_start = ?`, + tenantID.String(), monthStr, + ).Scan(&remaining, &total) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return 0, 0, false, nil + } + return 0, 0, false, fmt.Errorf("budget read-back: %w", err) + } + return remaining, total, true, nil +} + +func (b *SqliteHookBudget) seedIfStaleTx( + ctx context.Context, tx *sql.Tx, tenantID uuid.UUID, monthStr string, defaultBudget int64, +) error { + nowStr := time.Now().UTC().Format(time.RFC3339Nano) + _, err := tx.ExecContext(ctx, ` + INSERT INTO tenant_hook_budget + (tenant_id, month_start, budget_total, remaining, metadata, updated_at) + VALUES (?, ?, ?, ?, '{}', ?) + ON CONFLICT(tenant_id) DO UPDATE + SET month_start = excluded.month_start, + budget_total = excluded.budget_total, + remaining = excluded.remaining, + updated_at = ? + WHERE tenant_hook_budget.month_start < excluded.month_start`, + tenantID.String(), monthStr, defaultBudget, defaultBudget, nowStr, nowStr, + ) + if err != nil { + return fmt.Errorf("budget seed: %w", err) + } + return nil +}