feat(unstable): Add support for request cancellation (#195)

* Add JSON-RPC request cancellation support

* Abort connection before rejecting pending responses

* Keep cancelled requests pending until response

* Handle abort errors in JSON-RPC requests

* Rename request cancellation signal option

* Queue already-aborted request cancellations

Author: benbrandt

scripts/generate.js

@@ -75,6 +75,8 @@ async function main() {
 
 export const CLIENT_METHODS = ${JSON.stringify(metadata.clientMethods, null, 2)} as const;
 
+export const PROTOCOL_METHODS = ${JSON.stringify(metadata.protocolMethods, null, 2)} as const;
+
 export const PROTOCOL_VERSION = ${metadata.version};
 `,
     { parser: "typescript" },

src/acp.test.ts

@@ -721,7 +721,8 @@ describe("Connection", () => {
     const appAgent = createAgent({ name: "app-agent" })
       .onRequest(AGENT_METHODS.initialize, (c) => {
         events.push(`initialize:${c.params.protocolVersion}`);
-        expect(Object.keys(c).sort()).toEqual(["client", "params"]);
+        expect(Object.keys(c).sort()).toEqual(["client", "params", "signal"]);
+        expect(c.signal.aborted).toBe(false);
 
         return {
           protocolVersion: c.params.protocolVersion,
@@ -742,7 +743,7 @@ describe("Connection", () => {
           },
         },
         (c) => {
-          expect(Object.keys(c).sort()).toEqual(["client", "params"]);
+          expect(Object.keys(c).sort()).toEqual(["client", "params", "signal"]);
           events.push(`agent-route:${String(c.params.message)}`);
         },
       )
@@ -777,7 +778,8 @@ describe("Connection", () => {
           return { message: String(message).toUpperCase() };
         },
         (c) => {
-          expect(Object.keys(c).sort()).toEqual(["agent", "params"]);
+          expect(Object.keys(c).sort()).toEqual(["agent", "params", "signal"]);
+          expect(c.signal.aborted).toBe(false);
           events.push(`client-route:${String(c.params.message)}`);
 
           return { message: c.params.message };
@@ -824,6 +826,78 @@ describe("Connection", () => {
     ]);
   });
 
+  it("aborts app request context signals for protocol cancellation", async () => {
+    const requestSignal = Promise.withResolvers<AbortSignal>();
+
+    const appAgent = createAgent({ name: "cancel-signal-agent" }).onRequest(
+      "vendor/slow",
+      (params) => params as Record<string, never>,
+      async (c) => {
+        requestSignal.resolve(c.signal);
+        await new Promise<void>((resolve) => {
+          c.signal.addEventListener("abort", () => resolve(), { once: true });
+        });
+
+        return { cancelled: c.signal.aborted };
+      },
+    );
+    const appClient = createClient({ name: "cancel-signal-client" });
+
+    const result = await appClient.connectWith(appAgent, async (agentCx) => {
+      const response = agentCx.request<{ cancelled: boolean }>(
+        "vendor/slow",
+        {},
+      );
+      const signal = await requestSignal.promise;
+
+      expect(signal.aborted).toBe(false);
+      await agentCx.notify(methods.protocol.cancelRequest, { requestId: 0 });
+
+      return response;
+    });
+
+    expect(result).toEqual({ cancelled: true });
+  });
+
+  it("maps app request abort errors to request cancellation", async () => {
+    const requestSignal = Promise.withResolvers<AbortSignal>();
+
+    const appAgent = createAgent({ name: "abort-error-agent" }).onRequest(
+      "vendor/abort",
+      (params) => params as Record<string, never>,
+      async (c) => {
+        requestSignal.resolve(c.signal);
+        await new Promise<void>((_, reject) => {
+          c.signal.addEventListener(
+            "abort",
+            () => {
+              const error = new Error("aborted");
+              error.name = "AbortError";
+              reject(error);
+            },
+            { once: true },
+          );
+        });
+
+        return {};
+      },
+    );
+    const appClient = createClient({ name: "abort-error-client" });
+
+    await appClient.connectWith(appAgent, async (agentCx) => {
+      const response = agentCx.request("vendor/abort", {});
+      const signal = await requestSignal.promise;
+
+      expect(signal.aborted).toBe(false);
+      await agentCx.notify(methods.protocol.cancelRequest, { requestId: 0 });
+
+      await expect(response).rejects.toMatchObject({
+        code: -32800,
+        message: "Request cancelled",
+      });
+    });
+  });
+
   it("returns peer contexts from app connection handles", async () => {
     const events: string[] = [];