byre/internal/terminal/terminal.go

package terminal

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/creack/pty"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/ids"
)

const (
	backlogSize = 256 << 10
	clientBuf   = 256
	MaxSize     = 1 << 15
	promptMark  = "\x1b]133;D;"
	lineMark    = "\x1b]7701;"
	carryMax    = 64 << 10
	keepSize    = 256 << 10
	Quiet       = 500 * time.Millisecond
	MaxWait     = 10 * time.Second
)

var lineMarks = [][]byte{[]byte(promptMark), []byte(lineMark)}

const zshenv = `if [[ -n $BYRE_ZDOTDIR ]]; then
  ZDOTDIR=$BYRE_ZDOTDIR
else
  unset ZDOTDIR
fi
unset BYRE_ZDOTDIR
[[ -f ${ZDOTDIR:-$HOME}/.zshenv ]] && source ${ZDOTDIR:-$HOME}/.zshenv
if [[ -o interactive ]]; then
  unsetopt PROMPT_SP
  _byre_precmd() { print -n "\e]133;D;$?\a" }
  precmd_functions+=(_byre_precmd)
  _byre_line() { local line=$PREBUFFER$BUFFER; print -rn -- $'\e]7701;L'"${line//[[:cntrl:]]/ }"$'\a' }
  _byre_line_end() { print -rn -- $'\e]7701;E\a' }
  zmodload zsh/zle
  autoload -Uz add-zle-hook-widget
  add-zle-hook-widget line-init _byre_line
  add-zle-hook-widget line-pre-redraw _byre_line
  add-zle-hook-widget line-finish _byre_line_end
fi
`

type Terminals struct {
	zdotdir  string
	changed  func(taskID string)
	prompted func(taskID string)

	mu     sync.Mutex
	byTask map[string][]*Terminal
}

type Terminal struct {
	ID       string
	Name     string
	ptmx     *os.File
	cmd      *exec.Cmd
	Prompts  bool
	prompted func()

	mu         sync.Mutex
	backlog    []byte
	clients    map[chan []byte]bool
	exit       string
	exited     bool
	started    time.Time
	lastOutput time.Time
	wake       *wake
	carry      []byte
	line       string
	atPrompt   bool
}

type wake struct {
	fn  func(exit string, out []byte)
	out []byte
}

func New(zdotdir string, changed, prompted func(taskID string)) (*Terminals, error) {
	err := os.MkdirAll(zdotdir, 0o755)
	if err != nil {
		return nil, err
	}
	err = os.WriteFile(filepath.Join(zdotdir, ".zshenv"), []byte(zshenv), 0o644)
	if err != nil {
		return nil, err
	}
	return &Terminals{zdotdir: zdotdir, changed: changed, prompted: prompted, byTask: map[string][]*Terminal{}}, nil
}

func (ts *Terminals) List(taskID string) []*Terminal {
	ts.mu.Lock()
	defer ts.mu.Unlock()
	return slices.Clone(ts.byTask[taskID])
}

func (ts *Terminals) All() map[string][]*Terminal {
	ts.mu.Lock()
	defer ts.mu.Unlock()
	out := map[string][]*Terminal{}
	for taskID, terms := range ts.byTask {
		out[taskID] = slices.Clone(terms)
	}
	return out
}

func (ts *Terminals) Get(taskID, id string) *Terminal {
	ts.mu.Lock()
	defer ts.mu.Unlock()
	for _, term := range ts.byTask[taskID] {
		if term.ID == id {
			return term
		}
	}
	return nil
}

func (ts *Terminals) First(taskID string) *Terminal {
	ts.mu.Lock()
	defer ts.mu.Unlock()
	if len(ts.byTask[taskID]) == 0 {
		return nil
	}
	return ts.byTask[taskID][0]
}

func (ts *Terminals) Start(t api.Task, name, command string, size api.TerminalSize) (*Terminal, error) {
	term, err := ts.start(t, name, size)
	if err != nil {
		return nil, err
	}
	ts.changed(t.Id)
	if command != "" {
		go term.typeWhenReady(command + "\r")
	}
	return term, nil
}

