agent-orchestrator/backend/internal/terminal/attachment.go

package terminal

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"sync"
	"time"

	"github.com/aoagents/agent-orchestrator/backend/internal/ports"
)

// PTYSource is what a terminal needs from the runtime: the argv that attaches a
// PTY to a session's pane (plus any env that argv needs but is not in the
// daemon's process env — e.g. a per-session ZELLIJ_SOCKET_DIR on Windows),
// and a liveness check used to decide whether a dropped PTY should be
// re-attached or treated as a clean exit. The Zellij runtime adapter
// satisfies this via AttachCommand/IsAlive; the interface lives here, next to
// its only consumer, so terminal does not depend on a concrete adapter.
type PTYSource interface {
	AttachCommand(handle ports.RuntimeHandle) (argv []string, env []string, err error)
	IsAlive(ctx context.Context, handle ports.RuntimeHandle) (bool, error)
}

// ptyProcess is a started PTY-backed attach process. It is the injection seam
// that keeps the attach loop testable without a real process: unit tests supply
// a scripted in-memory implementation; production uses a creack/pty-backed one
// (see pty_unix.go).
type ptyProcess interface {
	io.ReadWriteCloser
	Resize(rows, cols uint16) error
}

// spawnFunc starts a PTY for argv, sized rows×cols from the start (zero means
// no size was recorded yet — kernel default). Spawning at the client's grid
// matters: the attach process reads the tty size once at startup, and sizing
// the PTY only after exec relies on SIGWINCH delivery that can race the
// process installing its handler — a missed signal leaves the zellij client
// laid out for a stale size. env, when non-nil, is the full env block for the
// child (mirrors exec.Cmd.Env: nil means inherit). ctx cancellation must
// terminate the process.
type spawnFunc func(ctx context.Context, argv []string, env []string, rows, cols uint16) (ptyProcess, error)

// reattach policy: a PTY that drops is re-attached while the underlying Zellij
// session is still alive, up to maxReattach consecutive failures. An attach that
// survived longer than reattachResetGrace before dropping resets the counter, so
// a long-lived pane that blips recovers but a tight crash-loop gives up.
const (
	defaultMaxReattach       = 5
	defaultReattachResetTime = 5 * time.Second
)

// attachment is ONE client's hold on a pane: a private `zellij attach` PTY
// spawned per mux open, streaming to a single sink. Zellij is the multiplexer —
// it owns the session's screen state and scrollback, and answers every fresh
// attach with its init handshake (alt screen, bracketed paste, and other terminal
// modes enabled by the embedded client options) followed by a faithful repaint.
// That handshake is why the PTY is per-client and there is no server-side replay
// buffer: a byte ring can replay recent output, but the one-time mode negotiation
// at the head of the stream scrolls out of any bounded buffer. A fresh attach per
// client makes Zellij re-send it, every time, by construction.
//
// onOpen fires once the attach PTY is actually ready to accept input. onData
// must not block: the WS layer funnels frames onto its own buffered writer.
// onExit fires at most once, when the attach loop gives up (runtime dead,
// attach failure cap) — never on close().
type attachment struct {
	id     string
	handle ports.RuntimeHandle
	src    PTYSource
	spawn  spawnFunc
	log    *slog.Logger
	onOpen func()
	onData func(data []byte)
	onExit func()

	maxReattach int
	resetGrace  time.Duration

	mu           sync.Mutex
	pty          ptyProcess
	cancel       context.CancelFunc
	rows         uint16 // last size the client asked for; re-applied on every attach
	cols         uint16
	closed       bool
	exited       bool
	opened       bool
	inputReady   bool
	pendingInput [][]byte
}

func newAttachment(id string, handle ports.RuntimeHandle, src PTYSource, spawn spawnFunc, onOpen func(), onData func([]byte), onExit func(), log *slog.Logger) *attachment {
	if log == nil {
		log = slog.Default()
	}
	if onData == nil {
		onData = func([]byte) {}
	}
	return &attachment{
		id:          id,
		handle:      handle,
		src:         src,
		spawn:       spawn,
		log:         log,
		onOpen:      onOpen,
		onData:      onData,
		onExit:      onExit,
		maxReattach: defaultMaxReattach,
		resetGrace:  defaultReattachResetTime,
	}
}

