From c82d633d8d3589f52090caf4ca06f148571afe91 Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Mon, 11 Nov 2024 22:08:46 +0000
Subject: [PATCH] Disconnect from client side if initialization fails

Resolves #24.
---
 src/client/index.test.ts |  2 ++
 src/client/index.ts      | 62 +++++++++++++++++++++++-----------------
 2 files changed, 37 insertions(+), 27 deletions(-)

diff --git a/src/client/index.test.ts b/src/client/index.test.ts
index a6dcc732..86b12ed7 100644
--- a/src/client/index.test.ts
+++ b/src/client/index.test.ts
@@ -122,6 +122,8 @@ test("should reject unsupported protocol version", async () => {
   await expect(client.connect(clientTransport)).rejects.toThrow(
     "Server's protocol version is not supported: invalid-version",
   );
+
+  expect(clientTransport.close).toHaveBeenCalled();
 });
 
 /*
diff --git a/src/client/index.ts b/src/client/index.ts
index c3662b29..3b194f21 100644
--- a/src/client/index.ts
+++ b/src/client/index.ts
@@ -32,7 +32,7 @@ import {
   ServerCapabilities,
   SubscribeRequest,
   SUPPORTED_PROTOCOL_VERSIONS,
-  UnsubscribeRequest
+  UnsubscribeRequest,
 } from "../types.js";
 
 /**
@@ -82,34 +82,40 @@ export class Client<
   override async connect(transport: Transport): Promise<void> {
     await super.connect(transport);
 
-    const result = await this.request(
-      {
-        method: "initialize",
-        params: {
-          protocolVersion: LATEST_PROTOCOL_VERSION,
-          capabilities: {},
-          clientInfo: this._clientInfo,
+    try {
+      const result = await this.request(
+        {
+          method: "initialize",
+          params: {
+            protocolVersion: LATEST_PROTOCOL_VERSION,
+            capabilities: {},
+            clientInfo: this._clientInfo,
+          },
         },
-      },
-      InitializeResultSchema,
-    );
-
-    if (result === undefined) {
-      throw new Error(`Server sent invalid initialize result: ${result}`);
-    }
-
-    if (!SUPPORTED_PROTOCOL_VERSIONS.includes(result.protocolVersion)) {
-      throw new Error(
-        `Server's protocol version is not supported: ${result.protocolVersion}`,
+        InitializeResultSchema,
       );
+
+      if (result === undefined) {
+        throw new Error(`Server sent invalid initialize result: ${result}`);
+      }
+
+      if (!SUPPORTED_PROTOCOL_VERSIONS.includes(result.protocolVersion)) {
+        throw new Error(
+          `Server's protocol version is not supported: ${result.protocolVersion}`,
+        );
+      }
+
+      this._serverCapabilities = result.capabilities;
+      this._serverVersion = result.serverInfo;
+
+      await this.notification({
+        method: "notifications/initialized",
+      });
+    } catch (error) {
+      // Disconnect if initialization fails.
+      void this.close();
+      throw error;
     }
-
-    this._serverCapabilities = result.capabilities;
-    this._serverVersion = result.serverInfo;
-
-    await this.notification({
-      method: "notifications/initialized",
-    });
   }
 
   /**
@@ -219,7 +225,9 @@ export class Client<
 
   async callTool(
     params: CallToolRequest["params"],
-    resultSchema: typeof CallToolResultSchema | typeof CompatibilityCallToolResultSchema = CallToolResultSchema,
+    resultSchema:
+      | typeof CallToolResultSchema
+      | typeof CompatibilityCallToolResultSchema = CallToolResultSchema,
     onprogress?: ProgressCallback,
   ) {
     return this.request(