orchard

  1// The key ring and the call log.
  2//
  3// One file, two tables, and the same reasoning chat's history uses: this holds
  4// the full text of every prompt and every completion the estate has produced,
  5// so it is its own database rather than tables beside anything else, and
  6// deleting it is one command that touches nothing else. secure_delete is on, so
  7// a deleted call is zeroed rather than left in the free pages.
  8package main
  9
 10import (
 11	"crypto/rand"
 12	"crypto/sha256"
 13	"crypto/subtle"
 14	"database/sql"
 15	"encoding/base64"
 16	"encoding/hex"
 17	"errors"
 18	"log/slog"
 19	"os"
 20	"path/filepath"
 21	"strings"
 22	"time"
 23
 24	_ "modernc.org/sqlite"
 25)
 26
 27type Store struct{ db *sql.DB }
 28
 29const schema = `
 30CREATE TABLE IF NOT EXISTS keys (
 31  id         INTEGER PRIMARY KEY,
 32  name       TEXT NOT NULL,
 33  -- SHA-256 and not Argon2id. A key is 32 bytes this service generated, so
 34  -- there is nothing to brute force and no reason to spend 100ms of every
 35  -- request proving it. Argon2id is for repos, where the secret is chosen by
 36  -- a person.
 37  hash       TEXT NOT NULL UNIQUE,
 38  -- The first characters, so a key is recognisable in a list without being
 39  -- recoverable from one.
 40  prefix     TEXT NOT NULL,
 41  created_at INTEGER NOT NULL,
 42  last_used  INTEGER NOT NULL DEFAULT 0,
 43  revoked_at INTEGER NOT NULL DEFAULT 0
 44);
 45CREATE TABLE IF NOT EXISTS calls (
 46  id          INTEGER PRIMARY KEY,
 47  key_id      INTEGER NOT NULL DEFAULT 0,
 48  caller      TEXT NOT NULL DEFAULT '',
 49  model       TEXT NOT NULL DEFAULT '',
 50  messages    TEXT NOT NULL DEFAULT '[]',
 51  completion  TEXT NOT NULL DEFAULT '',
 52  tools       TEXT NOT NULL DEFAULT '',
 53  prompt_tok  INTEGER NOT NULL DEFAULT 0,
 54  output_tok  INTEGER NOT NULL DEFAULT 0,
 55  decode_tps  REAL NOT NULL DEFAULT 0,
 56  ms          INTEGER NOT NULL DEFAULT 0,
 57  status      INTEGER NOT NULL DEFAULT 0,
 58  err         TEXT NOT NULL DEFAULT '',
 59  at          INTEGER NOT NULL
 60);
 61CREATE INDEX IF NOT EXISTS calls_at ON calls(at DESC);
 62CREATE INDEX IF NOT EXISTS calls_caller ON calls(caller, at DESC);
 63`
 64
 65func OpenStore(path string) (*Store, error) {
 66	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
 67		return nil, err
 68	}
 69	db, err := sql.Open("sqlite", path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=secure_delete(ON)&_pragma=foreign_keys(ON)")
 70	if err != nil {
 71		return nil, err
 72	}
 73	if _, err := db.Exec(schema); err != nil {
 74		return nil, err
 75	}
 76	return &Store{db: db}, nil
 77}
 78
 79// Close checkpoints the write ahead log first, so a container stopped with
 80// SIGKILL rather than SIGTERM loses at most the call in flight.
 81func (s *Store) Close() error {
 82	if _, err := s.db.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil {
 83		slog.Warn("wal checkpoint", "err", err)
 84	}
 85	return s.db.Close()
 86}
 87
 88func (s *Store) Checkpoint() {
 89	if _, err := s.db.Exec(`PRAGMA wal_checkpoint(PASSIVE)`); err != nil {
 90		slog.Debug("wal checkpoint", "err", err)
 91	}
 92}
 93
 94type Key struct {
 95	ID      int64     `json:"id"`
 96	Name    string    `json:"name"`
 97	Prefix  string    `json:"prefix"`
 98	Created time.Time `json:"created"`
 99	Used    time.Time `json:"used,omitempty"`
100	Revoked bool      `json:"revoked"`
101}
102
103// keyPrefix is on every key this service issues, so one found in a config file
104// is identifiable as belonging here rather than to some other provider.
105const keyPrefix = "orch-"
106
107// NewKey mints a key and returns the plaintext exactly once. Nothing keeps it,
108// so a lost key is reissued rather than recovered.
109func (s *Store) NewKey(name string) (string, Key, error) {
110	name = strings.TrimSpace(name)
111	if name == "" {
112		return "", Key{}, errors.New("a key needs a name saying what will use it")
113	}
114	raw := make([]byte, 32)
115	if _, err := rand.Read(raw); err != nil {
116		return "", Key{}, err
117	}
118	secret := keyPrefix + base64.RawURLEncoding.EncodeToString(raw)
119	sum := hashKey(secret)
120	now := time.Now().Unix()
121	r, err := s.db.Exec(`INSERT INTO keys(name, hash, prefix, created_at) VALUES(?,?,?,?)`,
122		name, sum, secret[:len(keyPrefix)+6], now)
123	if err != nil {
124		return "", Key{}, err
125	}
126	id, _ := r.LastInsertId()
127	return secret, Key{ID: id, Name: name, Prefix: secret[:len(keyPrefix)+6], Created: time.Unix(now, 0)}, nil
128}
129
130func hashKey(secret string) string {
131	sum := sha256.Sum256([]byte(secret))
132	return hex.EncodeToString(sum[:])
133}
134
135// Authenticate resolves a presented key. A revoked key is not a match, and the
136// comparison is constant time even though the lookup is by hash, because the
137// row is fetched by hash and then confirmed rather than trusted.
138func (s *Store) Authenticate(secret string) (Key, bool) {
139	if !strings.HasPrefix(secret, keyPrefix) {
140		return Key{}, false
141	}
142	sum := hashKey(secret)
143	var k Key
144	var created, used, revoked int64
145	var stored string
146	err := s.db.QueryRow(`SELECT id, name, prefix, hash, created_at, last_used, revoked_at FROM keys WHERE hash=?`, sum).
147		Scan(&k.ID, &k.Name, &k.Prefix, &stored, &created, &used, &revoked)
148	if err != nil {
149		return Key{}, false
150	}
151	if subtle.ConstantTimeCompare([]byte(stored), []byte(sum)) != 1 {
152		return Key{}, false
153	}
154	if revoked != 0 {
155		return Key{}, false
156	}
157	k.Created, k.Used = time.Unix(created, 0), time.Unix(used, 0)
158	return k, true
159}
160
161func (s *Store) TouchKey(id int64) {
162	if _, err := s.db.Exec(`UPDATE keys SET last_used=? WHERE id=?`, time.Now().Unix(), id); err != nil {
163		slog.Debug("touch key", "err", err)
164	}
165}
166
167func (s *Store) Keys() ([]Key, error) {
168	rows, err := s.db.Query(`SELECT id, name, prefix, created_at, last_used, revoked_at FROM keys ORDER BY revoked_at, id DESC`)
169	if err != nil {
170		return nil, err
171	}
172	defer rows.Close()
173	out := []Key{}
174	for rows.Next() {
175		var k Key
176		var created, used, revoked int64
177		if err := rows.Scan(&k.ID, &k.Name, &k.Prefix, &created, &used, &revoked); err != nil {
178			return nil, err
179		}
180		k.Created, k.Revoked = time.Unix(created, 0), revoked != 0
181		if used > 0 {
182			k.Used = time.Unix(used, 0)
183		}
184		out = append(out, k)
185	}
186	return out, rows.Err()
187}
188
189// Revoke is not a delete, so the call log keeps pointing at a real key and the
190// history of what a retired service asked for stays readable.
191func (s *Store) Revoke(id int64) error {
192	_, err := s.db.Exec(`UPDATE keys SET revoked_at=? WHERE id=? AND revoked_at=0`, time.Now().Unix(), id)
193	return err
194}
195
196func (s *Store) DeleteKey(id int64) error {
197	_, err := s.db.Exec(`DELETE FROM keys WHERE id=?`, id)
198	return err
199}
200
201// Call is one request through the gateway, prompt and completion included.
202type Call struct {
203	ID         int64     `json:"id"`
204	KeyID      int64     `json:"key_id"`
205	Caller     string    `json:"caller"`
206	Model      string    `json:"model"`
207	Messages   string    `json:"messages"`
208	Completion string    `json:"completion"`
209	Tools      string    `json:"tools,omitempty"`
210	PromptTok  int       `json:"prompt_tokens"`
211	OutputTok  int       `json:"output_tokens"`
212	DecodeTPS  float64   `json:"decode_tps"`
213	MS         int64     `json:"ms"`
214	Status     int       `json:"status"`
215	Err        string    `json:"err,omitempty"`
216	At         time.Time `json:"at"`
217}
218
219func (s *Store) LogCall(c Call) {
220	_, err := s.db.Exec(`INSERT INTO calls
221		(key_id, caller, model, messages, completion, tools, prompt_tok, output_tok, decode_tps, ms, status, err, at)
222		VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)`,
223		c.KeyID, c.Caller, c.Model, c.Messages, c.Completion, c.Tools,
224		c.PromptTok, c.OutputTok, c.DecodeTPS, c.MS, c.Status, c.Err, time.Now().Unix())
225	if err != nil {
226		slog.Error("logging a call failed", "err", err)
227	}
228}
229
230func (s *Store) Calls(caller string, limit int) ([]Call, error) {
231	q := `SELECT id, key_id, caller, model, messages, completion, tools,
232	             prompt_tok, output_tok, decode_tps, ms, status, err, at
233	      FROM calls`
234	args := []any{}
235	if caller != "" {
236		q += ` WHERE caller=?`
237		args = append(args, caller)
238	}
239	q += ` ORDER BY id DESC LIMIT ?`
240	args = append(args, limit)
241
242	rows, err := s.db.Query(q, args...)
243	if err != nil {
244		return nil, err
245	}
246	defer rows.Close()
247	out := []Call{}
248	for rows.Next() {
249		var c Call
250		var at int64
251		if err := rows.Scan(&c.ID, &c.KeyID, &c.Caller, &c.Model, &c.Messages, &c.Completion,
252			&c.Tools, &c.PromptTok, &c.OutputTok, &c.DecodeTPS, &c.MS, &c.Status, &c.Err, &at); err != nil {
253			return nil, err
254		}
255		c.At = time.Unix(at, 0)
256		out = append(out, c)
257	}
258	return out, rows.Err()
259}
260
261// Usage is the per caller roll up the overview page shows.
262type Usage struct {
263	Caller    string  `json:"caller"`
264	Calls     int     `json:"calls"`
265	PromptTok int     `json:"prompt_tokens"`
266	OutputTok int     `json:"output_tokens"`
267	Errors    int     `json:"errors"`
268	AvgTPS    float64 `json:"avg_tps"`
269}
270
271func (s *Store) Usage(since time.Time) ([]Usage, error) {
272	rows, err := s.db.Query(`
273		SELECT caller, COUNT(*), COALESCE(SUM(prompt_tok),0), COALESCE(SUM(output_tok),0),
274		       COALESCE(SUM(CASE WHEN status >= 400 OR err <> '' THEN 1 ELSE 0 END),0),
275		       COALESCE(AVG(NULLIF(decode_tps,0)),0)
276		FROM calls WHERE at >= ? GROUP BY caller ORDER BY COUNT(*) DESC`, since.Unix())
277	if err != nil {
278		return nil, err
279	}
280	defer rows.Close()
281	out := []Usage{}
282	for rows.Next() {
283		var u Usage
284		if err := rows.Scan(&u.Caller, &u.Calls, &u.PromptTok, &u.OutputTok, &u.Errors, &u.AvgTPS); err != nil {
285			return nil, err
286		}
287		out = append(out, u)
288	}
289	return out, rows.Err()
290}
291
292// Prune drops calls older than the retention window. The prompts are the whole
293// of what anyone asked this estate, so they do not accumulate forever by
294// default.
295func (s *Store) Prune(keep time.Duration) (int64, error) {
296	r, err := s.db.Exec(`DELETE FROM calls WHERE at < ?`, time.Now().Add(-keep).Unix())
297	if err != nil {
298		return 0, err
299	}
300	return r.RowsAffected()
301}
302
303func (s *Store) Counts() (keys, calls int) {
304	_ = s.db.QueryRow(`SELECT COUNT(*) FROM keys WHERE revoked_at=0`).Scan(&keys)
305	_ = s.db.QueryRow(`SELECT COUNT(*) FROM calls`).Scan(&calls)
306	return
307}