diff --git a/packages/coding-agent/src/system-prompt.test.ts b/packages/coding-agent/src/system-prompt.test.ts index a55f77208..456d1964f 100644 --- a/packages/coding-agent/src/system-prompt.test.ts +++ b/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 { 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); }); diff --git a/packages/coding-agent/src/system-prompt.ts b/packages/coding-agent/src/system-prompt.ts index d409d0e3a..3c410b396 100644 --- a/packages/coding-agent/src/system-prompt.ts +++ b/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 { try { @@ -138,13 +140,18 @@ async function runGpuProbe(cmd: string[]): Promise { 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) {