llmgw/internal/store/memory_test.go

package store

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"
)

func memoryProject(t *testing.T, s *DB) string {
	t.Helper()
	proj := tree(t, s)
	if err := s.SetMemoryRetention(context.Background(), proj, 24); err != nil {
		t.Fatal(err)
	}
	return proj
}

func TestAProjectKeepsNoNotesUnlessItAsks(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := tree(t, s)

	err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "secret"})
	if !errors.Is(err, ErrMemoryNotEnabled) {
		t.Fatalf("writing into a project that never asked returned %v", err)
	}
	if _, err := s.GetNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) {
		t.Errorf("a refused write left a note behind: %v", err)
	}
}

func TestANoteIsWrittenReadBackAndReplaced(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)

	if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Ankara", Tags: []string{"Geo", "geo", "facts"}}); err != nil {
		t.Fatal(err)
	}
	got, err := s.GetNote(ctx, proj, "a", "city")
	if err != nil {
		t.Fatal(err)
	}
	if got.Text != "Ankara" || len(got.Tags) != 2 || got.Tags[0] != "geo" || got.Tags[1] != "facts" {
		t.Fatalf("note came back as %+v: tags should be lowercased and deduplicated", got)
	}

	if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Istanbul"}); err != nil {
		t.Fatal(err)
	}
	got, _ = s.GetNote(ctx, proj, "a", "city")
	if got.Text != "Istanbul" || len(got.Tags) != 0 {
		t.Errorf("rewriting a key should replace text and tags: %+v", got)
	}
	notes, _ := s.SearchNotes(ctx, proj, "a", "", "", 0)
	if len(notes) != 1 {
		t.Errorf("rewriting a key created %d notes, want 1", len(notes))
	}
}

func TestNotesAreFoundByWordsAndTagsAndSpecialCharactersAreHarmless(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)
	put := func(key, text string, tags ...string) {
		t.Helper()
		if err := s.PutNote(ctx, proj, "a", Note{Key: key, Text: text, Tags: tags}); err != nil {
			t.Fatal(err)
		}
	}
	put("deploy", "the staging cluster is on port 8443", "ops")
	put("food", "she prefers green tea", "personal")
	put("old", "deploy notes from last year", "ops", "archive")

	byWord, err := s.SearchNotes(ctx, proj, "a", "staging cluster", "", 0)
	if err != nil || len(byWord) != 1 || byWord[0].Key != "deploy" {
		t.Fatalf("search by words: %+v (err %v)", byWord, err)
	}
	byTag, _ := s.SearchNotes(ctx, proj, "a", "", "ops", 0)
	if len(byTag) != 2 {
		t.Errorf("tag ops matched %d notes, want 2", len(byTag))
	}
	both, _ := s.SearchNotes(ctx, proj, "a", "deploy", "archive", 0)
	if len(both) != 1 || both[0].Key != "old" {
		t.Errorf("words and tag together: %+v", both)
	}
	// A tag is matched whole: "op" must not match "ops".
	if part, _ := s.SearchNotes(ctx, proj, "a", "", "op", 0); len(part) != 0 {
		t.Errorf("a partial tag matched %d notes", len(part))
	}
	// FTS operators in an agent's words must not become a syntax error.
	for _, q := range []string{`"`, `a -b`, `x:y`, `foo*`, `AND OR NOT`, `(`} {
		if _, err := s.SearchNotes(ctx, proj, "a", q, "", 0); err != nil {
			t.Errorf("query %q failed: %v", q, err)
		}
	}
}

func TestOneAgentOrProjectCannotReadAnothersNotes(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)
	if err := s.PutOrg(ctx, Org{ID: "o2", Name: "Other"}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutTeam(ctx, Team{ID: "t2", OrgID: "o2", Name: "T"}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "t2", Name: "P"}); err != nil {
		t.Fatal(err)
	}
	if err := s.SetMemoryRetention(ctx, "proj2", 24); err != nil {
		t.Fatal(err)
	}

	if err := s.PutNote(ctx, proj, "agent1", Note{Key: "k", Text: "private"}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.GetNote(ctx, proj, "agent2", "k"); !errors.Is(err, ErrNoNote) {
		t.Errorf("another agent in the same project read the note: %v", err)
	}
	if _, err := s.GetNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) {
		t.Errorf("another project read the note: %v", err)
	}
	if got, _ := s.SearchNotes(ctx, "proj2", "agent1", "private", "", 0); len(got) != 0 {
		t.Errorf("another project found the note by search: %+v", got)
	}
	if err := s.DeleteNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) {
		t.Errorf("another project deleted the note: %v", err)
	}
}

func TestNotesAreBoundedAndBadInputIsRefused(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)

	bad := []Note{
		{Key: "", Text: "x"},
		{Key: "k", Text: ""},
		{Key: strings.Repeat("k", MaxNoteKey+1), Text: "x"},
		{Key: "k", Text: strings.Repeat("x", MaxNoteText+1)},
		{Key: "k", Text: "x", Tags: []string{"has space"}},
	}
	for i, n := range bad {
		if err := s.PutNote(ctx, proj, "a", n); err == nil {
			t.Errorf("bad note %d was accepted", i)
		}
	}

	for i := 0; i < MaxNotesPerAgent; i++ {
		if err := s.PutNote(ctx, proj, "a", Note{Key: fmt.Sprintf("note-%d", i), Text: "x"}); err != nil {
			t.Fatalf("note %d: %v", i, err)
		}
	}
	if err := s.PutNote(ctx, proj, "a", Note{Key: "one-too-many", Text: "x"}); !errors.Is(err, ErrMemoryFull) {
		t.Errorf("exceeding the per-agent limit returned %v", err)
	}
	// Replacing an existing note is not growth and must still work when full.
	if err := s.PutNote(ctx, proj, "a", Note{Key: "note-0", Text: "y"}); err != nil {
		t.Errorf("replacing a note in a full agent failed: %v", err)
	}
	// Another agent has its own allowance.
	if err := s.PutNote(ctx, proj, "b", Note{Key: "note-0", Text: "x"}); err != nil {
		t.Errorf("a second agent was blocked by the first one's notes: %v", err)
	}
}

func TestAForgottenNoteIsGoneFromSearchToo(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)

	if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "needle in a haystack"}); err != nil {
		t.Fatal(err)
	}
	if err := s.DeleteNote(ctx, proj, "a", "k"); err != nil {
		t.Fatal(err)
	}
	if got, _ := s.SearchNotes(ctx, proj, "a", "needle", "", 0); len(got) != 0 {
		t.Errorf("a deleted note is still searchable: %+v", got)
	}
	if err := s.DeleteNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) {
		t.Errorf("deleting twice returned %v", err)
	}
}

func TestNotesNotRefreshedWithinRetentionArePruned(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	proj := memoryProject(t, s)

	if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "x"}); err != nil {
		t.Fatal(err)
	}
	if n, err := s.PruneMemory(ctx, time.Now()); err != nil || n != 0 {
		t.Fatalf("a fresh note was pruned: %d, %v", n, err)
	}
	if n, err := s.PruneMemory(ctx, time.Now().Add(25*time.Hour)); err != nil || n != 1 {
		t.Fatalf("an expired note survived: %d, %v", n, err)
	}
	if got, _ := s.SearchNotes(ctx, proj, "a", "", "", 0); len(got) != 0 {
		t.Errorf("pruned note still listed: %+v", got)
	}
}