Vectrum
Sign In
409 lines · 14.3 KB · Go
package tui

import (
	"bytes"
	"context"
	_ "embed"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unicode"

	tea "github.com/charmbracelet/bubbletea"

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

// defaultScript is what is written to ~/.config/codey/worktree.sh the first
// time /worktree is used. It is a file on disk rather than a command built
// here so that a tree that needs a .env copied or dependencies installed can
// be taught to do it without touching Codey.
//
//go:embed worktree.sh
var defaultScript string

// namingModel is asked to turn a task into a directory name. The job is to
// compress a sentence into three words, so the cheapest model that can read
// is the right one — and it is cloned off the session's client, so it
// authenticates exactly as the session does.
const namingModel = "claude-haiku-4-5-20251001"

// namingPrompt is deliberately narrow. The reply is slugified and validated
// whatever it says, but asking for one word on one line makes the common case
// need no repair.
const namingPrompt = `Name a git worktree for this task.

Reply with ONLY a short kebab-case slug: two to four words, lowercase, hyphens
between words, no punctuation, no quotes, no explanation. Name it after what
the work is, not what the code currently does.

The task:
`

// nameMax is how long a generated name may be. Long enough to stay readable
// in a directory listing, short enough not to hide the interesting part of a
// path behind a wrapped line.
const nameMax = 48

// worktreeReady carries a finished /worktree back to the update loop: the
// tree exists and the work can move into it. task is what the agent should
// start on once it has, and is empty when /worktree was given a bare name and
// so has nothing to do afterwards.
type worktreeReady struct {
	name string
	task string
	path string
	err  error
}

// startWorktree answers /worktree. The argument is read two ways, and which
// one you get is visible in what you typed: a single word is the name, and
// anything with a space in it is a description of the work, which the model
// names and the agent then starts on.
func (m *model) startWorktree(arg string) tea.Cmd {
	arg = strings.TrimSpace(arg)

	// "resume" is the one word that is not taken as a name: it goes to the
	// other tree rather than making one. It is checked before anything else so
	// that /worktree resume reads the same whether or not a tree could be made
	// right now.
	if arg == resumeWord || strings.HasPrefix(arg, resumeWord+" ") {
		return m.resumeWorktree(strings.TrimSpace(arg[len(resumeWord):]))
	}

	switch {
	case m.opts.NewTools == nil:
		m.tr.addError("/worktree is not available in this session — it has no way to move the tools")
		return m.sync()
	case arg == "":
		m.tr.addError("/worktree needs a name or a task — /worktree fix-login, /worktree fix the login redirect, or /worktree resume to open one that exists")
		return m.sync()
	case m.phase == phaseRunning:
		m.tr.addError("/worktree would move the ground under the running turn — interrupt first")
		return m.sync()
	case m.compacting || m.worktreeing:
		return m.sync()
	}

	// A bare name is taken at its word: no model call, and no task to start
	// on afterwards. That is the "just give me a tree" path, and it has to be
	// quick and offline.
	name, verbatim := verbatimName(arg)
	task := arg
	if verbatim {
		task = ""
	}

	m.worktreeing = true
	m.started = time.Now()
	m.closeMenu()

	ctx, cancel := context.WithCancel(m.parent)
	m.cancel = cancel

	client, prefix := m.opts.Client, m.opts.SystemPrefix
	home := m.opts.HomeRoot
	run := func() tea.Msg {
		defer cancel()
		if !verbatim {
			name = worktreeName(ctx, client, prefix, task)
		}
		// Naming is the slow half, so this is where an interrupt usually
		// lands. Checking here means esc during the naming call leaves no
		// tree behind rather than making one nobody asked for.
		if err := ctx.Err(); err != nil {
			return worktreeReady{name: name, task: task, err: err}
		}
		path, err := runWorktreeScript(ctx, home, name)
		return worktreeReady{name: name, task: task, path: path, err: err}
	}
	return tea.Batch(m.sync(), m.spin.Tick, m.tag(run))
}

// finishWorktree lands the result: the work moves, and if there was a task it
// starts. The turn is begun from here rather than from the goroutine so that
// the conversation is only ever touched by the update loop.
func (m *model) finishWorktree(msg worktreeReady) tea.Cmd {
	m.worktreeing = false
	m.interrupting = false
	m.cancel = nil

	if msg.err != nil {
		// An interrupt is something you did, not something that went wrong.
		if errors.Is(msg.err, context.Canceled) {
			m.tr.addNotice("/worktree interrupted — still working in " + shortenHome(m.opts.Root))
		} else {
			m.tr.addError("/worktree: " + msg.err.Error())
		}
		return m.sync()
	}

	m.switchRoot(msg.path)
	m.tr.addNotice("working in " + shortenHome(msg.path) + " — the session stays filed in " +
		shortenHome(m.opts.HomeRoot) + ", and reopening it comes back here")

	if msg.task == "" {
		// No turn follows, so nothing else is going to write the session —
		// and a worktree that was never recorded is one that reopening the
		// session will not come back to.
		if m.opts.Agent != nil && m.opts.SessionID != "" {
			if err := m.saveSession(); err != nil {
				m.tr.addError("could not record the worktree in the session: " + err.Error())
			}
		}
		return m.sync()
	}
	// The model is told where it is before it is told what to do. It has just
	// been handed tools pointing somewhere new, and nothing in the
	// conversation so far would explain why the paths moved.
	return m.startTurn(msg.task, fmt.Sprintf(
		"I have moved your working directory into a fresh git worktree at %s, on a new branch %q, made from the branch we were on. Everything you read, write and run now happens there. Work on this:\n\n%s",
		msg.path, msg.name, msg.task))
}

// switchRoot moves the work without moving the session. The tools are rebuilt
// because each one captured the old root when it was made; the session
// directory is untouched, so the conversation stays filed in the checkout
// Codey was started in rather than following the work around.
func (m *model) switchRoot(path string) {
	m.opts.Root = path
	// A session records the tree its work moved into so that reopening it
	// comes back here. Moving back to the checkout Codey was started in is not
	// moving into a tree, so the field is cleared rather than made to point at
	// the home root — which resuming would then treat as a worktree.
	m.opts.Worktree = path
	if samePath(path, m.opts.HomeRoot) {
		m.opts.Worktree = ""
	}

	if m.opts.NewTools == nil {
		return
	}
	// Building is what re-points the subagent supervisor, so it happens even
	// in the odd case of there being no agent to hand the result to.
	reg := m.opts.NewTools(path)
	if m.opts.Agent != nil {
		m.opts.Agent.SetTools(reg)
	}
}

// worktreeActivity is the line above the prompt while /worktree works. Naming
// is a model call and the script may be doing anything from a checkout to an
// install, so the wait is worth a spinner.
func (m *model) worktreeActivity() string {
	detail := fmt.Sprintf("%ds · esc to interrupt", int(time.Since(m.started).Seconds()))
	verb := "Making a worktree"
	if m.interrupting {
		verb = "Interrupting"
	}
	return m.spin.View() + " " + styCoral.Render(verb+"…") + styDim.Render(" ("+detail+")")
}

// worktreeName asks the model for a name and repairs whatever comes back. A
// name is not worth failing over: if the call errors, is interrupted, or
// returns prose, the task's own opening words make a serviceable name, and
// the work carries on.
func worktreeName(ctx context.Context, client *llm.Client, prefix, task string) string {
	fallback := slugify(task)
	if fallback == "" {
		fallback = "worktree"
	}
	if client != nil {
		req := llm.Request{
			MaxTokens: 64,
			Messages: []llm.Message{{
				Role:    llm.RoleUser,
				Content: []llm.ContentBlock{llm.TextBlock(namingPrompt + task)},
			}},
		}
		// An OAuth session must carry the Claude Code identity as the FIRST
		// system block or the API answers 429 rather than 401 — which reads
		// as a transient failure and burns the retries before surfacing. The
		// same warning is on the compaction request, for the same reason.
		if prefix != "" {
			req.System = []llm.SystemBlock{{Type: "text", Text: prefix}}
		}
		if resp, err := client.Clone(namingModel).CreateMessage(ctx, req); err == nil {
			if name := slugify(resp.Text()); name != "" {
				return name
			}
		}
	}
	return fallback
}

// verbatimName decides whether the argument is a name or a description of
// work. The rule is the whitespace: one word is a name, and a phrase is a task
// for the model to name. That is a rule you can see in what you typed, which
// matters more here than cleverness would — /worktree dark-mode should never
// pause to ask an API what to call itself.
//
// The word is still slugified, so a capital or a stray dot is repaired rather
// than refused.
func verbatimName(arg string) (string, bool) {
	if strings.ContainsFunc(arg, unicode.IsSpace) {
		return "", false
	}
	return slugify(arg), true
}

// slugify turns anything into something safe to use as a directory name:
// lowercase, words joined by single hyphens, and nothing that could escape the
// worktrees directory or be read by git as a flag. It is applied to the
// model's answer as well as to the fallback, because a model asked for a slug
// will now and then reply with a sentence about one.
func slugify(s string) string {
	// Only the first line is considered: a model that explains itself does so
	// underneath, and a first line is all a name can come from anyway.
	if i := strings.IndexByte(s, '\n'); i >= 0 {
		s = s[:i]
	}
	var b strings.Builder
	var pendingDash bool
	for _, r := range strings.ToLower(strings.TrimSpace(s)) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			if pendingDash && b.Len() > 0 {
				b.WriteByte('-')
			}
			pendingDash = false
			b.WriteRune(r)
		default:
			// Everything else — spaces, punctuation, dots, slashes, anything
			// non-ASCII — becomes a word boundary rather than a character.
			// That is what makes "../etc" and "feat/x" harmless here.
			pendingDash = true
		}
		if b.Len() >= nameMax {
			break
		}
	}
	return strings.Trim(b.String(), "-")
}

