Vectrum
Sign In
246 lines · 7.3 KB · Go
package tui

import (
	"context"
	"fmt"
	"os/exec"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"

	"harness/internal/llm"
	"harness/internal/tools"
)

// localMode is which interpreter the prompt box is pointed at. Off is the
// normal case: what you type goes to the model.
type localMode int

const (
	localOff localMode = iota
	localShell
	localPython
)

// localSpec is one mode's whole definition — the character that opens it, the
// binary it needs, and how a typed line becomes an argv. Keeping the three
// together is what stops the key, the prompt marker and the interpreter from
// drifting apart as modes are added.
type localSpec struct {
	marker string
	bin    string
	hint   string
	argv   func(line string) []string
}

var localSpecs = map[localMode]localSpec{
	localShell: {
		marker: "!",
		bin:    "bash",
		hint:   "shell — enter runs it here, esc to leave",
		argv:   func(line string) []string { return []string{"bash", "-c", line} },
	},
	localPython: {
		marker: "$",
		bin:    "python3",
		hint:   "python — enter runs it here, esc to leave",
		argv:   func(line string) []string { return []string{"python3", "-c", line} },
	},
}

// localFor maps an opening keystroke to its mode.
func localFor(key string) localMode {
	for mode, spec := range localSpecs {
		if key == spec.marker {
			return mode
		}
	}
	return localOff
}

// localRun is the command in flight. Only ever one: two at once would land in
// whatever order they finished, and the earlier one's output would be attached
// to the later one's transcript entry.
type localRun struct {
	spec    localSpec
	command string
	started time.Time
	cancel  context.CancelFunc
}

// localDone carries a finished local run back to the update loop. The command
// runs on its own goroutine but lands here, so the transcript and the
// conversation are only ever touched from the single-threaded loop.
type localDone struct {
	spec    localSpec
	command string
	report  string
	failed  bool
}

// onLocalKey handles the keys that belong to the local-command modes and says
// whether it took the key. Everything it does not take falls through to the
// prompt, which is what keeps typing a command rather than a set of controls.
func (m *model) onLocalKey(msg tea.KeyMsg) (tea.Cmd, bool) {
	key := msg.String()

	if m.local == localOff {
		// The marker opens a mode only from an empty prompt that is yours to
		// type into. Mid-sentence it is prose, and mid-turn there is nothing
		// to type into at all.
		if m.phase != phaseIdle || m.running != nil || m.input.Value() != "" {
			return nil, false
		}
		mode := localFor(key)
		if mode == localOff {
			return nil, false
		}
		m.setLocal(mode)
		m.closeMenu()
		return nil, true
	}

	switch key {
	case "esc":
		m.setLocal(localOff)
		return nil, true

	case "backspace":
		// Backspacing over the marker undoes it, the way it looks like it
		// should. With a command half-typed it is still just backspace.
		if m.input.Value() == "" {
			m.setLocal(localOff)
			return nil, true
		}
		return nil, false

	case "enter":
		line := strings.TrimSpace(m.input.Value())
		if line == "" || m.running != nil {
			return nil, true
		}
		mode := m.local
		m.input.Reset()
		m.growInput()
		return m.runLocal(mode, line), true
	}
	return nil, false
}

// setLocal switches the prompt box into a mode, or out of every mode. The box
// is emptied either way: a command typed at the shell prompt is not a sentence
// you meant to send to the model.
func (m *model) setLocal(mode localMode) {
	m.local = mode
	m.input.Reset()
	m.growInput()

	if mode == localOff {
		m.input.Prompt = "› "
		m.input.Placeholder = defaultPlaceholder
		m.input.FocusedStyle.Prompt = styCoral
		return
	}
	spec := localSpecs[mode]
	m.input.Prompt = spec.marker + " "
	m.input.Placeholder = spec.hint
	m.input.FocusedStyle.Prompt = localStyle(mode)
}

// localStyle is the mode's colour, carried by the prompt marker and the
// transcript entry so the two read as the same thing.
func localStyle(mode localMode) lipgloss.Style {
	if mode == localPython {
		return styLocalPython
	}
	return styLocalShell
}

// runLocal starts one command. It returns to the normal prompt immediately:
// the mode is per-command, like a shell prompt coming back, so a run is never
// followed by a box you thought you had left.
func (m *model) runLocal(mode localMode, line string) tea.Cmd {
	spec := localSpecs[mode]
	m.setLocal(localOff)

	// Looking the binary up first turns "python3 is not installed" into a
	// sentence rather than an exec error the shell spells its own way.
	if _, err := exec.LookPath(spec.bin); err != nil {
		m.tr.addError(spec.bin + " is not on your PATH, so " + spec.marker + " has nothing to run")
		return m.sync()
	}

	m.tr.startLocal(spec, line)
	ctx, cancel := context.WithCancel(m.parent)
	m.running = &localRun{spec: spec, command: line, started: time.Now(), cancel: cancel}

	root, argv := m.opts.Root, spec.argv(line)
	return tea.Batch(m.sync(), m.spin.Tick, func() tea.Msg {
		report, failed := tools.RunCommand(ctx, root, argv, tools.DefaultCommandTimeout)
		cancel()
		return localDone{spec: spec, command: line, report: report, failed: failed}
	})
}

// interruptLocal stops the command in flight. Its output still comes back and
// still reaches the model: what a command managed to print before you stopped
// it is evidence, not noise.
func (m *model) interruptLocal() tea.Cmd {
	if m.running != nil {
		m.running.cancel()
	}
	return nil
}

// localActivity is the line above the prompt while a command runs: which one,
// and for how long. Without it a slow command looks like a frozen prompt.
func (m *model) localActivity() string {
	r := m.running
	detail := fmt.Sprintf("%ds · esc to interrupt", int(time.Since(r.started).Seconds()))
	return m.spin.View() + " " +
		localStyle(localOf(r.spec)).Render("Running "+r.spec.marker+" "+r.command) +
		styDim.Render(" ("+detail+")")
}

// localOf recovers a mode from its spec, so the running line can be painted in
// the same colour the prompt marker was.
func localOf(spec localSpec) localMode {
	for mode, s := range localSpecs {
		if s.marker == spec.marker {
			return mode
		}
	}
	return localShell
}

// finishLocal lands a completed run: into the scrollback for you, and into the
// conversation for the model. The second half is the point of the feature —
// running `!go test` and then saying "fix that" only works if the agent was
// shown the same output you were.
//
// Nothing is sent now. The message sits in history and goes up with the next
// real turn, so several commands can be run before saying anything about them.
func (m *model) finishLocal(msg localDone) tea.Cmd {
	if m.running != nil {
		m.running.cancel()
		m.running = nil
	}
	m.tr.finishLocal(msg.report, msg.failed)
	if m.opts.Agent != nil {
		m.opts.Agent.State().Append(llm.Message{
			Role:    llm.RoleUser,
			Content: []llm.ContentBlock{llm.TextBlock(localMessage(msg))},
		})
	}
	return m.sync()
}

// localMessage is what the model is told. It is the capped report the tool
// would have produced, not the raw output: a noisy command should not be able
// to cost more context from the prompt box than it would through the tool.
func localMessage(msg localDone) string {
	return fmt.Sprintf("<local-command shell=%q>%s</local-command>\n<local-command-output>\n%s\n</local-command-output>",
		msg.spec.bin, msg.command, msg.report)
}