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 }