feat: add abort handling with raceAbort helper

index.ts

@@ -201,23 +201,38 @@ function streamCommandCode(
     };
 
     const controller = new AbortController();
+    let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
+
+    const abortUpstream = () => {
+      if (!controller.signal.aborted) controller.abort();
+      try { reader?.cancel().catch(() => undefined); } catch { /* best-effort */ }
+    };
+
+    if (options?.signal?.aborted) {
+      abortUpstream();
+    } else {
+      options?.signal?.addEventListener("abort", abortUpstream, { once: true });
+    }
+
+    // Helper: race a promise against the abort signal.
+    const raceAbort = <T>(promise: Promise<T>): Promise<T> => {
+      if (controller.signal.aborted) {
+        return Promise.reject(new DOMException("The operation was aborted", "AbortError"));
+      }
+      return new Promise<T>((resolve, reject) => {
+        const onAbort = () => reject(new DOMException("The operation was aborted", "AbortError"));
+        controller.signal.addEventListener("abort", onAbort, { once: true });
+        promise.then(
+          (v) => { controller.signal.removeEventListener("abort", onAbort); resolve(v); },
+          (e) => { controller.signal.removeEventListener("abort", onAbort); reject(e); },
+        );
+      });
+    };
 
     try {
       stream.push({ type: "start", partial: output });
 
-      const response = await fetch(`${API_BASE}/alpha/generate`, {
+      let body: unknown = {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
-          Authorization: `Bearer ${apiKey}`,
-          "x-command-code-version": "0.24.1",
-          "x-cli-environment": "production",
-          "x-project-slug": "pi-cc",
-          "x-taste-learning": "false",
-          "x-co-flag": "false",
-          "x-session-id": uuid(),
-        },
-        body: JSON.stringify({
         config: {
           workingDir: process.cwd(),
           date: new Date().toISOString().split("T")[0],
@@ -239,16 +254,34 @@ function streamCommandCode(
           max_tokens: Math.min(options?.maxTokens ?? model.maxTokens, 200_000),
           stream: true,
         },
-        }),
+      };
+
+      const nextBody = await raceAbort(Promise.resolve(options?.onPayload?.(body, model)));
+      if (nextBody !== undefined) body = nextBody;
+
+      const response = await raceAbort(fetch(`${API_BASE}/alpha/generate`, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${apiKey}`,
+          "x-command-code-version": "0.24.1",
+          "x-cli-environment": "production",
+          "x-project-slug": "pi-cc",
+          "x-taste-learning": "false",
+          "x-co-flag": "false",
+          "x-session-id": uuid(),
+          ...options?.headers,
+        },
+        body: JSON.stringify(body),
         signal: controller.signal,
-      });
+      }));
 
       if (!response.ok) {
-        const errBody = await response.text().catch(() => "");
+        const errBody = await raceAbort(response.text().catch(() => ""));
         throw new Error(`Command Code API error ${response.status}: ${errBody.slice(0, 500)}`);
       }
 
-      const reader = response.body?.getReader();
+      reader = response.body?.getReader();
       if (!reader) throw new Error("No response body");
 
       const decoder = new TextDecoder();
@@ -260,13 +293,16 @@ function streamCommandCode(
       let finished = false;
 
       readLoop: for (;;) {
-        const { done, value } = await reader.read();
+        if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError");
+        const { done, value } = await raceAbort(reader.read());
         if (done) break;
+        if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError");
         buffer += decoder.decode(value, { stream: true });
         const lines = buffer.split("\n");
         buffer = lines.pop() ?? "";
 
         for (const line of lines) {
+          if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError");
           const event = parseStreamEventLine(line);
           if (!event) continue;
 
@@ -344,10 +380,19 @@ function streamCommandCode(
       stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output });
       stream.end();
     } catch (error: any) {
+      if (controller.signal.aborted) {
+        output.stopReason = "aborted";
+        output.errorMessage = "Request aborted";
+      } else {
         output.stopReason = "error";
         output.errorMessage = error?.message ?? String(error);
-      stream.push({ type: "error", reason: "error", error: output });
+      }
+      stream.push({ type: "error", reason: output.stopReason, error: output });
       stream.end();
+    } finally {
+      options?.signal?.removeEventListener("abort", abortUpstream);
+      try { await reader?.cancel(); } catch { /* best-effort */ }
+      try { reader?.releaseLock(); } catch { /* may already be released */ }
     }
   })();