feat(hooks/budget): per-tenant token spend tracking subsystem

Atomic budget store for prompt-handler LLM spend, with PG and SQLite
implementations. Enforces per-hook token caps to prevent runaway cost
from matcher-bypass scenarios.

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
+}

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)")
+	}
+}

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
+}

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
+}