Store tool-result content compressed in its own column

internal/store/content.go

@@ -0,0 +1,69 @@
+package store
+
+import (
+	"database/sql/driver"
+	"encoding/json"
+	"fmt"
+	"maps"
+
+	"github.com/klauspost/compress/zstd"
+	"modernc.org/sqlite"
+
+	"g.fc.run/ai/byre/internal/api"
+)
+
+var (
+	zstdEncoder = must(zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedBetterCompression)))
+	zstdDecoder = must(zstd.NewReader(nil))
+)
+
+func must[T any](v T, err error) T {
+	if err != nil {
+		panic(err)
+	}
+	return v
+}
+
+func init() {
+	sqlite.MustRegisterDeterministicScalarFunction("zstd", 1, func(ctx *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
+		switch v := args[0].(type) {
+		case nil:
+			return nil, nil
+		case string:
+			return zstdEncoder.EncodeAll([]byte(v), nil), nil
+		default:
+			return nil, fmt.Errorf("zstd: want text, got %T", v)
+		}
+	})
+}
+
+func splitContent(ev api.Event) (map[string]any, []byte, error) {
+	c, ok := ev.Payload["content"]
+	if ev.Kind != "agent.tool_result" || !ok {
+		return ev.Payload, nil, nil
+	}
+	raw, err := json.Marshal(c)
+	if err != nil {
+		return nil, nil, err
+	}
+	payload := maps.Clone(ev.Payload)
+	delete(payload, "content")
+	return payload, zstdEncoder.EncodeAll(raw, nil), nil
+}
+
+func joinContent(payload map[string]any, content []byte) error {
+	if content == nil {
+		return nil
+	}
+	raw, err := zstdDecoder.DecodeAll(content, nil)
+	if err != nil {
+		return err
+	}
+	var c any
+	err = json.Unmarshal(raw, &c)
+	if err != nil {
+		return err
+	}
+	payload["content"] = c
+	return nil
+}

internal/store/migrations/113-compress-tool-result-content.sql

@@ -0,0 +1,4 @@
+ALTER TABLE events ADD COLUMN content BLOB;
+
+UPDATE events SET content = zstd(payload -> '$.content'), payload = json_remove(payload, '$.content')
+WHERE kind = 'agent.tool_result' AND json_type(payload, '$.content') IS NOT NULL;

internal/store/store.go

@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"log/slog"
 	"slices"
 	"strconv"
 	"strings"
@@ -87,6 +88,12 @@ func Open(path string, onEvent func([]api.Event)) (*Store, error) {
 		return nil, err
 	}
 
+	err = vacuumIfSparse(db)
+	if err != nil {
+		db.Close()
+		return nil, err
+	}
+
 	schema, err := schemaFingerprint(db)
 	if err != nil {
 		db.Close()
@@ -160,6 +167,24 @@ func migrate(db *sql.DB) error {
 	return nil
 }
 
+func vacuumIfSparse(db *sql.DB) error {
+	free, pages := 0, 0
+	err := db.QueryRow(`SELECT freelist_count, page_count FROM pragma_freelist_count, pragma_page_count`).Scan(&free, &pages)
+	if err != nil {
+		return err
+	}
+	if free*4 < pages {
+		return nil
+	}
+	slog.Info("vacuuming database", "free_pages", free, "pages", pages)
+	_, err = db.Exec(`VACUUM`)
+	if err != nil {
+		return err
+	}
+	_, err = db.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`)
+	return err
+}
+
 func appliedMigrations(db *sql.DB) (map[string]bool, error) {
 	rows, err := db.Query(`SELECT name FROM migrations`)
 	if err != nil {
@@ -928,7 +953,7 @@ func participants(ctx context.Context, q querier, taskID string) ([]api.Particip
 	return participants, rows.Err()
 }
 
-const eventColumns = `seq, kind, project_id, task_id, agent_session_id, owner, actor, payload, created_at`
+const eventColumns = `seq, kind, project_id, task_id, agent_session_id, owner, actor, payload, content, created_at`
 
 func (s *Store) queryEvents(ctx context.Context, query string, args ...any) ([]api.Event, error) {
 	rows, err := s.db.QueryContext(ctx, query, args...)
@@ -959,7 +984,7 @@ func (s *Store) ForkEvents(ctx context.Context, taskID, sessionID string) ([]api
 		SELECT seq, kind, project_id, task_id, agent_session_id, owner, actor,
 			json_object('vendor_uuid', json_extract(payload, '$.vendor_uuid'),
 				'id', json_extract(payload, '$.id'), 'tool_use_id', json_extract(payload, '$.tool_use_id')),
-			created_at
+			NULL, created_at
 		FROM events WHERE task_id = ? AND agent_session_id = ? ORDER BY seq
 	`, taskID, sessionID)
 }
@@ -1110,8 +1135,9 @@ func scanEvents(rows *sql.Rows) ([]api.Event, error) {
 	for rows.Next() {
 		ev := api.Event{}
 		payload := ""
+		content := []byte(nil)
 		createdAt := ""
-		err := rows.Scan(&ev.Seq, &ev.Kind, &ev.ProjectId, &ev.TaskId, &ev.AgentSessionId, &ev.Owner, &ev.Actor, &payload, &createdAt)
+		err := rows.Scan(&ev.Seq, &ev.Kind, &ev.ProjectId, &ev.TaskId, &ev.AgentSessionId, &ev.Owner, &ev.Actor, &payload, &content, &createdAt)
 		if err != nil {
 			return nil, err
 		}
@@ -1119,6 +1145,10 @@ func scanEvents(rows *sql.Rows) ([]api.Event, error) {
 		if err != nil {
 			return nil, err
 		}
+		err = joinContent(ev.Payload, content)
+		if err != nil {
+			return nil, err
+		}
 		ev.CreatedAt, err = parseTime(createdAt)
 		if err != nil {
 			return nil, err
@@ -1487,14 +1517,18 @@ func (s *Store) write(ctx context.Context, fn func(tx *sql.Tx, emit emitFunc) er
 	events := []api.Event{}
 	emit := func(ev api.Event) error {
 		ev.CreatedAt = time.Now().UTC()
-		payload, err := json.Marshal(ev.Payload)
+		fields, content, err := splitContent(ev)
+		if err != nil {
+			return err
+		}
+		payload, err := json.Marshal(fields)
 		if err != nil {
 			return err
 		}
 		err = tx.QueryRowContext(ctx, `
-			INSERT INTO events (kind, project_id, task_id, agent_session_id, owner, actor, payload, created_at)
-			VALUES (?, ?, ?, ?, ?, ?, ?, ?) RETURNING seq
-		`, ev.Kind, ev.ProjectId, ev.TaskId, ev.AgentSessionId, ev.Owner, ev.Actor, string(payload), formatTime(ev.CreatedAt)).Scan(&ev.Seq)
+			INSERT INTO events (kind, project_id, task_id, agent_session_id, owner, actor, payload, content, created_at)
+			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING seq
+		`, ev.Kind, ev.ProjectId, ev.TaskId, ev.AgentSessionId, ev.Owner, ev.Actor, string(payload), content, formatTime(ev.CreatedAt)).Scan(&ev.Seq)
 		if err != nil {
 			return err
 		}