Add web socket protocol support for IDE MCP server (#987)
Co-authored-by: matt korwel <matt.korwel@gmail.com>
This commit is contained in:
parent
ff478781ad
commit
1fcbdef994
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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"
|
||||
},
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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));
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue