Vectrum
Sign In
353 lines · 9.3 KB · Go
package tui

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/charmbracelet/bubbles/spinner"
	tea "github.com/charmbracelet/bubbletea"

	"harness/internal/agent"
	"harness/internal/approval"
	"harness/internal/llm"
	"harness/internal/tools"
)

// newLocalModel builds a model with a real agent behind it, so a local
// command has somewhere to leave its transcript for the next turn.
func newLocalModel(t *testing.T) *model {
	t.Helper()
	root := t.TempDir()
	client := llm.New(llm.WithAPIKey("x"))
	a := agent.New(client, tools.NewRegistry(nil), agent.NewState(), nil, agent.Options{})
	m := newModel(context.Background(), Options{
		Agent:  a,
		Client: client,
		Root:   root,
		Policy: approval.New(nil),
	})
	m.Update(tea.WindowSizeMsg{Width: 80, Height: 24})
	return m
}

// drain runs a command the way Bubble Tea would and feeds every message it
// produces back into the model, flattening batches. Returns what was seen.
//
// Spinner ticks are dropped rather than followed: each one schedules the next,
// so a real program stops animating when the frame does and a test that
// followed them would never return.
func drain(t *testing.T, m *model, c tea.Cmd) []tea.Msg {
	t.Helper()
	if c == nil {
		return nil
	}
	msg := c()
	if batch, ok := msg.(tea.BatchMsg); ok {
		var out []tea.Msg
		for _, sub := range batch {
			out = append(out, drain(t, m, sub)...)
		}
		return out
	}
	if msg == nil {
		return nil
	}
	if _, tick := msg.(spinner.TickMsg); tick {
		return nil
	}
	_, next := m.Update(msg)
	return append([]tea.Msg{msg}, drain(t, m, next)...)
}

func TestBangOnAnEmptyPromptEntersShellMode(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("!"))

	if m.local != localShell {
		t.Errorf("local = %v, want shell mode", m.local)
	}
	if v := m.input.Value(); v != "" {
		t.Errorf("prompt = %q: the ! is the mode, not the first character of the command", v)
	}
}

func TestDollarOnAnEmptyPromptEntersPythonMode(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("$"))

	if m.local != localPython {
		t.Errorf("local = %v, want python mode", m.local)
	}
	if v := m.input.Value(); v != "" {
		t.Errorf("prompt = %q, want empty", v)
	}
}

func TestBangAfterTextIsJustAnExclamationMark(t *testing.T) {
	m := newTestModel(t)
	typeIn(m, "it works!")

	if m.local != localOff {
		t.Errorf("local = %v: a ! mid-sentence is prose, not a mode", m.local)
	}
	if v := m.input.Value(); v != "it works!" {
		t.Errorf("prompt = %q, want %q", v, "it works!")
	}
}

func TestLocalModeIsNotOfferedWhileTheAgentIsWorking(t *testing.T) {
	m := newTestModel(t)
	m.phase = phaseRunning
	m.Update(key("!"))

	if m.local != localOff {
		t.Errorf("local = %v: a turn is running, so there is nothing to type into", m.local)
	}
}

func TestEscLeavesLocalMode(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("!"))
	typeIn(m, "rm -rf /")
	m.Update(key("esc"))

	if m.local != localOff {
		t.Errorf("local = %v, want off", m.local)
	}
	if v := m.input.Value(); v != "" {
		t.Errorf("prompt = %q: leaving the mode drops what was typed into it", v)
	}
}

func TestBackspaceOnAnEmptyPromptLeavesLocalMode(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("!"))
	m.Update(key("backspace"))

	if m.local != localOff {
		t.Errorf("local = %v: backspacing over the ! undoes it", m.local)
	}
}

func TestBackspaceWithTextTypedStaysInLocalMode(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("!"))
	typeIn(m, "ls")
	m.Update(key("backspace"))

	if m.local != localShell {
		t.Errorf("local = %v: backspace was editing the command, not leaving the mode", m.local)
	}
	if v := m.input.Value(); v != "l" {
		t.Errorf("prompt = %q, want %q", v, "l")
	}
}

func TestShellModeMarksThePromptSoTheModeIsVisible(t *testing.T) {
	m := newTestModel(t)
	m.Update(key("!"))
	if !strings.HasPrefix(m.input.Prompt, "!") {
		t.Errorf("prompt marker = %q, want it to start with !", m.input.Prompt)
	}

	m.Update(key("esc"))
	if strings.HasPrefix(m.input.Prompt, "!") {
		t.Errorf("prompt marker = %q, want the normal one back", m.input.Prompt)
	}
}

func TestShellCommandRunsAndItsOutputReachesTheAgent(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "echo hello-from-the-shell")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	history := m.opts.Agent.State().History()
	if len(history) != 1 {
		t.Fatalf("history has %d messages, want the one local run", len(history))
	}
	msg := history[0]
	if msg.Role != llm.RoleUser {
		t.Errorf("role = %q, want user: it is something the human did", msg.Role)
	}
	text := msg.Content[0].Text
	if !strings.Contains(text, "echo hello-from-the-shell") {
		t.Errorf("the command itself is missing:\n%s", text)
	}
	if !strings.Contains(text, "hello-from-the-shell") {
		t.Errorf("the output is missing:\n%s", text)
	}
	if !strings.Contains(text, "exit code: 0") {
		t.Errorf("the exit code is missing:\n%s", text)
	}
}

