heaven/vojo
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(ai-bot): move the operational store off SQLite onto a dedicated Postgres database (vojo_ai) via pgx

Browse source
This commit is contained in:
heaven 2026-06-01 02:40:26 +03:00
parent ebb2363d9d
commit f7f6984d18
11 changed files with 567 additions and 144 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
apps/ai-bot/.env.example
View file

@ -33,6 +33,14 @@ XAI_PRICE_INPUT_PER_M=1.25 # fallback per-1M prices that bound t
XAI_PRICE_CACHED_PER_M=0.20 XAI_PRICE_CACHED_PER_M=0.20
XAI_PRICE_OUTPUT_PER_M=2.50 XAI_PRICE_OUTPUT_PER_M=2.50
# --- Database (vojo_ai Postgres) ---
# Operational store (txn/event dedup, the daily spend ledger, the encrypted-warned
# set) — NOT message content (that lives in Synapse). A dedicated database+role on
# the shared Postgres, like each mautrix bridge. Inside the docker network the host
# is the `postgres` service. The DSN embeds the role password, so treat ai-bot.env
# as sensitive (chmod 600). Required.
AI_BOT_DATABASE_URL=postgres://vojo_ai:CHANGE_ME@postgres:5432/vojo_ai?sslmode=disable
# --- Paths (non-secret) --- # --- Paths (non-secret) ---
SYSTEM_PROMPT_PATH=prompts/system_ru.txt SYSTEM_PROMPT_PATH=prompts/system_ru.txt
STATE_DIR=/state STATE_DIR=/state

5
apps/ai-bot/Dockerfile
View file

@ -1,4 +1,4 @@
# Multi-stage: static CGO-free build (pure-Go SQLite) → distroless runtime. # Multi-stage: static CGO-free build (pure-Go pgx driver) → distroless runtime.
FROM golang:1.25 AS build FROM golang:1.25 AS build
WORKDIR /src WORKDIR /src
@ -15,7 +15,8 @@ WORKDIR /app
COPY --from=build /out/ai-bot /app/ai-bot COPY --from=build /out/ai-bot /app/ai-bot
# System prompt(s) ship with the image; override via SYSTEM_PROMPT_PATH + a mount. # System prompt(s) ship with the image; override via SYSTEM_PROMPT_PATH + a mount.
COPY --from=build /src/prompts /app/prompts COPY --from=build /src/prompts /app/prompts
# STATE_DIR (SQLite: spend ledger + txn dedup) is a mounted volume in compose. # The operational store now lives in Postgres (AI_BOT_DATABASE_URL → the vojo_ai
# database). STATE_DIR remains the runtime dir (registration.yaml etc.); no DB here.
ENV STATE_DIR=/state ENV STATE_DIR=/state
# Appservice transaction-push port (Synapse → bot). Match AS_ADDR / the # Appservice transaction-push port (Synapse → bot). Match AS_ADDR / the
# registration `url`. # registration `url`.

54
apps/ai-bot/README.md
View file

@ -34,7 +34,7 @@ apps/ai-bot/
├── mentions.go # m.mentions + pill/reply fallbacks (F29/F30) ├── mentions.go # m.mentions + pill/reply fallbacks (F29/F30)
├── context.go # xAI message-window assembly (trigger + bot replies) ├── context.go # xAI message-window assembly (trigger + bot replies)
├── xai.go # chat/completions client + retry (F6) ├── xai.go # chat/completions client + retry (F6)
├── store.go # SQLite: spend ledger, txn dedup, encrypted-warned set ├── store.go # Postgres (vojo_ai): spend ledger, txn/event dedup, encrypted-warned set
├── messages.go # bot-authored RU notices ├── messages.go # bot-authored RU notices
├── markdown.go # markdown → org.matrix.custom.html for the reply's formatted_body ├── markdown.go # markdown → org.matrix.custom.html for the reply's formatted_body
├── util.go # bounded dedup set ├── util.go # bounded dedup set
@ -46,11 +46,36 @@ apps/ai-bot/
## Configuration ## Configuration
All via environment (see `.env.example`). Required: `HOMESERVER_URL`, `BOT_MXID`, All via environment (see `.env.example`). Required: `HOMESERVER_URL`, `BOT_MXID`,
`AS_TOKEN`, `HS_TOKEN`, `XAI_API_KEY`, `ALLOWED_SERVERS`. `AS_ADDR` (default `AS_TOKEN`, `HS_TOKEN`, `XAI_API_KEY`, `ALLOWED_SERVERS`, `AI_BOT_DATABASE_URL`.
`:8009`) is the transaction-push listen address — it must match the `url` port in `AS_ADDR` (default `:8009`) is the transaction-push listen address — it must match
the registration. The model is env-configurable (`XAI_MODEL`, default the `url` port in the registration. The model is env-configurable (`XAI_MODEL`,
`grok-4.20-0309-non-reasoning`; `grok-4.3` is an alternative — **re-verify the id default `grok-4.20-0309-non-reasoning`; `grok-4.3` is an alternative — **re-verify
+ price on docs.x.ai before deploy**). the id + price on docs.x.ai before deploy**).
### Database
The bot keeps its **operational state** — appservice transaction + event dedup, the
daily spend ledger, and the encrypted-room warned set — in a dedicated Postgres
database `vojo_ai` on the shared server, mirroring the per-service bridge databases
(each bridge owns its own role + DB). It stores **no message content**: the room
timeline is canonical in Synapse, and the bot's xAI context window is the in-memory
buffer in `bot.go`. The schema is created/migrated on startup (a `schema_version`
table + idempotent `CREATE TABLE IF NOT EXISTS`), so a fresh `vojo_ai` needs no
manual DDL — just the role + database:
```sql
-- once, as the Postgres superuser (e.g. `docker exec vojo-postgres-1 psql -U synapse -d postgres`):
CREATE ROLE vojo_ai LOGIN PASSWORD '<32-char secret>'; -- least privilege; NOT a superuser
CREATE DATABASE vojo_ai OWNER vojo_ai;
```
Point the bot at it with `AI_BOT_DATABASE_URL` (libpq/pgx DSN). Inside the docker
network the host is the `postgres` service; `sslmode=disable` matches Synapse and
the bridges on the internal network:
```
AI_BOT_DATABASE_URL=postgres://vojo_ai:<secret>@postgres:5432/vojo_ai?sslmode=disable
```
The hard USD ceiling is priced from the **API-returned token usage** times the The hard USD ceiling is priced from the **API-returned token usage** times the
configured `XAI_PRICE_*_PER_M` fallbacks, so a price change only needs those configured `XAI_PRICE_*_PER_M` fallbacks, so a price change only needs those
@ -108,11 +133,11 @@ ai-bot:
image: ai-bot:custom image: ai-bot:custom
container_name: vojo-ai-bot container_name: vojo-ai-bot
restart: unless-stopped restart: unless-stopped
depends_on: [synapse] depends_on: [synapse, postgres] # needs both up before it starts
env_file: ./ai-bot/ai-bot.env # NON-SECRET config (chmod 600) env_file: ./ai-bot/ai-bot.env # config incl. AI_BOT_DATABASE_URL (chmod 600 — embeds the DB password)
environment: environment:
REGISTRATION_PATH: /data/registration.yaml # tokens (generated; shared with Synapse) REGISTRATION_PATH: /data/registration.yaml # tokens (generated; shared with Synapse)
STATE_DIR: /data/state # SQLite: spend ledger + txn dedup STATE_DIR: /data/state # runtime dir (the operational store is now in Postgres)
XAI_API_KEY_FILE: /data/secrets/xai_api_key # the one standalone secret XAI_API_KEY_FILE: /data/secrets/xai_api_key # the one standalone secret
volumes: volumes:
- ./ai-bot:/data # owned by uid 65532 (see setup) - ./ai-bot:/data # owned by uid 65532 (see setup)
@ -144,6 +169,17 @@ Compile-level + unit-tested locally:
safe-URL allowlist, false-positive guards, oversize/adversarial fallbacks). safe-URL allowlist, false-positive guards, oversize/adversarial fallbacks).
- ✅ `check-config` reads env + loads the system prompt. - ✅ `check-config` reads env + loads the system prompt.
The store-backed tests (appservice transaction handling + the dedup/limiter/warned
store in `store_test.go`, including the concurrent per-user-cap guarantee and
restart-durability) need a throwaway Postgres via `AI_BOT_TEST_DATABASE_URL`; they
**skip** when it is unset, so `go test ./...` stays green without one. To run them:
```bash
docker run -d --name pg -e POSTGRES_PASSWORD=p -p 5432:5432 postgres:16
# … create role+db vojo_ai, then:
AI_BOT_TEST_DATABASE_URL=postgres://vojo_ai:…@localhost:5432/vojo_ai?sslmode=disable go test ./...
```
Deferred to a live homeserver + xAI key + a loaded registration (runtime ✔): Deferred to a live homeserver + xAI key + a loaded registration (runtime ✔):
- Synapse pushes transactions → bot replies (`authenticated as @ai:vojo.chat` in logs); - Synapse pushes transactions → bot replies (`authenticated as @ai:vojo.chat` in logs);

