dddvchang-mcp-proxy
Version:
Smart MCP proxy with automatic JetBrains IDE discovery, WebSocket support, and intelligent connection naming
221 lines (220 loc) • 7.51 kB
JavaScript
import WebSocket from 'ws';
import { log } from './index.js';
/**
* WebSocket client for MCP protocol
* Provides persistent connection with better performance
*/
export class WebSocketMCPClient {
ws = null;
endpoint = '';
connected = false;
messageHandlers = new Map();
messageIdCounter = 0;
reconnectAttempts = 0;
maxReconnectAttempts = 5;
reconnectDelay = 1000; // Start with 1 second
constructor() {
log('WebSocket client created');
}
/**
* Connect to WebSocket server
*/
async connect(endpoint) {
if (endpoint) {
this.endpoint = endpoint.replace('http://', 'ws://').replace('/api', '/api/mcp-ws/');
}
if (!this.endpoint) {
throw new Error('No endpoint specified for WebSocket connection');
}
return new Promise((resolve, reject) => {
try {
log(`Attempting WebSocket connection to ${this.endpoint}`);
this.ws = new WebSocket(this.endpoint);
this.ws.on('open', () => {
log('WebSocket connection established');
this.connected = true;
this.reconnectAttempts = 0;
this.reconnectDelay = 1000;
resolve();
});
this.ws.on('message', (data) => {
try {
const message = JSON.parse(data.toString());
log('WebSocket message received:', message);
this.handleMessage(message);
}
catch (error) {
log('Error parsing WebSocket message:', error);
}
});
this.ws.on('close', (code, reason) => {
log(`WebSocket connection closed: ${code} - ${reason}`);
this.connected = false;
this.handleDisconnect();
});
this.ws.on('error', (error) => {
log('WebSocket error:', error);
this.connected = false;
reject(error);
});
}
catch (error) {
log('Failed to create WebSocket connection:', error);
reject(error);
}
});
}
/**
* Send a tool request over WebSocket
*/
async sendToolRequest(toolName, args) {
if (!this.connected || !this.ws) {
throw new Error('WebSocket not connected');
}
const messageId = `msg-${++this.messageIdCounter}`;
const message = {
tool: toolName,
args: typeof args === 'string' ? args : JSON.stringify(args),
id: messageId
};
return new Promise((resolve, reject) => {
// Set up response handler
this.messageHandlers.set(messageId, (response) => {
this.messageHandlers.delete(messageId);
if (response.error) {
reject(new Error(response.error));
}
else {
resolve(response);
}
});
// Send message
try {
log(`Sending WebSocket message: ${JSON.stringify(message)}`);
this.ws.send(JSON.stringify(message));
// Set timeout for response
setTimeout(() => {
if (this.messageHandlers.has(messageId)) {
this.messageHandlers.delete(messageId);
reject(new Error(`Timeout waiting for response to ${toolName}`));
}
}, 30000); // 30 second timeout
}
catch (error) {
this.messageHandlers.delete(messageId);
reject(error);
}
});
}
/**
* Handle incoming WebSocket messages
*/
handleMessage(message) {
// Check if this is a response to a request
const requestId = message.requestId || message.id;
if (requestId && this.messageHandlers.has(requestId)) {
const handler = this.messageHandlers.get(requestId);
handler(message);
}
else {
// Handle other message types (e.g., notifications)
log('Received unsolicited message:', message);
}
}
/**
* Handle WebSocket disconnection
*/
handleDisconnect() {
// Clear all pending handlers
this.messageHandlers.forEach((handler, id) => {
handler({ error: 'WebSocket disconnected' });
});
this.messageHandlers.clear();
// Attempt to reconnect
if (this.reconnectAttempts < this.maxReconnectAttempts) {
this.reconnectAttempts++;
log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts}) in ${this.reconnectDelay}ms`);
setTimeout(() => {
this.connect().catch(error => {
log('Reconnection failed:', error);
});
}, this.reconnectDelay);
// Exponential backoff
this.reconnectDelay = Math.min(this.reconnectDelay * 2, 30000);
}
else {
log('Max reconnection attempts reached, giving up');
}
}
/**
* Check if WebSocket is connected
*/
isConnected() {
return this.connected && this.ws !== null && this.ws.readyState === WebSocket.OPEN;
}
/**
* Close WebSocket connection
*/
disconnect() {
if (this.ws) {
log('Closing WebSocket connection');
this.ws.close();
this.ws = null;
this.connected = false;
}
}
/**
* Test connection to WebSocket endpoint
*/
async testConnection(endpoint) {
const wsEndpoint = endpoint.replace('http://', 'ws://').replace('/api', '/api/mcp-ws/');
return new Promise((resolve) => {
try {
const ws = new WebSocket(wsEndpoint);
const timeout = setTimeout(() => {
ws.close();
resolve(false);
}, 5000); // 5 second timeout
ws.on('open', () => {
clearTimeout(timeout);
ws.close();
resolve(true);
});
ws.on('error', () => {
clearTimeout(timeout);
resolve(false);
});
}
catch (error) {
resolve(false);
}
});
}
}
/**
* Test WebSocket endpoint availability
*/
export async function testWebSocketEndpoint(endpoint) {
const wsEndpoint = endpoint.replace('http://', 'ws://').replace('/api', '/api/mcp-ws');
return new Promise((resolve) => {
try {
const ws = new WebSocket(wsEndpoint);
const timeout = setTimeout(() => {
ws.close();
resolve(false);
}, 5000); // 5 second timeout
ws.on('open', () => {
clearTimeout(timeout);
ws.close();
resolve(true);
});
ws.on('error', () => {
clearTimeout(timeout);
resolve(false);
});
}
catch (error) {
resolve(false);
}
});
}