fix: Handle ACP connection transport failures cleanly (#103)

* Handle ACP connection transport failures cleanly

* Harden a few more cases

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

Author: advait

src/acp.test.ts

@@ -44,6 +44,7 @@ import {
   ResumeSessionRequest,
   ResumeSessionResponse,
 } from "./acp.js";
+import type { AnyMessage } from "./acp.js";
 
 describe("Connection", () => {
   let clientToAgent: TransformStream<Uint8Array, Uint8Array>;
@@ -988,6 +989,107 @@ describe("Connection", () => {
     expect(closeLog).toContain("client connection closed (signal)");
   });
 
+  class MinimalTestClient implements Client {
+    async writeTextFile(
+      _: WriteTextFileRequest,
+    ): Promise<WriteTextFileResponse> {
+      return {};
+    }
+    async readTextFile(_: ReadTextFileRequest): Promise<ReadTextFileResponse> {
+      return { content: "test" };
+    }
+    async requestPermission(
+      _: RequestPermissionRequest,
+    ): Promise<RequestPermissionResponse> {
+      return {
+        outcome: {
+          outcome: "selected",
+          optionId: "allow",
+        },
+      };
+    }
+    async sessionUpdate(_: SessionNotification): Promise<void> {
+      // no-op
+    }
+  }
+
+  it("rejects pending requests when the stream errors", async () => {
+    let readableController!: ReadableStreamDefaultController<AnyMessage>;
+
+    const connection = new ClientSideConnection(() => new MinimalTestClient(), {
+      readable: new ReadableStream<AnyMessage>({
+        start(controller) {
+          readableController = controller;
+        },
+      }),
+      writable: new WritableStream<AnyMessage>({
+        async write() {
+          // no-op
+        },
+      }),
+    });
+
+    const requestPromise = connection.newSession({
+      cwd: "/test",
+      mcpServers: [],
+    });
+    const error = new Error("stream exploded");
+
+    readableController.error(error);
+
+    await expect(requestPromise).rejects.toThrow("stream exploded");
+    await expect(connection.closed).resolves.toBeUndefined();
+    expect(connection.signal.aborted).toBe(true);
+  });
+
+  it("rejects pending requests when the writable stream errors", async () => {
+    const writeError = new Error("write failed");
+
+    const connection = new ClientSideConnection(() => new MinimalTestClient(), {
+      readable: new ReadableStream<AnyMessage>({
+        // Never produces messages; stays open.
+        start() {},
+      }),
+      writable: new WritableStream<AnyMessage>({
+        async write() {
+          throw writeError;
+        },
+      }),
+    });
+
+    const requestPromise = connection.newSession({
+      cwd: "/test",
+      mcpServers: [],
+    });
+
+    await expect(requestPromise).rejects.toThrow("write failed");
+    await expect(connection.closed).resolves.toBeUndefined();
+    expect(connection.signal.aborted).toBe(true);
+  });
+
+  it("rejects requests issued after the connection is closed", async () => {
+    const connection = new ClientSideConnection(() => new MinimalTestClient(), {
+      readable: new ReadableStream<AnyMessage>({
+        start(controller) {
+          // Close the readable stream immediately so the connection closes.
+          controller.close();
+        },
+      }),
+      writable: new WritableStream<AnyMessage>({
+        async write() {
+          // no-op
+        },
+      }),
+    });
+
+    await connection.closed;
+    expect(connection.signal.aborted).toBe(true);
+
+    await expect(
+      connection.newSession({ cwd: "/test", mcpServers: [] }),
+    ).rejects.toThrow("ACP connection closed");
+  });
+
   it("supports removing signal event listeners", async () => {
     const closeLog: string[] = [];
 

src/acp.ts

@@ -11,11 +11,15 @@ import type {
   AnyResponse,
   Result,
   ErrorResponse,
-  PendingResponse,
   RequestHandler,
   NotificationHandler,
 } from "./jsonrpc.js";
 
+type ConnectionPendingResponse = {
+  resolve: (response: unknown) => void;
+  reject: (error: unknown) => void;
+};
+
 /**
  * An agent-side connection to a client.
  *
@@ -1167,7 +1171,8 @@ export class ClientSideConnection implements Agent {
 export type { AnyMessage } from "./jsonrpc.js";
 
 class Connection {
-  #pendingResponses: Map<string | number | null, PendingResponse> = new Map();
+  #pendingResponses: Map<string | number | null, ConnectionPendingResponse> =
+    new Map();
   #nextRequestId: number = 0;
   #requestHandler: RequestHandler;
   #notificationHandler: NotificationHandler;
@@ -1187,7 +1192,7 @@ class Connection {
     this.#closedPromise = new Promise((resolve) => {
       this.#abortController.signal.addEventListener("abort", () => resolve());
     });
-    this.#receive();
+    void this.#receive();
   }
 
   /**
@@ -1221,44 +1226,64 @@ class Connection {
   }
 
   async #receive() {
-    const reader = this.#stream.readable.getReader();
+    let closeError: unknown = undefined;
+
     try {
-      while (true) {
-        const { value: message, done } = await reader.read();
-        if (done) {
-          break;
-        }
-        if (!message) {
-          continue;
-        }
+      const reader = this.#stream.readable.getReader();
+      try {
+        while (!this.#abortController.signal.aborted) {
+          const { value: message, done } = await reader.read();
+          if (done) {
+            break;
+          }
+          if (!message) {
+            continue;
+          }
 
-        try {
-          this.#processMessage(message);
-        } catch (err) {
-          console.error(
-            "Unexpected error during message processing:",
-            message,
-            err,
-          );
-          // Only send error response if the message had an id (was a request)
-          if ("id" in message && message.id !== undefined) {
-            this.#sendMessage({
-              jsonrpc: "2.0",
-              id: message.id,
-              error: {
-                code: -32700,
-                message: "Parse error",
-              },
-            });
+          try {
+            this.#processMessage(message);
+          } catch (err) {
+            console.error(
+              "Unexpected error during message processing:",
+              message,
+              err,
+            );
+            // Only send error response if the message had an id (was a request)
+            if ("id" in message && message.id !== undefined) {
+              this.#sendMessage({
+                jsonrpc: "2.0",
+                id: message.id,
+                error: {
+                  code: -32700,
+                  message: "Parse error",
+                },
+              });
+            }
           }
         }
+      } finally {
+        reader.releaseLock();
       }
+    } catch (error) {
+      closeError = error;
     } finally {
-      reader.releaseLock();
-      this.#abortController.abort();
+      this.#close(closeError);
     }
   }
 
+  #close(error?: unknown) {
+    if (this.#abortController.signal.aborted) {
+      return;
+    }
+
+    const closeError: unknown = error ?? new Error("ACP connection closed");
+    for (const pendingResponse of this.#pendingResponses.values()) {
+      pendingResponse.reject(closeError);
+    }
+    this.#pendingResponses.clear();
+    this.#abortController.abort(closeError);
+  }
+
   async #processMessage(message: AnyMessage) {
     if ("method" in message && "id" in message) {
       // It's a request
@@ -1376,7 +1401,8 @@ class Connection {
       if ("result" in response) {
         pendingResponse.resolve(response.result);
       } else if ("error" in response) {
-        pendingResponse.reject(response.error);
+        const { code, message, data } = response.error;
+        pendingResponse.reject(new RequestError(code, message, data));
       }
       this.#pendingResponses.delete(response.id);
     } else {
@@ -1385,6 +1411,7 @@ class Connection {
   }
 
   async sendRequest<Req, Resp>(method: string, params?: Req): Promise<Resp> {
+    this.#throwIfClosed();
     const id = this.#nextRequestId++;
     const responsePromise = new Promise((resolve, reject) => {
       this.#pendingResponses.set(id, { resolve, reject });
@@ -1394,9 +1421,19 @@ class Connection {
   }
 
   async sendNotification<N>(method: string, params?: N): Promise<void> {
+    this.#throwIfClosed();
     await this.#sendMessage({ jsonrpc: "2.0", method, params });
   }
 
+  #throwIfClosed() {
+    if (this.#abortController.signal.aborted) {
+      throw (
+        this.#abortController.signal.reason ??
+        new Error("ACP connection closed")
+      );
+    }
+  }
+
   async #sendMessage(message: AnyMessage) {
     this.#writeQueue = this.#writeQueue
       .then(async () => {
@@ -1408,8 +1445,7 @@ class Connection {
         }
       })
       .catch((error) => {
-        // Continue processing writes on error
-        console.error("ACP write error:", error);
+        this.#close(error);
       });
     return this.#writeQueue;
   }

src/jsonrpc.ts

@@ -36,11 +36,6 @@ export type ErrorResponse = {
   data?: unknown;
 };
 
-export type PendingResponse = {
-  resolve: (response: unknown) => void;
-  reject: (error: ErrorResponse) => void;
-};
-
 export type RequestHandler = (
   method: string,
   params: unknown,