Add web socket protocol support for IDE MCP server (#987)

Co-authored-by: matt korwel <matt.korwel@gmail.com>

packages/core/package.json

@@ -39,12 +39,14 @@
     "shell-quote": "^1.8.2",
     "simple-git": "^3.28.0",
     "strip-ansi": "^7.1.0",
-    "undici": "^7.10.0"
+    "undici": "^7.10.0",
+    "ws": "^8.18.0"
   },
   "devDependencies": {
     "@types/diff": "^7.0.2",
     "@types/dotenv": "^6.1.1",
     "@types/minimatch": "^5.1.2",
+    "@types/ws": "^8.5.10",
     "typescript": "^5.3.3",
     "vitest": "^3.1.1"
   },

packages/core/src/config/config.ts

@@ -45,6 +45,8 @@ export class MCPServerConfig {
     readonly cwd?: string,
     // For sse transport
     readonly url?: string,
+    // For websocket transport
+    readonly tcp?: string,
     // Common
     readonly timeout?: number,
     readonly trust?: boolean,

packages/core/src/tools/mcp-client.ts

@@ -12,6 +12,7 @@ import { MCPServerConfig } from '../config/config.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
 import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
 import { ToolRegistry } from './tool-registry.js';
+import { WebSocketClientTransport } from './websocket-client-transport.js';
 
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 
@@ -164,6 +165,8 @@ async function connectAndDiscover(
   let transport;
   if (mcpServerConfig.url) {
     transport = new SSEClientTransport(new URL(mcpServerConfig.url));
+  } else if (mcpServerConfig.tcp) {
+    transport = new WebSocketClientTransport(new URL(mcpServerConfig.tcp));
   } else if (mcpServerConfig.command) {
     transport = new StdioClientTransport({
       command: mcpServerConfig.command,
@@ -177,7 +180,7 @@ async function connectAndDiscover(
     });
   } else {
     console.error(
-      `MCP server '${mcpServerName}' has invalid configuration: missing both url (for SSE) and command (for stdio). Skipping.`,
+      `MCP server '${mcpServerName}' has invalid configuration: missing url (for SSE), tcp (for websocket), and command (for stdio). Skipping.`,
     );
     // Update status to disconnected
     updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
@@ -254,9 +257,11 @@ async function connectAndDiscover(
       console.error(
         `MCP server '${mcpServerName}' did not return valid tool function declarations. Skipping.`,
       );
-      if (transport instanceof StdioClientTransport) {
-        await transport.close();
-      } else if (transport instanceof SSEClientTransport) {
+      if (
+        transport instanceof StdioClientTransport ||
+        transport instanceof SSEClientTransport ||
+        transport instanceof WebSocketClientTransport
+      ) {
         await transport.close();
       }
       // Update status to disconnected
@@ -315,7 +320,8 @@ async function connectAndDiscover(
     // Ensure transport is cleaned up on error too
     if (
       transport instanceof StdioClientTransport ||
-      transport instanceof SSEClientTransport
+      transport instanceof SSEClientTransport ||
+      transport instanceof WebSocketClientTransport
     ) {
       await transport.close();
     }
@@ -334,7 +340,8 @@ async function connectAndDiscover(
     );
     if (
       transport instanceof StdioClientTransport ||
-      transport instanceof SSEClientTransport
+      transport instanceof SSEClientTransport ||
+      transport instanceof WebSocketClientTransport
     ) {
       await transport.close();
       // Update status to disconnected

packages/core/src/tools/websocket-client-transport.ts

@@ -0,0 +1,97 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import WebSocket from 'ws';
+import {
+  Transport,
+  TransportSendOptions,
+} from '@modelcontextprotocol/sdk/shared/transport.js';
+import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
+import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
+
+export class WebSocketClientTransport implements Transport {
+  private socket: WebSocket | null = null;
+  onclose?: () => void;
+  onerror?: (error: Error) => void;
+  onmessage?: (
+    message: JSONRPCMessage,
+    extra?: { authInfo?: AuthInfo },
+  ) => void;
+
+  constructor(private readonly url: URL) {}
+
+  async start(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      const handshakeTimeoutDuration = 10000;
+      let connectionTimeout: NodeJS.Timeout | null = null;
+
+      try {
+        this.socket = new WebSocket(this.url.toString(), {
+          handshakeTimeout: handshakeTimeoutDuration,
+        });
+
+        connectionTimeout = setTimeout(() => {
+          this.socket?.close();
+          reject(
+            new Error(
+              `WebSocket connection timed out after ${handshakeTimeoutDuration}ms`,
+            ),
+          );
+        }, handshakeTimeoutDuration);
+
+        this.socket.on('open', () => {
+          clearTimeout(connectionTimeout!);
+          resolve();
+        });
+
+        this.socket.on('message', (data) => {
+          try {
+            const parsedMessage: JSONRPCMessage = JSON.parse(data.toString());
+            this.onmessage?.(parsedMessage, { authInfo: undefined }); // Auth unsupported currently
+          } catch (error: unknown) {
+            this.onerror?.(
+              error instanceof Error ? error : new Error(String(error)),
+            );
+          }
+        });
+
+        this.socket.on('error', (error) => {
+          clearTimeout(connectionTimeout!);
+          this.onerror?.(error);
+          reject(error);
+        });
+
+        this.socket.on('close', () => {
+          clearTimeout(connectionTimeout!);
+          this.onclose?.();
+          this.socket = null;
+        });
+      } catch (error: unknown) {
+        clearTimeout(connectionTimeout!);
+        reject(error instanceof Error ? error : new Error(String(error)));
+      }
+    });
+  }
+
+  async close(): Promise<void> {
+    if (this.socket) {
+      this.socket.close();
+      this.socket = null;
+    }
+  }
+
+  async send(
+    message: JSONRPCMessage,
+    _options?: TransportSendOptions,
+  ): Promise<void> {
+    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
+      throw new Error(
+        'WebSocket is not connected or not open. Cannot send message.',
+      );
+    }
+    this.socket.send(JSON.stringify(message));
+  }
+}