Fix PTY attach output ordering

server/cmd/api/api/process.go

@@ -38,6 +38,9 @@ type processHandle struct {
 	stdout     io.ReadCloser
 	stderr     io.ReadCloser
 	ptyFile    *os.File
+	ptyOut     ptyOutput
+	ptyDrained chan struct{}
+	ptyStop    chan struct{}
 	isTTY      bool
 	output     *processOutput
 	reason     string
@@ -710,32 +713,28 @@ func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Reques
 		}
 	}()
 
-	// Goroutine: read from PTY (stdout), write to WebSocket
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		err := ptyio.ReadPTYToWriter(h.ptyFile, func(data []byte) error {
-			select {
-			case writeCh <- wsWriteOp{msgType: websocket.MessageBinary, data: data}:
-				return nil
-			case <-done:
-				return io.EOF // Signal to stop reading
-			}
-		}, done)
-		if err != nil {
-			log.Error("pty read error", "err", err)
+	// Replay output already drained from the PTY, then forward live output. The
+	// ptyOutput lock keeps replay ahead of any writes that arrive concurrently.
+	h.ptyOut.attach(func(data []byte) error {
+		select {
+		case writeCh <- wsWriteOp{msgType: websocket.MessageBinary, data: data}:
+			return nil
+		case <-done:
+			return io.EOF
 		}
-		shutdown()
-	}()
+	})
 
-	// Goroutine: watch for process exit and send exit code.
+	// Goroutine: watch for process exit and send exit code after PTY output drains.
 	// This must send the exit code BEFORE triggering shutdown so the writer can deliver it.
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		select {
 		case <-h.doneCh:
-			// Process exited - send exit code before shutdown
+			// The process waiter closes doneCh only after PTY output has drained.
+			// Flush the attached sink before enqueueing the exit event so replayed
+			// and live output stay ahead of it even when the socket was backpressured.
+			h.ptyOut.flush()
 			h.mu.RLock()
 			exitCode := h.exitCode
 			h.mu.RUnlock()
@@ -760,6 +759,8 @@ func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Reques
 	// Wait for shutdown signal
 	<-done
 
+	h.ptyOut.detach()
+
 	// Wait for all goroutines to complete before clearing attachActive.
 	// This prevents a race where a new client attaches while goroutines are still running.
 	wg.Wait()

server/cmd/api/api/process_attach_test.go

@@ -0,0 +1,73 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	oapi "github.com/kernel/kernel-images/server/lib/oapi"
+	"github.com/kernel/kernel-images/server/lib/scaletozero"
+	"github.com/kernel/kernel-images/server/lib/wsdrain"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProcessAttachReplaysOutputThenExits(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()}
+
+	// The first line is written before any client attaches, so it must come back as replay.
+	args := []string{"-c", "echo first; sleep 0.3; echo second"}
+	allocateTTY := true
+	body := &oapi.ProcessSpawnRequest{Command: "sh", Args: &args, AllocateTty: &allocateTTY}
+	spawnResp, err := svc.ProcessSpawn(ctx, oapi.ProcessSpawnRequestObject{Body: body})
+	require.NoError(t, err)
+	s200, ok := spawnResp.(oapi.ProcessSpawn200JSONResponse)
+	require.True(t, ok, "unexpected spawn resp: %T", spawnResp)
+	id := s200.ProcessId.String()
+
+	time.Sleep(100 * time.Millisecond)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		svc.HandleProcessAttachWS(w, r, id, wsdrain.New())
+	}))
+	defer srv.Close()
+
+	conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(srv.URL, "http"), nil)
+	require.NoError(t, err)
+	defer conn.CloseNow()
+
+	var out strings.Builder
+	exitCode := -1
+	readCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	for {
+		msgType, data, err := conn.Read(readCtx)
+		require.NoError(t, err, "attach ended before exit message; output so far %q", out.String())
+		if msgType == websocket.MessageText {
+			var ctrl struct {
+				Type     string `json:"type"`
+				ExitCode *int   `json:"exitCode"`
+			}
+			require.NoError(t, json.Unmarshal(data, &ctrl))
+			if ctrl.Type == "exit" {
+				require.NotNil(t, ctrl.ExitCode)
+				exitCode = *ctrl.ExitCode
+				break
+			}
+			continue
+		}
+		out.Write(data)
+	}
+
+	require.Equal(t, 0, exitCode)
+	got := out.String()
+	require.Contains(t, got, "first")
+	require.Contains(t, got, "second")
+	require.Less(t, strings.Index(got, "first"), strings.Index(got, "second"))
+}

