llmgw/internal/store/threads.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// ErrThreadsNotEnabled means the project has not asked llmgw to keep
// conversations.
//
// Returned rather than silently ignoring the thread. Ignoring it would send
// the model a conversation with no history, and it would answer confidently
// about something it was never told — which reads as the model being wrong
// rather than as a setting being off.
var ErrThreadsNotEnabled = errors.New("store: this project does not keep conversations")

// ErrNoThread means the id does not exist, or belongs to another project.
//
// One error for both on purpose. Distinguishing them would turn thread ids
// into an oracle: a caller could learn which ids exist in projects they cannot
// read by watching which error comes back.
var ErrNoThread = errors.New("store: no such thread")

// Message is one turn of a stored conversation.
type Message struct {
	Role    string
	Content string
	// Model is which model produced an assistant turn, empty for a user one. A
	// thread that moved between models otherwise reads as a single voice.
	Model string
}

// EnsureAnonymousProject creates the single project that owns everything kept
// by a gateway with no authentication, and sets how long it keeps it.
//
// A real project rather than a special case in every query: threads and notes
// are scoped by project id and enforce it in SQL, and an empty id would have to
// be an exception in each of those places.
func (s *DB) EnsureAnonymousProject(ctx context.Context, id string, retainHours int) error {
	if err := s.PutOrg(ctx, Org{ID: id, Name: id}); err != nil {
		return err
	}
	if err := s.PutTeam(ctx, Team{ID: id, OrgID: id, Name: id}); err != nil {
		return err
	}
	if err := s.PutProject(ctx, Project{ID: id, TeamID: id, Name: id}); err != nil {
		return err
	}
	if err := s.SetThreadRetention(ctx, id, retainHours); err != nil {
		return err
	}
	return s.SetMemoryRetention(ctx, id, retainHours)
}

// ThreadRetention reports how long a project keeps conversations.
//
// Zero means it keeps none, which is the default and was the only behaviour
// before threads existed.
func (s *DB) ThreadRetention(ctx context.Context, projectID string) (time.Duration, error) {
	var hours sql.NullInt64
	err := s.db.QueryRowContext(ctx,
		`SELECT thread_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours)
	if errors.Is(err, sql.ErrNoRows) {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	if !hours.Valid || hours.Int64 <= 0 {
		return 0, nil
	}
	return time.Duration(hours.Int64) * time.Hour, nil
}

// SetThreadRetention turns conversation storage on for a project, or off.
//
// Hours rather than a boolean: a thread store with no expiry is a transcript
// of everything a team ever typed, kept forever because nobody chose a
// number. Turning it on means saying for how long.
func (s *DB) SetThreadRetention(ctx context.Context, projectID string, hours int) error {
	var v any
	if hours > 0 {
		v = hours
	}
	res, err := s.db.ExecContext(ctx,
		`UPDATE projects SET thread_retention_hours = ? WHERE id = ?`, v, projectID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("store: no project %q", projectID)
	}
	return nil
}

// Thread returns a conversation's messages in order.
//
// The project is part of the lookup rather than checked afterwards, so there
// is no path that reads another project's conversation even briefly.
func (s *DB) Thread(ctx context.Context, projectID, threadID string) ([]Message, error) {
	var owner string
	err := s.db.QueryRowContext(ctx,
		`SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
	if errors.Is(err, sql.ErrNoRows) || (err == nil && owner != projectID) {
		return nil, ErrNoThread
	}
	if err != nil {
		return nil, err
	}

	return queryAll(ctx, s.db,
		`SELECT role, content, COALESCE(model, '') FROM thread_messages
		 WHERE thread_id = ? ORDER BY seq`,
		[]any{threadID}, func(rows *sql.Rows) (Message, error) {
			var m Message
			return m, rows.Scan(&m.Role, &m.Content, &m.Model)
		})
}

// AppendThread adds turns to a conversation, creating it if it is new.
//
// One transaction, because a half-written exchange is worse than a lost one:
// a stored question with no answer would be replayed to the next model as
// though the assistant had ignored it.
func (s *DB) AppendThread(ctx context.Context, projectID, threadID string, msgs []Message) error {
	if len(msgs) == 0 {
		return nil
	}
	retention, err := s.ThreadRetention(ctx, projectID)
	if err != nil {
		return err
	}
	if retention == 0 {
		return ErrThreadsNotEnabled
	}

	q, tx, err := s.begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()

	now := time.Now().UnixMilli()
	var owner string
	err = q.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		if _, err := q.ExecContext(ctx,
			`INSERT INTO threads (id, project_id, created_at, updated_at) VALUES (?, ?, ?, ?)`,
			threadID, projectID, now, now); err != nil {
			return err
		}
	case err != nil:
		return err
	case owner != projectID:
		// Appending to a thread another project owns would let one team write
		// into another's conversation, and then read it back through their
		// own replies.
		return ErrNoThread
	}

	var next int64
	if err := q.QueryRowContext(ctx,
		`SELECT COALESCE(MAX(seq), -1) + 1 FROM thread_messages WHERE thread_id = ?`,
		threadID).Scan(&next); err != nil {
		return err
	}
	for i, m := range msgs {
		var model any
		if m.Model != "" {
			model = m.Model
		}
		if _, err := q.ExecContext(ctx,
			`INSERT INTO thread_messages (thread_id, seq, role, content, model) VALUES (?, ?, ?, ?, ?)`,
			threadID, next+int64(i), m.Role, m.Content, model); err != nil {
			return err
		}
	}
	// Touched on every append, so retention measures idleness rather than age.
	if _, err := q.ExecContext(ctx,
		`UPDATE threads SET updated_at = ? WHERE id = ?`, now, threadID); err != nil {
		return err
	}
	return tx.Commit()
}

// DeleteThread removes a conversation and everything in it.
//
// Somebody has to be able to say "forget that", and a retention window is not
// an answer to a message sent by mistake.
func (s *DB) DeleteThread(ctx context.Context, projectID, threadID string) error {
	res, err := s.db.ExecContext(ctx,
		`DELETE FROM threads WHERE id = ? AND project_id = ?`, threadID, projectID)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNoThread
	}
	return nil
}

// PruneThreads deletes conversations past their project's retention.
//
// Per project, because the window is per project. A single sweep with one
// cutoff would apply the shortest retention to everyone or the longest to
// nobody, and the number a project chose is the promise made to it.
func (s *DB) PruneThreads(ctx context.Context, now time.Time) (int, error) {
	type window struct {
		project string
		hours   int64
	}
	windows, err := queryAll(ctx, s.db,
		`SELECT id, thread_retention_hours FROM projects WHERE thread_retention_hours > 0`,
		nil, func(rows *sql.Rows) (window, error) {
			var w window
			return w, rows.Scan(&w.project, &w.hours)
		})
	if err != nil {
		return 0, err
	}

	var removed int
	for _, w := range windows {
		cutoff := now.Add(-time.Duration(w.hours) * time.Hour).UnixMilli()
		res, err := s.db.ExecContext(ctx,
			`DELETE FROM threads WHERE project_id = ? AND updated_at < ?`, w.project, cutoff)
		if err != nil {
			return removed, err
		}
		n, _ := res.RowsAffected()
		removed += int(n)
	}
	return removed, nil
}