Vectrum
Sign In
297 lines · 9.5 KB · Go
package tui

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

	tea "github.com/charmbracelet/bubbletea"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/approval"
	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/metrics"
	"github.com/johannSo/codey/internal/tools"
)

// saveSession writes one session file directly, which is the same thing
// Agent.Save produces and lets a test say exactly whose root it was in.
func saveSession(t *testing.T, dir string, s agent.Session) {
	t.Helper()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	b, err := json.Marshal(s)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, s.ID+".json"), b, 0o644); err != nil {
		t.Fatal(err)
	}
}

func user(text string) llm.Message {
	return llm.Message{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock(text)}}
}

func assistant(text string) llm.Message {
	return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock(text)}}
}

// resumeModel is a model with a real agent behind it and a session directory
// holding three sessions: two from this root and one from another.
func resumeModel(t *testing.T) *model {
	t.Helper()
	root := t.TempDir()
	dir := filepath.Join(root, agent.SessionDir)

	saveSession(t, dir, agent.Session{
		ID: "20260101T090000Z", Root: root, Model: "claude-opus-5", Turns: 3,
		Updated: time.Now().Add(-3 * time.Hour),
		Usage:   llm.Usage{InputTokens: 900, OutputTokens: 90},
		Messages: []llm.Message{
			user("rename the widget"),
			assistant("renamed it"),
		},
	})
	saveSession(t, dir, agent.Session{
		ID: "20260102T090000Z", Root: root, Model: "claude-opus-5", Turns: 2,
		Updated: time.Now().Add(-10 * time.Minute),
		Usage:   llm.Usage{InputTokens: 400, OutputTokens: 40},
		Messages: []llm.Message{
			user("fix the parser"),
			{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
				llm.TextBlock("looking at it"),
				{Type: "tool_use", ID: "tu_1", Name: "read_file", Input: json.RawMessage(`{"path":"parse.go"}`)},
			}},
			{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.ToolResultBlock("tu_1", "package parse", false)}},
			assistant("the parser is fixed"),
		},
	})
	// Another checkout's work, filed here as if it had been copied in.
	saveSession(t, dir, agent.Session{
		ID: "20260103T090000Z", Root: "/home/joni/test2", Turns: 1,
		Updated:  time.Now(),
		Messages: []llm.Message{user("a different checkout")},
	})

	log := metrics.NewWriter(&bytes.Buffer{}, "test")
	client := llm.New(llm.WithAPIKey("k"))
	state := agent.NewState()
	a := agent.New(client, tools.NewRegistry(log), state, log, agent.Options{})

	m := newModel(context.Background(), Options{
		Agent: a, Client: client, Policy: approval.New(nil),
		Root: root, SessionDir: dir, SessionID: "20260104T090000Z",
	})
	m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
	return m
}

// The whole point of the chooser: it offers this directory's sessions, by the
// name of what was first asked of them, and never another directory's.
func TestResumeListsOnlyThisDirectorysSessions(t *testing.T) {
	m := resumeModel(t)
	m.submit("/resume")

	if m.phase != phaseSessions {
		t.Fatalf("phase = %v, want the session chooser open", m.phase)
	}
	view := m.sessionView()
	for _, want := range []string{"Resume session", "fix the parser", "rename the widget"} {
		if !strings.Contains(view, want) {
			t.Errorf("%q missing from the chooser:\n%s", want, view)
		}
	}
	if strings.Contains(view, "a different checkout") {
		t.Errorf("the chooser offered another directory's session:\n%s", view)
	}
	// Newest first, so the one you were just in is the one under the cursor.
	if m.sessions[0].ID != "20260102T090000Z" {
		t.Errorf("first row is %q, want the most recently updated", m.sessions[0].ID)
	}
}

// Choosing a session replaces the conversation: the agent's history, the file
// the next turn saves into, and the scrollback, which is replayed so you can
// read what you are continuing.
func TestResumeSwitchesTheConversation(t *testing.T) {
	m := resumeModel(t)
	m.submit("/resume")
	pressKey(t, m, key("enter"))

	if m.phase != phaseIdle {
		t.Errorf("phase = %v, want the chooser closed", m.phase)
	}
	if m.opts.SessionID != "20260102T090000Z" {
		t.Errorf("session id = %q, want the chosen one", m.opts.SessionID)
	}

	history := m.opts.Agent.State().History()
	if len(history) != 4 {
		t.Fatalf("restored %d messages, want the session's 4", len(history))
	}
	if history[0].Content[0].Text != "fix the parser" {
		t.Errorf("history starts with %q", history[0].Content[0].Text)
	}
	if m.turns != 2 || m.usage.InputTokens != 400 {
		t.Errorf("turn and usage totals = %d, %+v; want the session's", m.turns, m.usage)
	}
	// Already paid for in an earlier run, so the next turn is billed on its
	// own delta rather than on the whole restored history.
	if m.priced != m.usage {
		t.Errorf("priced = %+v, want the restored usage %+v", m.priced, m.usage)
	}

	out := m.tr.render()
	for _, want := range []string{"fix the parser", "looking at it", "ReadFile", "the parser is fixed", "resumed 20260102T090000Z"} {
		if !strings.Contains(out, want) {
			t.Errorf("%q missing from the replayed transcript:\n%s", want, out)
		}
	}
}