server/cmd/api/api/process_runner.go

@@ -13,6 +13,7 @@ import (
 	"github.com/creack/pty"
 	"github.com/google/uuid"
 	"github.com/kernel/kernel-images/server/lib/logger"
+	"github.com/kernel/kernel-images/server/lib/ptyio"
 )
 
 const processOutputDrainTimeout = time.Second
@@ -57,6 +58,18 @@ func (s *ApiService) startProcess(ctx context.Context, cmd *exec.Cmd, tty, stdin
 		return nil, err
 	}
 	h.pid = cmd.Process.Pid
+	if tty {
+		h.ptyDrained = make(chan struct{})
+		h.ptyStop = make(chan struct{})
+		// PTY output is drained for the process lifetime so output remains available
+		// when a client attaches after the process has written it.
+		go func() {
+			defer close(h.ptyDrained)
+			if err := ptyio.ReadPTYToWriter(h.ptyFile, h.ptyOut.write, h.ptyStop); err != nil {
+				logger.FromContext(ctx).Error("pty read error", "err", err)
+			}
+		}()
+	}
 	s.procMu.Lock()
 	if s.procs == nil {
 		s.procs = make(map[string]*processHandle)
@@ -190,6 +203,17 @@ func (h *processHandle) run(ctx, lifecycle context.Context, timeoutSec *int, rel
 			drain = nil
 		}
 	}
+	if h.isTTY {
+		// The reader polls the raw fd, so stop it before closing the master.
+		// Closing first leaves it polling a dead descriptor forever. A background
+		// child can hold the slave open after the shell exits, so the wait is bounded.
+		select {
+		case <-h.ptyDrained:
+		case <-time.After(ptyDrainWait):
+			close(h.ptyStop)
+			<-h.ptyDrained
+		}
+	}
 	h.closeOutputPipes()
 	if h.stdin != nil {
 		h.stdin.Close()

server/cmd/api/api/process_test.go

@@ -286,3 +286,30 @@ func TestBuildCmd_AsRootAndAsUserConflict(t *testing.T) {
 	_, err := buildCmd(body)
 	require.Error(t, err, "expected error when both as_root and as_user are set")
 }
+
+func TestProcessTTYExitWithBackgroundChildHoldingPTYDoesNotBlock(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()}
+
+	// The background sleep ignores SIGHUP, so it keeps the PTY slave open after the
+	// shell exits and the reader never sees EOF.
+	args := []string{"-c", "trap \"\" HUP; sleep 30 & exit 0"}
+	allocateTTY := true
+	body := &oapi.ProcessSpawnRequest{Command: "sh", Args: &args, AllocateTty: &allocateTTY}
+	spawnResp, err := svc.ProcessSpawn(ctx, oapi.ProcessSpawnRequestObject{Body: body})
+	require.NoError(t, err)
+	s200, ok := spawnResp.(oapi.ProcessSpawn200JSONResponse)
+	require.True(t, ok, "unexpected spawn resp: %T", spawnResp)
+
+	svc.procMu.Lock()
+	h := svc.procs[s200.ProcessId.String()]
+	svc.procMu.Unlock()
+	require.NotNil(t, h)
+
+	select {
+	case <-h.doneCh:
+	case <-time.After(10 * time.Second):
+		t.Fatal("process waiter did not finish after the PTY drain deadline")
+	}
+}

server/cmd/api/api/pty_output.go

