byre/internal/terminal/terminal.go

package terminal

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

	"github.com/creack/pty"

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

const (
	backlogSize = 256 << 10
	clientBuf   = 256
	MaxSize     = 1 << 15
	promptMark  = "\x1b]133;D;"
)

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
  _byre_precmd() { print -n "\e]133;D;$?\a" }
  precmd_functions+=(_byre_precmd)
fi
`

type Terminals struct {
	zdotdir string

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

type Terminal struct {
	ptmx    *os.File
	cmd     *exec.Cmd
	Prompts bool

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

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

func New(zdotdir 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, byTask: map[string]*Terminal{}}, nil
}

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

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

	term := ts.byTask[t.Id]
	if term != nil {
		return term, nil
	}

	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{ptmx: ptmx, cmd: cmd, Prompts: prompts, clients: map[chan []byte]bool{}, started: time.Now()}
	ts.byTask[t.Id] = term
	slog.Info("terminal started", "task", t.Id, "shell", shell, "pid", cmd.Process.Pid)
	go func() {
		term.pump()
		ts.mu.Lock()
		if ts.byTask[t.Id] == term {
			delete(ts.byTask, t.Id)
		}
		ts.mu.Unlock()
		slog.Info("terminal ended", "task", t.Id, "exit", term.exit)
	}()
	return term, nil
}

func (ts *Terminals) Close(taskID string) {
	ts.mu.Lock()
	term := ts.byTask[taskID]
	delete(ts.byTask, taskID)
	ts.mu.Unlock()
	if term != nil {
		term.kill()
	}
}

func (ts *Terminals) CloseAll() {
	ts.mu.Lock()
	all := ts.byTask
	ts.byTask = map[string]*Terminal{}
	ts.mu.Unlock()
	for _, term := range all {
		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)
	}
	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) WakeOnPrompt(fn func(exit string, out []byte)) {
	term.mu.Lock()
	defer term.mu.Unlock()
	term.wake = &wake{fn: fn}
}

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) 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()
}