func (ts *Terminals) start(t api.Task, name string, size api.TerminalSize) (*Terminal, error) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	shell := os.Getenv("SHELL")
	if shell == "" {
		return nil, errors.New("SHELL is not set in byre's environment")
	}
	cmd := exec.Command(shell, "-l")
	cmd.Dir = t.WorktreePath
	cmd.Env = append(os.Environ(), "TERM=xterm-256color", "COLORTERM=truecolor", fmt.Sprintf("PORT=%d", t.PortBase))
	prompts := filepath.Base(shell) == "zsh"
	if prompts {
		cmd.Env = append(cmd.Env, "BYRE_ZDOTDIR="+os.Getenv("ZDOTDIR"), "ZDOTDIR="+ts.zdotdir)
	}
	ptmx, err := pty.StartWithSize(cmd, winsize(size))
	if err != nil {
		return nil, fmt.Errorf("start %s: %w", shell, err)
	}

	term := &Terminal{ID: ids.New(ids.Terminal), Name: name, ptmx: ptmx, cmd: cmd, Prompts: prompts, prompted: func() { ts.prompted(t.Id) }, clients: map[chan []byte]bool{}, started: time.Now()}
	ts.byTask[t.Id] = append(ts.byTask[t.Id], term)
	slog.Info("terminal started", "task", t.Id, "terminal", term.ID, "name", name, "shell", shell, "pid", cmd.Process.Pid)
	go func() {
		term.pump()
		if ts.remove(t.Id, term) {
			ts.changed(t.Id)
		}
		slog.Info("terminal ended", "task", t.Id, "terminal", term.ID, "exit", term.exit)
	}()
	return term, nil
}

func (ts *Terminals) remove(taskID string, term *Terminal) bool {
	ts.mu.Lock()
	defer ts.mu.Unlock()
	i := slices.Index(ts.byTask[taskID], term)
	if i < 0 {
		return false
	}
	ts.byTask[taskID] = slices.Delete(ts.byTask[taskID], i, i+1)
	if len(ts.byTask[taskID]) == 0 {
		delete(ts.byTask, taskID)
	}
	return true
}

func (ts *Terminals) Kill(taskID, id string) bool {
	term := ts.Get(taskID, id)
	if term == nil || !ts.remove(taskID, term) {
		return false
	}
	term.kill()
	ts.changed(taskID)
	return true
}

func (ts *Terminals) Close(taskID string) {
	ts.mu.Lock()
	terms := ts.byTask[taskID]
	delete(ts.byTask, taskID)
	ts.mu.Unlock()
	for _, term := range terms {
		term.kill()
	}
}

func (ts *Terminals) CloseAll() {
	ts.mu.Lock()
	all := ts.byTask
	ts.byTask = map[string][]*Terminal{}
	ts.mu.Unlock()
	for _, terms := range all {
		for _, term := range terms {
			term.kill()
		}
	}
}

func winsize(size api.TerminalSize) *pty.Winsize {
	return &pty.Winsize{Cols: uint16(size.Cols), Rows: uint16(size.Rows)}
}

func ValidSize(size api.TerminalSize) bool {
	return size.Cols >= 1 && size.Rows >= 1 && size.Cols <= MaxSize && size.Rows <= MaxSize
}

func (term *Terminal) pump() {
	buf := make([]byte, 32<<10)
	for {
		n, err := term.ptmx.Read(buf)
		if n > 0 {
			term.output(bytes.Clone(buf[:n]))
		}
		if err != nil {
			break
		}
	}

	err := term.cmd.Wait()
	term.ptmx.Close()
	exit := fmt.Sprint(err)
	if term.cmd.ProcessState != nil {
		exit = term.cmd.ProcessState.String()
	}

	term.mu.Lock()
	defer term.mu.Unlock()
	term.exit = exit
	term.exited = true
	for c := range term.clients {
		close(c)
	}
	term.clients = nil
}

func (term *Terminal) output(b []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	term.lastOutput = time.Now()
	if term.wake != nil {
		term.checkWake(b)
	}
	if term.Prompts {
		term.trackLine(b)
	}
	term.backlog = append(term.backlog, b...)
	if len(term.backlog) > 2*backlogSize {
		term.backlog = bytes.Clone(term.backlog[len(term.backlog)-backlogSize:])
	}
	for c := range term.clients {
		select {
		case c <- b:
		default:
			delete(term.clients, c)
			close(c)
		}
	}
}

func (term *Terminal) checkWake(b []byte) {
	w := term.wake
	from := max(0, len(w.out)-len(promptMark)-8)
	w.out = append(w.out, b...)
	i := bytes.Index(w.out[from:], []byte(promptMark))
	if i < 0 {
		if len(w.out) > 2*backlogSize {
			w.out = bytes.Clone(w.out[len(w.out)-backlogSize:])
		}
		return
	}
	start := from + i
	end := bytes.IndexByte(w.out[start:], '\a')
	if end < 0 {
		return
	}
	term.wake = nil
	go w.fn(string(w.out[start+len(promptMark):start+end]), w.out[:start])
}

