dddvchang-mcp-proxy
Version:
Smart MCP proxy with automatic JetBrains IDE discovery, WebSocket support, and intelligent connection naming
374 lines (373 loc) • 14.6 kB
JavaScript
import { WebSocketMCPClient } from './websocket-client.js';
import { clusterDiscovery } from './cluster-discovery.js';
// Simple logging function to avoid circular imports
const log = (message) => {
const LOG_ENABLED = process.env.LOG_ENABLED === 'true';
if (LOG_ENABLED) {
console.error(message);
}
};
/**
* Load balancing strategies for cluster connections
*/
export var LoadBalancingStrategy;
(function (LoadBalancingStrategy) {
LoadBalancingStrategy["ROUND_ROBIN"] = "round_robin";
LoadBalancingStrategy["LEAST_CONNECTIONS"] = "least_connections";
LoadBalancingStrategy["PROJECT_AFFINITY"] = "project_affinity";
LoadBalancingStrategy["RANDOM"] = "random";
})(LoadBalancingStrategy || (LoadBalancingStrategy = {}));
/**
* Enhanced WebSocket manager with cluster awareness and load balancing
*/
export class ClusterWebSocketManager {
connectionPool = new Map();
currentRoundRobinIndex = 0;
strategy = LoadBalancingStrategy.PROJECT_AFFINITY;
maxConnectionsPerNode = 10;
connectionHealthCheckInterval = null;
isStarted = false;
constructor(strategy = LoadBalancingStrategy.PROJECT_AFFINITY) {
this.strategy = strategy;
log(`Cluster WebSocket manager created with strategy: ${strategy}`);
}
/**
* Start the cluster WebSocket manager
*/
async start() {
if (this.isStarted) {
return;
}
// Set up cluster discovery callbacks
clusterDiscovery.onNodeDiscovered((node) => {
this.handleNodeDiscovered(node);
});
clusterDiscovery.onNodeRemoved((nodeId) => {
this.handleNodeRemoved(nodeId);
});
// Start cluster discovery
await clusterDiscovery.startDiscovery();
// Start health check interval
this.startHealthChecks();
this.isStarted = true;
log('Cluster WebSocket manager started');
}
/**
* Stop the cluster WebSocket manager
*/
async stop() {
if (!this.isStarted) {
return;
}
// Stop health checks
if (this.connectionHealthCheckInterval) {
clearInterval(this.connectionHealthCheckInterval);
this.connectionHealthCheckInterval = null;
}
// Close all connections
for (const entry of this.connectionPool.values()) {
entry.client.disconnect();
}
this.connectionPool.clear();
// Stop cluster discovery
await clusterDiscovery.stopDiscovery();
this.isStarted = false;
log('Cluster WebSocket manager stopped');
}
/**
* Execute a tool request with load balancing
*/
async executeToolRequest(toolName, args, projectPath) {
const selectedNode = this.selectNode(toolName, projectPath);
if (!selectedNode) {
throw new Error('No healthy cluster nodes available for WebSocket connection');
}
const entry = this.connectionPool.get(selectedNode.id);
if (!entry || !entry.isConnected) {
throw new Error(`Selected node ${selectedNode.id} is not connected`);
}
// Track active request
entry.activeRequests++;
entry.lastUsed = new Date();
try {
log(`Executing tool ${toolName} on node ${selectedNode.id} (${selectedNode.httpEndpoint})`);
const result = await entry.client.sendToolRequest(toolName, args);
// Reset failure count on success
entry.consecutiveFailures = 0;
return result;
}
catch (error) {
// Track failure
entry.consecutiveFailures++;
log(`Tool execution failed on node ${selectedNode.id}: ${error}`);
// If this node has too many failures, try another node
if (entry.consecutiveFailures >= 3) {
log(`Node ${selectedNode.id} has ${entry.consecutiveFailures} consecutive failures, trying alternative`);
return this.executeToolRequestWithFailover(toolName, args, projectPath, selectedNode.id);
}
throw error;
}
finally {
entry.activeRequests = Math.max(0, entry.activeRequests - 1);
}
}
/**
* Execute tool request with failover to another node
*/
async executeToolRequestWithFailover(toolName, args, projectPath, excludeNodeId) {
const healthyNodes = clusterDiscovery.getHealthyNodes()
.filter(node => node.id !== excludeNodeId);
if (healthyNodes.length === 0) {
throw new Error('No alternative healthy nodes available for failover');
}
// Try the first available healthy node
const fallbackNode = healthyNodes[0];
const entry = this.connectionPool.get(fallbackNode.id);
if (!entry || !entry.isConnected) {
throw new Error('Failover node is not connected');
}
entry.activeRequests++;
entry.lastUsed = new Date();
try {
log(`Executing tool ${toolName} on failover node ${fallbackNode.id}`);
const result = await entry.client.sendToolRequest(toolName, args);
entry.consecutiveFailures = 0;
return result;
}
finally {
entry.activeRequests = Math.max(0, entry.activeRequests - 1);
}
}
/**
* Select the best node for a request based on the load balancing strategy
*/
selectNode(toolName, projectPath) {
const healthyConnections = Array.from(this.connectionPool.values())
.filter(entry => entry.isConnected && entry.consecutiveFailures < 3);
if (healthyConnections.length === 0) {
return null;
}
// If there's a preferred node set (for IDE switching), use it if available and healthy
if (this.preferredNodeId) {
const preferredEntry = healthyConnections.find(entry => entry.node.id === this.preferredNodeId);
if (preferredEntry) {
return preferredEntry.node;
}
}
switch (this.strategy) {
case LoadBalancingStrategy.PROJECT_AFFINITY:
return this.selectNodeByProjectAffinity(healthyConnections, projectPath);
case LoadBalancingStrategy.LEAST_CONNECTIONS:
return this.selectNodeByLeastConnections(healthyConnections);
case LoadBalancingStrategy.ROUND_ROBIN:
return this.selectNodeByRoundRobin(healthyConnections);
case LoadBalancingStrategy.RANDOM:
return this.selectNodeRandomly(healthyConnections);
default:
return healthyConnections[0].node;
}
}
/**
* Select node based on project affinity
*/
selectNodeByProjectAffinity(connections, projectPath) {
if (projectPath) {
// First try exact path match
let nodeWithProject = connections.find(entry => entry.node.projectPaths.includes(projectPath));
if (nodeWithProject) {
log(`Found exact project path match: ${projectPath}`);
return nodeWithProject.node;
}
// Try fuzzy matching based on project name
nodeWithProject = connections.find(entry => entry.node.projectPaths.some(path => {
// Extract project name from path and compare
const pathSegments = path.split('/');
const projectName = pathSegments[pathSegments.length - 1];
return projectName.toLowerCase().includes(projectPath.toLowerCase()) ||
projectPath.toLowerCase().includes(projectName.toLowerCase());
}));
if (nodeWithProject) {
log(`Found fuzzy project name match: ${projectPath}`);
return nodeWithProject.node;
}
// Try partial path matching
nodeWithProject = connections.find(entry => entry.node.projectPaths.some(path => path.toLowerCase().includes(projectPath.toLowerCase()) ||
projectPath.toLowerCase().includes(path.toLowerCase())));
if (nodeWithProject) {
log(`Found partial project path match: ${projectPath}`);
return nodeWithProject.node;
}
log(`No project match found for: ${projectPath}, falling back to least connections`);
}
// Fall back to least connections
return this.selectNodeByLeastConnections(connections);
}
/**
* Select node with least active connections
*/
selectNodeByLeastConnections(connections) {
return connections.reduce((min, current) => current.activeRequests < min.activeRequests ? current : min).node;
}
/**
* Select node using round-robin
*/
selectNodeByRoundRobin(connections) {
const selectedEntry = connections[this.currentRoundRobinIndex % connections.length];
this.currentRoundRobinIndex = (this.currentRoundRobinIndex + 1) % connections.length;
return selectedEntry.node;
}
/**
* Select node randomly
*/
selectNodeRandomly(connections) {
const randomIndex = Math.floor(Math.random() * connections.length);
return connections[randomIndex].node;
}
/**
* Handle discovery of a new cluster node
*/
async handleNodeDiscovered(node) {
if (this.connectionPool.has(node.id)) {
// Update existing node info
const entry = this.connectionPool.get(node.id);
entry.node = node;
return;
}
// Create new connection for the node
try {
const client = new WebSocketMCPClient();
await client.connect(node.websocketEndpoint);
const entry = {
node,
client,
isConnected: true,
activeRequests: 0,
lastUsed: new Date(),
consecutiveFailures: 0
};
this.connectionPool.set(node.id, entry);
log(`Added WebSocket connection to cluster node: ${node.id} (${node.websocketEndpoint})`);
}
catch (error) {
log(`Failed to establish WebSocket connection to node ${node.id}: ${error}`);
}
}
/**
* Handle removal of a cluster node
*/
handleNodeRemoved(nodeId) {
const entry = this.connectionPool.get(nodeId);
if (entry) {
entry.client.disconnect();
this.connectionPool.delete(nodeId);
log(`Removed WebSocket connection to cluster node: ${nodeId}`);
}
}
/**
* Start periodic health checks for connections
*/
startHealthChecks() {
this.connectionHealthCheckInterval = setInterval(() => {
this.performHealthChecks();
}, 30000); // Check every 30 seconds
}
/**
* Perform health checks on all connections
*/
async performHealthChecks() {
const checks = Array.from(this.connectionPool.entries()).map(async ([nodeId, entry]) => {
if (!entry.client.isConnected()) {
log(`Node ${nodeId} WebSocket connection is not healthy, attempting reconnection`);
try {
await entry.client.connect(entry.node.websocketEndpoint);
entry.isConnected = true;
entry.consecutiveFailures = 0;
log(`Reconnected to node ${nodeId}`);
}
catch (error) {
entry.isConnected = false;
entry.consecutiveFailures++;
log(`Failed to reconnect to node ${nodeId}: ${error}`);
}
}
});
await Promise.allSettled(checks);
}
/**
* Get connection statistics
*/
getConnectionStats() {
const stats = {
totalNodes: this.connectionPool.size,
connectedNodes: Array.from(this.connectionPool.values()).filter(e => e.isConnected).length,
totalActiveRequests: Array.from(this.connectionPool.values()).reduce((sum, e) => sum + e.activeRequests, 0),
strategy: this.strategy,
nodeDetails: Array.from(this.connectionPool.entries()).map(([nodeId, entry]) => ({
nodeId,
endpoint: entry.node.httpEndpoint,
isConnected: entry.isConnected,
activeRequests: entry.activeRequests,
consecutiveFailures: entry.consecutiveFailures,
lastUsed: entry.lastUsed,
projectPaths: entry.node.projectPaths
}))
};
return stats;
}
/**
* Change load balancing strategy
*/
setLoadBalancingStrategy(strategy) {
this.strategy = strategy;
this.currentRoundRobinIndex = 0; // Reset round-robin counter
log(`Load balancing strategy changed to: ${strategy}`);
}
/**
* Check if cluster manager has any healthy connections
*/
hasHealthyConnections() {
return Array.from(this.connectionPool.values())
.some(entry => entry.isConnected && entry.consecutiveFailures < 3);
}
/**
* Set a preferred node for routing (used for IDE switching)
*/
preferredNodeId = null;
async setPreferredNode(nodeId) {
const entry = this.connectionPool.get(nodeId);
if (!entry) {
throw new Error(`Node ${nodeId} not found in connection pool`);
}
if (!entry.isConnected) {
// Try to reconnect if not connected
try {
await entry.client.connect(entry.node.websocketEndpoint);
entry.isConnected = true;
entry.consecutiveFailures = 0;
log(`Reconnected to preferred node: ${nodeId}`);
}
catch (error) {
throw new Error(`Failed to connect to node ${nodeId}: ${error}`);
}
}
this.preferredNodeId = nodeId;
log(`Set preferred node to: ${nodeId}`);
}
/**
* Clear preferred node setting
*/
clearPreferredNode() {
this.preferredNodeId = null;
log("Cleared preferred node setting");
}
/**
* Get currently preferred node ID
*/
getPreferredNodeId() {
return this.preferredNodeId;
}
}
/**
* Singleton instance for cluster WebSocket management
*/
export const clusterWebSocketManager = new ClusterWebSocketManager();