From 9a9c655c75ead06738974d0b7fad08e184d68ea8 Mon Sep 17 00:00:00 2001 From: Amit Kumar Date: Sun, 26 Jul 2026 04:47:31 +0000 Subject: [PATCH] fix(host): flush queued output before closing clients on shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pumpPTY has already returned by the time shutdown runs, so everything the agent ever wrote is sitting in each client's queue. Closing the connection there truncated the agent's last screen, and cutting a frame in half reached the viewer as "attach output: unexpected EOF" with a non-zero exit instead of a clean "[uam: session ended]". Clients now carry a second signal: done still means stop now, flush means finish first. Shutdown asks every client to flush, then waits up to shutdownFlushWindow for each writer to drain its queue and close the connection itself. A viewer that has stopped reading cannot hold teardown open — the same window bounds both the socket write deadline and the wait — and a client with no writer running is force-closed when the window expires. Verified end to end against the built binary: an agent that prints 60 lines and exits while a client is attached now delivers every line, prints the session-ended note, and the attach client exits 0. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01X3PrstXFvYdemQrz8Vb1b1 --- internal/session/client_registry.go | 22 ++- internal/session/host.go | 67 +++++++- internal/session/host_shutdown_flush_test.go | 153 +++++++++++++++++++ 3 files changed, 232 insertions(+), 10 deletions(-) create mode 100644 internal/session/host_shutdown_flush_test.go diff --git a/internal/session/client_registry.go b/internal/session/client_registry.go index 7441de7..f95eba7 100644 --- a/internal/session/client_registry.go +++ b/internal/session/client_registry.go @@ -24,9 +24,14 @@ type serverMessage struct { } type attachClient struct { - conn net.Conn - out chan serverMessage - done chan struct{} + conn net.Conn + out chan serverMessage + done chan struct{} + // flush is closed by shutdown to ask this client's writer to drain what is + // queued and then close the connection itself. done means "stop now"; + // flush means "finish, then stop". + flush chan struct{} + flushOnce sync.Once version protocolVersion id string requestedRole clientRole @@ -49,6 +54,17 @@ func (client *attachClient) drop() { }) } +// requestFlush asks the client's writer to drain its queue and close. It is a +// no-op for a client that has no writer running; shutdown falls back to drop +// once its flush window expires. +func (client *attachClient) requestFlush() { + client.flushOnce.Do(func() { + if client.flush != nil { + close(client.flush) + } + }) +} + type clientRegistration struct { requestedRole clientRole hello clientHello diff --git a/internal/session/host.go b/internal/session/host.go index 03c49b0..f1024a5 100644 --- a/internal/session/host.go +++ b/internal/session/host.go @@ -55,6 +55,12 @@ const killGrace = 1500 * time.Millisecond // disconnected rather than allowed to stall the session. const attachBufFrames = 512 +// shutdownFlushWindow bounds how long teardown waits for attached viewers to +// receive the agent's final output. Long enough for a queued screen on a slow +// link, short enough that a viewer which stopped reading cannot hold the host +// open. +const shutdownFlushWindow = 250 * time.Millisecond + const ( markClosedRetryWindow = 2 * time.Second markClosedRetryBase = 25 * time.Millisecond @@ -464,7 +470,8 @@ func (h *host) handleAttach(conn net.Conn, br *bufio.Reader, req request) { return } client := &attachClient{ - conn: conn, out: make(chan serverMessage, attachBufFrames), done: make(chan struct{}), version: version, + conn: conn, out: make(chan serverMessage, attachBufFrames), done: make(chan struct{}), + flush: make(chan struct{}), version: version, fallback: version == protocolV1 && !req.versionPresent, } attachResponse, err := h.registerAttachClient(client, registration) @@ -491,6 +498,9 @@ func (h *host) attachWriter(client *attachClient) { select { case <-client.done: return + case <-client.flush: + h.flushClient(client) + return case message := <-client.out: if err := h.writeServerMessage(client, message); err != nil { h.dropClientReason(client, "connection_write") @@ -500,6 +510,28 @@ func (h *host) attachWriter(client *attachClient) { } } +// flushClient writes what the shutting-down client has already queued and then +// closes the connection itself. The write deadline bounds a viewer that has +// stopped reading, so a stalled socket cannot hold up host teardown. +func (h *host) flushClient(client *attachClient) { + defer client.drop() + if client.conn != nil { + if err := client.conn.SetWriteDeadline(time.Now().Add(shutdownFlushWindow)); err != nil { + return + } + } + for { + select { + case message := <-client.out: + if err := h.writeServerMessage(client, message); err != nil { + return + } + default: + return + } + } +} + func (h *host) writeServerMessage(client *attachClient, message serverMessage) error { if client.version == protocolV1 { if message.kind != serverFramePTY { @@ -659,6 +691,17 @@ func (h *host) signalChild(sig syscall.Signal) { // closed (the native replacement for the tmux session-closed hook), tell any // attached clients, and remove the runtime files. func (h *host) shutdown(exitCode int) { + h.shutdownClients() + providerID := readProviderIdentityHandoff(h.dir, h.name, h.providerIdentityFile) + if err := removeSessionFiles(h.dir, h.name); err != nil { + log.Warn("remove session files failed", "session", h.name, "error", err) + } + h.recordExit(exitCode, providerID) +} + +// shutdownClients releases every attached client, giving each writer a bounded +// window to put the agent's last output on the wire first. +func (h *host) shutdownClients() { h.mu.Lock() clients := h.registry.drain() h.mu.Unlock() @@ -667,13 +710,23 @@ func (h *host) shutdown(exitCode int) { Event: "attach.lifecycle", Session: h.name, ClientID: client.id, Protocol: int(client.version), Role: string(client.assignedRole), Reason: "host_shutdown", }) - client.drop() - } - providerID := readProviderIdentityHandoff(h.dir, h.name, h.providerIdentityFile) - if err := removeSessionFiles(h.dir, h.name); err != nil { - log.Warn("remove session files failed", "session", h.name, "error", err) + client.requestFlush() + } + // pumpPTY has already returned, so everything the agent ever wrote is + // queued; each writer only needs the chance to put it on the wire. Closing + // the connection first — as this used to — truncated the agent's last + // screen, and cutting a frame in half turned a clean exit into + // "attach output: unexpected EOF" on the viewer. + flushBy := time.Now().Add(shutdownFlushWindow) + for _, client := range clients { + select { + case <-client.done: + case <-time.After(time.Until(flushBy)): + // No writer is running for this client, or it is wedged behind a + // stalled socket. Either way the deadline is the whole budget. + client.drop() + } } - h.recordExit(exitCode, providerID) } func (h *host) recordExit(exitCode int, providerID string) { diff --git a/internal/session/host_shutdown_flush_test.go b/internal/session/host_shutdown_flush_test.go new file mode 100644 index 0000000..f66fedd --- /dev/null +++ b/internal/session/host_shutdown_flush_test.go @@ -0,0 +1,153 @@ +package session + +import ( + "bufio" + "bytes" + "net" + "testing" + "time" +) + +// Output still queued when the host tears down must reach the viewer. The +// distinction the writer has to honour: `done` means stop now, `flush` means +// finish first. Dropping the client outright — as shutdown used to — truncated +// the agent's last screen, and a half-written frame surfaced on the viewer as +// "attach output: unexpected EOF" instead of a clean session end. +func TestWriterFlushesQueuedOutputOnShutdownSignal(t *testing.T) { + // Given three frames queued and a shutdown already requested. + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := newFlushTestClient(t, h, server) + payloads := [][]byte{[]byte("first\r\n"), []byte("second\r\n"), []byte("agent final line\r\n")} + for _, payload := range payloads { + if !h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: payload}) { + t.Fatal("queueing the agent's output failed") + } + } + client.requestFlush() + + // When the writer runs. + go h.attachWriter(client) + + // Then every queued frame arrives before the connection closes. + if err := viewer.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil { + t.Fatal(err) + } + reader := bufio.NewReader(viewer) + for i, want := range payloads { + kind, payload, err := readFrame(reader) + if err != nil { + t.Fatalf("frame %d lost on shutdown: %v", i, err) + } + if kind != serverFramePTY || !bytes.Equal(payload, want) { + t.Fatalf("frame %d = %d %q, want %d %q", i, kind, payload, serverFramePTY, want) + } + } + waitClosed(t, client) +} + +// The contrast that makes the distinction meaningful: a client told to stop +// now discards what is queued. +func TestWriterStopsImmediatelyOnDone(t *testing.T) { + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := newFlushTestClient(t, h, server) + if err := viewer.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: []byte("discarded")}) + client.drop() + + go h.attachWriter(client) + + if _, _, err := readFrame(bufio.NewReader(viewer)); err == nil { + t.Fatal("a dropped client still wrote its queue") + } +} + +func newFlushTestClient(t *testing.T, h *host, conn net.Conn) *attachClient { + t.Helper() + client := &attachClient{ + conn: conn, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + client.ready = true + return client +} + +func waitClosed(t *testing.T, client *attachClient) { + t.Helper() + select { + case <-client.done: + case <-time.After(3 * time.Second): + t.Fatal("client was not closed after the flush") + } +} + +// A viewer that has stopped reading must not hold teardown open. +func TestShutdownDoesNotWaitForeverOnAStalledViewer(t *testing.T) { + // Given a client whose peer never reads. + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := &attachClient{ + conn: server, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + client.ready = true + go h.attachWriter(client) + h.enqueueClient(client, serverMessage{kind: serverFramePTY, payload: bytes.Repeat([]byte("x"), 4096)}) + + // When the host tears down. + start := time.Now() + h.shutdownClients() + + // Then it returns within the flush budget and the client is closed. + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Fatalf("shutdown blocked for %s on a stalled viewer", elapsed) + } + select { + case <-client.done: + default: + t.Fatal("stalled client was not force-closed") + } +} + +// A registered client whose writer never started still has to be closed. +func TestShutdownClosesClientsWithoutAWriter(t *testing.T) { + server, viewer := net.Pipe() + t.Cleanup(func() { _ = viewer.Close() }) + h := &host{name: "uam-fake-11112222", registry: newClientRegistry()} + client := &attachClient{ + conn: server, out: make(chan serverMessage, attachBufFrames), + done: make(chan struct{}), flush: make(chan struct{}), version: protocolV2, + } + if err := h.registry.register(client, clientRegistration{ + requestedRole: roleController, hello: validTestHello(), + size: terminalSize{cols: 80, rows: 24}, + }); err != nil { + t.Fatal(err) + } + + h.shutdownClients() + + select { + case <-client.done: + default: + t.Fatal("client with no writer was left open") + } +}