// run drives attach → read-loop → re-attach until the pane exits cleanly, the
// attachment is closed, or ctx is cancelled. It is started once per attachment.
func (a *attachment) run(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	if !a.setRunCancel(cancel) {
		cancel()
		return
	}
	defer a.clearRunCancel(cancel)

	failures := 0
	for {
		if a.shouldStop(ctx) {
			return
		}

		// Gate EVERY attach (including the first) on the runtime actually
		// being alive. `zellij attach` resurrects EXITED sessions — re-running
		// the serialized agent command — so attaching to a dead handle would
		// re-create a runtime the daemon already destroyed, outside lifecycle
		// control. A definitive "not alive" is a clean exit. A probe ERROR is
		// not proof of death: it retries with backoff up to the same
		// consecutive-failure cap as attach failures.
		alive, err := a.src.IsAlive(ctx, a.handle)
		if a.shouldStop(ctx) {
			return
		}
		if err != nil {
			failures++
			if failures > a.maxReattach {
				a.fail("liveness probe: " + err.Error())
				return
			}
			if !a.backoff(ctx, failures) {
				return
			}
			continue
		}
		if !alive {
			a.markExited()
			return
		}

		argv, env, err := a.src.AttachCommand(a.handle)
		if a.shouldStop(ctx) {
			return
		}
		if err != nil {
			a.fail("attach command: " + err.Error())
			return
		}
		rows, cols := a.size()
		if a.shouldStop(ctx) {
			return
		}
		p, err := a.spawn(ctx, argv, env, rows, cols)
		if a.shouldStop(ctx) {
			if p != nil {
				_ = p.Close()
			}
			return
		}
		if err != nil {
			failures++
			if failures > a.maxReattach {
				a.fail("spawn pty: " + err.Error())
				return
			}
			if !a.backoff(ctx, failures) {
				return
			}
			continue
		}

		if !a.setPTY(p) {
			_ = p.Close()
			return
		}
		start := time.Now()
		a.copyOut(p)
		a.clearPTY(p)
		_ = p.Close()
		if a.shouldStop(ctx) {
			return
		}

		if time.Since(start) >= a.resetGrace {
			failures = 0
		}
		failures++

		if failures > a.maxReattach {
			a.markExited()
			return
		}
		if !a.backoff(ctx, failures) {
			return
		}
		a.log.Debug("terminal re-attaching", "id", a.id, "failures", failures)
	}
}

// copyOut pumps PTY output to the sink until the PTY closes or errors.
func (a *attachment) copyOut(p ptyProcess) {
	buf := make([]byte, 32*1024)
	for {
		n, err := p.Read(buf)
		if n > 0 {
			chunk := make([]byte, n)
			copy(chunk, buf[:n])
			a.onData(chunk)
		}
		if err != nil {
			return
		}
	}
}

// backoff sleeps between attach attempts; false means ctx was cancelled.
// Whether another attempt is warranted at all (liveness, failure cap) is
// decided at the top of the run loop, so a re-attach and a first attach share
// one gate.
func (a *attachment) backoff(ctx context.Context, failures int) bool {
	select {
	case <-ctx.Done():
		return false
	case <-time.After(reattachBackoff(failures)):
		return true
	}
}

func reattachBackoff(failures int) time.Duration {
	d := time.Duration(failures) * 200 * time.Millisecond
	if d > time.Second {
		d = time.Second
	}
	return d
}

// write sends client keystrokes to the PTY. Input that arrives after open but
// before the attach PTY is published is buffered and flushed as soon as setPTY
// runs, so a fast user cannot type into the attach race and lose bytes.
func (a *attachment) write(p []byte) error {
	if len(p) == 0 {
		return nil
	}
	chunk := append([]byte(nil), p...)

	a.mu.Lock()
	if a.closed || a.exited {
		a.mu.Unlock()
		return errors.New("terminal: attachment closed")
	}
	pty := a.pty
	if pty == nil || !a.inputReady {
		a.pendingInput = append(a.pendingInput, chunk)
		a.mu.Unlock()
		return nil
	}
	a.mu.Unlock()
	_, err := pty.Write(chunk)
	return err
}

