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

415 lines
12 KiB
Go

package terminal
import (
"context"
"io"
"log/slog"
"sync"
"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 Source, onData func([]byte), onExit func()) *attachment {
return newTestAttachmentWithOpen(src, nil, onData, onExit)
}
func newTestAttachmentWithOpen(src Source, onOpen func(), onData func([]byte), onExit func()) *attachment {
return newAttachment("t1", ports.RuntimeHandle{ID: "t1"}, src, onOpen, onData, onExit, testLogger())
}
func currentPTY(a *attachment) ports.Stream {
a.mu.Lock()
defer a.mu.Unlock()
return a.pty
}
func TestAttachmentStreamsOutputToSink(t *testing.T) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp}
var sink safeBytes
a := newTestAttachment(src, 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) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp}
a := newTestAttachment(src, 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}
})
}
func TestAttachmentSignalsOpenOnlyAfterPTYIsPublished(t *testing.T) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp}
opened := make(chan bool, 1)
var a *attachment
a = newTestAttachmentWithOpen(src, func() {
opened <- currentPTY(a) == pty
}, nil, nil)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go a.run(ctx)
select {
case sawPTY := <-opened:
if !sawPTY {
t.Fatal("open callback fired before the PTY was published")
}
case <-time.After(time.Second):
t.Fatal("open callback did not fire")
}
}
func TestAttachmentBuffersInputUntilPTYReady(t *testing.T) {
pty := newFakePTY()
spawnStarted := make(chan struct{})
releaseSpawn := make(chan struct{})
src := &fakeSource{alive: true, attachFn: func(context.Context, uint16, uint16) (ports.Stream, error) {
close(spawnStarted)
<-releaseSpawn
return pty, nil
}}
a := newTestAttachment(src, nil, nil)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go a.run(ctx)
select {
case <-spawnStarted:
case <-time.After(time.Second):
t.Fatal("spawn was not reached")
}
if err := a.write([]byte("hello\n")); err != nil {
t.Fatalf("write before PTY ready: %v", err)
}
close(releaseSpawn)
eventually(t, time.Second, func() bool { return string(pty.writtenBytes()) == "hello\n" })
}
// 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) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp}
a := newTestAttachment(src, 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 Stream must be OPENED at the recorded grid, not sized after attach: the
// attach client reads the tty size once at startup, and a post-attach resize
// 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 attach, not the first grid
// forever.
func TestAttachmentSpawnsPTYAtRecordedSize(t *testing.T) {
first := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{first}}
src := &fakeSource{alive: true, spawner: sp}
a := newTestAttachment(src, 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) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp} // alive for the first attach
exited := make(chan struct{})
a := newTestAttachment(src, 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) // runtime 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: a mux
// 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 opening an attach Stream — even on the very first
// open (the original code only checked liveness on re-attach).
func TestAttachmentNeverAttachesToDeadRuntime(t *testing.T) {
sp := &fakeSpawner{}
src := &fakeSource{alive: false, spawner: sp}
exited := make(chan struct{})
a := newTestAttachment(src, 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 attaches", 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) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{aliveErr: io.ErrUnexpectedEOF, spawner: sp}
a := newTestAttachment(src, 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 attaches", 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) {
p1, p2 := newFakePTY(), newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{p1, p2}}
src := &fakeSource{alive: true, spawner: sp} // session still alive -> re-attach on drop
a := newTestAttachment(src, 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() })
}
// A persistent Attach error is treated like the old spawn-error path: it backs
// off and retries up to the failure cap, then reports exited. (The old
// AttachCommand-error path failed immediately; folding the argv build into
// Attach means an attach failure now shares the spawn-failure retry policy,
// which is the correct behavior for a transient dial/exec failure.)
func TestAttachmentFailsWhenAttachErrors(t *testing.T) {
src := &fakeSource{alive: true, attachErr: io.ErrUnexpectedEOF}
exited := make(chan struct{})
a := newTestAttachment(src, nil, func() { close(exited) })
a.maxReattach = 2 // keep the retry budget small so the test is fast
go a.run(context.Background())
select {
case <-exited:
case <-time.After(3 * time.Second):
t.Fatal("expected exit after attach errors exhaust the retry budget")
}
}
// close() is a detach, not a pane death: it must stop the attach loop and kill
// the client's PTY without firing onExit — the runtime session is still alive
// and an exited frame would wrongly flip the client UI to its terminal state.
func TestAttachmentCloseDoesNotFireExit(t *testing.T) {
pty := newFakePTY()
sp := &fakeSpawner{ptys: []*fakePTY{pty}}
src := &fakeSource{alive: true, spawner: sp}
exited := make(chan struct{})
a := newTestAttachment(src, 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 attaches", got)
}
}
type closeOrderPTY struct {
*fakePTY
ctx context.Context
before chan struct{}
after chan struct{}
once sync.Once
}
func (p *closeOrderPTY) Close() error {
p.once.Do(func() {
select {
case <-p.ctx.Done():
close(p.after)
default:
close(p.before)
}
})
return p.fakePTY.Close()
}
func TestAttachmentCloseClosesPTYBeforeCancel(t *testing.T) {
beforeCancel := make(chan struct{})
afterCancel := make(chan struct{})
var spawnCtx context.Context
src := &fakeSource{alive: true, attachFn: func(ctx context.Context, _, _ uint16) (ports.Stream, error) {
spawnCtx = ctx
return &closeOrderPTY{
fakePTY: newFakePTY(),
ctx: ctx,
before: beforeCancel,
after: afterCancel,
}, nil
}}
a := newTestAttachment(src, nil, nil)
done := make(chan struct{})
go func() {
a.run(context.Background())
close(done)
}()
eventually(t, time.Second, func() bool { return currentPTY(a) != nil })
a.close()
select {
case <-beforeCancel:
case <-afterCancel:
t.Fatal("attachment cancelled the run context before closing the PTY")
case <-time.After(time.Second):
t.Fatal("PTY was not closed")
}
select {
case <-spawnCtx.Done():
case <-time.After(time.Second):
t.Fatal("close must still cancel the attach context")
}
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("run must return after close")
}
}