// ensureScript returns the path of the worktree script, writing the default
// one if there is nothing there. An existing script is never touched: once it
// is on disk it belongs to whoever edited it.
func ensureScript() (string, error) {
	path, err := config.ScriptPath()
	if err != nil {
		return "", err
	}
	if _, err := os.Stat(path); err == nil {
		return path, nil
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return "", err
	}
	if err := os.WriteFile(path, []byte(defaultScript), 0o755); err != nil {
		return "", err
	}
	return path, nil
}

// runWorktreeScript pipes the name in and reads the path back out. The script
// decides where the tree lands — it is allowed to put it somewhere else
// entirely — so the last line of its stdout is taken as the answer rather than
// the path Codey would have guessed.
func runWorktreeScript(ctx context.Context, repoRoot, name string) (string, error) {
	if name == "" {
		return "", fmt.Errorf("could not work out a name for the worktree")
	}
	script, err := ensureScript()
	if err != nil {
		return "", err
	}
	trees, err := config.WorktreesDir()
	if err != nil {
		return "", err
	}

	env := append(os.Environ(),
		"CODEY_REPO_ROOT="+repoRoot,
		"CODEY_WORKTREE_NAME="+name,
		"CODEY_WORKTREES_DIR="+trees,
	)
	stdout, stderr, err := runScript(ctx, repoRoot, script, name, env, worktreeTimeout)
	if err != nil {
		return "", scriptError(script, err, stderr)
	}

	path := lastLine(stdout)
	if path == "" {
		return "", fmt.Errorf("%s printed no path — it should echo the worktree's directory as its last line", shortenHome(script))
	}
	if !filepath.IsAbs(path) {
		path = filepath.Join(repoRoot, path)
	}
	if fi, err := os.Stat(path); err != nil || !fi.IsDir() {
		return "", fmt.Errorf("%s said the worktree is at %s, but there is no directory there", shortenHome(script), path)
	}
	return path, nil
}

