package store import ( "context" "errors" "testing" "time" ) func TestAProjectKeepsNoConversationsUnlessItAsks(t *testing.T) { // The property being given up here is worth protecting in a test: until a // project turns this on, llmgw stores no message content anywhere. Every // project that existed before threads did, and every one created since // without asking, must stay that way — a gateway that quietly starts // keeping what a team types is a much harder thing to have deployed. ctx := context.Background() s := openDB(t) proj := tree(t, s) if d, err := s.ThreadRetention(ctx, proj); err != nil || d != 0 { t.Fatalf("a new project retains %v (err %v), want none", d, err) } err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "secret"}}) if !errors.Is(err, ErrThreadsNotEnabled) { t.Fatalf("storing into a project that never asked returned %v", err) } // And nothing was written on the way to refusing. if _, err := s.Thread(ctx, proj, "th1"); !errors.Is(err, ErrNoThread) { t.Errorf("a refused append left a thread behind: %v", err) } } func TestATurnedOnProjectStoresAndReplaysInOrder(t *testing.T) { ctx := context.Background() s := openDB(t) proj := tree(t, s) if err := s.SetThreadRetention(ctx, proj, 24); err != nil { t.Fatal(err) } turns := []Message{ {Role: "user", Content: "first"}, {Role: "assistant", Content: "answer one", Model: "small"}, {Role: "user", Content: "second"}, {Role: "assistant", Content: "answer two", Model: "large"}, } for _, m := range turns { if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil { t.Fatal(err) } } got, err := s.Thread(ctx, proj, "th1") if err != nil { t.Fatal(err) } if len(got) != len(turns) { t.Fatalf("stored %d turns, read back %d", len(turns), len(got)) } for i := range turns { if got[i].Role != turns[i].Role || got[i].Content != turns[i].Content { t.Errorf("turn %d came back as %+v", i, got[i]) } } // Which model answered, because a thread that moved between models // otherwise reads as one voice — and when the answers disagree that is the // first thing anyone asks. if got[1].Model != "small" || got[3].Model != "large" { t.Errorf("the models that answered were lost: %q and %q", got[1].Model, got[3].Model) } } func TestAnotherProjectsThreadIsNotReadable(t *testing.T) { // Thread ids are opaque, but an unguessable id is a secret, and secrets end // up in logs, shell history and bug reports. The project scope is what // makes a leaked id useless to anyone outside it. ctx := context.Background() s := openDB(t) mine := tree(t, s) if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil { t.Fatal(err) } for _, p := range []string{mine, "proj2"} { if err := s.SetThreadRetention(ctx, p, 24); err != nil { t.Fatal(err) } } if err := s.AppendThread(ctx, mine, "th1", []Message{{Role: "user", Content: "private"}}); err != nil { t.Fatal(err) } // The same error as a thread that does not exist. Two errors would let a // caller learn which ids are real in projects they cannot read. if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) { t.Errorf("reading across projects returned %v", err) } if _, err := s.Thread(ctx, "proj2", "never-existed"); !errors.Is(err, ErrNoThread) { t.Errorf("a missing thread returned a different error: %v", err) } // Writing across projects is the same hole in the other direction: append, // then read your own reply back. if err := s.AppendThread(ctx, "proj2", "th1", []Message{{Role: "user", Content: "x"}}); !errors.Is(err, ErrNoThread) { t.Errorf("writing into another project's thread returned %v", err) } } func TestDeletingAThreadTakesItsMessagesWithIt(t *testing.T) { // Somebody has to be able to say "forget that", and a retention window is // no answer to a message sent by mistake. ctx := context.Background() s := openDB(t) proj := tree(t, s) if err := s.SetThreadRetention(ctx, proj, 24); err != nil { t.Fatal(err) } if err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "oops"}}); err != nil { t.Fatal(err) } if err := s.DeleteThread(ctx, proj, "th1"); err != nil { t.Fatal(err) } if _, err := s.Thread(ctx, proj, "th1"); !errors.Is(err, ErrNoThread) { t.Errorf("the thread survived deletion: %v", err) } // The rows, not just the row. Without the cascade the content would still // be in the database after someone was told it was gone. var left int if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM thread_messages WHERE thread_id = ?`, "th1").Scan(&left); err != nil { t.Fatal(err) } if left != 0 { t.Errorf("%d messages survived the thread being deleted", left) } } func TestRetentionIsMeasuredFromTheLastTurnNotTheFirst(t *testing.T) { // Expiring a conversation someone is still having would be a strange way // to keep a promise about how long things are kept. ctx := context.Background() s := openDB(t) proj := tree(t, s) if err := s.SetThreadRetention(ctx, proj, 1); err != nil { t.Fatal(err) } if err := s.AppendThread(ctx, proj, "live", []Message{{Role: "user", Content: "hello"}}); err != nil { t.Fatal(err) } // Two hours on, a thread touched now is still inside a one-hour window. if n, err := s.PruneThreads(ctx, time.Now().Add(2*time.Hour)); err != nil { t.Fatal(err) } else if n != 1 { t.Fatalf("an idle thread past its window was not pruned: %d", n) } if err := s.AppendThread(ctx, proj, "live2", []Message{{Role: "user", Content: "hello"}}); err != nil { t.Fatal(err) } if n, err := s.PruneThreads(ctx, time.Now()); err != nil { t.Fatal(err) } else if n != 0 { t.Errorf("a thread inside its window was pruned: %d", n) } } func TestPruningLeavesProjectsThatStoreNothingAlone(t *testing.T) { // They have nothing to prune, and a sweep that treated "no retention" as // "retain for zero hours" would be indistinguishable from one that worked. ctx := context.Background() s := openDB(t) tree(t, s) if n, err := s.PruneThreads(ctx, time.Now()); err != nil || n != 0 { t.Errorf("pruning removed %d rows in a store with no threads (err %v)", n, err) } } func TestTurningRetentionOffAgain(t *testing.T) { ctx := context.Background() s := openDB(t) proj := tree(t, s) if err := s.SetThreadRetention(ctx, proj, 24); err != nil { t.Fatal(err) } if err := s.SetThreadRetention(ctx, proj, 0); err != nil { t.Fatal(err) } if d, err := s.ThreadRetention(ctx, proj); err != nil || d != 0 { t.Fatalf("retention is %v after being turned off (err %v)", d, err) } if err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "x"}}); !errors.Is(err, ErrThreadsNotEnabled) { t.Errorf("storing continued after being turned off: %v", err) } }