llmgw/internal/store/memory.go

package store

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

// ErrMemoryNotEnabled means the project has not asked llmgw to keep notes.
//
// Returned rather than ignoring the write: an agent that believes it saved
// something and finds nothing later has been lied to.
var ErrMemoryNotEnabled = errors.New("store: this project does not keep memory notes")

// ErrNoNote means the key does not exist, or belongs to another project or agent.
var ErrNoNote = errors.New("store: no such note")

// ErrMemoryFull means the agent already holds as many notes as one may.
var ErrMemoryFull = errors.New("store: this agent has reached its limit of notes")

// Limits on what one note and one agent may hold. Notes are memory held on
// behalf of a caller, so each is bounded rather than trusted.
const (
	MaxNoteKey       = 200
	MaxNoteText      = 16 << 10
	MaxNoteTags      = 16
	MaxNotesPerAgent = 1000
	MaxNoteResults   = 100
)

var tagPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,63}$`)

// Note is one thing an agent chose to remember.
type Note struct {
	Key       string
	Text      string
	Tags      []string
	CreatedAt time.Time
	UpdatedAt time.Time
}

// NormalizeTag lowercases and validates a tag, so "Work" and "work" are one tag
// and nothing outside a small alphabet can reach the stored form.
func NormalizeTag(tag string) (string, error) {
	t := strings.ToLower(strings.TrimSpace(tag))
	if !tagPattern.MatchString(t) {
		return "", fmt.Errorf("invalid tag %q: use lowercase letters, digits, - and _ (up to 64)", tag)
	}
	return t, nil
}

func encodeTags(tags []string) (string, error) {
	seen := map[string]bool{}
	var out []string
	for _, tag := range tags {
		t, err := NormalizeTag(tag)
		if err != nil {
			return "", err
		}
		if !seen[t] {
			seen[t] = true
			out = append(out, t)
		}
	}
	if len(out) > MaxNoteTags {
		return "", fmt.Errorf("too many tags: at most %d", MaxNoteTags)
	}
	if len(out) == 0 {
		return "", nil
	}
	return "," + strings.Join(out, ",") + ",", nil
}

func decodeTags(s string) []string {
	s = strings.Trim(s, ",")
	if s == "" {
		return nil
	}
	return strings.Split(s, ",")
}

// MemoryRetention reports how long a project keeps notes. Zero means none.
func (s *DB) MemoryRetention(ctx context.Context, projectID string) (time.Duration, error) {
	var hours sql.NullInt64
	err := s.db.QueryRowContext(ctx,
		`SELECT memory_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
}

// SetMemoryRetention turns note storage on for a project, or off.
func (s *DB) SetMemoryRetention(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 memory_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
}

