test(tools): poll for the exec server's address instead of racing its startup

internal/tools/bash_tool_test.go

@@ -70,6 +70,9 @@ func runBashToolHelper(command string) {
 		fmt.Println("listening", listener.Addr().String())
 		server := &http.Server{
 			Handler: http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+				// Printed per request, so a test can prove the session delivers output
+				// mid-flight: this line cannot exist before the first read has happened.
+				fmt.Println("served", request.URL.Path)
 				_, _ = response.Write([]byte("zero-server-ok"))
 			}),
 		}
@@ -522,7 +525,14 @@ func TestBashToolTimeoutKillsBackgroundChildren(t *testing.T) {
 	if result.Status != StatusError {
 		t.Fatalf("expected timeout error status, got %s: %q", result.Status, result.Output)
 	}
-	if elapsed > time.Second {
+	// The regression this guards is Run blocking until the background child's own
+	// sleep ends, so the bound only has to sit under that. One second sat almost
+	// exactly on it: the call legitimately costs a sandbox plan, a fork and exec of
+	// a shell plus a subshell, the 300ms timeout itself and the post-kill drain,
+	// and spawn alone was measured here at 720ms to 1.44s with every core busy. A
+	// merely slow runner therefore failed a correct system. Three seconds is still
+	// far below the child's lifetime, so nothing is given up.
+	if elapsed > 3*time.Second {
 		t.Fatalf("Run blocked %s past the 300ms timeout; background child held the pipes", elapsed)
 	}
 

internal/tools/exec_command_test.go

@@ -3,6 +3,7 @@ package tools
 import (
 	"context"
 	"io"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -373,18 +374,19 @@ func TestExecCommandForegroundServerReturnsSessionAndServesHTTP(t *testing.T) {
 	if err != nil {
 		t.Fatalf("foreground server should return session_id, meta=%#v output=%q", start.Meta, start.Output)
 	}
-	addr := parseListeningAddress(start.Output)
-	if addr == "" {
-		t.Fatalf("server output did not include listening address: %q", start.Output)
-	}
 	t.Cleanup(func() {
 		writeTool.Run(context.Background(), map[string]any{
 			"session_id": sessionID,
 			"chars":      "\u0003",
 		})
 	})
 
-	response, err := http.Get("http://" + addr)
+	addr := waitForListeningAddress(t, writeTool, sessionID, start.Output)
+
+	// Bounded, so a server that accepts a connection and never answers fails here
+	// with its address instead of hanging until the package timeout kills the run.
+	client := &http.Client{Timeout: 30 * time.Second}
+	response, err := client.Get("http://" + addr)
 	if err != nil {
 		t.Fatalf("foreground exec server was not reachable at %s: %v; output=%q", addr, err, start.Output)
 	}
@@ -396,6 +398,153 @@ func TestExecCommandForegroundServerReturnsSessionAndServesHTTP(t *testing.T) {
 	if string(bytes) != "zero-server-ok" {
 		t.Fatalf("server response = %q", string(bytes))
 	}
+
+	// AND THE SESSION DELIVERS OUTPUT MID-FLIGHT, PINNED DETERMINISTICALLY.
+	//
+	// The helper prints this line only when it serves a request, so it cannot have
+	// been in the first read: the request had not happened yet. It can only arrive
+	// through a poll of a process that is still running, which is the delivery the
+	// address parse used to cover by accident and only on a machine slow enough to
+	// miss the first read. A regression where a live session buffers everything
+	// until it exits leaves the rest of this test green and fails here.
+	deadline := time.Now().Add(60 * time.Second)
+	var served strings.Builder
+	for !strings.Contains(served.String(), "served") && time.Now().Before(deadline) {
+		poll := writeTool.Run(context.Background(), map[string]any{
+			"session_id":    sessionID,
+			"chars":         "",
+			"yield_time_ms": 250,
+		})
+		served.WriteString(poll.Output)
+		if poll.Status != StatusOK {
+			t.Fatalf("the exec session was gone before it reported serving the request:\n%s", served.String())
+		}
+	}
+	if !strings.Contains(served.String(), "served") {
+		t.Fatalf("a still-running session never delivered the line the server printed while serving:\n%s", served.String())
+	}
+
+	// AND A READ DRAINS WHAT IT RETURNS.
+	//
+	// One more poll, because "the line showed up" is satisfied just as well by
+	// output that is re-delivered forever. drain copies and then nils the buffer,
+	// so this read carries nothing; if it stopped clearing, every session in
+	// production would repeat its whole output on every poll while every other
+	// assertion here stayed green. The manager-level test that asserts the same
+	// count is skipped on Windows, so this is the only place guarding it there.
+	again := writeTool.Run(context.Background(), map[string]any{
+		"session_id":    sessionID,
+		"chars":         "",
+		"yield_time_ms": 250,
+	})
+	if strings.Contains(again.Output, "served") {
+		t.Errorf("a second read returned output the first had already delivered, so reads are no longer draining:\n%s", again.Output)
+	}
+}
+
+// waitForListeningAddress polls the exec session until the helper says where it
+// is listening.
+//
+// ONE YIELD IS AN ALLOWANCE, NOT A GUARANTEE, AND THIS TEST TREATED IT AS ONE.
+//
+// The helper is this test binary re-executed. The first read used to be its only
+// chance: 500ms to be spawned through the process manager, reach net.Listen and
+// have its line drained back. A Windows CI runner under load does not get there
+// in time, so the read returned "Command is still running." with no address and
+// the test failed as though the server were broken. It failed that way on
+// unrelated contributor PRs, which is worse than a slow test: it reports somebody
+// else's change as the fault.
+//
+// Polling is what the banner in that very output tells the caller to do, so it is
+// what this does, up to a bound far above any plausible start.
+//
+// AND IT STILL FAILS FAST WHEN THE FAILURE IS REAL. A session that has gone away
+// answers with StatusError, and a process that exited before listening answers
+// with an exit_code line. Either is a genuine defect, so neither is worth waiting
+// out: they end the loop immediately with what the session actually said.
+//
+// Each read DRAINS: managedProcess.collect returns only the bytes produced since
+// the previous call, so these chunks are incremental and none of them is the whole
+// output. Each is parsed on its own, which is sound because the helper writes its
+// address with a single Println and a write that small does not tear. The whole
+// transcript is kept anyway, so a failure shows every chunk rather than the last.
+func waitForListeningAddress(t *testing.T, writeTool Tool, sessionID int, first string) string {
+	t.Helper()
+	if addr := validListeningAddress(first); addr != "" {
+		return addr
+	}
+	// AND WHEN IT DOES NOT, THIS TEST STOPS COVERING THE FIRST READ.
+	//
+	// That parse used to be the only assertion in the repo that a still-running
+	// session's first read carries the child's stdout at all: every other test
+	// either uses a tiny yield against a helper that prints later, or uses a large
+	// one against a helper that exits, which is the collect-on-exit branch. The
+	// manager-level equivalent is skipped on Windows because it shells out to
+	// /bin/sh. The coverage cannot be pinned deterministically without
+	// reintroducing the race, since collect() blocks the whole window either way,
+	// so it is given up deliberately here and paid back after the request below,
+	// where a line that cannot exist before the first read proves the poll path.
+	transcript := []string{first}
+	// Sixty seconds, not thirty. Starting this helper was measured here at 0.74s
+	// to 4.79s with every core saturated, on a box faster than a CI runner, and
+	// windows-latest additionally scans a freshly linked test binary on its first
+	// execution. A deadline of the same order as the latency it tolerates is the
+	// mistake that produced the 500ms budget; the point is a bound, and every
+	// genuine failure below leaves before reaching it anyway.
+	deadline := time.Now().Add(60 * time.Second)
+	for time.Now().Before(deadline) {
+		poll := writeTool.Run(context.Background(), map[string]any{
+			"session_id":    sessionID,
+			"chars":         "",
+			"yield_time_ms": 250,
+		})
+		transcript = append(transcript, poll.Output)
+		if addr := validListeningAddress(poll.Output); addr != "" {
+			return addr
+		}
+		// FAST-FAIL, BECAUSE THE LOOP HAS NO PACING OF ITS OWN. All the spacing
+		// comes from collect() waiting out its yield inside a LIVE session; every
+		// error return happens in microseconds. Continue removes the process the
+		// moment it observes the exit, so without these two the loop free-runs on
+		// a dead id: measured at 3028 polls in 5 seconds, which on a 4-vCPU runner
+		// would peg a core for the whole deadline and hand t.Fatalf hundreds of
+		// kilobytes of the same sentence. The exit result is also delivered
+		// exactly once, so folding it into the blob and continuing loses the one
+		// poll that says what went wrong.
+		if poll.Status != StatusOK {
+			t.Fatalf("the exec session was gone before the server reported an address:\n%s", strings.Join(transcript, "\n--- poll ---\n"))
+		}
+		if strings.Contains(poll.Output, "exit_code:") {
+			t.Fatalf("the server process exited before it reported an address:\n%s", strings.Join(transcript, "\n--- poll ---\n"))
+		}
+	}
+	t.Fatalf("the server never reported a listening address within the deadline:\n%s", strings.Join(transcript, "\n--- poll ---\n"))
+	return ""
+}
+
+// validListeningAddress is parseListeningAddress plus the check that what it
+// found is actually an address.
+//
+// A read DRAINS, and the formatter wraps each drain in its own body, so a line
+// split across two reads arrives as two fragments with banner text wedged between
+// them and no concatenation can rejoin it. The formatter also terminates the
+// partial chunk, so the fragment looks like a whole line and the bare parse
+// returns a truncated address: the test would then fail on an unreachable port
+// and point the next reader at the network instead of at the parse. Rejecting
+// anything that is not host:port keeps polling instead.
+func validListeningAddress(output string) string {
+	addr := parseListeningAddress(output)
+	if addr == "" {
+		return ""
+	}
+	host, port, err := net.SplitHostPort(addr)
+	if err != nil || host == "" {
+		return ""
+	}
+	if number, err := strconv.Atoi(port); err != nil || number <= 0 {
+		return ""
+	}
+	return addr
 }
 
 func parseListeningAddress(output string) string {
@@ -427,7 +576,12 @@ func TestExecCommandReapsFinishedUnpolledSession(t *testing.T) {
 		t.Fatalf("session_id is not numeric: %v", err)
 	}
 
-	deadline := time.Now().Add(2 * time.Second)
+	// Ten seconds, not two. Inside this the helper must cold-start, run its 250ms
+	// sleep, exit and pass the reap, and a cold start alone was measured at up to
+	// 1.44s under load: two seconds was a budget with almost nothing left in it. The
+	// loop returns as soon as the session is gone, so the bound only costs anything
+	// when the reap genuinely never happens.
+	deadline := time.Now().Add(10 * time.Second)
 	for {
 		if _, ok := manager.Snapshot(sessionID); !ok {
 			return
@@ -739,6 +893,13 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
 		t.Fatalf("session_id is not numeric: %v", err)
 	}
 
+	// SAME SHAPE AS THE SERVER TEST ABOVE, one window narrower.
+	//
+	// The shell is started with a 10ms yield, so it has almost certainly not
+	// reached `read` yet, and the echo then had exactly one fixed 1000ms window to
+	// travel back through the PTY. Writing early is fine, the terminal buffers it,
+	// but a loaded runner can spend that whole second still getting the shell up.
+	// So the echo is polled for rather than demanded within one window.
 	result := writeTool.Run(context.Background(), map[string]any{
 		"session_id":    sessionID,
 		"chars":         "hello\n",
@@ -747,11 +908,39 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
 	if result.Status != StatusOK {
 		t.Fatalf("write_stdin status = %s: %s", result.Status, result.Output)
 	}
-	if !strings.Contains(result.Output, "got:hello") {
-		t.Fatalf("expected PTY input output, got %q", result.Output)
+	var transcript strings.Builder
+	transcript.WriteString(result.Output)
+	exitCode := result.Meta["exit_code"]
+	// BOTH THE ECHO AND THE EXIT, because the test asserts both.
+	//
+	// collect can return the drained echo on a yield timer before markDone has
+	// run, so a poll can carry got:hello with no exit_code meta yet. Stopping on
+	// the echo alone would then assert an exit that had not been observed. The
+	// condition deliberately keeps the echo term as well: gating on exit_code
+	// alone would stop polling before a slow shell start delivers the echo, which
+	// is the race this whole change exists to remove.
+	deadline := time.Now().Add(30 * time.Second)
+	for (!strings.Contains(transcript.String(), "got:hello") || exitCode == "") && time.Now().Before(deadline) {
+		poll := writeTool.Run(context.Background(), map[string]any{
+			"session_id":    sessionID,
+			"chars":         "",
+			"yield_time_ms": 250,
+		})
+		transcript.WriteString(poll.Output)
+		if code := poll.Meta["exit_code"]; code != "" {
+			exitCode = code
+		}
+		if poll.Status != StatusOK {
+			// The session is gone. Either the echo already arrived, which the loop
+			// condition will see, or it never will and the assertions below say so.
+			break
+		}
 	}
-	if result.Meta["exit_code"] != "0" {
-		t.Fatalf("expected exited session, got meta=%#v output=%q", result.Meta, result.Output)
+	if !strings.Contains(transcript.String(), "got:hello") {
+		t.Fatalf("expected PTY input output, got %q", transcript.String())
+	}
+	if exitCode != "0" {
+		t.Fatalf("expected exited session, got exit_code=%q output=%q", exitCode, transcript.String())
 	}
 }