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

package terminal

import (
	"context"
	"io"
	"log/slog"
	"testing"
	"time"

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

func testLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }

func newTestAttachment(src PTYSource, spawn spawnFunc, onData func([]byte), onExit func()) *attachment {
	return newAttachment("t1", ports.RuntimeHandle{ID: "t1"}, src, spawn, onData, onExit, testLogger())
}

func TestAttachmentStreamsOutputToSink(t *testing.T) {
	src := &fakeSource{alive: true}
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}

	var sink safeBytes
	a := newTestAttachment(src, sp.spawn, sink.add, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	pty.push([]byte("hello"))
	eventually(t, time.Second, func() bool { return sink.string() == "hello" })
}

func TestAttachmentWriteAndResizeReachPTY(t *testing.T) {
	src := &fakeSource{alive: true}
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}
	a := newTestAttachment(src, sp.spawn, nil, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	eventually(t, time.Second, func() bool { return a.write([]byte("ls\n")) == nil })
	eventually(t, time.Second, func() bool { return string(pty.writtenBytes()) == "ls\n" })

	if err := a.resize(24, 80); err != nil {
		t.Fatalf("resize: %v", err)
	}
	eventually(t, time.Second, func() bool {
		rs := pty.resizeCalls()
		return len(rs) == 1 && rs[0] == [2]uint16{24, 80}
	})
}

// A size requested before the PTY exists (the open frame's cols/rows, or a
// resize racing the attach) must not be lost: the attach applies it the moment
// the PTY is up, instead of leaving the pane at the kernel default grid.
func TestAttachmentAppliesRequestedSizeOnAttach(t *testing.T) {
	src := &fakeSource{alive: true}
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}
	a := newTestAttachment(src, sp.spawn, nil, nil)

	if err := a.resize(30, 100); err != nil {
		t.Fatalf("resize before attach: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	eventually(t, time.Second, func() bool {
		rs := pty.resizeCalls()
		return len(rs) == 1 && rs[0] == [2]uint16{30, 100}
	})
}

// The PTY must be SPAWNED at the recorded grid, not sized after exec: the
// attach process (zellij) reads the tty size once at startup, and a post-spawn
// TIOCSWINSZ depends on SIGWINCH delivery that can race the client installing
// its handler — a missed signal left the session laid out for the previous
// client's size (the "terminal doesn't repaint after a resize" desync). Also
// covers re-attach: a later resize must reach the NEXT spawn, not the first
// grid forever.
func TestAttachmentSpawnsPTYAtRecordedSize(t *testing.T) {
	src := &fakeSource{alive: true}
	first := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{first}}
	a := newTestAttachment(src, sp.spawn, nil, nil)
	a.resetGrace = time.Hour // keep the failure counter deterministic

	if err := a.resize(37, 115); err != nil {
		t.Fatalf("resize before attach: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	eventually(t, time.Second, func() bool {
		ss := sp.spawnSizes()
		return len(ss) == 1 && ss[0] == [2]uint16{37, 115}
	})

	// Client resized, then the PTY dropped: the re-attach spawn must start at
	// the latest grid.
	if err := a.resize(40, 148); err != nil {
		t.Fatalf("resize while attached: %v", err)
	}
	_ = first.Close()

	eventually(t, time.Second, func() bool {
		ss := sp.spawnSizes()
		return len(ss) == 2 && ss[1] == [2]uint16{40, 148}
	})
}

func TestAttachmentSkipsReattachOnCleanExit(t *testing.T) {
	src := &fakeSource{alive: true} // alive for the first attach
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}

	exited := make(chan struct{})
	a := newTestAttachment(src, sp.spawn, nil, func() { close(exited) })

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	eventually(t, time.Second, func() bool { return sp.calls() == 1 })
	src.setAlive(false) // Zellij session gone -> no re-attach
	pty.Close()         // pane ends
	select {
	case <-exited:
	case <-time.After(time.Second):
		t.Fatal("expected exit callback after clean pane exit")
	}
	if got := sp.calls(); got != 1 {
		t.Fatalf("expected exactly one attach, got %d", got)
	}
}

