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

2026-10-18 15:05:28 -07:00
parent 2658dcbfc4
commit 1ef3594016
8 changed files with 480 additions and 201 deletions
@@ -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()
+}
@@ -0,0 +1,87 @@
+package terminal
+
+import (
+	"strings"
+	"unicode/utf8"
+)
+
+const textLimit = 64 << 10
+
+func Text(b []byte) string {
+	out := strings.Builder{}
+	line := []rune{}
+	col := 0
+	flush := func() {
+		out.WriteString(strings.TrimRight(string(line), " "))
+		out.WriteByte('\n')
+		line = line[:0]
+		col = 0
+	}
+	for i := 0; i < len(b); {
+		c := b[i]
+		switch {
+		case c == 0x1b:
+			end, final := skipEscape(b, i)
+			if final == 'K' {
+				line = line[:min(col, len(line))]
+			}
+			i = end + 1
+			continue
+		case c == '\n':
+			flush()
+		case c == '\r':
+			col = 0
+		case c == '\b':
+			col = max(0, col-1)
+		case c == '\t' || c >= 0x20 && c != 0x7f:
+			r, n := utf8.DecodeRune(b[i:])
+			for len(line) < col {
+				line = append(line, ' ')
+			}
+			if col < len(line) {
+				line[col] = r
+			} else {
+				line = append(line, r)
+			}
+			col++
+			i += n
+			continue
+		}
+		i++
+	}
+	s := out.String() + strings.TrimRight(string(line), " ")
+	if len(s) > textLimit {
+		s = s[len(s)-textLimit:]
+		s = s[strings.IndexByte(s, '\n')+1:]
+	}
+	return s
+}
+
+func skipEscape(b []byte, i int) (int, byte) {
+	if i+1 >= len(b) {
+		return i, 0
+	}
+	i++
+	switch b[i] {
+	case '[':
+		for i++; i < len(b); i++ {
+			if b[i] >= 0x40 && b[i] <= 0x7e {
+				return i, b[i]
+			}
+		}
+		return len(b) - 1, 0
+	case ']', 'P', '_', '^':
+		for i++; i < len(b); i++ {
+			if b[i] == 0x07 {
+				return i, 0
+			}
+			if b[i] == 0x1b && i+1 < len(b) && b[i+1] == '\\' {
+				return i + 1, 0
+			}
+		}
+		return len(b) - 1, 0
+	case '(', ')', '*', '+', '#', '%':
+		return min(i+1, len(b)-1), 0
+	}
+	return i, 0
+}