fix(process): don't drop process output and the exit event from streams

server/cmd/api/api/process.go

@@ -30,6 +30,11 @@ import (
 	openapi_types "github.com/oapi-codegen/runtime/types"
 )
 
+// processOutputDrainTimeout bounds how long a spawned process's exit is held
+// back while its output readers drain the pipes. A child the process started
+// in the background can keep the pipes open after the process itself exits.
+const processOutputDrainTimeout = 2 * time.Second
+
 type processHandle struct {
 	id       openapi_types.UUID
 	pid      int
@@ -298,19 +303,12 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 		stdin = ptyFile
 		isTTY = true
 	} else {
-		stdout, err = cmd.StdoutPipe()
-		if err != nil {
-			return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stdout"}}, nil
-		}
-		stderr, err = cmd.StderrPipe()
-		if err != nil {
-			return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stderr"}}, nil
-		}
 		stdin, err = cmd.StdinPipe()
 		if err != nil {
 			return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to open stdin"}}, nil
 		}
-		if err := cmd.Start(); err != nil {
+		stdout, stderr, err = startWithOutputPipes(cmd)
+		if err != nil {
 			if isUserCmdError(err) {
 				return oapi.ProcessSpawn400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: err.Error()}}, nil
 			}
@@ -354,8 +352,11 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 	// Reader goroutines
 	// In PTY mode, do NOT read from the PTY here to avoid competing with the /attach endpoint.
 	// In non‑PTY mode, stdout and stderr are separate pipes, so we run two readers and tag chunks accordingly.
+	var readers sync.WaitGroup
 	if !isTTY {
+		readers.Add(2)
 		go func() {
+			defer readers.Done()
 			reader := bufio.NewReader(stdout)
 			buf := make([]byte, 4096)
 			for {
@@ -371,6 +372,7 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 			}
 		}()
 		go func() {
+			defer readers.Done()
 			reader := bufio.NewReader(stderr)
 			buf := make([]byte, 4096)
 			for {
@@ -405,6 +407,17 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 			}
 		}
 		h.setExited(code)
+		// Let the readers drain what the process wrote before it exited, so
+		// the exit event follows all of its output.
+		readersDone := make(chan struct{})
+		go func() {
+			readers.Wait()
+			close(readersDone)
+		}()
+		select {
+		case <-readersDone:
+		case <-time.After(processOutputDrainTimeout):
+		}
 		// Ensure all related FDs are closed to avoid leaking descriptors.
 		// In PTY mode, close the PTY master; in non-PTY mode, close individual pipes.
 		if h.isTTY {
@@ -431,6 +444,11 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 			}
 		}
 
+		// Closing the pipes above ends any reader still reading them. A reader
+		// blocked on a full outCh finishes once a stream takes its event, just
+		// as the exit event below would have waited for one.
+		<-readersDone
+
 		// Send exit event
 		evt := oapi.ProcessStreamEventEvent("exit")
 		h.outCh <- oapi.ProcessStreamEvent{Event: &evt, ExitCode: &code}
@@ -457,6 +475,36 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn
 	}, nil
 }
 
