orchard

  1// Conversations on disk, and the mode that keeps them off it.
  2//
  3// This is a second database rather than more tables next to anything else, for
  4// the reason search's history is: deleting every conversation has to be
  5// possible without touching whatever else the site caches, and a separate file
  6// is one line to exclude from a backup. Deletion is forward only, so the honest
  7// way to keep something out of a backup is never to have put it in one.
  8package main
  9
 10import (
 11	"crypto/rand"
 12	"database/sql"
 13	"encoding/json"
 14	"fmt"
 15	"log/slog"
 16	"os"
 17	"path/filepath"
 18	"strings"
 19	"time"
 20
 21	_ "modernc.org/sqlite"
 22)
 23
 24type Store struct{ db *sql.DB }
 25
 26const schema = `
 27CREATE TABLE IF NOT EXISTS conversations (
 28  -- A uuid rather than a counter, so an address carries no ordering and says
 29  -- nothing about how many conversations there are.
 30  id         TEXT PRIMARY KEY,
 31  title      TEXT NOT NULL DEFAULT '',
 32  created_at INTEGER NOT NULL,
 33  updated_at INTEGER NOT NULL,
 34  -- The running summary that replaces the oldest turns once a conversation
 35  -- outgrows the window. Empty until the first compaction.
 36  summary    TEXT NOT NULL DEFAULT '',
 37  -- How many messages the summary already covers, so compaction knows where
 38  -- to resume rather than summarising the same turns again.
 39  summarized INTEGER NOT NULL DEFAULT 0
 40);
 41CREATE TABLE IF NOT EXISTS messages (
 42  id      INTEGER PRIMARY KEY,
 43  conv_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
 44  role    TEXT NOT NULL,
 45  content TEXT NOT NULL,
 46  tools   TEXT NOT NULL DEFAULT '[]',
 47  -- What the user typed, when content also carries the text of their
 48  -- attachments. Empty when the two are the same.
 49  display TEXT NOT NULL DEFAULT '',
 50  files   TEXT NOT NULL DEFAULT '[]',
 51  -- The numbered pages an answer cites, so a reload links the same way the
 52  -- turn did. The page text they were matched against is not kept.
 53  sources TEXT NOT NULL DEFAULT '[]',
 54  -- The charts an answer drew, as their subjects rather than their readings, so
 55  -- reopening a conversation fetches today's numbers rather than replaying old
 56  -- ones.
 57  widgets TEXT NOT NULL DEFAULT '[]',
 58  steps TEXT NOT NULL DEFAULT '[]',
 59  at      INTEGER NOT NULL
 60);
 61CREATE INDEX IF NOT EXISTS messages_conv ON messages(conv_id, id);
 62
 63-- The rate limit penalty box, kept here so a deploy does not clear it. An
 64-- in-process map meant every restart asked a host that was already refusing,
 65-- which is the surest way to keep a ban alive.
 66CREATE TABLE IF NOT EXISTS penalties (
 67  host  TEXT PRIMARY KEY,
 68  till  INTEGER NOT NULL,
 69  trips INTEGER NOT NULL DEFAULT 1
 70);
 71
 72-- One row per outbound call to a budgeted host, so a deploy does not hand the
 73-- model a fresh day's allowance. Nothing older than a day is kept.
 74CREATE TABLE IF NOT EXISTS spend (
 75  host TEXT    NOT NULL,
 76  at   INTEGER NOT NULL
 77);
 78CREATE INDEX IF NOT EXISTS spend_host_at ON spend(host, at);
 79`
 80
 81func OpenStore(path string) (*Store, error) {
 82	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
 83		return nil, err
 84	}
 85	// secure_delete zeroes a deleted row rather than leaving it in the free
 86	// pages, which is the difference between deleting a conversation and
 87	// deleting the text of one.
 88	db, err := sql.Open("sqlite", path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=secure_delete(ON)&_pragma=foreign_keys(ON)")
 89	if err != nil {
 90		return nil, err
 91	}
 92	if _, err := db.Exec(schema); err != nil {
 93		return nil, err
 94	}
 95	migrate(db)
 96	migrateIDs(db)
 97	st := &Store{db: db}
 98	if err := st.initFacts(); err != nil {
 99		return nil, err
100	}
101	return st, nil
102}
103
104// migrate adds what a database written by an older build does not have. SQLite
105// has no ADD COLUMN IF NOT EXISTS, so each one is attempted and a duplicate is
106// the expected answer on every run after the first.
107func migrate(db *sql.DB) {
108	for _, stmt := range []string{
109		`ALTER TABLE messages ADD COLUMN display TEXT NOT NULL DEFAULT ''`,
110		`ALTER TABLE messages ADD COLUMN files TEXT NOT NULL DEFAULT '[]'`,
111		`ALTER TABLE messages ADD COLUMN sources TEXT NOT NULL DEFAULT '[]'`,
112		`ALTER TABLE messages ADD COLUMN widgets TEXT NOT NULL DEFAULT '[]'`,
113		`ALTER TABLE messages ADD COLUMN steps TEXT NOT NULL DEFAULT '[]'`,
114	} {
115		if _, err := db.Exec(stmt); err != nil && !strings.Contains(err.Error(), "duplicate column") {
116			slog.Warn("migrate", "stmt", stmt, "err", err)
117		}
118	}
119}
120
121// migrateIDs rebuilds both tables when the conversation id is still a counter.
122// SQLite cannot change a column's type, and the ids have to be handed out
123// before anything can point at them, so it is a copy rather than an update. The
124// old addresses stop resolving, which was accepted when the change was asked
125// for.
126func migrateIDs(db *sql.DB) {
127	rows, err := db.Query(`PRAGMA table_info(conversations)`)
128	if err != nil {
129		return
130	}
131	kind := ""
132	for rows.Next() {
133		var cid int
134		var name, typ, dflt any
135		var notnull, pk int
136		if rows.Scan(&cid, &name, &typ, &notnull, &dflt, &pk) == nil && name == "id" {
137			kind, _ = typ.(string)
138		}
139	}
140	rows.Close()
141	if !strings.EqualFold(kind, "INTEGER") {
142		return
143	}
144
145	tx, err := db.Begin()
146	if err != nil {
147		slog.Error("could not start the id migration", "err", err)
148		return
149	}
150	defer tx.Rollback()
151
152	fail := func(step string, err error) bool {
153		if err == nil {
154			return false
155		}
156		slog.Error("the id migration stopped", "step", step, "err", err)
157		return true
158	}
159
160	var old []int64
161	ids, err := tx.Query(`SELECT id FROM conversations ORDER BY id`)
162	if fail("reading the conversations", err) {
163		return
164	}
165	for ids.Next() {
166		var id int64
167		if ids.Scan(&id) == nil {
168			old = append(old, id)
169		}
170	}
171	ids.Close()
172
173	for _, stmt := range []string{
174		`CREATE TABLE conversations_new (
175		  id TEXT PRIMARY KEY, title TEXT NOT NULL DEFAULT '',
176		  created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL,
177		  summary TEXT NOT NULL DEFAULT '', summarized INTEGER NOT NULL DEFAULT 0)`,
178		`CREATE TABLE messages_new (
179		  id INTEGER PRIMARY KEY,
180		  conv_id TEXT NOT NULL REFERENCES conversations_new(id) ON DELETE CASCADE,
181		  role TEXT NOT NULL, content TEXT NOT NULL,
182		  tools TEXT NOT NULL DEFAULT '[]', display TEXT NOT NULL DEFAULT '',
183		  files TEXT NOT NULL DEFAULT '[]', sources TEXT NOT NULL DEFAULT '[]',
184		  widgets TEXT NOT NULL DEFAULT '[]',
185		  steps TEXT NOT NULL DEFAULT '[]',
186		  at INTEGER NOT NULL)`,
187	} {
188		if fail("creating the new tables", run(tx, stmt)) {
189			return
190		}
191	}
192
193	for _, o := range old {
194		id := newID()
195		if fail("copying a conversation", run(tx,
196			`INSERT INTO conversations_new SELECT ?, title, created_at, updated_at, summary, summarized FROM conversations WHERE id=?`,
197			id, o)) {
198			return
199		}
200		if fail("copying its messages", run(tx,
201			`INSERT INTO messages_new(conv_id, role, content, tools, display, files, sources, widgets, steps, at)
202			 SELECT ?, role, content, tools, display, files, sources, widgets, steps, at FROM messages WHERE conv_id=? ORDER BY id`,
203			id, o)) {
204			return
205		}
206	}
207
208	// The child goes first, so dropping the parent never fires a cascade over
209	// rows that are still the only copy.
210	for _, stmt := range []string{
211		`DROP TABLE messages`,
212		`DROP TABLE conversations`,
213		`ALTER TABLE conversations_new RENAME TO conversations`,
214		`ALTER TABLE messages_new RENAME TO messages`,
215		`CREATE INDEX IF NOT EXISTS messages_conv ON messages(conv_id, id)`,
216	} {
217		if fail("swapping the tables", run(tx, stmt)) {
218			return
219		}
220	}
221	if fail("committing", tx.Commit()) {
222		return
223	}
224	slog.Info("conversation ids are uuids now", "conversations", len(old))
225}
226
227func run(tx *sql.Tx, stmt string, args ...any) error {
228	_, err := tx.Exec(stmt, args...)
229	return err
230}
231
232// Close checkpoints the write ahead log into the database file before closing.
233// Without this the newest conversations live only in the -wal, and a container
234// stopped with SIGKILL rather than SIGTERM can lose them.
235func (s *Store) Close() error {
236	if _, err := s.db.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil {
237		slog.Warn("wal checkpoint", "err", err)
238	}
239	return s.db.Close()
240}
241
242// Checkpoint folds the log into the file without closing. It runs after every
243// turn, so the cost of a hard kill is at most the turn in flight rather than
244// everything since the process started.
245func (s *Store) Checkpoint() {
246	if _, err := s.db.Exec(`PRAGMA wal_checkpoint(PASSIVE)`); err != nil {
247		slog.Debug("wal checkpoint", "err", err)
248	}
249}
250
251type Conversation struct {
252	ID        string    `json:"id"`
253	Title     string    `json:"title"`
254	Updated   time.Time `json:"updated"`
255	Preview   string    `json:"preview,omitempty"`
256	Summary   string    `json:"-"`
257	Summarize int       `json:"-"`
258}
259
260type Stored struct {
261	Role    Role          `json:"role"`
262	Content string        `json:"content"`
263	Display string        `json:"display,omitempty"`
264	Files   []Attachment  `json:"files,omitempty"`
265	Tools   []ToolSummary `json:"tools,omitempty"`
266	Sources []Source      `json:"sources,omitempty"`
267	Widgets []Widget      `json:"widgets,omitempty"`
268	Steps   []Step        `json:"steps,omitempty"`
269	At      time.Time     `json:"at"`
270}
271
272// Shown is what the user typed, which is the whole message unless files were
273// attached and their text was folded into it.
274func (m Stored) Shown() string {
275	if m.Display != "" {
276		return m.Display
277	}
278	return m.Content
279}
280
281// ToolSummary is what the UI shows under a message: which tools ran and how
282// they went, not their whole payload.
283type ToolSummary struct {
284	Name string `json:"name"`
285	Args string `json:"args,omitempty"`
286	MS   int64  `json:"ms"`
287	OK   bool   `json:"ok"`
288	Err  string `json:"err,omitempty"`
289	// How old the data behind this call is, for a tool that reads a snapshot
290	// rather than the live thing. Shown on the chip, because an answer off a
291	// local corpus looks exactly like one off the web otherwise.
292	Age string `json:"age,omitempty"`
293}
294
295func (s *Store) NewConversation(title string) (string, error) {
296	now := time.Now().Unix()
297	id := newID()
298	if _, err := s.db.Exec(`INSERT INTO conversations(id, title, created_at, updated_at) VALUES(?,?,?,?)`,
299		id, title, now, now); err != nil {
300		return "", err
301	}
302	return id, nil
303}
304
305// newID is a version 4 uuid. Nothing in the repo needed one before this, and a
306// dependency for sixteen random bytes and a format string is not a trade worth
307// making.
308func newID() string {
309	var b [16]byte
310	if _, err := rand.Read(b[:]); err != nil {
311		// The only way this fails is a broken kernel, and carrying on with a
312		// predictable id would be worse than saying so.
313		panic("no randomness for a conversation id: " + err.Error())
314	}
315	b[6] = b[6]&0x0f | 0x40
316	b[8] = b[8]&0x3f | 0x80
317	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
318}
319
320func (s *Store) Append(convID string, m Stored) error {
321	b, _ := json.Marshal(m.Tools)
322	f, _ := json.Marshal(m.Files)
323	src, _ := json.Marshal(m.Sources)
324	wid, _ := json.Marshal(m.Widgets)
325	stp, _ := json.Marshal(m.Steps)
326	if _, err := s.db.Exec(`INSERT INTO messages(conv_id, role, content, tools, display, files, sources, widgets, steps, at) VALUES(?,?,?,?,?,?,?,?,?,?)`,
327		convID, string(m.Role), m.Content, string(b), m.Display, string(f), string(src), string(wid), string(stp), time.Now().Unix()); err != nil {
328		return err
329	}
330	_, err := s.db.Exec(`UPDATE conversations SET updated_at=? WHERE id=?`, time.Now().Unix(), convID)
331	return err
332}
333
334func (s *Store) Messages(convID string) ([]Stored, error) {
335	rows, err := s.db.Query(`SELECT role, content, tools, display, files, sources, widgets, steps, at FROM messages WHERE conv_id=? ORDER BY id`, convID)
336	if err != nil {
337		return nil, err
338	}
339	defer rows.Close()
340	var out []Stored
341	for rows.Next() {
342		var m Stored
343		var role, tools, files, sources, widgets, steps string
344		var at int64
345		if err := rows.Scan(&role, &m.Content, &tools, &m.Display, &files, &sources, &widgets, &steps, &at); err != nil {
346			return nil, err
347		}
348		m.Role, m.At = Role(role), time.Unix(at, 0)
349		_ = json.Unmarshal([]byte(tools), &m.Tools)
350		_ = json.Unmarshal([]byte(files), &m.Files)
351		_ = json.Unmarshal([]byte(sources), &m.Sources)
352		_ = json.Unmarshal([]byte(widgets), &m.Widgets)
353		_ = json.Unmarshal([]byte(steps), &m.Steps)
354		out = append(out, m)
355	}
356	return out, rows.Err()
357}
358
359func (s *Store) List(limit int) ([]Conversation, error) {
360	rows, err := s.db.Query(`
361		SELECT c.id, c.title, c.updated_at,
362		       COALESCE((SELECT COALESCE(NULLIF(display, ''), content) FROM messages WHERE conv_id=c.id AND role='user' ORDER BY id LIMIT 1), '')
363		FROM conversations c
364		WHERE EXISTS (SELECT 1 FROM messages WHERE conv_id=c.id)
365		ORDER BY c.updated_at DESC LIMIT ?`, limit)
366	if err != nil {
367		return nil, err
368	}
369	defer rows.Close()
370	var out []Conversation
371	for rows.Next() {
372		var c Conversation
373		var up int64
374		var first string
375		if err := rows.Scan(&c.ID, &c.Title, &up, &first); err != nil {
376			return nil, err
377		}
378		c.Updated = time.Unix(up, 0)
379		if c.Title == "" {
380			c.Title = titleFrom(first)
381		}
382		c.Preview = trim(first, 90)
383		out = append(out, c)
384	}
385	return out, rows.Err()
386}
387
388func (s *Store) Get(convID string) (Conversation, error) {
389	var c Conversation
390	var up int64
391	err := s.db.QueryRow(`SELECT id, title, updated_at, summary, summarized FROM conversations WHERE id=?`,
392		convID).Scan(&c.ID, &c.Title, &up, &c.Summary, &c.Summarize)
393	c.Updated = time.Unix(up, 0)
394	return c, err
395}
396
397func (s *Store) SetSummary(convID string, summary string, covered int) error {
398	_, err := s.db.Exec(`UPDATE conversations SET summary=?, summarized=? WHERE id=?`, summary, covered, convID)
399	return err
400}
401
402func (s *Store) SetTitle(convID string, title string) error {
403	_, err := s.db.Exec(`UPDATE conversations SET title=? WHERE id=? AND title=''`, title, convID)
404	return err
405}
406
407func (s *Store) Delete(convID string) error {
408	_, err := s.db.Exec(`DELETE FROM conversations WHERE id=?`, convID)
409	return err
410}
411
412// DeleteAll drops every conversation. VACUUM afterwards so the pages actually
413// come back rather than sitting in the file as free space.
414func (s *Store) DeleteAll() error {
415	if _, err := s.db.Exec(`DELETE FROM conversations`); err != nil {
416		return err
417	}
418	_, err := s.db.Exec(`VACUUM`)
419	return err
420}
421
422func (s *Store) Count() (convs, msgs int) {
423	_ = s.db.QueryRow(`SELECT COUNT(*) FROM conversations`).Scan(&convs)
424	_ = s.db.QueryRow(`SELECT COUNT(*) FROM messages`).Scan(&msgs)
425	return
426}
427
428func titleFrom(first string) string {
429	t := trim(strings.TrimSpace(first), 48)
430	if t == "" {
431		return "New chat"
432	}
433	return t
434}
435
436func trim(s string, n int) string {
437	s = strings.Join(strings.Fields(s), " ")
438	if len(s) <= n {
439		return s
440	}
441	return strings.TrimSpace(s[:n-1]) + "…"
442}
443
444func fmtWhen(t time.Time) string {
445	d := time.Since(t)
446	switch {
447	case d < time.Minute:
448		return "just now"
449	case d < time.Hour:
450		return fmt.Sprintf("%dm ago", int(d.Minutes()))
451	case d < 24*time.Hour:
452		return fmt.Sprintf("%dh ago", int(d.Hours()))
453	default:
454		return t.Format("2 Jan")
455	}
456}
457
458// Penalties reads back the rate limit boxes that outlived the last process.
459// Anything already expired is left behind rather than deleted here, since the
460// next Trip overwrites it and a read should not write.
461func (s *Store) Penalties() map[string][2]int64 {
462	out := map[string][2]int64{}
463	rows, err := s.db.Query(`SELECT host, till, trips FROM penalties WHERE till > ?`,
464		time.Now().UnixMilli())
465	if err != nil {
466		return out
467	}
468	defer rows.Close()
469	for rows.Next() {
470		var host string
471		var till, trips int64
472		if err := rows.Scan(&host, &till, &trips); err != nil {
473			return out
474		}
475		out[host] = [2]int64{till, trips}
476	}
477	return out
478}
479
480// SavePenalty records one host's box so it survives a restart.
481func (s *Store) SavePenalty(host string, till time.Time, trips int) {
482	_, _ = s.db.Exec(`INSERT INTO penalties (host, till, trips) VALUES (?,?,?)
483		ON CONFLICT(host) DO UPDATE SET till = excluded.till, trips = excluded.trips`,
484		host, till.UnixMilli(), trips)
485}
486
487// ClearPenalty forgets a host that answered, so one bad afternoon does not
488// leave it on a long backoff for good.
489func (s *Store) ClearPenalty(host string) {
490	_, _ = s.db.Exec(`DELETE FROM penalties WHERE host = ?`, host)
491}
492
493// Spend reads back a host's calls from the last day.
494func (s *Store) Spend(host string) []time.Time {
495	rows, err := s.db.Query(`SELECT at FROM spend WHERE host = ? AND at > ? ORDER BY at`,
496		host, time.Now().Add(-24*time.Hour).UnixMilli())
497	if err != nil {
498		return nil
499	}
500	defer rows.Close()
501	var out []time.Time
502	for rows.Next() {
503		var at int64
504		if err := rows.Scan(&at); err != nil {
505			return out
506		}
507		out = append(out, time.UnixMilli(at))
508	}
509	return out
510}
511
512// SaveSpend replaces a host's record with what the budget currently holds, and
513// drops what has aged out in the same statement. Called after a turn rather
514// than per request, since losing the last turn's counts to a hard kill costs
515// less than a write on the request path.
516func (s *Store) SaveSpend(host string, at []time.Time) {
517	tx, err := s.db.Begin()
518	if err != nil {
519		return
520	}
521	defer func() { _ = tx.Rollback() }()
522	if _, err := tx.Exec(`DELETE FROM spend WHERE host = ?`, host); err != nil {
523		return
524	}
525	for _, t := range at {
526		if _, err := tx.Exec(`INSERT INTO spend (host, at) VALUES (?, ?)`, host, t.UnixMilli()); err != nil {
527			return
528		}
529	}
530	_ = tx.Commit()
531}