Cancel idle commands without racing process startup

internal/headless/headless.go

@@ -5,6 +5,7 @@ package headless
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -219,20 +220,19 @@ func (t *trackingWriter) Write(p []byte) (n int, err error) {
 // runCommandAndCapture shells out the given command and intercepts both standard streams.
 // It kills the process if it produces no output for IdleTimeout.
 func runCommandAndCapture(shell, command string) (string, string, error) {
-	cmd := exec.Command(shell, "-c", command)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cmd := exec.CommandContext(ctx, shell, "-c", command)
 	cmd.Env = os.Environ()
 
 	var stdoutBuf, stderrBuf bytes.Buffer
 
-	timer := time.AfterFunc(IdleTimeout, func() {
-		if cmd.Process != nil {
-			_ = cmd.Process.Kill()
-		}
-	})
+	timeout := IdleTimeout
+	timer := time.AfterFunc(timeout, cancel)
 	defer timer.Stop()
 
 	resetTimer := func() {
-		timer.Reset(IdleTimeout)
+		timer.Reset(timeout)
 	}
 
 	cmd.Stdout = &trackingWriter{w: &stdoutBuf, onWrite: resetTimer}

internal/headless/idle_timeout_test.go

@@ -0,0 +1,29 @@
+package headless
+
+import (
+	"testing"
+	"time"
+)
+
+func TestCaptureIdleTimeout(t *testing.T) {
+	previous := IdleTimeout
+	IdleTimeout = 100 * time.Millisecond
+	t.Cleanup(func() { IdleTimeout = previous })
+
+	for _, tc := range []struct {
+		name, command, stdout string
+	}{
+		{"silent", "exec sleep 2", ""},
+		{"output before idle", "printf ready; exec sleep 2", "ready"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			stdout, stderr, err := runCommandAndCapture("sh", tc.command)
+			if err == nil {
+				t.Fatal("expected the idle command to be terminated")
+			}
+			if stdout != tc.stdout || stderr != "" {
+				t.Fatalf("captured stdout=%q stderr=%q, want stdout=%q and empty stderr", stdout, stderr, tc.stdout)
+			}
+		})
+	}
+}