34
apps/ai-bot/appservice_test.go
View file

@ -6,22 +6,46 @@ import (
"log/slog" "log/slog"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"path/filepath" "os"
"strings" "strings"
"testing" "testing"
"time" "time"
) )
// testDSN is the throwaway Postgres the store-backed tests run against. When unset,
// those tests skip rather than fail, so `go test ./...` stays green on a machine
// without a Postgres (the build/vet gates still cover the package).
func testDSN() string { return os.Getenv("AI_BOT_TEST_DATABASE_URL") }
// openTestStore opens the store against the test database with a clean slate, so a
// shared/persistent test database doesn't leak rows between tests or runs. Skips the
// test when AI_BOT_TEST_DATABASE_URL is unset.
func openTestStore(t *testing.T) *Store {
t.Helper()
dsn := testDSN()
if dsn == "" {
t.Skip("set AI_BOT_TEST_DATABASE_URL (a throwaway Postgres) to run store-backed tests")
}
st, err := OpenStore(dsn)
if err != nil {
t.Fatalf("open store: %v", err)
}
ctx, cancel := opContext()
defer cancel()
if _, err := st.pool.Exec(ctx, `TRUNCATE processed_txn, processed_event, spend, warned_encrypted`); err != nil {
st.Close()
t.Fatalf("truncate test tables: %v", err)
}
return st
}
// newTestAS wires an AppService whose handler pushes each dispatched batch onto a // newTestAS wires an AppService whose handler pushes each dispatched batch onto a
// channel. Transactions are now processed asynchronously (the 200 is returned before // channel. Transactions are now processed asynchronously (the 200 is returned before
// the handler runs), so tests read from the channel with a timeout instead of // the handler runs), so tests read from the channel with a timeout instead of
// inspecting a slice immediately after the call. // inspecting a slice immediately after the call.
func newTestAS(t *testing.T) (*AppService, *Store, chan []Event) { func newTestAS(t *testing.T) (*AppService, *Store, chan []Event) {
t.Helper() t.Helper()
st, err := OpenStore(filepath.Join(t.TempDir(), "t.db")) st := openTestStore(t)
if err != nil {
t.Fatalf("open store: %v", err)
}
dispatched := make(chan []Event, 8) dispatched := make(chan []Event, 8)
as := NewAppService( as := NewAppService(
&Config{HSToken: "secret", BotMXID: "@ai:vojo.chat"}, &Config{HSToken: "secret", BotMXID: "@ai:vojo.chat"},

6
apps/ai-bot/bot.go
View file

@ -51,7 +51,7 @@ func NewBot(ctx context.Context, cfg *Config, logger *slog.Logger) (*Bot, error)
mx := NewMatrixClient(cfg.HomeserverURL, cfg.ASToken, cfg.BotMXID) mx := NewMatrixClient(cfg.HomeserverURL, cfg.ASToken, cfg.BotMXID)
xai := NewXAIClient(cfg.XAIBaseURL, cfg.XAIAPIKey, logger) xai := NewXAIClient(cfg.XAIBaseURL, cfg.XAIAPIKey, logger)
st, err := OpenStore(cfg.statePath("ai-bot.db")) st, err := OpenStore(cfg.DatabaseURL)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -164,8 +164,8 @@ func (b *Bot) handleEvent(ctx context.Context, ev *Event) {
// markSeen records an event id in both the in-memory set and the durable store and // markSeen records an event id in both the in-memory set and the durable store and
// reports whether it is NEW (first time). The in-memory Add is atomic, and SeenEvent // reports whether it is NEW (first time). The in-memory Add is atomic, and SeenEvent
// is an atomic INSERT OR IGNORE, so two racing goroutines for the same event can never // is an atomic INSERT … ON CONFLICT DO NOTHING, so two racing goroutines for the same
// both proceed. On a durable-store error we fall through (the in-memory set still // event can never both proceed. On a durable-store error we fall through (the in-memory set still
// guards this session). // guards this session).
func (b *Bot) markSeen(eventID string) bool { func (b *Bot) markSeen(eventID string) bool {
if !b.seen.Add(eventID) { if !b.seen.Add(eventID) {

9
apps/ai-bot/config.go
View file

@ -56,6 +56,12 @@ type Config struct {
SystemPromptPath string SystemPromptPath string
SystemPrompt string SystemPrompt string
StateDir string StateDir string
// DatabaseURL is the libpq/pgx DSN of the bot's dedicated Postgres database
// (`vojo_ai`), e.g. postgres://vojo_ai:***@postgres:5432/vojo_ai?sslmode=disable.
// It holds only operational state (txn/event dedup, the daily spend ledger, the
// encrypted-warned set) — never message content. Required.
DatabaseURL string
} }
func getenv(key, def string) string { func getenv(key, def string) string {
@ -130,6 +136,7 @@ func LoadConfig() (*Config, error) {
XAIModel: getenv("XAI_MODEL", "grok-4.20-0309-non-reasoning"), XAIModel: getenv("XAI_MODEL", "grok-4.20-0309-non-reasoning"),
SystemPromptPath: getenv("SYSTEM_PROMPT_PATH", "prompts/system_ru.txt"), SystemPromptPath: getenv("SYSTEM_PROMPT_PATH", "prompts/system_ru.txt"),
StateDir: strings.TrimRight(getenv("STATE_DIR", "/state"), "/"), StateDir: strings.TrimRight(getenv("STATE_DIR", "/state"), "/"),
DatabaseURL: getenv("AI_BOT_DATABASE_URL", ""),
AllowedServers: parseServerSet(getenv("ALLOWED_SERVERS", "")), AllowedServers: parseServerSet(getenv("ALLOWED_SERVERS", "")),
UnlimitedUsers: parseServerSet(getenv("UNLIMITED_USERS", "")), UnlimitedUsers: parseServerSet(getenv("UNLIMITED_USERS", "")),
} }
@ -179,6 +186,7 @@ func LoadConfig() (*Config, error) {
req("AS_TOKEN", cfg.ASToken) req("AS_TOKEN", cfg.ASToken)
req("HS_TOKEN", cfg.HSToken) req("HS_TOKEN", cfg.HSToken)
req("XAI_API_KEY", cfg.XAIAPIKey) req("XAI_API_KEY", cfg.XAIAPIKey)
req("AI_BOT_DATABASE_URL", cfg.DatabaseURL)
if len(cfg.AllowedServers) == 0 { if len(cfg.AllowedServers) == 0 {
problems = append(problems, "ALLOWED_SERVERS is required (comma-separated homeserver allowlist)") problems = append(problems, "ALLOWED_SERVERS is required (comma-separated homeserver allowlist)")
} }
@ -259,5 +267,6 @@ func (c *Config) Summary() string {
c.PriceInputPerM, c.PriceCachedPerM, c.PriceOutputPerM), c.PriceInputPerM, c.PriceCachedPerM, c.PriceOutputPerM),
" SYSTEM_PROMPT_PATH = " + c.SystemPromptPath, " SYSTEM_PROMPT_PATH = " + c.SystemPromptPath,
" STATE_DIR = " + c.StateDir, " STATE_DIR = " + c.StateDir,
" AI_BOT_DATABASE_URL= " + redact(c.DatabaseURL),
}, "\n") }, "\n")
} }

18
apps/ai-bot/go.mod
View file

@ -3,23 +3,21 @@ module vojo.chat/ai-bot
go 1.25.0 go 1.25.0
require ( require (
github.com/jackc/pgx/v5 v5.9.2
github.com/microcosm-cc/bluemonday v1.0.27 github.com/microcosm-cc/bluemonday v1.0.27
github.com/yuin/goldmark v1.8.2 github.com/yuin/goldmark v1.8.2
gopkg.in/yaml.v3 v3.0.1 gopkg.in/yaml.v3 v3.0.1
modernc.org/sqlite v1.51.0
) )
require ( require (
github.com/aymerick/douceur v0.2.0 // indirect github.com/aymerick/douceur v0.2.0 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/gorilla/css v1.0.1 // indirect github.com/gorilla/css v1.0.1 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect
github.com/kr/text v0.2.0 // indirect
github.com/rogpeppe/go-internal v1.15.0 // indirect
golang.org/x/net v0.26.0 // indirect golang.org/x/net v0.26.0 // indirect
golang.org/x/sys v0.42.0 // indirect golang.org/x/sync v0.20.0 // indirect
modernc.org/libc v1.72.3 // indirect golang.org/x/text v0.29.0 // indirect
modernc.org/mathutil v1.7.1 // indirect
modernc.org/memory v1.11.0 // indirect
) )

80
apps/ai-bot/go.sum
View file

@ -1,65 +1,45 @@
github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8=
github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0=
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw=
github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk=
github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/yuin/goldmark v1.8.2 h1:kEGpgqJXdgbkhcOgBxkC0X0PmoPG1ZyoZ117rDVp4zE= github.com/yuin/goldmark v1.8.2 h1:kEGpgqJXdgbkhcOgBxkC0X0PmoPG1ZyoZ117rDVp4zE=
github.com/yuin/goldmark v1.8.2/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= github.com/yuin/goldmark v1.8.2/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg=
golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk=
golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4=
golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY=
modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ=
modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A=
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo=
modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU=
modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
modernc.org/sqlite v1.51.0 h1:aH/MMSoayAIhozZ7uJbVTT9QO/VhzBf0J9tymmmuC/U=
modernc.org/sqlite v1.51.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM=
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=

280
apps/ai-bot/store.go
View file

@ -1,11 +1,13 @@
package main package main
import ( import (
"database/sql" "context"
"errors"
"fmt" "fmt"
"time" "time"
_ "modernc.org/sqlite" "github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
) )
// reserveResult is the outcome of a pre-call limiter reservation. // reserveResult is the outcome of a pre-call limiter reservation.
@ -17,61 +19,169 @@ const (
reserveDeniedGlobal // global daily USD ceiling hit (⏳ rate-limit reaction, F24) reserveDeniedGlobal // global daily USD ceiling hit (⏳ rate-limit reaction, F24)
) )
// LRU bounds for the dedup tables (unchanged from the former SQLite store): keep
// only the most recent ids so the tables don't grow without limit.
const (
maxProcessedTxn = 5000
maxProcessedEvent = 20000
)
// opTimeout bounds every store operation. SQLite (a local file) effectively never
// blocked; Postgres is over the docker network, so a cap keeps a stalled DB from
// hanging a per-room handler goroutine forever.
const opTimeout = 10 * time.Second
// Store is the durable bot state: transaction + event dedup, the daily spend // 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 // ledger, and the encrypted-room warned set. It holds ONLY operational data — no
// stays static for a distroless/scratch image. // message content (the room timeline lives in Synapse). Backed by a dedicated
// Postgres database (`vojo_ai`), in line with the per-service bridge databases, so
// the spend ledger, dedup state and warned set share the server's backup/restore.
type Store struct { type Store struct {
db *sql.DB pool *pgxpool.Pool
} }
func OpenStore(path string) (*Store, error) { // OpenStore connects to the `vojo_ai` Postgres database via the AI_BOT_DATABASE_URL
db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)") // DSN, applies pending migrations, and returns a ready Store. A small pool suffices:
// the bot processes transactions serially and every statement here is short.
func OpenStore(dsn string) (*Store, error) {
cfg, err := pgxpool.ParseConfig(dsn)
if err != nil { if err != nil {
return nil, fmt.Errorf("parse AI_BOT_DATABASE_URL: %w", err)
}
// The former SQLite store pinned a single connection to serialize all callers;
// pgx gives us a real pool. Keep it small — the per-room handler goroutines only
// ever issue brief statements, and the shared server runs many other databases.
cfg.MaxConns = 4
cfg.MinConns = 1
ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
defer cancel()
pool, err := pgxpool.NewWithConfig(ctx, cfg)
if err != nil {
return nil, fmt.Errorf("connect vojo_ai: %w", err)
}
if err := pool.Ping(ctx); err != nil {
pool.Close()
return nil, fmt.Errorf("ping vojo_ai: %w", err)
}
s := &Store{pool: pool}
if err := s.migrate(ctx); err != nil {
pool.Close()
return nil, err return nil, err
} }
// One connection: database/sql serializes all callers onto it, which keeps the return s, nil
// 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 (s *Store) Close() error {
s.pool.Close()
return nil
}
// migrationLockKey namespaces the advisory lock that guards the migration runner so
// two starting instances (the bot is single-instance, but be robust) can't race the
// version check. Arbitrary fixed constant.
const migrationLockKey = 0x76_6f_6a_6f // "vojo"
// migrations are applied in order; schema_version records the highest applied
// version so re-runs are no-ops. Every step is also idempotent (CREATE TABLE IF NOT
// EXISTS) so a half-applied database still converges.
var migrations = []string{
// v1: the operational schema — a 1:1 port of the former SQLite tables.
// processed_* carry a surrogate identity column because Postgres has no rowid:
// the LRU trim orders by it, and txn_id/event_id stay UNIQUE for the upsert.
`CREATE TABLE IF NOT EXISTS processed_txn (
id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
txn_id TEXT UNIQUE NOT NULL
);
CREATE TABLE IF NOT EXISTS processed_event (
id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
event_id TEXT UNIQUE NOT NULL
);
CREATE TABLE IF NOT EXISTS spend (
date TEXT NOT NULL,
mxid TEXT NOT NULL,
requests INTEGER NOT NULL DEFAULT 0,
usd DOUBLE PRECISION NOT NULL DEFAULT 0,
PRIMARY KEY (date, mxid)
);
CREATE TABLE IF NOT EXISTS warned_encrypted (room_id TEXT PRIMARY KEY);`,
}
// migrate runs all pending migrations on a single connection under a session
// advisory lock, recording each in schema_version.
func (s *Store) migrate(ctx context.Context) error {
conn, err := s.pool.Acquire(ctx)
if err != nil {
return fmt.Errorf("migrate: acquire: %w", err)
}
defer conn.Release()
if _, err := conn.Exec(ctx, `SELECT pg_advisory_lock($1)`, int64(migrationLockKey)); err != nil {
return fmt.Errorf("migrate: lock: %w", err)
}
defer func() {
_, _ = conn.Exec(ctx, `SELECT pg_advisory_unlock($1)`, int64(migrationLockKey))
}()
if _, err := conn.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_version (version INTEGER PRIMARY KEY)`); err != nil {
return fmt.Errorf("migrate: schema_version: %w", err)
}
var current int
if err := conn.QueryRow(ctx, `SELECT COALESCE(MAX(version), 0) FROM schema_version`).Scan(&current); err != nil {
return fmt.Errorf("migrate: read version: %w", err)
}
for v := current; v < len(migrations); v++ {
tx, err := conn.Begin(ctx)
if err != nil {
return fmt.Errorf("migrate: begin %d: %w", v+1, err)
}
if _, err := tx.Exec(ctx, migrations[v]); err != nil {
_ = tx.Rollback(ctx)
return fmt.Errorf("migrate: apply %d: %w", v+1, err)
}
if _, err := tx.Exec(ctx, `INSERT INTO schema_version (version) VALUES ($1)`, v+1); err != nil {
_ = tx.Rollback(ctx)
return fmt.Errorf("migrate: record %d: %w", v+1, err)
}
if err := tx.Commit(ctx); err != nil {
return fmt.Errorf("migrate: commit %d: %w", v+1, err)
}
}
return nil
}
func todayUTC() string { return time.Now().UTC().Format("2006-01-02") } func todayUTC() string { return time.Now().UTC().Format("2006-01-02") }
// opContext derives a bounded context for a single store operation.
func opContext() (context.Context, context.CancelFunc) {
return context.WithTimeout(context.Background(), opTimeout)
}
// HasTxn / MarkTxn give appservice transactions idempotency across restarts: a // HasTxn / MarkTxn give appservice transactions idempotency across restarts: a
// transaction Synapse retries (because our 200 was lost) is processed at most // transaction Synapse retries (because our 200 was lost) is processed at most
// once. The table is bounded to the most recent ids. // once. The table is bounded to the most recent ids.
func (s *Store) HasTxn(txnID string) (bool, error) { func (s *Store) HasTxn(txnID string) (bool, error) {
ctx, cancel := opContext()
defer cancel()
var one int var one int
err := s.db.QueryRow(`SELECT 1 FROM processed_txn WHERE txn_id = ?`, txnID).Scan(&one) err := s.pool.QueryRow(ctx, `SELECT 1 FROM processed_txn WHERE txn_id = $1`, txnID).Scan(&one)
if err == sql.ErrNoRows { if errors.Is(err, pgx.ErrNoRows) {
return false, nil return false, nil
} }
return err == nil, err return err == nil, err
} }
func (s *Store) MarkTxn(txnID string) error { func (s *Store) MarkTxn(txnID string) error {
if _, err := s.db.Exec(`INSERT OR IGNORE INTO processed_txn (txn_id) VALUES (?)`, txnID); err != nil { ctx, cancel := opContext()
defer cancel()
if _, err := s.pool.Exec(ctx,
`INSERT INTO processed_txn (txn_id) VALUES ($1) ON CONFLICT DO NOTHING`, txnID); err != nil {
return err return err
} }
_, err := s.db.Exec(`DELETE FROM processed_txn WHERE rowid NOT IN _, err := s.pool.Exec(ctx, `DELETE FROM processed_txn WHERE id NOT IN
(SELECT rowid FROM processed_txn ORDER BY rowid DESC LIMIT 5000)`) (SELECT id FROM processed_txn ORDER BY id DESC LIMIT $1)`, maxProcessedTxn)
return err return err
} }
@ -79,28 +189,37 @@ func (s *Store) MarkTxn(txnID string) error {
// or already seen (false) — the DURABLE equivalent of the in-memory dedup set, so // 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 // 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 // the bot reprocess it (dup answer + double-bill + cap inflation). Bounded to the
// most recent ids. // most recent ids. INSERT … ON CONFLICT DO NOTHING affects 1 row on insert and 0 on
// conflict, so RowsAffected distinguishes new from already-seen.
func (s *Store) SeenEvent(eventID string) (bool, error) { func (s *Store) SeenEvent(eventID string) (bool, error) {
res, err := s.db.Exec(`INSERT OR IGNORE INTO processed_event (event_id) VALUES (?)`, eventID) ctx, cancel := opContext()
defer cancel()
tag, err := s.pool.Exec(ctx,
`INSERT INTO processed_event (event_id) VALUES ($1) ON CONFLICT DO NOTHING`, eventID)
if err != nil { if err != nil {
return false, err return false, err
} }
if n, _ := res.RowsAffected(); n == 0 { if tag.RowsAffected() == 0 {
return false, nil // already recorded → not new return false, nil // already recorded → not new
} }
_, err = s.db.Exec(`DELETE FROM processed_event WHERE rowid NOT IN _, err = s.pool.Exec(ctx, `DELETE FROM processed_event WHERE id NOT IN
(SELECT rowid FROM processed_event ORDER BY rowid DESC LIMIT 20000)`) (SELECT id FROM processed_event ORDER BY id DESC LIMIT $1)`, maxProcessedEvent)
return true, err return true, err
} }
// SpentTodayUSD sums all spend for the current UTC day. // SpentTodayUSD sums all spend for the current UTC day. SUM over no rows is NULL,
// which scans into a nil *float64 → treated as 0.
func (s *Store) SpentTodayUSD() (float64, error) { func (s *Store) SpentTodayUSD() (float64, error) {
var v sql.NullFloat64 ctx, cancel := opContext()
err := s.db.QueryRow(`SELECT SUM(usd) FROM spend WHERE date = ?`, todayUTC()).Scan(&v) defer cancel()
if err != nil { var v *float64
if err := s.pool.QueryRow(ctx, `SELECT SUM(usd) FROM spend WHERE date = $1`, todayUTC()).Scan(&v); err != nil {
return 0, err return 0, err
} }
return v.Float64, nil if v == nil {
return 0, nil
}
return *v, nil
} }
// Reserve runs the two independent gates in one transaction, BEFORE the xAI call // Reserve runs the two independent gates in one transaction, BEFORE the xAI call
@ -108,38 +227,62 @@ func (s *Store) SpentTodayUSD() (float64, error) {
// anti-abuse. It increments the per-user request count on success; the USD 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 // reconciled after the response. Order: global first (cheapest to deny), then
// per-user. // per-user.
//
// A transaction-scoped advisory lock on (date, mxid) serializes concurrent
// reservations for the SAME user+day, so the per-user check-then-increment stays
// atomic. The former SQLite store got this for free (one connection serialized all
// callers); the pgx pool is concurrent, and the same user messaging from two rooms
// at once would otherwise be able to slip past the per-user cap. Different users
// never contend.
func (s *Store) Reserve(mxid string, perUserCap int, dailyUSDCeiling float64) (reserveResult, error) { func (s *Store) Reserve(mxid string, perUserCap int, dailyUSDCeiling float64) (reserveResult, error) {
ctx, cancel := opContext()
defer cancel()
day := todayUTC() day := todayUTC()
tx, err := s.db.Begin()
tx, err := s.pool.Begin(ctx)
if err != nil { if err != nil {
return reserveOK, err return reserveOK, err
} }
defer tx.Rollback() defer tx.Rollback(ctx)
var global sql.NullFloat64 // Key on date|mxid. The separator only needs to avoid cross-key ambiguity; a
if err := tx.QueryRow(`SELECT SUM(usd) FROM spend WHERE date = ?`, day).Scan(&global); err != nil { // hash collision would merely over-serialize two unrelated users, never corrupt a
// count. (NUL is rejected by Postgres text, so use a printable separator.)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, day+"|"+mxid); err != nil {
return reserveOK, err return reserveOK, err
} }
if global.Float64 >= dailyUSDCeiling {
// SUM over zero rows is NULL → nil pointer → treat as 0.0, exactly as the SQLite
// store's sql.NullFloat64 did (and as SpentTodayUSD does). This keeps the gate 1:1
// even at the degenerate dailyUSDCeiling == 0 (deny everything), where 0 >= 0.
var global *float64
if err := tx.QueryRow(ctx, `SELECT SUM(usd) FROM spend WHERE date = $1`, day).Scan(&global); err != nil {
return reserveOK, err
}
spentToday := 0.0
if global != nil {
spentToday = *global
}
if spentToday >= dailyUSDCeiling {
return reserveDeniedGlobal, nil return reserveDeniedGlobal, nil
} }
var requests int var requests int
err = tx.QueryRow(`SELECT requests FROM spend WHERE date = ? AND mxid = ?`, day, mxid).Scan(&requests) err = tx.QueryRow(ctx, `SELECT requests FROM spend WHERE date = $1 AND mxid = $2`, day, mxid).Scan(&requests)
if err != nil && err != sql.ErrNoRows { if err != nil && !errors.Is(err, pgx.ErrNoRows) {
return reserveOK, err return reserveOK, err
} }
if requests >= perUserCap { if requests >= perUserCap {
return reserveDeniedUser, nil return reserveDeniedUser, nil
} }
if _, err := tx.Exec( if _, err := tx.Exec(ctx,
`INSERT INTO spend (date, mxid, requests, usd) VALUES (?, ?, 1, 0) `INSERT INTO spend (date, mxid, requests, usd) VALUES ($1, $2, 1, 0)
ON CONFLICT(date, mxid) DO UPDATE SET requests = requests + 1`, ON CONFLICT (date, mxid) DO UPDATE SET requests = spend.requests + 1`,
day, mxid); err != nil { day, mxid); err != nil {
return reserveOK, err return reserveOK, err
} }
if err := tx.Commit(); err != nil { if err := tx.Commit(ctx); err != nil {
return reserveOK, err return reserveOK, err
} }
return reserveOK, nil return reserveOK, nil
@ -147,20 +290,26 @@ func (s *Store) Reserve(mxid string, perUserCap int, dailyUSDCeiling float64) (r
// RefundRequest gives back a reserved request slot when the call ultimately // 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 // failed (e.g. an xAI outage), so a transient failure doesn't burn the user's
// daily cap. Never drops below zero. // daily cap. Never drops below zero. A single UPDATE is atomic, so concurrent
// refunds settle correctly without extra locking.
func (s *Store) RefundRequest(mxid string) error { func (s *Store) RefundRequest(mxid string) error {
_, err := s.db.Exec( ctx, cancel := opContext()
`UPDATE spend SET requests = MAX(0, requests - 1) WHERE date = ? AND mxid = ?`, defer cancel()
_, err := s.pool.Exec(ctx,
`UPDATE spend SET requests = GREATEST(0, requests - 1) WHERE date = $1 AND mxid = $2`,
todayUTC(), mxid) todayUTC(), mxid)
return err return err
} }
// Reconcile books the actual USD cost of a completed call against the user's // Reconcile books the actual USD cost of a completed call against the user's
// daily row (and thus the global total). // daily row (and thus the global total). The accumulating upsert is atomic and
// commutative, so concurrent reconciles for the same user sum correctly.
func (s *Store) Reconcile(mxid string, usd float64) error { func (s *Store) Reconcile(mxid string, usd float64) error {
_, err := s.db.Exec( ctx, cancel := opContext()
`INSERT INTO spend (date, mxid, requests, usd) VALUES (?, ?, 0, ?) defer cancel()
ON CONFLICT(date, mxid) DO UPDATE SET usd = usd + excluded.usd`, _, err := s.pool.Exec(ctx,
`INSERT INTO spend (date, mxid, requests, usd) VALUES ($1, $2, 0, $3)
ON CONFLICT (date, mxid) DO UPDATE SET usd = spend.usd + excluded.usd`,
todayUTC(), mxid, usd) todayUTC(), mxid, usd)
return err return err
} }
@ -170,15 +319,20 @@ func (s *Store) Reconcile(mxid string, usd float64) error {
// message (F5). The bot never reacts to its own events: m.reaction is not an // message (F5). The bot never reacts to its own events: m.reaction is not an
// m.room.message, so it never re-enters handleMessage. // m.room.message, so it never re-enters handleMessage.
func (s *Store) HasWarnedEncrypted(roomID string) (bool, error) { func (s *Store) HasWarnedEncrypted(roomID string) (bool, error) {
ctx, cancel := opContext()
defer cancel()
var one int var one int
err := s.db.QueryRow(`SELECT 1 FROM warned_encrypted WHERE room_id = ?`, roomID).Scan(&one) err := s.pool.QueryRow(ctx, `SELECT 1 FROM warned_encrypted WHERE room_id = $1`, roomID).Scan(&one)
if err == sql.ErrNoRows { if errors.Is(err, pgx.ErrNoRows) {
return false, nil return false, nil
} }
return err == nil, err return err == nil, err
} }
func (s *Store) SetWarnedEncrypted(roomID string) error { func (s *Store) SetWarnedEncrypted(roomID string) error {
_, err := s.db.Exec(`INSERT OR IGNORE INTO warned_encrypted (room_id) VALUES (?)`, roomID) ctx, cancel := opContext()
defer cancel()
_, err := s.pool.Exec(ctx,
`INSERT INTO warned_encrypted (room_id) VALUES ($1) ON CONFLICT DO NOTHING`, roomID)
return err return err
} }

211
apps/ai-bot/store_test.go Normal file
View file

@ -0,0 +1,211 @@
package main
import (
"sync"
"sync/atomic"
"testing"
)
// These tests exercise the Postgres-backed store directly. They run only when
// AI_BOT_TEST_DATABASE_URL points at a throwaway database (openTestStore skips
// otherwise) and start from a clean slate (openTestStore truncates).
func TestStoreTxnDedup(t *testing.T) {
st := openTestStore(t)
defer st.Close()
if got, err := st.HasTxn("txn-1"); err != nil || got {
t.Fatalf("fresh txn: got (%v,%v), want (false,nil)", got, err)
}
if err := st.MarkTxn("txn-1"); err != nil {
t.Fatalf("mark: %v", err)
}
if got, err := st.HasTxn("txn-1"); err != nil || !got {
t.Fatalf("marked txn: got (%v,%v), want (true,nil)", got, err)
}
// Re-marking is idempotent (a retried transaction).
if err := st.MarkTxn("txn-1"); err != nil {
t.Fatalf("re-mark: %v", err)
}
if got, _ := st.HasTxn("txn-2"); got {
t.Fatalf("unrelated txn must be unseen")
}
}
func TestStoreSeenEvent(t *testing.T) {
st := openTestStore(t)
defer st.Close()
first, err := st.SeenEvent("$ev1")
if err != nil || !first {
t.Fatalf("first SeenEvent: got (%v,%v), want (true,nil)", first, err)
}
again, err := st.SeenEvent("$ev1")
if err != nil || again {
t.Fatalf("repeat SeenEvent: got (%v,%v), want (false,nil)", again, err)
}
other, err := st.SeenEvent("$ev2")
if err != nil || !other {
t.Fatalf("new SeenEvent: got (%v,%v), want (true,nil)", other, err)
}
}
// Dedup state must survive a process restart — the whole point of the durable store.
func TestStoreDedupSurvivesRestart(t *testing.T) {
st := openTestStore(t)
if _, err := st.SeenEvent("$ev-restart"); err != nil {
t.Fatalf("seen: %v", err)
}
if err := st.MarkTxn("txn-restart"); err != nil {
t.Fatalf("mark: %v", err)
}
st.Close()
// Reopen the same database WITHOUT truncating: simulates a container restart.
st2, err := OpenStore(testDSN())
if err != nil {
t.Fatalf("reopen: %v", err)
}
defer st2.Close()
if isNew, err := st2.SeenEvent("$ev-restart"); err != nil || isNew {
t.Fatalf("event after restart must be already-seen: got (%v,%v)", isNew, err)
}
if seen, err := st2.HasTxn("txn-restart"); err != nil || !seen {
t.Fatalf("txn after restart must be seen: got (%v,%v)", seen, err)
}
}
func TestStoreLimiterPerUserCap(t *testing.T) {
st := openTestStore(t)
defer st.Close()
const user = "@u:vojo.chat"
const cap, ceiling = 2, 100.0
for i := 0; i < cap; i++ {
if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveOK {
t.Fatalf("reserve %d: got (%v,%v), want reserveOK", i, res, err)
}
}
// The (cap+1)th request is denied per-user.
if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveDeniedUser {
t.Fatalf("over-cap reserve: got (%v,%v), want reserveDeniedUser", res, err)
}
// A different user is unaffected.
if res, err := st.Reserve("@v:vojo.chat", cap, ceiling); err != nil || res != reserveOK {
t.Fatalf("other user reserve: got (%v,%v), want reserveOK", res, err)
}
// Refund returns a slot, so the first user can reserve once more.
if err := st.RefundRequest(user); err != nil {
t.Fatalf("refund: %v", err)
}
if res, err := st.Reserve(user, cap, ceiling); err != nil || res != reserveOK {
t.Fatalf("post-refund reserve: got (%v,%v), want reserveOK", res, err)
}
}
// A zero per-user cap denies even the first request — the SQLite store's
// requests(0) >= cap(0) behaviour, preserved.
func TestStoreLimiterZeroCap(t *testing.T) {
st := openTestStore(t)
defer st.Close()
if res, err := st.Reserve("@u:vojo.chat", 0, 100.0); err != nil || res != reserveDeniedUser {
t.Fatalf("zero-cap reserve: got (%v,%v), want reserveDeniedUser", res, err)
}
}
// A zero ceiling denies the very first request of the day even before any spend row
// exists — the SQLite store treated SUM(NULL) as 0.0 (0 >= 0), and the PG store must
// match (SUM over zero rows is NULL).
func TestStoreLimiterZeroCeiling(t *testing.T) {
st := openTestStore(t)
defer st.Close()
if res, err := st.Reserve("@u:vojo.chat", 1_000_000, 0); err != nil || res != reserveDeniedGlobal {
t.Fatalf("zero-ceiling reserve on empty store: got (%v,%v), want reserveDeniedGlobal", res, err)
}
}
func TestStoreLimiterGlobalCeiling(t *testing.T) {
st := openTestStore(t)
defer st.Close()
const ceiling = 1.0
// Book spend up to the ceiling (Reconcile is what feeds the global gate).
if err := st.Reconcile("@a:vojo.chat", 0.6); err != nil {
t.Fatalf("reconcile a: %v", err)
}
if err := st.Reconcile("@b:vojo.chat", 0.5); err != nil {
t.Fatalf("reconcile b: %v", err)
}
if spent, err := st.SpentTodayUSD(); err != nil || spent < 1.1 {
t.Fatalf("spent today: got (%v,%v), want >= 1.1", spent, err)
}
// Now any reservation is denied globally, regardless of the per-user cap.
if res, err := st.Reserve("@c:vojo.chat", 1_000_000, ceiling); err != nil || res != reserveDeniedGlobal {
t.Fatalf("over-ceiling reserve: got (%v,%v), want reserveDeniedGlobal", res, err)
}
}
// The pgx pool is concurrent (the SQLite store serialized on one connection). The
// advisory lock in Reserve must still admit EXACTLY perUserCap requests when many
// arrive at once for the same user — the same user messaging from several rooms
// simultaneously must not slip past the cap.
func TestStoreReserveConcurrentRespectsCap(t *testing.T) {
st := openTestStore(t)
defer st.Close()
const user = "@race:vojo.chat"
const cap = 10
const goroutines = 50
var ok int64
var wg sync.WaitGroup
for i := 0; i < goroutines; i++ {
wg.Add(1)
go func() {
defer wg.Done()
res, err := st.Reserve(user, cap, 1e9)
if err != nil {
t.Errorf("reserve: %v", err)
return
}
if res == reserveOK {
atomic.AddInt64(&ok, 1)
}
}()
}
wg.Wait()
if ok != cap {
t.Fatalf("concurrent reserves admitted %d, want exactly %d (the per-user cap)", ok, cap)
}
}
func TestStoreWarnedEncrypted(t *testing.T) {
st := openTestStore(t)
const room = "!enc:vojo.chat"
if warned, err := st.HasWarnedEncrypted(room); err != nil || warned {
t.Fatalf("fresh room: got (%v,%v), want (false,nil)", warned, err)
}
if err := st.SetWarnedEncrypted(room); err != nil {
t.Fatalf("set: %v", err)
}
// Setting twice is idempotent.
if err := st.SetWarnedEncrypted(room); err != nil {
t.Fatalf("re-set: %v", err)
}
if warned, err := st.HasWarnedEncrypted(room); err != nil || !warned {
t.Fatalf("warned room: got (%v,%v), want (true,nil)", warned, err)
}
st.Close()
// The one-shot flag must outlive a restart (F5: no re-react after restart).
st2, err := OpenStore(testDSN())
if err != nil {
t.Fatalf("reopen: %v", err)
}
defer st2.Close()
if warned, err := st2.HasWarnedEncrypted(room); err != nil || !warned {
t.Fatalf("warned after restart: got (%v,%v), want (true,nil)", warned, err)
}
}

6
docs/ai/server-side.md
View file

@ -40,7 +40,8 @@ you're touching server config, SSH to the box and read the live file. Last refre
## Postgres roles & databases ## Postgres roles & databases
`synapse` is the historical superuser. Each bridge has its own role + DB with a 32-char `synapse` is the historical superuser. Each bridge — and the Vojo AI bot — has its own
non-superuser role + DB with a 32-char
`openssl rand -base64 32 | tr -d '/+=' | head -c 32` password. `openssl rand -base64 32 | tr -d '/+=' | head -c 32` password.
| Role | Database | Used by | | Role | Database | Used by |
@ -49,6 +50,7 @@ you're touching server config, SSH to the box and read the live file. Last refre
| `mautrix_telegram` | `mautrix_telegram` | Telegram bridge | | `mautrix_telegram` | `mautrix_telegram` | Telegram bridge |
| `mautrix_discord` | `mautrix_discord` | Discord bridge | | `mautrix_discord` | `mautrix_discord` | Discord bridge |
| `mautrix_whatsapp` | `mautrix_whatsapp` | WhatsApp bridge | | `mautrix_whatsapp` | `mautrix_whatsapp` | WhatsApp bridge |
| `vojo_ai` | `vojo_ai` | Vojo AI bot — operational store (txn/event dedup, daily spend ledger, encrypted-warned set; **no** message content). DSN via `AI_BOT_DATABASE_URL`. |
## `~/vojo/synapse/homeserver.yaml` — relevant slices ## `~/vojo/synapse/homeserver.yaml` — relevant slices
@ -130,7 +132,7 @@ in doubt.
| `telegram-bridge` | `dock.mau.dev/mautrix/telegram:<v26.04 bridgev2 tag>` | `./bridges/telegram:/data` | | `telegram-bridge` | `dock.mau.dev/mautrix/telegram:<v26.04 bridgev2 tag>` | `./bridges/telegram:/data` |
| `discord-bridge` | `dock.mau.dev/mautrix/discord:v0.7.5` | `./bridges/discord:/data` (legacy bridge — runtime reports `0.7.6+dev`) | | `discord-bridge` | `dock.mau.dev/mautrix/discord:v0.7.5` | `./bridges/discord:/data` (legacy bridge — runtime reports `0.7.6+dev`) |
| `whatsapp-bridge` | `dock.mau.dev/mautrix/whatsapp:v0.12.4` | `./bridges/whatsapp:/data` | | `whatsapp-bridge` | `dock.mau.dev/mautrix/whatsapp:v0.12.4` | `./bridges/whatsapp:/data` |
| `ai-bot` | `ai-bot:custom` (built locally from [`apps/ai-bot/`](../../apps/ai-bot/), shipped via `docker save \| ssh docker load` — VS Code task **Deploy AI bot**) | **Vojo AI** = `@ai:vojo.chat`, an xAI-Grok-backed **application service** (NOT a normal bot user). Answers `@`-mentions in groups + everything in 1:1s; the Grok reply (markdown) is rendered to `org.matrix.custom.html` and sent as `formatted_body` (in-bot `markdown.go`, zero deps; emits only tags Cinny's sanitizer keeps, escapes all model text), falling back to the plain `body` when there's no formatting. Mounts `./ai-bot:/data` (owned **uid 65532**, distroless nonroot) holding `registration.yaml` (self-generated, `generate-registration`), `state/` (SQLite spend ledger + txn dedup) and `secrets/xai_api_key`. Push port `:8009` (registration `url: http://ai-bot:8009`). Secrets via env/`*_FILE`; `as_token`/`hs_token` read from `registration.yaml` (no rotation). See [`apps/ai-bot/README.md`](../../apps/ai-bot/README.md). | | `ai-bot` | `ai-bot:custom` (built locally from [`apps/ai-bot/`](../../apps/ai-bot/), shipped via `docker save \| ssh docker load` — VS Code task **Deploy AI bot**) | **Vojo AI** = `@ai:vojo.chat`, an xAI-Grok-backed **application service** (NOT a normal bot user). Answers `@`-mentions in groups + everything in 1:1s; the Grok reply (markdown) is rendered to `org.matrix.custom.html` and sent as `formatted_body` (in-bot `markdown.go`, zero deps; emits only tags Cinny's sanitizer keeps, escapes all model text), falling back to the plain `body` when there's no formatting. Mounts `./ai-bot:/data` (owned **uid 65532**, distroless nonroot) holding `registration.yaml` (self-generated, `generate-registration`), `state/` (runtime dir) and `secrets/xai_api_key`. Its **operational store** (txn/event dedup, daily spend ledger, encrypted-warned set) lives in the dedicated `vojo_ai` Postgres DB via `AI_BOT_DATABASE_URL` — `depends_on: [synapse, postgres]`. Push port `:8009` (registration `url: http://ai-bot:8009`). Secrets via env/`*_FILE`; `as_token`/`hs_token` read from `registration.yaml` (no rotation). See [`apps/ai-bot/README.md`](../../apps/ai-bot/README.md). |
### Bridge service stanza (template) ### Bridge service stanza (template)