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=

272
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 := ` func (s *Store) Close() error {
CREATE TABLE IF NOT EXISTS processed_txn (txn_id TEXT PRIMARY KEY); 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 ( CREATE TABLE IF NOT EXISTS spend (
date TEXT NOT NULL, mxid TEXT NOT NULL, date TEXT NOT NULL,
requests INTEGER NOT NULL DEFAULT 0, usd REAL NOT NULL DEFAULT 0, mxid TEXT NOT NULL,
requests INTEGER NOT NULL DEFAULT 0,
usd DOUBLE PRECISION NOT NULL DEFAULT 0,
PRIMARY KEY (date, mxid) PRIMARY KEY (date, mxid)
); );
CREATE TABLE IF NOT EXISTS warned_encrypted (room_id TEXT PRIMARY KEY); 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() } // 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)