func (term *Terminal) trackLine(b []byte) {
	s := append(term.carry, b...)
	term.carry = nil
	for {
		i := bytes.Index(s, []byte("\x1b]"))
		if i < 0 {
			if len(s) > 0 && s[len(s)-1] == 0x1b {
				term.carry = []byte{0x1b}
			}
			return
		}
		rest := s[i:]
		var prefix []byte
		for _, m := range lineMarks {
			if bytes.HasPrefix(rest, m) {
				prefix = m
				break
			}
			if len(rest) < len(m) && bytes.HasPrefix(m, rest) {
				term.carry = bytes.Clone(rest)
				return
			}
		}
		if prefix == nil {
			s = rest[2:]
			continue
		}
		end := bytes.IndexByte(rest, '\a')
		if end < 0 {
			if len(rest) < carryMax {
				term.carry = bytes.Clone(rest)
			}
			return
		}
		body := string(rest[len(prefix):end])
		switch {
		case string(prefix) == promptMark:
			term.atPrompt, term.line = true, ""
			go term.prompted()
		case body == "E":
			term.atPrompt, term.line = false, ""
		case strings.HasPrefix(body, "L"):
			term.atPrompt, term.line = true, body[1:]
		}
		s = rest[end+1:]
	}
}

func (term *Terminal) InputLine() (string, bool) {
	term.mu.Lock()
	defer term.mu.Unlock()
	return term.line, term.atPrompt
}

func (term *Terminal) WakeOnPrompt(fn func(exit string, out []byte), mayRun bool) bool {
	term.mu.Lock()
	defer term.mu.Unlock()
	if term.atPrompt && term.line == "" || !term.atPrompt && !mayRun {
		return false
	}
	term.wake = &wake{fn: fn}
	return true
}

func (term *Terminal) CancelWake() {
	term.mu.Lock()
	defer term.mu.Unlock()
	term.wake = nil
}

func (term *Terminal) Backlog() []byte {
	term.mu.Lock()
	defer term.mu.Unlock()
	return bytes.Clone(term.backlog[max(0, len(term.backlog)-backlogSize):])
}

func (term *Terminal) Subscribe() ([]byte, chan []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	if term.exited {
		return nil, nil
	}
	c := make(chan []byte, clientBuf)
	term.clients[c] = true
	return bytes.Clone(term.backlog[max(0, len(term.backlog)-backlogSize):]), c
}

func (term *Terminal) Unsubscribe(c chan []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	if term.clients[c] {
		delete(term.clients, c)
		close(c)
	}
}

func (term *Terminal) ExitStatus() (string, bool) {
	term.mu.Lock()
	defer term.mu.Unlock()
	return term.exit, term.exited
}

func (term *Terminal) Activity() (started, lastOutput time.Time) {
	term.mu.Lock()
	defer term.mu.Unlock()
	return term.started, term.lastOutput
}

func (term *Terminal) Started() time.Time {
	term.mu.Lock()
	defer term.mu.Unlock()
	return term.started
}

func (term *Terminal) typeWhenReady(text string) {
	backlog, out := term.Subscribe()
	if out == nil {
		return
	}
	_, err := term.AwaitPrompt(context.Background(), out, backlog)
	term.Unsubscribe(out)
	if err != nil {
		return
	}
	err = term.Write([]byte(text))
	if err != nil {
		slog.Warn("type saved command", "terminal", term.ID, "err", err)
	}
}

func (term *Terminal) Settling() bool {
	return term.promptWait() > 0
}

func (term *Terminal) promptWait() time.Duration {
	started, last := term.Activity()
	wait := time.Until(started.Add(MaxWait))
	if !last.IsZero() {
		wait = min(wait, time.Until(last.Add(Quiet)))
	}
	return wait
}

func (term *Terminal) AwaitPrompt(ctx context.Context, out chan []byte, got []byte) ([]byte, error) {
	for {
		wait := term.promptWait()
		if wait <= 0 {
			return got, nil
		}
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		case b, open := <-out:
			timer.Stop()
			if !open {
				return got, nil
			}
			got = KeepTail(append(got, b...))
		}
	}
}

func KeepTail(b []byte) []byte {
	if len(b) > 2*keepSize {
		return bytes.Clone(b[len(b)-keepSize:])
	}
	return b
}

func (term *Terminal) Write(b []byte) error {
	_, err := term.ptmx.Write(b)
	return err
}

func (term *Terminal) Resize(size api.TerminalSize) {
	err := pty.Setsize(term.ptmx, winsize(size))
	if err != nil {
		slog.Debug("terminal resize", "err", err)
	}
}

func (term *Terminal) kill() {
	err := syscall.Kill(-term.cmd.Process.Pid, syscall.SIGHUP)
	if err != nil && !errors.Is(err, syscall.ESRCH) {
		slog.Warn("hang up terminal", "pid", term.cmd.Process.Pid, "err", err)
	}
	term.ptmx.Close()
}