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

Co-authored-by: matt korwel <matt.korwel@gmail.com>
This commit is contained in:
Shreya Keshive 2025-06-13 13:30:44 +00:00 committed by GitHub
parent ff478781ad
commit 1fcbdef994
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 149 additions and 8 deletions

35
package-lock.json generated
View File

@ -2321,6 +2321,16 @@
"integrity": "sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==",
"license": "MIT"
},
"node_modules/@types/ws": {
"version": "8.18.1",
"resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.18.1.tgz",
"integrity": "sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==",
"dev": true,
"license": "MIT",
"dependencies": {
"@types/node": "*"
}
},
"node_modules/@types/yargs": {
"version": "17.0.33",
"resolved": "https://registry.npmjs.org/@types/yargs/-/yargs-17.0.33.tgz",
@ -10816,12 +10826,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"
},
@ -10837,6 +10849,27 @@
"engines": {
"node": ">= 4"
}
},
"packages/core/node_modules/ws": {
"version": "8.18.2",
"resolved": "https://registry.npmjs.org/ws/-/ws-8.18.2.tgz",
"integrity": "sha512-DMricUmwGZUVr++AEAe2uiVM7UoO9MAVZMDu05UQOaUII0lp+zOzLLU4Xqh/JvTqklB1T4uELaaPBKyjE1r4fQ==",
"license": "MIT",
"engines": {
"node": ">=10.0.0"
},
"peerDependencies": {
"bufferutil": "^4.0.1",
"utf-8-validate": ">=5.0.2"
},
"peerDependenciesMeta": {
"bufferutil": {
"optional": true
},
"utf-8-validate": {
"optional": true
}
}
}
}
}

View File

@ -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"
},

View File

@ -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,

View File

@ -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

View File

@ -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));
}
}