// TestAttachmentNeverAttachesToDeadRuntime covers the resurrection bug: `zellij
// attach` on a killed-but-cached session resurrects it, re-running the agent
// command. An attachment whose runtime probes definitively dead must therefore
// report exited WITHOUT ever spawning an attach PTY — even on the very first
// open (the original code only checked liveness on re-attach).
func TestAttachmentNeverAttachesToDeadRuntime(t *testing.T) {
	src := &fakeSource{alive: false}
	sp := &fakeSpawner{}

	exited := make(chan struct{})
	a := newTestAttachment(src, sp.spawn, nil, func() { close(exited) })

	go a.run(context.Background())
	select {
	case <-exited:
	case <-time.After(time.Second):
		t.Fatal("expected exit when runtime is dead before first attach")
	}
	if got := sp.calls(); got != 0 {
		t.Fatalf("attach must never run against a dead runtime, got %d spawns", got)
	}
}

// TestAttachmentRetriesProbeErrorsBeforeAttaching pins the hard rule that a
// failed liveness probe is NOT proof of death: a transient probe error must
// not flip the terminal to exited, and the attach proceeds once the probe
// recovers.
func TestAttachmentRetriesProbeErrorsBeforeAttaching(t *testing.T) {
	src := &fakeSource{aliveErr: io.ErrUnexpectedEOF}
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}
	a := newTestAttachment(src, sp.spawn, nil, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	// While probes error the attachment must neither exit nor attach.
	time.Sleep(50 * time.Millisecond)
	if a.isExited() {
		t.Fatal("probe error must not be treated as runtime death")
	}
	if got := sp.calls(); got != 0 {
		t.Fatalf("attach must wait for a successful probe, got %d spawns", got)
	}

	// Probe recovers -> the attach goes through.
	src.setAliveResult(true, nil)
	eventually(t, 2*time.Second, func() bool { return sp.calls() == 1 })
	if a.isExited() {
		t.Fatal("attachment exited despite a live runtime")
	}
}

func TestAttachmentReattachesWhileSessionAlive(t *testing.T) {
	src := &fakeSource{alive: true} // session still alive -> re-attach on drop
	p1, p2 := newFakePTY(), newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{p1, p2}}
	a := newTestAttachment(src, sp.spawn, nil, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go a.run(ctx)

	eventually(t, time.Second, func() bool { return sp.calls() >= 1 })
	if err := a.resize(24, 80); err != nil {
		t.Fatalf("resize: %v", err)
	}
	p1.Close() // first attach drops
	eventually(t, 2*time.Second, func() bool { return sp.calls() >= 2 })

	// The client's grid survives the re-attach: the fresh PTY is sized to the
	// last requested grid without waiting for the client to resize again.
	eventually(t, time.Second, func() bool {
		for _, rs := range p2.resizeCalls() {
			if rs == [2]uint16{24, 80} {
				return true
			}
		}
		return false
	})

	// Now the session is gone: the next drop must not re-attach.
	src.setAlive(false)
	p2.Close()
	eventually(t, 2*time.Second, func() bool { return a.isExited() })
}

func TestAttachmentFailsWhenAttachCommandErrors(t *testing.T) {
	src := &fakeSource{alive: true, attachErr: io.ErrUnexpectedEOF}
	sp := &fakeSpawner{}

	exited := make(chan struct{})
	a := newTestAttachment(src, sp.spawn, nil, func() { close(exited) })

	go a.run(context.Background())
	select {
	case <-exited:
	case <-time.After(time.Second):
		t.Fatal("expected exit when attach command fails")
	}
	if sp.calls() != 0 {
		t.Fatalf("spawn should not run when attach command errors, got %d calls", sp.calls())
	}
}

// close() is a detach, not a pane death: it must stop the attach loop and kill
// the client's PTY without firing onExit — the Zellij session is still alive
// and an exited frame would wrongly flip the client UI to its terminal state.
func TestAttachmentCloseDoesNotFireExit(t *testing.T) {
	src := &fakeSource{alive: true}
	pty := newFakePTY()
	sp := &fakeSpawner{ptys: []*fakePTY{pty}}

	exited := make(chan struct{})
	a := newTestAttachment(src, sp.spawn, nil, func() { close(exited) })

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan struct{})
	go func() {
		a.run(ctx)
		close(done)
	}()

	eventually(t, time.Second, func() bool { return sp.calls() == 1 })
	a.close()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("run must return after close")
	}
	select {
	case <-exited:
		t.Fatal("close must not fire onExit")
	default:
	}
	if got := sp.calls(); got != 1 {
		t.Fatalf("close must stop re-attaching, got %d spawns", got)
	}
}