Vectrum
Sign In
251 lines · 8.4 KB · Go
package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/johannSo/codey/internal/llm"
)

func TestSaveAndResumeSession(t *testing.T) {
	dir := t.TempDir()
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "boom", `{}`)),
		msg("end_turn", textBlock("first answer")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "the original task"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if err := a.Save(dir, "sess-1", "/repo", "", "claude-opus-5"); err != nil {
		t.Fatalf("Save: %v", err)
	}

	loaded, err := ResolveSession(dir, "sess-1")
	if err != nil {
		t.Fatalf("ResolveSession: %v", err)
	}
	if loaded.Model != "claude-opus-5" || loaded.Root != "/repo" || loaded.Turns != 2 {
		t.Errorf("session = %+v", loaded)
	}
	if len(loaded.Messages) != len(a.State().History()) {
		t.Errorf("%d messages saved, want %d", len(loaded.Messages), len(a.State().History()))
	}

	// A resumed agent continues the same conversation.
	f2 := newFakeAPI(t, msg("end_turn", textBlock("second answer")))
	b := newTestAgent(t, f2, t.TempDir(), &bytes.Buffer{})
	b.Restore(loaded)
	if _, err := b.Run(context.Background(), "and now this"); err != nil {
		t.Fatalf("Run after restore: %v", err)
	}
	sent := f2.rawRequest(0)
	if !strings.Contains(sent, "the original task") {
		t.Errorf("resumed request lost the earlier conversation:\n%s", sent)
	}
	if !strings.Contains(sent, "and now this") {
		t.Error("resumed request is missing the new prompt")
	}
	if _, turns := b.State().Usage(); turns != 3 {
		t.Errorf("turns = %d, want the restored 2 plus 1", turns)
	}
}

func TestSavePreservesCreatedTime(t *testing.T) {
	dir := t.TempDir()
	f := newFakeAPI(t, msg("end_turn", textBlock("x")), msg("end_turn", textBlock("y")))
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})

	a.Run(context.Background(), "one")
	a.Save(dir, "s", "/r", "", "m")
	first, _ := LoadSession(filepath.Join(dir, "s.json"))

	time.Sleep(5 * time.Millisecond)
	a.Run(context.Background(), "two")
	a.Save(dir, "s", "/r", "", "m")
	second, _ := LoadSession(filepath.Join(dir, "s.json"))

	if !second.Created.Equal(first.Created) {
		t.Errorf("created moved from %v to %v across saves", first.Created, second.Created)
	}
	if !second.Updated.After(first.Updated) {
		t.Error("updated should move forward")
	}
}

func TestResolveLatestSession(t *testing.T) {
	dir := t.TempDir()
	for _, id := range []string{"older", "newer"} {
		s := Session{ID: id, Messages: []llm.Message{{Role: llm.RoleUser}}}
		writeSession(t, dir, s)
		time.Sleep(10 * time.Millisecond)
	}
	got, err := ResolveSession(dir, "latest")
	if err != nil {
		t.Fatalf("ResolveSession: %v", err)
	}
	if got.ID != "newer" {
		t.Errorf("resolved %q, want the most recently updated", got.ID)
	}
	if _, err := ResolveSession(t.TempDir(), "latest"); err == nil {
		t.Error("an empty session directory should be an error, not an empty session")
	}
	if _, err := ResolveSession(dir, "nonexistent"); err == nil {
		t.Error("an unknown id should be an error")
	}
}

func writeSession(t *testing.T, dir string, s Session) {
	t.Helper()
	a := &Agent{state: NewState()}
	a.state.Replace(s.Messages)
	if err := a.Save(dir, s.ID, "/r", "", "m"); err != nil {
		t.Fatal(err)
	}
}

func TestLoadSessionRejectsGarbage(t *testing.T) {
	dir := t.TempDir()
	os.WriteFile(filepath.Join(dir, "broken.json"), []byte("{not json"), 0o644)
	if _, err := ResolveSession(dir, "broken"); err == nil {
		t.Error("want an error for a corrupt session file")
	}
}

