diff options
| -rw-r--r-- | package-lock.json | 35 | ||||
| -rw-r--r-- | packages/core/package.json | 4 | ||||
| -rw-r--r-- | packages/core/src/config/config.ts | 2 | ||||
| -rw-r--r-- | packages/core/src/tools/mcp-client.ts | 19 | ||||
| -rw-r--r-- | packages/core/src/tools/websocket-client-transport.ts | 97 |
5 files changed, 149 insertions, 8 deletions
diff --git a/package-lock.json b/package-lock.json index 845de105..1d3c475f 100644 --- a/package-lock.json +++ b/package-lock.json @@ -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 + } + } } } } diff --git a/packages/core/package.json b/packages/core/package.json index 32f1462a..c216905b 100644 --- a/packages/core/package.json +++ b/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" }, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index abc2240b..5e13241d 100644 --- a/packages/core/src/config/config.ts +++ b/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, diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index a7d6e00c..6f498730 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/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 diff --git a/packages/core/src/tools/websocket-client-transport.ts b/packages/core/src/tools/websocket-client-transport.ts new file mode 100644 index 00000000..ff754c0a --- /dev/null +++ b/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)); + } +} |