From c39f8084d739c99e2d2d383e79654ee4f52e8f3d Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Fri, 11 Oct 2024 16:13:12 +0100
Subject: [PATCH] Update client transports

---
 src/client/sse.ts       | 13 +++++++++----
 src/client/stdio.ts     | 28 ++++++++++++++++++++--------
 src/client/websocket.ts |  9 +++++++--
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/client/sse.ts b/src/client/sse.ts
index dd06b77f..2ebf23e8 100644
--- a/src/client/sse.ts
+++ b/src/client/sse.ts
@@ -11,14 +11,19 @@ export class SSEClientTransport implements Transport {
   private _eventSource?: EventSource;
   private _endpoint?: URL;
   private _abortController?: AbortController;
+  private _url: URL;
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
   onmessage?: (message: JSONRPCMessage) => void;
 
-  connect(url: URL): Promise<void> {
+  constructor(url: URL) {
+    this._url = url;
+  }
+
+  start(): Promise<void> {
     return new Promise((resolve, reject) => {
-      this._eventSource = new EventSource(url.href);
+      this._eventSource = new EventSource(this._url.href);
       this._abortController = new AbortController();
 
       this._eventSource.onerror = (event) => {
@@ -35,8 +40,8 @@ export class SSEClientTransport implements Transport {
         const messageEvent = event as MessageEvent;
 
         try {
-          this._endpoint = new URL(messageEvent.data, url);
-          if (this._endpoint.origin !== url.origin) {
+          this._endpoint = new URL(messageEvent.data, this._url);
+          if (this._endpoint.origin !== this._url.origin) {
             throw new Error(
               `Endpoint origin does not match connection origin: ${this._endpoint.origin}`,
             );
diff --git a/src/client/stdio.ts b/src/client/stdio.ts
index 30da0799..36d4f206 100644
--- a/src/client/stdio.ts
+++ b/src/client/stdio.ts
@@ -31,22 +31,34 @@ export class StdioClientTransport implements Transport {
   private _process?: ChildProcess;
   private _abortController: AbortController = new AbortController();
   private _readBuffer: ReadBuffer = new ReadBuffer();
+  private _serverParams: StdioServerParameters;
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
   onmessage?: (message: JSONRPCMessage) => void;
 
+  constructor(server: StdioServerParameters) {
+    this._serverParams = server;
+  }
+
   /**
-   * Spawns the server process and prepare to communicate with it.
+   * Starts the server process and prepares to communicate with it.
    */
-  spawn(server: StdioServerParameters): Promise<void> {
+  async start(): Promise<void> {
     return new Promise((resolve, reject) => {
-      this._process = spawn(server.command, server.args ?? [], {
-        // The parent process may have sensitive secrets in its env, so don't inherit it automatically.
-        env: server.env === undefined ? {} : { ...server.env },
-        stdio: ["pipe", "pipe", "inherit"],
-        signal: this._abortController.signal,
-      });
+      this._process = spawn(
+        this._serverParams.command,
+        this._serverParams.args ?? [],
+        {
+          // The parent process may have sensitive secrets in its env, so don't inherit it automatically.
+          env:
+            this._serverParams.env === undefined
+              ? {}
+              : { ...this._serverParams.env },
+          stdio: ["pipe", "pipe", "inherit"],
+          signal: this._abortController.signal,
+        },
+      );
 
       this._process.on("error", (error) => {
         if (error.name === "AbortError") {
diff --git a/src/client/websocket.ts b/src/client/websocket.ts
index 47856cc7..97dc5529 100644
--- a/src/client/websocket.ts
+++ b/src/client/websocket.ts
@@ -8,14 +8,19 @@ const SUBPROTOCOL = "mcp";
  */
 export class WebSocketClientTransport implements Transport {
   private _socket?: WebSocket;
+  private _url: URL;
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
   onmessage?: (message: JSONRPCMessage) => void;
 
-  connect(url: URL): Promise<void> {
+  constructor(url: URL) {
+    this._url = url;
+  }
+
+  start(): Promise<void> {
     return new Promise((resolve, reject) => {
-      this._socket = new WebSocket(url, SUBPROTOCOL);
+      this._socket = new WebSocket(this._url, SUBPROTOCOL);
 
       this._socket.onerror = (event) => {
         const error =