// PutNote writes a note, replacing any note the agent already holds under the
// same key.
func (s *DB) PutNote(ctx context.Context, projectID, agentID string, n Note) error {
	if n.Key == "" || len(n.Key) > MaxNoteKey {
		return fmt.Errorf("key must be 1 to %d bytes", MaxNoteKey)
	}
	if n.Text == "" || len(n.Text) > MaxNoteText {
		return fmt.Errorf("text must be 1 to %d bytes", MaxNoteText)
	}
	tags, err := encodeTags(n.Tags)
	if err != nil {
		return err
	}
	retention, err := s.MemoryRetention(ctx, projectID)
	if err != nil {
		return err
	}
	if retention == 0 {
		return ErrMemoryNotEnabled
	}

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

	now := time.Now().UnixMilli()
	res, err := q.ExecContext(ctx,
		`UPDATE memory_notes SET text = ?, tags = ?, updated_at = ?
		 WHERE project_id = ? AND agent_id = ? AND key = ?`,
		n.Text, tags, now, projectID, agentID, n.Key)
	if err != nil {
		return err
	}
	if updated, _ := res.RowsAffected(); updated == 0 {
		var count int
		if err := q.QueryRowContext(ctx,
			`SELECT COUNT(*) FROM memory_notes WHERE project_id = ? AND agent_id = ?`,
			projectID, agentID).Scan(&count); err != nil {
			return err
		}
		if count >= MaxNotesPerAgent {
			return ErrMemoryFull
		}
		if _, err := q.ExecContext(ctx,
			`INSERT INTO memory_notes (project_id, agent_id, key, text, tags, created_at, updated_at)
			 VALUES (?, ?, ?, ?, ?, ?, ?)`,
			projectID, agentID, n.Key, n.Text, tags, now, now); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func scanNote(rows interface{ Scan(...any) error }) (Note, error) {
	var (
		n          Note
		tags       string
		created    int64
		updatedRaw int64
	)
	if err := rows.Scan(&n.Key, &n.Text, &tags, &created, &updatedRaw); err != nil {
		return Note{}, err
	}
	n.Tags = decodeTags(tags)
	n.CreatedAt = time.UnixMilli(created)
	n.UpdatedAt = time.UnixMilli(updatedRaw)
	return n, nil
}

const noteColumns = `key, text, tags, created_at, updated_at`

// GetNote reads one note. Scoped to the project and agent in the query itself.
func (s *DB) GetNote(ctx context.Context, projectID, agentID, key string) (Note, error) {
	n, err := scanNote(s.db.QueryRowContext(ctx,
		`SELECT `+noteColumns+` FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`,
		projectID, agentID, key))
	if errors.Is(err, sql.ErrNoRows) {
		return Note{}, ErrNoNote
	}
	return n, err
}

// ftsQuery turns free text into a query that cannot be a syntax error: every
// word becomes a quoted term, and they must all match. FTS5 treats characters
// such as - : * and " as operators, and an agent's search words are not written
// for that syntax.
func ftsQuery(q string) string {
	words := strings.Fields(q)
	for i, w := range words {
		words[i] = `"` + strings.ReplaceAll(w, `"`, `""`) + `"`
	}
	return strings.Join(words, " ")
}

// SearchNotes finds an agent's notes by words and/or tag, newest first. With
// neither it lists the most recent.
func (s *DB) SearchNotes(ctx context.Context, projectID, agentID, query, tag string, limit int) ([]Note, error) {
	if limit <= 0 || limit > MaxNoteResults {
		limit = MaxNoteResults
	}
	q := `SELECT ` + noteColumns + ` FROM memory_notes WHERE project_id = ? AND agent_id = ?`
	args := []any{projectID, agentID}

	if match := ftsQuery(query); match != "" {
		// The one query in this store that is not the same sentence in both
		// dialects. SQLite searches an FTS5 table kept in step by triggers;
		// PostgreSQL matches a tsvector derived from the row. Writing it out
		// twice is honest about that — the alternative was a shared spelling
		// that means something subtly different in each.
		if s.dialect == PostgresDialect {
			q += ` AND search @@ plainto_tsquery('simple', ?)`
			// plainto_tsquery takes words, not FTS5 syntax, so the raw query
			// goes in rather than the quoted form SQLite needs.
			args = append(args, query)
		} else {
			q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)`
			args = append(args, match)
		}
	}
	if tag != "" {
		t, err := NormalizeTag(tag)
		if err != nil {
			return nil, err
		}
		q += ` AND tags LIKE ?`
		args = append(args, "%,"+t+",%")
	}
	q += ` ORDER BY updated_at DESC, id DESC LIMIT ?`
	args = append(args, limit)

	return queryAll(ctx, s.db, q, args, func(rows *sql.Rows) (Note, error) { return scanNote(rows) })
}

// DeleteNote forgets one note.
func (s *DB) DeleteNote(ctx context.Context, projectID, agentID, key string) error {
	res, err := s.db.ExecContext(ctx,
		`DELETE FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`,
		projectID, agentID, key)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNoNote
	}
	return nil
}

// PruneMemory deletes notes not refreshed within their project's retention.
func (s *DB) PruneMemory(ctx context.Context, now time.Time) (int, error) {
	type window struct {
		project string
		hours   int64
	}
	windows, err := queryAll(ctx, s.db,
		`SELECT id, memory_retention_hours FROM projects WHERE memory_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 memory_notes WHERE project_id = ? AND updated_at < ?`, w.project, cutoff)
		if err != nil {
			return removed, err
		}
		n, _ := res.RowsAffected()
		removed += int(n)
	}
	return removed, nil
}