// A session recorded somewhere else is refused by id too, not just hidden from
// the list — otherwise the rule would only be a display convention.
func TestResumeRefusesASessionFromAnotherDirectory(t *testing.T) {
	m := resumeModel(t)
	before := m.opts.SessionID
	m.submit("/resume 20260103T090000Z")

	out := m.tr.render()
	if !strings.Contains(out, "belongs to") {
		t.Errorf("resuming another directory's session was not refused:\n%s", out)
	}
	if m.opts.SessionID != before {
		t.Errorf("session id moved to %q", m.opts.SessionID)
	}
	if len(m.opts.Agent.State().History()) != 0 {
		t.Error("another directory's conversation was restored anyway")
	}

	m.submit("/resume nonexistent")
	if !strings.Contains(m.tr.render(), "no session nonexistent") {
		t.Errorf("an unknown id was not reported:\n%s", m.tr.render())
	}
}

// Switching mid-turn would pull the history out from under the request in
// flight, so it is refused while anything is running.
func TestResumeIsRefusedMidTurn(t *testing.T) {
	m := resumeModel(t)
	m.phase = phaseRunning
	m.submit("/resume")

	if m.phase != phaseRunning {
		t.Errorf("phase = %v, want the turn left alone", m.phase)
	}
	if !strings.Contains(m.tr.render(), "interrupt first") {
		t.Errorf("a mid-turn /resume was not explained:\n%s", m.tr.render())
	}
}

// A directory nobody has worked in yet has nothing to offer, and says so
// rather than opening an empty chooser.
func TestResumeWithNothingSavedSaysSo(t *testing.T) {
	m := resumeModel(t)
	m.opts.SessionDir = filepath.Join(t.TempDir(), "sessions")
	m.submit("/resume")

	if m.phase == phaseSessions {
		t.Error("an empty chooser was opened")
	}
	if !strings.Contains(m.tr.render(), "no saved sessions") {
		t.Errorf("an empty session directory was not reported:\n%s", m.tr.render())
	}
}

// The session being left is written out before the swap, so nothing that
// happened since the last turn is lost by switching away.
func TestResumeSavesTheSessionBeingLeft(t *testing.T) {
	m := resumeModel(t)
	m.opts.Agent.State().Replace([]llm.Message{user("work in the session being left")})
	m.turns = 1

	m.submit("/resume 20260101T090000Z")

	left, err := agent.LoadSession(filepath.Join(m.opts.SessionDir, "20260104T090000Z.json"))
	if err != nil {
		t.Fatalf("the session being left was not saved: %v", err)
	}
	if got := left.Title(); got != "work in the session being left" {
		t.Errorf("saved session title = %q", got)
	}
	if left.Root != m.opts.Root {
		t.Errorf("saved root = %q, want %q", left.Root, m.opts.Root)
	}
}

func TestAgoReadsAsAList(t *testing.T) {
	cases := []struct {
		d    time.Duration
		want string
	}{
		{30 * time.Second, "just now"},
		{90 * time.Second, "1m ago"},
		{3 * time.Hour, "3h ago"},
		{50 * time.Hour, "2d ago"},
		{3 * 7 * 24 * time.Hour, "3w ago"},
	}
	for _, c := range cases {
		if got := ago(c.d); got != c.want {
			t.Errorf("ago(%v) = %q, want %q", c.d, got, c.want)
		}
	}
}

// A click on a row is a choice here too, and it has to land on the session it
// looks like it is on — the heading rows above the list are the easy thing to
// get wrong.
func TestClickingASessionRowResumesIt(t *testing.T) {
	m := resumeModel(t)
	m.mouseOn = true
	m.submit("/resume")
	m.View()

	// Two rows of heading come before the list; the second row is the older
	// of this directory's two sessions.
	row, ok := m.rowOf(secSessions, sessionHeadRows+1)
	if !ok {
		t.Fatal("the session chooser is not on screen anywhere")
	}
	m.Update(press(4, row))

	if m.opts.SessionID != "20260101T090000Z" {
		t.Errorf("clicking the second row resumed %q", m.opts.SessionID)
	}
	if m.phase != phaseIdle {
		t.Error("the chooser stayed open after a session was clicked")
	}
	if !strings.Contains(m.tr.render(), "rename the widget") {
		t.Errorf("the clicked session was not replayed:\n%s", m.tr.render())
	}
}

// The wheel moves through the chooser rather than scrolling the transcript
// behind it, the way it does over every other open picker.
func TestTheWheelMovesThroughTheSessionChooser(t *testing.T) {
	m := resumeModel(t)
	m.mouseOn = true
	m.submit("/resume")
	m.View()

	row, _ := m.rowOf(secSessions, sessionHeadRows)
	m.Update(wheel(4, row, false))
	if m.sessionIdx != 1 {
		t.Errorf("wheel down over the chooser moved to %d, want 1", m.sessionIdx)
	}
	m.Update(wheel(4, row, true))
	if m.sessionIdx != 0 {
		t.Errorf("wheel up over the chooser moved to %d, want 0", m.sessionIdx)
	}
}