llmgw/internal/store/threads_test.go

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)
	}
}