vojo/apps/ai-bot/store.go

package main

import (
	"database/sql"
	"fmt"
	"time"

	_ "modernc.org/sqlite"
)

// reserveResult is the outcome of a pre-call limiter reservation.
type reserveResult int

const (
	reserveOK           reserveResult = iota
	reserveDeniedUser                 // per-user daily request cap hit (⏳ rate-limit reaction, F24)
	reserveDeniedGlobal               // global daily USD ceiling hit (⏳ rate-limit reaction, F24)
)

// Store is the durable bot state: transaction + event dedup, the daily spend
// ledger, and the encrypted-room warned set. Pure-Go SQLite (no cgo) so the binary
// stays static for a distroless/scratch image.
type Store struct {
	db *sql.DB
}

func OpenStore(path string) (*Store, error) {
	db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
	if err != nil {
		return nil, err
	}
	// One connection: database/sql serializes all callers onto it, which keeps the
	// now-concurrent handler goroutines from contending on SQLite write locks. All
	// statements here are short, so callers block only briefly and never deadlock.
	db.SetMaxOpenConns(1)

	schema := `
	CREATE TABLE IF NOT EXISTS processed_txn (txn_id TEXT PRIMARY KEY);
	CREATE TABLE IF NOT EXISTS spend (
		date TEXT NOT NULL, mxid TEXT NOT NULL,
		requests INTEGER NOT NULL DEFAULT 0, usd REAL NOT NULL DEFAULT 0,
		PRIMARY KEY (date, mxid)
	);
	CREATE TABLE IF NOT EXISTS warned_encrypted (room_id TEXT PRIMARY KEY);
	CREATE TABLE IF NOT EXISTS processed_event (event_id TEXT PRIMARY KEY);`
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("init schema: %w", err)
	}
	return &Store{db: db}, nil
}

func (s *Store) Close() error { return s.db.Close() }

func todayUTC() string { return time.Now().UTC().Format("2006-01-02") }

// HasTxn / MarkTxn give appservice transactions idempotency across restarts: a
// transaction Synapse retries (because our 200 was lost) is processed at most
// once. The table is bounded to the most recent ids.
func (s *Store) HasTxn(txnID string) (bool, error) {
	var one int
	err := s.db.QueryRow(`SELECT 1 FROM processed_txn WHERE txn_id = ?`, txnID).Scan(&one)
	if err == sql.ErrNoRows {
		return false, nil
	}
	return err == nil, err
}

func (s *Store) MarkTxn(txnID string) error {
	if _, err := s.db.Exec(`INSERT OR IGNORE INTO processed_txn (txn_id) VALUES (?)`, txnID); err != nil {
		return err
	}
	_, err := s.db.Exec(`DELETE FROM processed_txn WHERE rowid NOT IN
		(SELECT rowid FROM processed_txn ORDER BY rowid DESC LIMIT 5000)`)
	return err
}

// SeenEvent records an event id as handled and reports whether it was NEW (true)
// or already seen (false) — the DURABLE equivalent of the in-memory dedup set, so
// a crash/restart between handling an event and acking its transaction can't make
// the bot reprocess it (dup answer + double-bill + cap inflation). Bounded to the
// most recent ids.
func (s *Store) SeenEvent(eventID string) (bool, error) {
	res, err := s.db.Exec(`INSERT OR IGNORE INTO processed_event (event_id) VALUES (?)`, eventID)
	if err != nil {
		return false, err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return false, nil // already recorded → not new
	}
	_, err = s.db.Exec(`DELETE FROM processed_event WHERE rowid NOT IN
		(SELECT rowid FROM processed_event ORDER BY rowid DESC LIMIT 20000)`)
	return true, err
}

// SpentTodayUSD sums all spend for the current UTC day.
func (s *Store) SpentTodayUSD() (float64, error) {
	var v sql.NullFloat64
	err := s.db.QueryRow(`SELECT SUM(usd) FROM spend WHERE date = ?`, todayUTC()).Scan(&v)
	if err != nil {
		return 0, err
	}
	return v.Float64, nil
}

// Reserve runs the two independent gates in one transaction, BEFORE the xAI call
// (F4): the global USD ceiling protects the wallet; the per-user request cap is
// anti-abuse. It increments the per-user request count on success; the USD is
// reconciled after the response. Order: global first (cheapest to deny), then
// per-user.
func (s *Store) Reserve(mxid string, perUserCap int, dailyUSDCeiling float64) (reserveResult, error) {
	day := todayUTC()
	tx, err := s.db.Begin()
	if err != nil {
		return reserveOK, err
	}
	defer tx.Rollback()

	var global sql.NullFloat64
	if err := tx.QueryRow(`SELECT SUM(usd) FROM spend WHERE date = ?`, day).Scan(&global); err != nil {
		return reserveOK, err
	}
	if global.Float64 >= dailyUSDCeiling {
		return reserveDeniedGlobal, nil
	}

	var requests int
	err = tx.QueryRow(`SELECT requests FROM spend WHERE date = ? AND mxid = ?`, day, mxid).Scan(&requests)
	if err != nil && err != sql.ErrNoRows {
		return reserveOK, err
	}
	if requests >= perUserCap {
		return reserveDeniedUser, nil
	}

	if _, err := tx.Exec(
		`INSERT INTO spend (date, mxid, requests, usd) VALUES (?, ?, 1, 0)
		 ON CONFLICT(date, mxid) DO UPDATE SET requests = requests + 1`,
		day, mxid); err != nil {
		return reserveOK, err
	}
	if err := tx.Commit(); err != nil {
		return reserveOK, err
	}
	return reserveOK, nil
}

// RefundRequest gives back a reserved request slot when the call ultimately
// failed (e.g. an xAI outage), so a transient failure doesn't burn the user's
// daily cap. Never drops below zero.
func (s *Store) RefundRequest(mxid string) error {
	_, err := s.db.Exec(
		`UPDATE spend SET requests = MAX(0, requests - 1) WHERE date = ? AND mxid = ?`,
		todayUTC(), mxid)
	return err
}

// Reconcile books the actual USD cost of a completed call against the user's
// daily row (and thus the global total).
func (s *Store) Reconcile(mxid string, usd float64) error {
	_, err := s.db.Exec(
		`INSERT INTO spend (date, mxid, requests, usd) VALUES (?, ?, 0, ?)
		 ON CONFLICT(date, mxid) DO UPDATE SET usd = usd + excluded.usd`,
		todayUTC(), mxid, usd)
	return err
}

// HasWarnedEncrypted / SetWarnedEncrypted persist the one-shot "reacted 🔒 to this
// room because I can't read encryption" flag so a restart doesn't re-react on every
// message (F5). The bot never reacts to its own events: m.reaction is not an
// m.room.message, so it never re-enters handleMessage.
func (s *Store) HasWarnedEncrypted(roomID string) (bool, error) {
	var one int
	err := s.db.QueryRow(`SELECT 1 FROM warned_encrypted WHERE room_id = ?`, roomID).Scan(&one)
	if err == sql.ErrNoRows {
		return false, nil
	}
	return err == nil, err
}

func (s *Store) SetWarnedEncrypted(roomID string) error {
	_, err := s.db.Exec(`INSERT OR IGNORE INTO warned_encrypted (room_id) VALUES (?)`, roomID)
	return err
}