From 6ec4c09325051d5296246ee6843dd8b5be69908e Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Thu, 14 Nov 2024 14:11:11 +0000
Subject: [PATCH] Add tests for issuing cancellation notifications

---
 src/client/index.test.ts | 55 ++++++++++++++++++++++++++++++++
 src/server/index.test.ts | 68 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+)

diff --git a/src/client/index.test.ts b/src/client/index.test.ts
index 5610a629..0d7eb941 100644
--- a/src/client/index.test.ts
+++ b/src/client/index.test.ts
@@ -436,3 +436,58 @@ test("should typecheck", () => {
       },
     });
 });
+
+test("should handle client cancelling a request", async () => {
+  const server = new Server(
+    {
+      name: "test server",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        resources: {},
+      },
+    },
+  );
+
+  // Set up server to delay responding to listResources
+  server.setRequestHandler(
+    ListResourcesRequestSchema,
+    async (request, extra) => {
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      return {
+        resources: [],
+      };
+    },
+  );
+
+  const [clientTransport, serverTransport] =
+    InMemoryTransport.createLinkedPair();
+
+  const client = new Client(
+    {
+      name: "test client",
+      version: "1.0",
+    },
+    {
+      capabilities: {},
+    },
+  );
+
+  await Promise.all([
+    client.connect(clientTransport),
+    server.connect(serverTransport),
+  ]);
+
+  // Set up abort controller
+  const controller = new AbortController();
+
+  // Issue request but cancel it immediately
+  const listResourcesPromise = client.listResources(undefined, {
+    signal: controller.signal,
+  });
+  controller.abort("Cancelled by test");
+
+  // Request should be rejected
+  await expect(listResourcesPromise).rejects.toBe("Cancelled by test");
+});
diff --git a/src/server/index.test.ts b/src/server/index.test.ts
index d30c670b..0697cc5c 100644
--- a/src/server/index.test.ts
+++ b/src/server/index.test.ts
@@ -407,3 +407,71 @@ test("should typecheck", () => {
     },
   );
 });
+
+test("should handle server cancelling a request", async () => {
+  const server = new Server(
+    {
+      name: "test server",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        sampling: {},
+      },
+    },
+  );
+
+  const client = new Client(
+    {
+      name: "test client",
+      version: "1.0",
+    },
+    {
+      capabilities: {
+        sampling: {},
+      },
+    },
+  );
+
+  // Set up client to delay responding to createMessage
+  client.setRequestHandler(
+    CreateMessageRequestSchema,
+    async (_request, extra) => {
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      return {
+        model: "test",
+        role: "assistant",
+        content: {
+          type: "text",
+          text: "Test response",
+        },
+      };
+    },
+  );
+
+  const [clientTransport, serverTransport] =
+    InMemoryTransport.createLinkedPair();
+
+  await Promise.all([
+    client.connect(clientTransport),
+    server.connect(serverTransport),
+  ]);
+
+  // Set up abort controller
+  const controller = new AbortController();
+
+  // Issue request but cancel it immediately
+  const createMessagePromise = server.createMessage(
+    {
+      messages: [],
+      maxTokens: 10,
+    },
+    {
+      signal: controller.signal,
+    },
+  );
+  controller.abort("Cancelled by test");
+
+  // Request should be rejected
+  await expect(createMessagePromise).rejects.toBe("Cancelled by test");
+});