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 (silent drop, F24) reserveDeniedGlobal // global daily USD ceiling hit (notice, F24) ) // Store is the durable bot state: /sync cursor, 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 } // Single writer — the loop is serial; one connection avoids lock churn. 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 "told this room I // can't read encryption" flag so a restart doesn't re-spam the notice (F5) — the // bot never sees its own notice (the sync filter drops nothing, but the loop skips // m.notice and self). 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 }