// scriptError puts the script's own complaint in front of the reader. A failed
// worktree is nearly always git explaining something specific — a branch that
// exists, a dirty index — and that sentence is worth more than the exit code.
func scriptError(script string, err error, stderr string) error {
	if detail := lastLine(stderr); detail != "" {
		return fmt.Errorf("%s: %s", shortenHome(script), detail)
	}
	return fmt.Errorf("%s: %w", shortenHome(script), err)
}

// lastLine is the final non-empty line, trimmed. Scripts are chatty and the
// interesting part is at the end of what they say.
func lastLine(s string) string {
	lines := strings.Split(s, "\n")
	for i := len(lines) - 1; i >= 0; i-- {
		if t := strings.TrimSpace(lines[i]); t != "" {
			return t
		}
	}
	return ""
}

// worktreeTimeout is how long the script gets. Generous, because a script
// taught to install dependencies is doing a real job, and a worktree is not
// something anyone asks for twice a minute.
const worktreeTimeout = 5 * time.Minute

// runScript runs one script with text on its stdin and its two streams kept
// apart. The bash tool's RunCommand would do most of this, but it folds stdout
// and stderr into one report and has nowhere to put stdin — and here the two
// streams mean different things: stdout is the answer, stderr is the
// explanation of a failure.
func runScript(ctx context.Context, dir, script, stdin string, env []string, timeout time.Duration) (string, string, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "bash", script)
	cmd.Dir = dir
	cmd.Env = env
	cmd.Stdin = strings.NewReader(stdin)
	// A process group, so a script that spawns a long install can be killed
	// whole rather than leaving the child holding the pipes open.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		if cmd.Process == nil {
			return nil
		}
		return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
	}
	cmd.WaitDelay = 2 * time.Second

	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr

	err := cmd.Run()
	if ctx.Err() == context.DeadlineExceeded {
		err = fmt.Errorf("timed out after %s", timeout)
	}
	return stdout.String(), stderr.String(), err
}