+// startWithOutputPipes starts cmd with stdout and stderr connected to pipes
+// whose read ends the caller owns and must close. cmd.StdoutPipe is not used
+// because cmd.Wait closes its read end as soon as the process exits, which
+// discards output the readers have not consumed yet.
+func startWithOutputPipes(cmd *exec.Cmd) (stdout, stderr *os.File, err error) {
+	stdoutR, stdoutW, err := os.Pipe()
+	if err != nil {
+		return nil, nil, err
+	}
+	stderrR, stderrW, err := os.Pipe()
+	if err != nil {
+		stdoutR.Close()
+		stdoutW.Close()
+		return nil, nil, err
+	}
+	cmd.Stdout = stdoutW
+	cmd.Stderr = stderrW
+	err = cmd.Start()
+	// The child holds its own copies of the write ends; closing ours lets the
+	// readers see EOF once the child, and anything it started, closes them.
+	stdoutW.Close()
+	stderrW.Close()
+	if err != nil {
+		stdoutR.Close()
+		stderrR.Close()
+		return nil, nil, err
+	}
+	return stdoutR, stderrR, nil
+}
+
 // Send signal to process
 // (POST /process/{process_id}/kill)
 func (s *ApiService) ProcessKill(ctx context.Context, request oapi.ProcessKillRequestObject) (oapi.ProcessKillResponseObject, error) {
@@ -581,27 +629,42 @@ func (s *ApiService) ProcessStdoutStream(ctx context.Context, request oapi.Proce
 	pr, pw := io.Pipe()
 	go func() {
 		defer pw.Close()
+		// writeEvent writes evt as an SSE frame: data: <json>\n\n
+		writeEvent := func(evt oapi.ProcessStreamEvent) error {
+			var buf bytes.Buffer
+			if err := json.NewEncoder(&buf).Encode(evt); err != nil {
+				log.Error("failed to marshal event", "err", err)
+				return err
+			}
+			line := bytes.TrimRight(buf.Bytes(), "\n")
+			if _, err := pw.Write([]byte("data: ")); err != nil {
+				return err
+			}
+			if _, err := pw.Write(line); err != nil {
+				return err
+			}
+			_, err := pw.Write([]byte("\n\n"))
+			return err
+		}
 		for {
 			select {
 			case evt := <-h.outCh:
-				// Write SSE: data: <json>\n\n
-				var buf bytes.Buffer
-				if err := json.NewEncoder(&buf).Encode(evt); err != nil {
-					log.Error("failed to marshal event", "err", err)
-					return
-				}
-				line := bytes.TrimRight(buf.Bytes(), "\n")
-				if _, err := pw.Write([]byte("data: ")); err != nil {
-					return
-				}
-				if _, err := pw.Write(line); err != nil {
-					return
-				}
-				if _, err := pw.Write([]byte("\n\n")); err != nil {
+				if writeEvent(evt) != nil {
 					return
 				}
 			case <-h.doneCh:
-				return
+				// The exit event and any output before it were queued before
+				// doneCh closed; deliver them before ending the stream.
+				for {
+					select {
+					case evt := <-h.outCh:
+						if writeEvent(evt) != nil {
+							return
+						}
+					default:
+						return
+					}
+				}
 			}
 		}
 	}()

server/cmd/api/api/process_test.go

@@ -9,6 +9,7 @@ import (
 	"os/user"
 	"strconv"
 	"strings"
+	"syscall"
 	"testing"
 	"time"
 
@@ -271,3 +272,138 @@ func TestBuildCmd_AsRootAndAsUserConflict(t *testing.T) {
 	_, err := buildCmd(body)
 	require.Error(t, err, "expected error when both as_root and as_user are set")
 }
+
+// spawnShForTest spawns `sh -c script` and returns its process ID.
+func spawnShForTest(t *testing.T, svc *ApiService, script string) openapi_types.UUID {
+	t.Helper()
+	args := []string{"-c", script}
+	resp, err := svc.ProcessSpawn(context.Background(), oapi.ProcessSpawnRequestObject{Body: &oapi.ProcessSpawnRequest{Command: "sh", Args: &args}})
+	require.NoError(t, err, "ProcessSpawn error")
+	s200, ok := resp.(oapi.ProcessSpawn200JSONResponse)
+	require.True(t, ok, "unexpected spawn resp type: %T", resp)
+	require.NotNil(t, s200.ProcessId, "missing ProcessId in spawn resp")
+	return *s200.ProcessId
+}
+
+// readProcessStream reads a process's SSE stream until it ends and returns
+// its stdout, its stderr and every event in the order received.
+func readProcessStream(t *testing.T, svc *ApiService, id openapi_types.UUID) (string, string, []oapi.ProcessStreamEvent) {
+	t.Helper()
+	resp, err := svc.ProcessStdoutStream(context.Background(), oapi.ProcessStdoutStreamRequestObject{ProcessId: id})
+	require.NoError(t, err, "StdoutStream error")
+	st200, ok := resp.(oapi.ProcessStdoutStream200TexteventStreamResponse)
+	require.True(t, ok, "unexpected stream resp: %T", resp)
+
+	type result struct {
+		stdout, stderr strings.Builder
+		events         []oapi.ProcessStreamEvent
+		err            error
+	}
+	done := make(chan *result, 1)
+	go func() {
+		r := &result{}
+		scanner := bufio.NewScanner(st200.Body)
+		for scanner.Scan() {
+			line := scanner.Text()
+			if !strings.HasPrefix(line, "data: ") {
+				continue
+			}
+			var evt oapi.ProcessStreamEvent
+			if r.err = json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &evt); r.err != nil {
+				break
+			}
+			r.events = append(r.events, evt)
+			if evt.DataB64 != nil {
+				b, _ := base64.StdEncoding.DecodeString(*evt.DataB64)
+				if evt.Stream != nil && *evt.Stream == "stderr" {
+					r.stderr.Write(b)
+				} else {
+					r.stdout.Write(b)
+				}
+			}
+		}
+		if r.err == nil {
+			r.err = scanner.Err()
+		}
+		done <- r
+	}()
+	select {
+	case r := <-done:
+		require.NoError(t, r.err, "read SSE stream")
+		return r.stdout.String(), r.stderr.String(), r.events
+	case <-time.After(10 * time.Second):
+		st200.Body.(io.Closer).Close()
+		t.Fatal("stream did not end within 10s")
+		return "", "", nil
+	}
+}
+
+// requireExitLast checks that the stream ended with exactly one exit event.
+func requireExitLast(t *testing.T, events []oapi.ProcessStreamEvent, wantCode int) {
+	t.Helper()
+	require.NotEmpty(t, events, "stream delivered no events")
+	last := events[len(events)-1]
+	require.True(t, last.Event != nil && *last.Event == "exit", "last event is not exit: %+v", last)
+	require.NotNil(t, last.ExitCode, "exit event has no exit code")
+	require.Equal(t, wantCode, *last.ExitCode, "exit code mismatch")
+	for _, evt := range events[:len(events)-1] {
+		require.Nil(t, evt.Event, "unexpected event before exit: %+v", evt)
+	}
+}
+
+func TestProcessStdoutStream_AfterExitDeliversQueuedEvents(t *testing.T) {
+	t.Parallel()
+	svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()}
+
+	// Opening the stream only after the process has exited is the common case
+	// for short commands. Everything queued for it must still be delivered.
+	for i := 0; i < 200; i++ {
+		id := spawnShForTest(t, svc, "printf one; printf two 1>&2; exit 7")
+		svc.procMu.RLock()
+		h := svc.procs[id.String()]
+		svc.procMu.RUnlock()
+		<-h.doneCh
+
+		stdout, stderr, events := readProcessStream(t, svc, id)
+		require.Equal(t, "one", stdout, "run %d: stdout mismatch", i)
+		require.Equal(t, "two", stderr, "run %d: stderr mismatch", i)
+		requireExitLast(t, events, 7)
+	}
+}
+
+func TestProcessStdoutStream_DeliversAllOutputBeforeExit(t *testing.T) {
+	t.Parallel()
+	svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()}
+
+	// The process exits right after writing more than a pipe buffer, so output
+	// is still unread in the pipe when it exits.
+	want := strings.Repeat("kernel\n", 50000)
+	for i := 0; i < 5; i++ {
+		id := spawnShForTest(t, svc, "yes kernel | head -n 50000; exit 3")
+		stdout, stderr, events := readProcessStream(t, svc, id)
+		require.Equal(t, len(want), len(stdout), "run %d: stdout length mismatch", i)
+		require.Equal(t, want, stdout, "run %d: stdout mismatch", i)
+		require.Empty(t, stderr, "run %d: unexpected stderr", i)
+		requireExitLast(t, events, 3)
+	}
+}
+
+func TestProcessSpawn_BackgroundChildDoesNotHoldExit(t *testing.T) {
+	t.Parallel()
+	svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()}
+
+	// The shell exits at once, but the background sleep keeps the output
+	// pipes open. The stream must still end with the shell's exit.
+	id := spawnShForTest(t, svc, "sleep 30 & echo $!")
+	stdout, _, events := readProcessStream(t, svc, id)
+	pid, err := strconv.Atoi(strings.TrimSpace(stdout))
+	require.NoError(t, err, "background pid not printed: %q", stdout)
+	t.Cleanup(func() { _ = syscall.Kill(pid, syscall.SIGKILL) })
+	requireExitLast(t, events, 0)
+
+	statusResp, err := svc.ProcessStatus(context.Background(), oapi.ProcessStatusRequestObject{ProcessId: id})
+	require.NoError(t, err, "ProcessStatus error")
+	sr, ok := statusResp.(oapi.ProcessStatus200JSONResponse)
+	require.True(t, ok, "unexpected status resp: %T", statusResp)
+	require.Equal(t, oapi.ProcessStatusState("exited"), *sr.State)
+}