@@ -0,0 +1,152 @@
+package api
+
+import (
+	"sync"
+	"time"
+)
+
+const (
+	ptyReplayBytes = 64 * 1024
+	ptyDrainWait   = 2 * time.Second
+)
+
+type ptyOutput struct {
+	mu        sync.Mutex
+	buf       []byte
+	start     uint64
+	end       uint64
+	sink      func([]byte) error
+	stop      chan struct{}
+	exited    chan struct{}
+	changed   chan struct{}
+	delivered uint64
+}
+
+func (o *ptyOutput) write(data []byte) error {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	if len(o.buf) == 0 {
+		o.buf = make([]byte, ptyReplayBytes)
+	}
+
+	if len(data) > ptyReplayBytes {
+		o.start = o.end + uint64(len(data)-ptyReplayBytes)
+		data = data[len(data)-ptyReplayBytes:]
+	}
+	for len(data) > 0 {
+		index := int(o.end % ptyReplayBytes)
+		n := min(len(data), ptyReplayBytes-index)
+		copy(o.buf[index:index+n], data[:n])
+		o.end += uint64(n)
+		data = data[n:]
+		if o.end-o.start > ptyReplayBytes {
+			o.start = o.end - ptyReplayBytes
+		}
+	}
+	o.signalLocked()
+	return nil
+}
+
+func (o *ptyOutput) attach(sink func([]byte) error) {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	if o.changed == nil {
+		o.changed = make(chan struct{})
+	}
+	o.sink = sink
+	o.stop = make(chan struct{})
+	o.exited = make(chan struct{})
+	o.delivered = o.start
+	go o.forward(o.stop, o.exited)
+}
+
+func (o *ptyOutput) forward(stop <-chan struct{}, exited chan<- struct{}) {
+	defer close(exited)
+	for {
+		o.mu.Lock()
+		if o.delivered < o.start {
+			o.delivered = o.start
+		}
+		if o.sink == nil {
+			o.mu.Unlock()
+			return
+		}
+		if o.delivered < o.end {
+			count := min(o.end-o.delivered, uint64(ptyReplayBytes))
+			data := o.copyLocked(o.delivered, count)
+			sink := o.sink
+			o.mu.Unlock()
+			if err := sink(data); err != nil {
+				o.mu.Lock()
+				o.sink = nil
+				o.signalLocked()
+				o.mu.Unlock()
+				return
+			}
+			o.mu.Lock()
+			if o.sink == nil {
+				o.signalLocked()
+				o.mu.Unlock()
+				return
+			}
+			o.delivered += uint64(len(data))
+			o.signalLocked()
+			o.mu.Unlock()
+			continue
+		}
+		changed := o.changed
+		o.mu.Unlock()
+		select {
+		case <-changed:
+		case <-stop:
+			return
+		}
+	}
+}
+
+func (o *ptyOutput) copyLocked(cursor, count uint64) []byte {
+	data := make([]byte, count)
+	index := int(cursor % ptyReplayBytes)
+	first := min(len(data), ptyReplayBytes-index)
+	copy(data, o.buf[index:index+first])
+	copy(data[first:], o.buf[:len(data)-first])
+	return data
+}
+
+func (o *ptyOutput) flush() {
+	o.mu.Lock()
+	target := o.end
+	for o.sink != nil && o.delivered < target {
+		changed := o.changed
+		o.mu.Unlock()
+		<-changed
+		o.mu.Lock()
+		target = o.end
+	}
+	o.mu.Unlock()
+}
+
+func (o *ptyOutput) detach() {
+	o.mu.Lock()
+	var exited chan struct{}
+	if o.stop != nil {
+		close(o.stop)
+		o.stop = nil
+		exited = o.exited
+	}
+	o.sink = nil
+	o.signalLocked()
+	o.mu.Unlock()
+	// Wait for the forwarder to leave any in-flight sink call so it cannot move
+	// the cursor of the next attachment.
+	if exited != nil {
+		<-exited
+	}
+}
+
+func (o *ptyOutput) signalLocked() {
+	if o.changed != nil {
+		close(o.changed)
+	}
+	o.changed = make(chan struct{})
+}