byre/internal/terminal/terminal.go

package terminal

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

	"github.com/creack/pty"

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

const (
	backlogSize = 256 << 10
	clientBuf   = 256
	MaxSize     = 1 << 15
)

type Terminals struct {
	mu     sync.Mutex
	byTask map[string]*Terminal
}

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

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

func New() *Terminals {
	return &Terminals{byTask: map[string]*Terminal{}}
}

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