Let task agents read and type into the task's terminal

internal/terminal/terminal.go

@@ -0,0 +1,213 @@
+package terminal
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"log/slog"
+	"os"
+	"os/exec"
+	"sync"
+	"syscall"
+
+	"github.com/creack/pty"
+
+	"g.fc.run/misc/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
+}
+
+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{}}
+	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.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) 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()
+}