feat(vscode): promote timed out commands to background jobs

packages/cli/src/lib/__tests__/background-job-manager.test.ts

@@ -1,6 +1,9 @@
+import type { ChildProcess } from "node:child_process";
+import { EventEmitter } from "node:events";
 import { mkdtemp, readFile, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
+import { PassThrough } from "node:stream";
 import { describe, expect, it } from "vitest";
 import { BackgroundJobManager } from "../background-job-manager";
 
@@ -44,6 +47,44 @@ describe("BackgroundJobManager", () => {
     expect(result2?.output).toBe("");
   });
 
+  it("captures live adopted output while replaying initial output", async () => {
+    const outputDir = await mkdtemp(join(tmpdir(), "pochi-bgjob-adopt-test-"));
+    try {
+      const manager = new BackgroundJobManager({ outputDir });
+      const stdout = new PassThrough();
+      const stderr = new PassThrough();
+      const child = Object.assign(new EventEmitter(), {
+        stdout,
+        stderr,
+        kill: () => true,
+      }) as unknown as ChildProcess;
+      let releaseReplay: () => void = () => {};
+      const replayGate = new Promise<void>((resolve) => {
+        releaseReplay = resolve;
+      });
+      const initialStdout = (async function* () {
+        yield Buffer.from("before");
+        await replayGate;
+      })();
+
+      const { outputFile } = manager.adopt(child, "test", {
+        stdout: initialStdout,
+        stderr: [],
+      });
+      stdout.end("after");
+      stderr.end();
+      stdout.destroy();
+      stderr.destroy();
+      releaseReplay();
+      child.emit("close", 0);
+
+      expect(await manager.waitForAllJobs(1000)).toBe("completed");
+      expect(await readFile(outputFile, "utf8")).toBe("beforeafter");
+    } finally {
+      await rm(outputDir, { recursive: true, force: true });
+    }
+  });
+
   it("should guide agents away from rapid empty reads", () => {
     const manager = new BackgroundJobManager();
     const { backgroundJobId } = manager.start("sleep 10", ".");

packages/cli/src/lib/__tests__/foreground-output-capture.test.ts

@@ -0,0 +1,69 @@
+import { once } from "node:events";
+import type { WriteStream } from "node:fs";
+import { stat } from "node:fs/promises";
+import { PassThrough } from "node:stream";
+import { describe, expect, it } from "vitest";
+import { ForegroundOutputCapture } from "../foreground-output-capture";
+
+describe("ForegroundOutputCapture", () => {
+  it("keeps only the bounded replay tail across stdout and stderr", async () => {
+    const stdout = new PassThrough();
+    const stderr = new PassThrough();
+    const capture = new ForegroundOutputCapture(stdout, stderr, 5);
+
+    stdout.write("1234");
+    stderr.write("abcdef");
+
+    await expect(capture.finish()).resolves.toEqual({
+      stdout: "",
+      stderr: "bcdef",
+    });
+  });
+
+  it("creates private spool files and removes them after promotion replay", async () => {
+    const stdout = new PassThrough();
+    const stderr = new PassThrough();
+    const capture = new ForegroundOutputCapture(stdout, stderr);
+    const internals = capture as unknown as {
+      stdoutCapture: { outputPath: string; writer: WriteStream };
+      stderrCapture: { outputPath: string; writer: WriteStream };
+    };
+    await Promise.all([
+      once(internals.stdoutCapture.writer, "open"),
+      once(internals.stderrCapture.writer, "open"),
+    ]);
+
+    if (process.platform !== "win32") {
+      const [stdoutStat, stderrStat] = await Promise.all([
+        stat(internals.stdoutCapture.outputPath),
+        stat(internals.stderrCapture.outputPath),
+      ]);
+      expect(stdoutStat.mode & 0o777).toBe(0o600);
+      expect(stderrStat.mode & 0o777).toBe(0o600);
+    }
+
+    stdout.write("stdout");
+    stderr.write("stderr");
+    const promoted = capture.promote();
+    const collect = async (
+      chunks: AsyncIterable<Buffer | string> | Iterable<Buffer | string>,
+    ) => {
+      const output: Buffer[] = [];
+      for await (const chunk of chunks) {
+        output.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
+      }
+      return Buffer.concat(output).toString("utf8");
+    };
+
+    await expect(
+      Promise.all([collect(promoted.stdout), collect(promoted.stderr)]),
+    ).resolves.toEqual(["stdout", "stderr"]);
+    await expect(stat(internals.stdoutCapture.outputPath)).rejects.toMatchObject(
+      { code: "ENOENT" },
+    );
+    await expect(stat(internals.stderrCapture.outputPath)).rejects.toMatchObject(
+      { code: "ENOENT" },
+    );
+    await promoted.dispose?.();
+  });
+});

packages/cli/src/lib/background-job-manager.ts

@@ -1,6 +1,7 @@
 import { type ChildProcess, spawn } from "node:child_process";
 import { tmpdir } from "node:os";
 import path from "node:path";
+import type { Readable } from "node:stream";
 import { StringDecoder } from "node:string_decoder";
 import type { BackgroundJobTerminalEvent } from "@getpochi/common";
 import { assertBackgroundJobReadInterval } from "@getpochi/common";
@@ -25,13 +26,24 @@ export interface BackgroundJob {
   lastReadAt?: number;
   stopRequested?: boolean;
   finalizing?: boolean;
+  disposeAbort?: () => void;
 }
 
 export interface BackgroundJobStartResult {
   backgroundJobId: string;
   outputFile: string;
 }
 
+export type BackgroundJobInitialOutputStream =
+  | Iterable<Buffer | string>
+  | AsyncIterable<Buffer | string>;
+
+export interface BackgroundJobInitialOutput {
+  stdout: BackgroundJobInitialOutputStream;
+  stderr: BackgroundJobInitialOutputStream;
+  dispose?: () => Promise<void>;
+}
+
 export interface BackgroundJobManagerOptions {
   taskId?: string;
   outputDir?: string;
@@ -50,6 +62,31 @@ export class BackgroundJobManager {
     command: string,
     cwd: string,
     envs?: Record<string, string>,
+  ): BackgroundJobStartResult {
+    const child = spawn(command, {
+      shell: getShellPath(),
+      cwd,
+      env: { ...process.env, ...getTerminalEnv(), ...envs },
+      stdio: ["ignore", "pipe", "pipe"],
+    });
+
+    return this.register(child, command);
+  }
+
+  adopt(
+    child: ChildProcess,
+    command: string,
+    initialOutput: BackgroundJobInitialOutput,
+    abortSignal?: AbortSignal,
+  ): BackgroundJobStartResult {
+    return this.register(child, command, initialOutput, abortSignal);
+  }
+
+  private register(
+    child: ChildProcess,
+    command: string,
+    initialOutput: BackgroundJobInitialOutput = { stdout: [], stderr: [] },
+    abortSignal?: AbortSignal,
   ): BackgroundJobStartResult {
     const id = createBackgroundJobId("command");
     const outputFile = this.options.outputDir
@@ -58,15 +95,6 @@ export class BackgroundJobManager {
         ? getBackgroundJobOutputPath(this.options.taskId, id)
         : path.join(tmpdir(), "pochi-background-jobs", `${id}.log`);
     const outputWriter = new BackgroundJobOutputFile(outputFile);
-
-    const shell = getShellPath();
-    const child = spawn(command, {
-      shell,
-      cwd,
-      env: { ...process.env, ...getTerminalEnv(), ...envs },
-      stdio: ["ignore", "pipe", "pipe"],
-    });
-
     const job: BackgroundJob = {
       id,
       command,
@@ -80,46 +108,113 @@ export class BackgroundJobManager {
 
     this.jobs.set(id, job);
 
-    const appendOutput = async (chunk: string) => {
-      if (chunk.length === 0) return;
-      await outputWriter.append(chunk);
-      if (job.output.length + chunk.length > this.maxOutputSize) {
-        const keep = this.maxOutputSize - chunk.length;
-        if (keep > 0) {
-          job.output = job.output.slice(-keep) + chunk;
+    let appendTail = Promise.resolve();
+    const appendOutput = (chunk: string): Promise<void> => {
+      appendTail = appendTail.then(async () => {
+        if (chunk.length === 0) return;
+        await outputWriter.append(chunk);
+        if (job.output.length + chunk.length > this.maxOutputSize) {
+          const keep = this.maxOutputSize - chunk.length;
+          if (keep > 0) {
+            job.output = job.output.slice(-keep) + chunk;
+          } else {
+            job.output = chunk.slice(-this.maxOutputSize);
+          }
         } else {
-          job.output = chunk.slice(-this.maxOutputSize);
+          job.output += chunk;
         }
-      } else {
-        job.output += chunk;
+      });
+      return appendTail;
+    };
+
+    const consumeOutput = async (
+      stream: Readable | null,
+      initialOutputStream: BackgroundJobInitialOutputStream,
+    ) => {
+      const decoder = new StringDecoder("utf8");
+      const sanitizer = new PlainOutputSanitizer();
+      const initialOutputFinished = (async () => {
+        for await (const chunk of initialOutputStream) {
+          await appendOutput(sanitizer.write(decoder.write(chunk)));
+        }
+      })();
+      const liveOutputFinished = stream
+        ? new Promise<void>((resolve, reject) => {
+            let liveOutputTail = initialOutputFinished;
+            let settled = false;
+            const cleanup = () => {
+              stream.removeListener("data", onData);
+              stream.removeListener("end", onFinished);
+              stream.removeListener("close", onFinished);
+              stream.removeListener("error", onError);
+            };
+            const settle = (error?: unknown) => {
+              if (settled) return;
+              settled = true;
+              cleanup();
+              liveOutputTail.then(
+                () => (error === undefined ? resolve() : reject(error)),
+                reject,
+              );
+            };
+            const onData = (chunk: Buffer | string) => {
+              stream.pause();
+              liveOutputTail = liveOutputTail
+                .then(() => appendOutput(sanitizer.write(decoder.write(chunk))))
+                .then(() => {
+                  if (!settled) stream.resume();
+                });
+              void liveOutputTail.catch(onError);
+            };
+            const onFinished = () => settle();
+            const onError = (error: unknown) => settle(error);
+
+            // Foreground capture pauses the child streams before handing them
+            // off. Install the live listener first, then resume. Pausing again
+            // on each chunk keeps the handoff bounded while initial output is
+            // replayed and retains the chunk even if the stream closes.
+            stream.on("data", onData);
+            stream.once("end", onFinished);
+            stream.once("close", onFinished);
+            stream.once("error", onError);
+            void initialOutputFinished.catch(onError);
+            stream.resume();
+          })
+        : Promise.resolve();
+
+      await Promise.all([initialOutputFinished, liveOutputFinished]);
+
+      // StringDecoder buffers an incomplete trailing UTF-8 sequence. A
+      // manually stopped process may end in the middle of a character, so
+      // discard that partial sequence instead of flushing it as U+FFFD.
+      if (!job.stopRequested) {
+        await appendOutput(sanitizer.write(decoder.end()));
       }
+      await appendOutput(sanitizer.end());
     };
 
     let outputError: unknown;
-    const outputFinished = Promise.all(
-      [child.stdout, child.stderr]
-        .filter((stream) => stream !== null)
-        .map(async (stream) => {
-          const decoder = new StringDecoder("utf8");
-          const sanitizer = new PlainOutputSanitizer();
-          for await (const chunk of stream) {
-            await appendOutput(sanitizer.write(decoder.write(chunk)));
-          }
-
-          // StringDecoder buffers an incomplete trailing UTF-8 sequence. A
-          // manually stopped process may end in the middle of a character, so
-          // discard that partial sequence instead of flushing it as U+FFFD.
-          // For a naturally completed process, preserve Node's usual behavior
-          // for genuinely malformed output by flushing the decoder.
-          if (!job.stopRequested) {
-            await appendOutput(sanitizer.write(decoder.end()));
-          }
-          await appendOutput(sanitizer.end());
-        }),
-    ).catch((error) => {
-      outputError = error;
-      child.kill();
-    });
+    const outputFinished = Promise.all([
+      consumeOutput(child.stdout, initialOutput.stdout),
+      consumeOutput(child.stderr, initialOutput.stderr),
+    ])
+      .finally(() => initialOutput.dispose?.())
+      .catch((error) => {
+        outputError = error;
+        child.kill();
+      });
+
+    if (abortSignal) {
+      const onAbort = () => {
+        if (job.status !== "running" || job.finalizing) return;
+        job.stopRequested = true;
+        child.kill();
+      };
+      abortSignal.addEventListener("abort", onAbort, { once: true });
+      job.disposeAbort = () =>
+        abortSignal.removeEventListener("abort", onAbort);
+      if (abortSignal.aborted) onAbort();
+    }
 
     child.on("close", async (code) => {
       const status = job.stopRequested
@@ -172,6 +267,7 @@ export class BackgroundJobManager {
       finalError =
         closeError instanceof Error ? closeError.message : String(closeError);
     }
+    job.disposeAbort?.();
     job.status = finalStatus;
     job.finalizing = false;