Merge remote-tracking branch 'origin/farm/11d4e46a/bound-gpu-probe'

packages/coding-agent/src/system-prompt.test.ts

@@ -14,6 +14,8 @@ async function runProbeScenario(options: {
 	runs: number;
 	sleepSeconds?: number;
 	holdStdoutOpen?: boolean;
+	descendantHoldsStdout?: boolean;
+	validOutput?: string;
 }): Promise<ProbeRunResult> {
 	const tempRoot = await fs.mkdtemp(path.join(os.tmpdir(), "omp-gpu-probe-"));
 	try {
@@ -25,7 +27,7 @@ async function runProbeScenario(options: {
 		const lspciPath = path.join(binDir, "lspci");
 		await Bun.write(
 			lspciPath,
-			'#!/usr/bin/env sh\nprintf x >> "$OMP_GPU_PROBE_COUNT"\nif [ "$OMP_GPU_PROBE_HOLD_STDOUT_OPEN" = "true" ]; then sleep "$OMP_GPU_PROBE_SLEEP" & wait "$!"; fi\nif [ -n "$OMP_GPU_PROBE_SLEEP" ]; then exec sleep "$OMP_GPU_PROBE_SLEEP"; fi\nexit 0\n',
+			'#!/usr/bin/env sh\nprintf x >> "$OMP_GPU_PROBE_COUNT"\nif [ -n "$OMP_GPU_PROBE_VALID_OUTPUT" ]; then printf "%s\\n" "$OMP_GPU_PROBE_VALID_OUTPUT"; fi\nif [ "$OMP_GPU_PROBE_DESCENDANT_HOLDS_STDOUT" = "true" ]; then sleep "$OMP_GPU_PROBE_SLEEP" & exit 0; fi\nif [ "$OMP_GPU_PROBE_HOLD_STDOUT_OPEN" = "true" ]; then sleep "$OMP_GPU_PROBE_SLEEP" & wait "$!"; fi\nif [ -n "$OMP_GPU_PROBE_SLEEP" ]; then exec sleep "$OMP_GPU_PROBE_SLEEP"; fi\nexit 0\n',
 		);
 		await fs.chmod(lspciPath, 0o755);
 
@@ -83,6 +85,16 @@ console.log(JSON.stringify({ elapsedMs: Math.round(performance.now() - startedAt
 		} else {
 			delete env.OMP_GPU_PROBE_HOLD_STDOUT_OPEN;
 		}
+		if (options.descendantHoldsStdout) {
+			env.OMP_GPU_PROBE_DESCENDANT_HOLDS_STDOUT = "true";
+		} else {
+			delete env.OMP_GPU_PROBE_DESCENDANT_HOLDS_STDOUT;
+		}
+		if (options.validOutput !== undefined) {
+			env.OMP_GPU_PROBE_VALID_OUTPUT = options.validOutput;
+		} else {
+			delete env.OMP_GPU_PROBE_VALID_OUTPUT;
+		}
 
 		const childStartedAt = performance.now();
 		const child = Bun.spawn([process.execPath, scenarioPath], { stdout: "pipe", stderr: "pipe", env });
@@ -118,4 +130,29 @@ describe.skipIf(process.platform !== "linux")("system prompt GPU probe", () => {
 		// not linger until a descendant holding stdout (sleep 7) exits on its own.
 		expect(result.childElapsedMs).toBeLessThan(6500);
 	}, 15_000);
+
+	it("does not wait on stdout held by a descendant after a successful probe", async () => {
+		const result = await runProbeScenario({ runs: 1, sleepSeconds: 3, descendantHoldsStdout: true });
+
+		expect(result.cached).toEqual({ gpu: null });
+		// Probe exits 0 immediately but leaves a backgrounded sleep holding the stdout
+		// pipe. The success path MUST bound the drain wait, not block until sleep exits.
+		expect(result.elapsedMs).toBeLessThan(2000);
+		expect(result.childElapsedMs).toBeLessThan(2000);
+	}, 15_000);
+
+	it("keeps probe output captured before a descendant delays EOF", async () => {
+		const result = await runProbeScenario({
+			runs: 1,
+			sleepSeconds: 3,
+			descendantHoldsStdout: true,
+			validOutput: "00:02.0 VGA compatible controller: NVIDIA TestGPU",
+		});
+
+		// Probe exited 0 with valid output before bg sleep held stdout open.
+		// Captured stdout MUST be cached, not discarded as if the probe failed.
+		expect(result.cached).toEqual({ gpu: "02.0 VGA compatible controller: NVIDIA TestGPU" });
+		expect(result.elapsedMs).toBeLessThan(2000);
+		expect(result.childElapsedMs).toBeLessThan(2000);
+	}, 15_000);
 });

packages/coding-agent/src/system-prompt.ts

@@ -113,6 +113,8 @@ function parseWmicTable(output: string, header: string): string | null {
 const SYSTEM_PROMPT_PREP_TIMEOUT_MS = 5000;
 /** Kept below prep timeout so timed-out probes can still write the null cache before fallback. */
 const GPU_PROBE_TIMEOUT_MS = SYSTEM_PROMPT_PREP_TIMEOUT_MS - 500;
+/** Drop stdout from a probe descendant that inherited the pipe after the probe exited. */
+const GPU_PROBE_STDOUT_DRAIN_MS = 250;
 
 async function runGpuProbe(cmd: string[]): Promise<string | null> {
 	try {
@@ -138,13 +140,18 @@ async function runGpuProbe(cmd: string[]): Promise<string | null> {
 			stdout += decoder.decode();
 		})();
 		const exitCode = await proc.exited;
-		if (exitCode !== 0) {
+		// Even on exit 0, a probe wrapper can leave a descendant holding stdout open.
+		// Bound the EOF wait so getCachedGpu cannot outlive the probe in either path;
+		// keep whatever bytes the reader already captured before cancelling.
+		const drained = await Promise.race([
+			stdoutDone.then(() => "ok" as const).catch(() => "err" as const),
+			Bun.sleep(GPU_PROBE_STDOUT_DRAIN_MS).then(() => "timeout" as const),
+		]);
+		if (drained !== "ok") {
 			await stdoutReader.cancel().catch(() => undefined);
 			await stdoutDone.catch(() => undefined);
-			return null;
 		}
-		await stdoutDone;
-		return stdout;
+		return exitCode === 0 ? stdout : null;
 	} catch {
 		return null;
 	}
@@ -531,7 +538,10 @@ export async function buildSystemPrompt(options: BuildSystemPromptOptions = {}):
 		gpu: undefined as string | undefined,
 	};
 
-	const deadline = Bun.sleep(SYSTEM_PROMPT_PREP_TIMEOUT_MS).then(() => "__timeout__" as const);
+	const { promise: deadline, resolve: fireDeadline } = Promise.withResolvers<"__timeout__">();
+	const deadlineTimer = setTimeout(() => fireDeadline("__timeout__"), SYSTEM_PROMPT_PREP_TIMEOUT_MS);
+	// Unref so a fast prep does not hold a one-shot CLI alive waiting for this timer.
+	deadlineTimer.unref();
 	const timedOut: string[] = [];
 	const failed: Array<{ name: string; error: unknown }> = [];
 
@@ -630,6 +640,7 @@ export async function buildSystemPrompt(options: BuildSystemPromptOptions = {}):
 		withDeadline("resolveActiveRepoContext", activeRepoContextPromise, prepDefaults.activeRepoContext),
 		withDeadline("getCachedGpu", gpuPromise, prepDefaults.gpu),
 	]);
+	clearTimeout(deadlineTimer);
 	const agentsMdFiles = Array.from(new Set(workspaceTree.agentsMdFiles)).sort().slice(0, AGENTS_MD_LIMIT);
 
 	if (timedOut.length > 0) {