From 58fc79653cadf44690a1765bc6bdf865fd6ffbc6 Mon Sep 17 00:00:00 2001 From: Tanish Gudise <174986239+TanishGudise@users.noreply.github.com> Date: Mon, 12 Oct 2026 02:12:35 -0400 Subject: [PATCH] fix(process): don't drop process output and the exit event from streams A /process/{id}/stdout/stream client could miss the end of a process's output and its exit event: - The stream goroutine selected on outCh and doneCh. Once the process exited both were ready, so it could return with output and the exit event still queued. - The waiter called cmd.Wait, which closes cmd.StdoutPipe/StderrPipe read ends as soon as the process exits, discarding output the readers had not read yet, and then queued the exit event ahead of the readers' last chunks. The API now owns the read ends of the output pipes, the waiter lets the readers drain them before closing them and queuing the exit event (the wait before closing is bounded by processOutputDrainTimeout, since a background child can hold the pipes open), and the stream delivers whatever is queued before it ends. Co-Authored-By: Claude Opus 5.5 --- server/cmd/api/api/process.go | 111 ++++++++++++++++++----- server/cmd/api/api/process_test.go | 136 +++++++++++++++++++++++++++++ 2 files changed, 223 insertions(+), 24 deletions(-) diff --git a/server/cmd/api/api/process.go b/server/cmd/api/api/process.go index 38a458f4e..b9b6ebb84 100644 --- a/server/cmd/api/api/process.go +++ b/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: \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: \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 + } + } } } }() diff --git a/server/cmd/api/api/process_test.go b/server/cmd/api/api/process_test.go index 741734d74..06ff82264 100644 --- a/server/cmd/api/api/process_test.go +++ b/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) +}