// A session belongs to the directory it was recorded in. The chooser is built
// on this: work done in /home/joni/test must not be reopenable in
// /home/joni/test2, where none of the paths it discusses exist.
func TestListSessionsOnlyOffersThisRoots(t *testing.T) {
	dir := t.TempDir()
	put := func(id, root, prompt string, age time.Duration) {
		t.Helper()
		s := Session{
			ID: id, Root: root, Model: "claude-opus-5", Turns: 2,
			Updated:  time.Now().Add(-age),
			Messages: []llm.Message{{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock(prompt)}}},
		}
		b, err := json.Marshal(s)
		if err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(filepath.Join(dir, id+".json"), b, 0o644); err != nil {
			t.Fatal(err)
		}
	}
	put("old", "/home/joni/test", "rename the widget", 2*time.Hour)
	put("new", "/home/joni/test", "fix the parser", time.Minute)
	put("elsewhere", "/home/joni/test2", "a different checkout", time.Second)
	// A directory is not a session, and neither is a file that will not parse:
	// one bad file must not hide the good ones.
	os.Mkdir(filepath.Join(dir, "subdir.json"), 0o755)
	os.WriteFile(filepath.Join(dir, "broken.json"), []byte("{not json"), 0o644)

	got, err := ListSessions(dir, "/home/joni/test/")
	if err != nil {
		t.Fatalf("ListSessions: %v", err)
	}
	var ids []string
	for _, s := range got {
		ids = append(ids, s.ID)
	}
	if strings.Join(ids, ",") != "new,old" {
		t.Errorf("listed %v, want this root's sessions, newest first", ids)
	}
	if got[0].Title != "fix the parser" {
		t.Errorf("title = %q, want the first prompt", got[0].Title)
	}
	if got[0].Turns != 2 || got[0].Messages != 1 || got[0].Model != "claude-opus-5" {
		t.Errorf("row = %+v, want the session's own numbers", got[0])
	}

	// The other checkout sees only its own, and a directory with nothing in it
	// is an empty list rather than an error: it means no session yet.
	other, err := ListSessions(dir, "/home/joni/test2")
	if err != nil || len(other) != 1 || other[0].ID != "elsewhere" {
		t.Errorf("ListSessions for the other root = %+v, %v", other, err)
	}
	if list, err := ListSessions(filepath.Join(t.TempDir(), "never-saved"), "/r"); err != nil || len(list) != 0 {
		t.Errorf("a missing session directory = %v, %v; want no sessions and no error", list, err)
	}
}

func TestSessionTitleIsTheFirstThingAsked(t *testing.T) {
	long := strings.Repeat("a", 200)
	cases := []struct {
		name string
		s    Session
		want string
	}{
		{"empty", Session{}, ""},
		{"first user line", Session{Messages: []llm.Message{
			{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("hello")}},
			{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("\n\n  add a resume command  \nand tests")}},
			{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("later")}},
		}}, "add a resume command"},
		{"tool results are not what anyone said", Session{Messages: []llm.Message{
			{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.ToolResultBlock("tu_1", "some output", false)}},
			{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("what went wrong?")}},
		}}, "what went wrong?"},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if got := c.s.Title(); got != c.want {
				t.Errorf("Title() = %q, want %q", got, c.want)
			}
		})
	}

	s := Session{Messages: []llm.Message{{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock(long)}}}}
	title := s.Title()
	if len([]rune(title)) != titleMax || !strings.HasSuffix(title, "…") {
		t.Errorf("a long prompt was not capped: %d runes, %q", len([]rune(title)), title)
	}
}

// The worktree a session's work had moved into is part of the session: it is
// what sends the work back there when the conversation is reopened. The
// directory the session is filed under stays the main checkout, so the
// chooser there still lists it.
func TestSessionRemembersItsWorktree(t *testing.T) {
	dir := t.TempDir()
	a := New(nil, nil, NewState(), nil, Options{})
	if err := a.Save(dir, "s", "/repo", "/trees/dark-mode", "m"); err != nil {
		t.Fatal(err)
	}
	got, err := LoadSession(filepath.Join(dir, "s.json"))
	if err != nil {
		t.Fatal(err)
	}
	if got.Worktree != "/trees/dark-mode" {
		t.Errorf("Worktree = %q, want the tree it moved into", got.Worktree)
	}
	if got.Root != "/repo" {
		t.Errorf("Root = %q, want the checkout it is filed under", got.Root)
	}
	if !got.BelongsTo("/repo") {
		t.Error("a session whose work moved no longer belongs to the checkout it is filed in")
	}
}

// A session that never left is unchanged on disk, so nothing that reads these
// files has to cope with a new key appearing where there was none.
func TestSessionWithoutAWorktreeOmitsTheField(t *testing.T) {
	dir := t.TempDir()
	a := New(nil, nil, NewState(), nil, Options{})
	if err := a.Save(dir, "s", "/repo", "", "m"); err != nil {
		t.Fatal(err)
	}
	b, err := os.ReadFile(filepath.Join(dir, "s.json"))
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(b), "worktree") {
		t.Errorf("session file mentions a worktree it never had:\n%s", b)
	}
}