fix(codex): make write_stdin polls abortable (#2627)

Co-authored-by: Letta Code <noreply@letta.com>

Author: cpacker

src/tools/exec-command.test.ts

@@ -97,6 +97,34 @@ describe.skipIf(isWindows)("Codex unified exec tools", () => {
     expect(second.output).toContain("Output:\ndone");
   });
 
+  test("empty write_stdin polls abort promptly", async () => {
+    const first = await exec_command({
+      cmd: "sleep 2",
+      yield_time_ms: 250,
+    });
+
+    const match = first.output.match(/Process running with session ID (\d+)/);
+    expect(match?.[1]).toBeDefined();
+
+    const controller = new AbortController();
+    const timer = setTimeout(() => controller.abort(), 50);
+    const startedAt = Date.now();
+    try {
+      await expect(
+        write_stdin({
+          session_id: Number(match?.[1]),
+          chars: "",
+          yield_time_ms: 30_000,
+          signal: controller.signal,
+        }),
+      ).rejects.toThrow("The operation was aborted");
+    } finally {
+      clearTimeout(timer);
+    }
+
+    expect(Date.now() - startedAt).toBeLessThan(1000);
+  });
+
   test("write_stdin sends input to tty-enabled sessions", async () => {
     const first = await exec_command({
       cmd: "cat",

src/tools/impl/exec-command.ts

@@ -52,6 +52,7 @@ interface WriteStdinArgs {
   chars?: string;
   yield_time_ms?: number;
   max_output_tokens?: number;
+  signal?: AbortSignal;
   onOutput?: (chunk: string, stream: "stdout" | "stderr") => void;
 }
 
@@ -129,8 +130,35 @@ type ExecOutputChunk = {
 
 const execSessions = new Map<string, ExecSession>();
 
-function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
+function createAbortError(): Error {
+  const error = new Error("The operation was aborted");
+  error.name = "AbortError";
+  return error;
+}
+
+function throwIfAborted(signal: AbortSignal | undefined): void {
+  if (signal?.aborted) {
+    throw createAbortError();
+  }
+}
+
+function sleep(ms: number, signal?: AbortSignal): Promise<void> {
+  if (!signal) {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+  throwIfAborted(signal);
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(() => {
+      signal.removeEventListener("abort", onAbort);
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      clearTimeout(timer);
+      signal.removeEventListener("abort", onAbort);
+      reject(createAbortError());
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+  });
 }
 
 function clampYieldTime(value: number | undefined, fallback: number): number {
@@ -568,12 +596,14 @@ async function waitForSessionOutput(params: {
   session: ExecSession;
   startOffset: number;
   yieldTimeMs: number;
+  signal?: AbortSignal;
   onOutput?: (chunk: string, stream: "stdout" | "stderr") => void;
 }): Promise<{ output: string; wallTimeMs: number }> {
   const startTime = Date.now();
   const deadline = startTime + params.yieldTimeMs;
   let emittedOffset = params.startOffset;
 
+  throwIfAborted(params.signal);
   while (Date.now() < deadline && params.session.status === "running") {
     if (params.onOutput && params.session.output.length > emittedOffset) {
       for (const chunk of getSessionOutputChunks(
@@ -585,8 +615,9 @@ async function waitForSessionOutput(params: {
       }
       emittedOffset = params.session.output.length;
     }
-    await sleep(25);
+    await sleep(25, params.signal);
   }
+  throwIfAborted(params.signal);
 
   if (params.onOutput && params.session.output.length > emittedOffset) {
     for (const chunk of getSessionOutputChunks(
@@ -690,6 +721,7 @@ export async function exec_command(
     session,
     startOffset: 0,
     yieldTimeMs,
+    signal: args.signal,
     onOutput: args.onOutput,
   });
 
@@ -731,7 +763,7 @@ export async function write_stdin(
   }
   if (chars) {
     (backgroundProcess.process as ProcessLauncher).write(chars);
-    await sleep(100);
+    await sleep(100, args.signal);
   }
 
   const startOffset = session.readOffset;
@@ -740,6 +772,7 @@ export async function write_stdin(
     session,
     startOffset,
     yieldTimeMs,
+    signal: args.signal,
     onOutput: args.onOutput,
   });