Wake agents when the human runs a command they left in the terminal

internal/terminal/terminal.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"sync"
 	"syscall"
 	"time"
@@ -20,16 +21,33 @@ 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
+	ptmx    *os.File
+	cmd     *exec.Cmd
+	Prompts bool
 
 	mu         sync.Mutex
 	backlog    []byte
@@ -38,10 +56,24 @@ type Terminal struct {
 	exited     bool
 	started    time.Time
 	lastOutput time.Time
+	wake       *wake
 }
 
-func New() *Terminals {
-	return &Terminals{byTask: map[string]*Terminal{}}
+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 {
@@ -66,12 +98,16 @@ func (ts *Terminals) Attach(t api.Task, size api.TerminalSize) (*Terminal, error
 	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, clients: map[chan []byte]bool{}, started: time.Now()}
+	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() {
@@ -148,6 +184,9 @@ func (term *Terminal) output(b []byte) {
 	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:])
@@ -162,6 +201,38 @@ func (term *Terminal) output(b []byte) {
 	}
 }
 
+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()