feat: WebSocket server transport implementation

package.json

@@ -69,15 +69,15 @@
     "@types/jest": "^29.5.12",
     "@types/node": "^22.0.2",
     "@types/supertest": "^6.0.2",
-    "@types/ws": "^8.5.12",
+    "@types/ws": "^8.18.0",
     "eslint": "^9.8.0",
     "jest": "^29.7.0",
     "supertest": "^7.0.0",
     "ts-jest": "^29.2.4",
     "tsx": "^4.16.5",
     "typescript": "^5.5.4",
     "typescript-eslint": "^8.0.0",
-    "ws": "^8.18.0"
+    "ws": "^8.18.1"
   },
   "resolutions": {
     "strip-ansi": "6.0.1"

src/server/index.ts

@@ -29,6 +29,10 @@ import {
   SUPPORTED_PROTOCOL_VERSIONS,
 } from "../types.js";
 
+// Export server transports
+export { StdioServerTransport } from "./stdio.js";
+export { WebSocketServerTransport } from "./websocket.js";
+
 export type ServerOptions = ProtocolOptions & {
   /**
    * Capabilities to advertise as being supported by this server.

src/server/websocket.ts

@@ -0,0 +1,108 @@
+import WebSocket from "ws";
+import { JSONRPCMessage } from "../types.js";
+import { Transport } from "../shared/transport.js";
+
+/**
+ * Server transport for WebSockets: this communicates with a MCP client over a single WebSocket connection.
+ *
+ * This transport is designed to be used with a WebSocket server implementation (like one built with `ws` or `express-ws`).
+ * You would typically create an instance of this transport for each incoming WebSocket connection.
+ */
+export class WebSocketServerTransport implements Transport {
+  private _started = false;
+
+  constructor(private _ws: WebSocket) {}
+
+  onclose?: () => void;
+  onerror?: (error: Error) => void;
+  onmessage?: (message: JSONRPCMessage) => void;
+
+  // Arrow functions to bind `this` properly
+  private _onMessageHandler = (data: WebSocket.RawData) => {
+    try {
+      const messageStr = data.toString("utf-8");
+      // TODO: Add robust JSON parsing and validation, potentially using zod
+      const message: JSONRPCMessage = JSON.parse(messageStr);
+      this.onmessage?.(message);
+    } catch (error) {
+      // Handle JSON parsing errors or other issues
+      this.onerror?.(
+        error instanceof Error ? error : new Error("Failed to process message"),
+      );
+    }
+  };
+
+  private _onErrorHandler = (error: Error) => {
+    this.onerror?.(error);
+  };
+
+  private _onCloseHandler = () => {
+    this.onclose?.();
+    // Clean up listeners after close
+    this._ws.off("message", this._onMessageHandler);
+    this._ws.off("error", this._onErrorHandler);
+    this._ws.off("close", this._onCloseHandler);
+  };
+
+  /**
+   * Starts listening for messages on the WebSocket.
+   */
+  async start(): Promise<void> {
+    if (this._started) {
+      throw new Error(
+        "WebSocketServerTransport already started! Ensure start() is called only once per connection.",
+      );
+    }
+    if (this._ws.readyState !== WebSocket.OPEN) {
+        throw new Error("WebSocket is not open. Cannot start transport.");
+    }
+
+    this._started = true;
+    this._ws.on("message", this._onMessageHandler);
+    this._ws.on("error", this._onErrorHandler);
+    this._ws.on("close", this._onCloseHandler);
+
+    // Unlike stdio, WebSocket connections are typically already established when the transport is created.
+    // No explicit connection action needed here, just attaching listeners.
+  }
+
+  /**
+   * Closes the WebSocket connection.
+   */
+  async close(): Promise<void> {
+    if (this._ws.readyState === WebSocket.OPEN || this._ws.readyState === WebSocket.CONNECTING) {
+      this._ws.close();
+    }
+    // Ensure listeners are removed even if close was called externally or connection was already closed
+    this._onCloseHandler();
+    this._started = false; // Mark as not started
+  }
+
+  /**
+   * Sends a JSON-RPC message over the WebSocket connection.
+   */
+  send(message: JSONRPCMessage): Promise<void> {
+    return new Promise((resolve, reject) => {
+      if (this._ws.readyState !== WebSocket.OPEN) {
+        return reject(new Error("WebSocket is not open. Cannot send message."));
+      }
+
+      try {
+        const json = JSON.stringify(message);
+        this._ws.send(json, (error) => {
+          if (error) {
+            this.onerror?.(error); // Notify via onerror
+            reject(error); // Reject the promise
+          } else {
+            resolve();
+          }
+        });
+      } catch (error) {
+        // Handle JSON stringification errors
+        const err = error instanceof Error ? error : new Error("Failed to serialize message");
+        this.onerror?.(err);
+        reject(err);
+      }
+    });
+  }
+}