// resize records the client's grid and applies it to the live PTY. The size is
// remembered so an attach that is still in flight (or a later re-attach) starts
// at the client's grid instead of the kernel default — the open frame's
// cols/rows land here before the PTY exists.
func (a *attachment) resize(rows, cols uint16) error {
	a.mu.Lock()
	a.rows, a.cols = rows, cols
	pty := a.pty
	a.mu.Unlock()
	if pty == nil {
		return nil
	}
	return pty.Resize(rows, cols)
}

// size returns the client's last requested grid (zero before the first
// open/resize recorded one). The spawn path reads it so the PTY starts at the
// client's grid instead of the kernel default.
func (a *attachment) size() (rows, cols uint16) {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.rows, a.cols
}

// setPTY publishes a freshly attached PTY and replays the client's last
// requested size onto it (see resize) — the spawn already started at the size
// read in run, but a resize frame can land between that read and registration
// here; the replay (Setsize + explicit WINCH) converges the late case.
func (a *attachment) setPTY(p ptyProcess) bool {
	a.mu.Lock()
	if a.closed || a.exited {
		a.mu.Unlock()
		return false
	}
	a.pty = p
	a.inputReady = false
	rows, cols := a.rows, a.cols
	shouldOpen := !a.opened
	if shouldOpen {
		a.opened = true
	}
	onOpen := a.onOpen
	a.mu.Unlock()
	if rows > 0 && cols > 0 {
		_ = p.Resize(rows, cols)
	}
	if shouldOpen && onOpen != nil {
		onOpen()
	}

	for {
		a.mu.Lock()
		pending := append([][]byte(nil), a.pendingInput...)
		a.pendingInput = nil
		if len(pending) == 0 {
			a.inputReady = true
			a.mu.Unlock()
			return true
		}
		a.mu.Unlock()

		for _, chunk := range pending {
			if _, err := p.Write(chunk); err != nil {
				a.fail("flush pending input: " + err.Error())
				return false
			}
		}
	}
}

func (a *attachment) clearPTY(p ptyProcess) {
	a.mu.Lock()
	if a.pty == p {
		a.pty = nil
		a.inputReady = false
	}
	a.mu.Unlock()
}

// close detaches this client: stop re-attaching and kill the attach PTY. It
// never touches the Zellij session itself, which the zellij server keeps alive
// for other clients.
func (a *attachment) close() {
	a.mu.Lock()
	if a.closed {
		a.mu.Unlock()
		return
	}
	a.closed = true
	pty := a.pty
	a.pty = nil
	a.inputReady = false
	a.pendingInput = nil
	cancel := a.cancel
	a.mu.Unlock()
	if pty != nil {
		_ = pty.Close()
	}
	if cancel != nil {
		cancel()
	}
}

func (a *attachment) setRunCancel(cancel context.CancelFunc) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.closed {
		return false
	}
	a.cancel = cancel
	return true
}

func (a *attachment) clearRunCancel(cancel context.CancelFunc) {
	a.mu.Lock()
	a.cancel = nil
	a.mu.Unlock()
	cancel()
}

func (a *attachment) isClosed() bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.closed
}

func (a *attachment) shouldStop(ctx context.Context) bool {
	return ctx.Err() != nil || a.isClosed()
}

func (a *attachment) isExited() bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.exited
}

// markExited flips the attachment to exited and fires onExit once.
func (a *attachment) markExited() {
	a.mu.Lock()
	if a.exited {
		a.mu.Unlock()
		return
	}
	a.exited = true
	a.mu.Unlock()
	if a.onExit != nil {
		a.onExit()
	}
}

// fail reports an unrecoverable attach error as an exit.
func (a *attachment) fail(reason string) {
	a.log.Warn("terminal attachment failed", "id", a.id, "reason", reason)
	a.markExited()
}