func TestAFailingShellCommandKeepsItsOutputAndExitCode(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "echo nope >&2; exit 2")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	text := m.opts.Agent.State().History()[0].Content[0].Text
	for _, want := range []string{"exit code: 2", "nope"} {
		if !strings.Contains(text, want) {
			t.Errorf("missing %q in:\n%s", want, text)
		}
	}
}

func TestPythonModeRunsThroughTheInterpreter(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("$"))
	typeIn(m, "print(6 * 7)")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	text := m.opts.Agent.State().History()[0].Content[0].Text
	if !strings.Contains(text, "42") {
		t.Errorf("want the interpreter's answer, got:\n%s", text)
	}
}

func TestAMissingInterpreterSaysSoRatherThanFailingObscurely(t *testing.T) {
	t.Setenv("PATH", t.TempDir())
	m := newLocalModel(t)
	m.Update(key("$"))
	typeIn(m, "print(1)")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	rendered := m.tr.render()
	if !strings.Contains(rendered, "python3") || !strings.Contains(rendered, "not on your PATH") {
		t.Errorf("want a plain explanation of the missing interpreter, got:\n%s", rendered)
	}
	if n := len(m.opts.Agent.State().History()); n != 0 {
		t.Errorf("history has %d messages: a command that never ran has nothing to tell the model", n)
	}
}

func TestARunReturnsThePromptToNormalSoTheNextThingTypedIsAMessage(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "true")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	if m.local != localOff {
		t.Errorf("local = %v: the mode is per-command, like a shell prompt returning", m.local)
	}
	if v := m.input.Value(); v != "" {
		t.Errorf("prompt = %q, want it emptied", v)
	}
}

func TestTheRunShowsUpInTheTranscriptWithItsOutput(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "echo visible-in-scrollback")
	_, cmd := m.Update(key("enter"))
	drain(t, m, cmd)

	rendered := m.tr.render()
	for _, want := range []string{"! echo visible-in-scrollback", "visible-in-scrollback"} {
		if !strings.Contains(rendered, want) {
			t.Errorf("missing %q in transcript:\n%s", want, rendered)
		}
	}
}

func TestAnEmptyLocalCommandRunsNothing(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	m.Update(key("enter"))

	if n := len(m.opts.Agent.State().History()); n != 0 {
		t.Errorf("history has %d messages, want none", n)
	}
	if m.local != localShell {
		t.Errorf("local = %v: an empty enter is a no-op, not a way out of the mode", m.local)
	}
}

// Two runs in flight at once would land out of order, and finishing one would
// attach its output to the other's transcript entry.
func TestASecondLocalCommandIsRefusedWhileOneIsStillRunning(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "sleep 5")
	m.Update(key("enter")) // deliberately not drained: still in flight

	m.Update(key("!"))
	typeIn(m, "echo second")
	m.Update(key("enter"))

	var runs int
	for _, e := range m.tr.entries {
		if e.kind == entryLocal {
			runs++
		}
	}
	if runs != 1 {
		t.Errorf("%d local entries, want only the one still running", runs)
	}
}

func TestAMessageToTheAgentWaitsForTheLocalCommandToFinish(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "sleep 5")
	m.Update(key("enter"))

	typeIn(m, "now fix that")
	m.Update(key("enter"))

	if m.phase == phaseRunning {
		t.Error("the turn started before the command it is about had landed in the history")
	}
	if v := m.input.Value(); v != "now fix that" {
		t.Errorf("prompt = %q: the message was dropped rather than held", v)
	}
}

func TestTheRunningCommandIsNamedWhileItWorks(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "sleep 5")
	m.Update(key("enter"))

	if !strings.Contains(m.activity(), "sleep 5") {
		t.Errorf("activity line = %q, want the running command named", m.activity())
	}
}

func TestEscInterruptsARunningLocalCommand(t *testing.T) {
	m := newLocalModel(t)
	m.Update(key("!"))
	typeIn(m, "sleep 30")
	_, cmd := m.Update(key("enter"))

	m.Update(key("esc"))

	start := time.Now()
	drain(t, m, cmd)
	if elapsed := time.Since(start); elapsed > 3*time.Second {
		t.Errorf("the command ran for %s: esc never reached it", elapsed)
	}
	if m.running != nil {
		t.Error("the run is still marked in flight after it came back")
	}
}

// A key with no command to complete it is invisible unless /help says so.
func TestHelpNamesTheLocalCommandKeys(t *testing.T) {
	help := newTestModel(t).helpText()
	for _, want := range []string{"!", "$", "shell", "python"} {
		if !strings.Contains(help, want) {
			t.Errorf("/help does not mention %q:\n%s", want, help)
		}
	}
}