merge: worker/65d6be2f - MCP server bug fixes

Author: echobt

mcp-server/src/client.ts

@@ -5,11 +5,13 @@
 import * as net from "net";
 
 interface SocketRequest {
+  id: number;
   command: string;
   payload: Record<string, unknown>;
 }
 
 interface SocketResponse {
+  id?: number;
   success: boolean;
   data?: unknown;
   error?: string;
@@ -24,10 +26,10 @@ interface PendingRequest {
 
 // Command-specific timeouts (in ms)
 const COMMAND_TIMEOUTS: Record<string, number> = {
-  takeScreenshot: 60000,  // Screenshots can take longer
-  getDom: 60000,          // DOM can be large
-  executeJs: 60000,       // JS execution varies
-  default: 30000,         // Default timeout
+  takeScreenshot: 60000,
+  getDom: 60000,
+  executeJs: 60000,
+  default: 30000,
 };
 
 export class CortexSocketClient {
@@ -37,9 +39,10 @@ export class CortexSocketClient {
   private connected: boolean = false;
   private connecting: boolean = false;
   private responseBuffer: string = "";
-  private pendingRequest: PendingRequest | null = null;
-  private reconnectAttempts: number = 0;
+  private pendingRequests: Map<number, PendingRequest> = new Map();
+  private nextRequestId: number = 1;
   private maxReconnectAttempts: number = 3;
+  private connectWaiters: Array<{ resolve: () => void; reject: (err: Error) => void }> = [];
 
   constructor(host: string = "127.0.0.1", port: number = 4000) {
     this.host = host;
@@ -52,34 +55,16 @@ export class CortexSocketClient {
     }
 
     if (this.connecting) {
-      // Wait for existing connection attempt
-      return new Promise((resolve, reject) => {
-        const checkInterval = setInterval(() => {
-          if (this.connected) {
-            clearInterval(checkInterval);
-            resolve();
-          } else if (!this.connecting) {
-            clearInterval(checkInterval);
-            reject(new Error("Connection failed"));
-          }
-        }, 100);
-        
-        // Timeout waiting for connection
-        setTimeout(() => {
-          clearInterval(checkInterval);
-          if (!this.connected) {
-            reject(new Error("Connection timeout"));
-          }
-        }, 10000);
+      return new Promise<void>((resolve, reject) => {
+        this.connectWaiters.push({ resolve, reject });
       });
     }
 
     this.connecting = true;
 
-    return new Promise((resolve, reject) => {
+    return new Promise<void>((resolve, reject) => {
       this.socket = new net.Socket();
-      
-      // Set socket timeout
+
       this.socket.setTimeout(60000);
 
       const connectTimeout = setTimeout(() => {
@@ -88,26 +73,31 @@ export class CortexSocketClient {
           this.socket.destroy();
           this.socket = null;
         }
-        reject(new Error("Connection timeout"));
+        const err = new Error("Connection timeout");
+        reject(err);
+        this.drainConnectWaiters(err);
       }, 10000);
 
+      let settled = false;
+
       this.socket.on("connect", () => {
         clearTimeout(connectTimeout);
         this.connected = true;
         this.connecting = false;
-        this.reconnectAttempts = 0;
+        settled = true;
         console.error(`[MCP Client] Connected to Cortex Desktop at ${this.host}:${this.port}`);
         resolve();
+        this.drainConnectWaiters(null);
       });
 
       this.socket.on("data", (data) => {
         this.handleData(data);
       });
 
       this.socket.on("timeout", () => {
-        console.error("[MCP Client] Socket timeout");
-        if (this.pendingRequest) {
-          this.clearPendingRequest(new Error("Socket timeout"));
+        console.error("[MCP Client] Socket idle timeout — destroying connection");
+        if (this.socket) {
+          this.socket.destroy();
         }
       });
 
@@ -116,149 +106,189 @@ export class CortexSocketClient {
         console.error("[MCP Client] Socket error:", err.message);
         this.connected = false;
         this.connecting = false;
-        
-        if (this.pendingRequest) {
-          this.clearPendingRequest(err);
+
+        this.rejectAllPending(err);
+
+        if (!settled) {
+          settled = true;
+          reject(err);
+          this.drainConnectWaiters(err);
         }
-        
-        reject(err);
       });
 
       this.socket.on("close", () => {
         console.error("[MCP Client] Connection closed");
+        const wasClosed = !this.connected;
         this.connected = false;
         this.connecting = false;
         this.socket = null;
-        
-        if (this.pendingRequest) {
-          this.clearPendingRequest(new Error("Connection closed"));
+        this.responseBuffer = "";
+
+        this.rejectAllPending(new Error("Connection closed"));
+
+        if (!settled && !wasClosed) {
+          settled = true;
+          const err = new Error("Connection closed before established");
+          reject(err);
+          this.drainConnectWaiters(err);
         }
       });
 
       this.socket.connect(this.port, this.host);
     });
   }
 
+  private drainConnectWaiters(err: Error | null): void {
+    const waiters = this.connectWaiters;
+    this.connectWaiters = [];
+    for (const w of waiters) {
+      if (err) {
+        w.reject(err);
+      } else {
+        w.resolve();
+      }
+    }
+  }
+
   private handleData(data: Buffer): void {
     this.responseBuffer += data.toString();
-    
-    // Process all complete JSON responses in buffer
+
     let newlineIndex: number;
     while ((newlineIndex = this.responseBuffer.indexOf("\n")) !== -1) {
       const jsonStr = this.responseBuffer.substring(0, newlineIndex);
       this.responseBuffer = this.responseBuffer.substring(newlineIndex + 1);
-      
+
       if (!jsonStr.trim()) continue;
-      
+
       try {
         const response = JSON.parse(jsonStr) as SocketResponse;
-        if (this.pendingRequest) {
-          const { resolve, timeoutId } = this.pendingRequest;
-          clearTimeout(timeoutId);
-          this.pendingRequest = null;
-          resolve(response);
+
+        if (response.id != null && this.pendingRequests.has(response.id)) {
+          const pending = this.pendingRequests.get(response.id)!;
+          this.pendingRequests.delete(response.id);
+          clearTimeout(pending.timeoutId);
+          pending.resolve(response);
+        } else {
+          const firstEntry = this.pendingRequests.entries().next();
+          if (!firstEntry.done) {
+            const [id, pending] = firstEntry.value;
+            this.pendingRequests.delete(id);
+            clearTimeout(pending.timeoutId);
+            pending.resolve(response);
+          }
         }
       } catch (e) {
         console.error("[MCP Client] Failed to parse response:", e, "Raw:", jsonStr.substring(0, 200));
-        if (this.pendingRequest) {
-          this.clearPendingRequest(new Error(`Failed to parse response: ${e}`));
+        const firstEntry = this.pendingRequests.entries().next();
+        if (!firstEntry.done) {
+          const [id, pending] = firstEntry.value;
+          this.pendingRequests.delete(id);
+          clearTimeout(pending.timeoutId);
+          pending.reject(new Error(`Failed to parse response: ${e}`));
         }
       }
     }
   }
 
-  private clearPendingRequest(error: Error): void {
-    if (this.pendingRequest) {
-      const { reject, timeoutId } = this.pendingRequest;
-      clearTimeout(timeoutId);
-      this.pendingRequest = null;
-      reject(error);
+  private rejectAllPending(error: Error): void {
+    const entries = Array.from(this.pendingRequests.entries());
+    this.pendingRequests.clear();
+    for (const [, pending] of entries) {
+      clearTimeout(pending.timeoutId);
+      pending.reject(error);
     }
   }
 
   private getTimeout(command: string): number {
-    return COMMAND_TIMEOUTS[command] || COMMAND_TIMEOUTS.default;
+    return COMMAND_TIMEOUTS[command] ?? COMMAND_TIMEOUTS["default"];
   }
 
   async sendCommand(command: string, payload: Record<string, unknown> = {}): Promise<SocketResponse> {
-    // Ensure we're connected
     if (!this.connected || !this.socket) {
-      try {
-        await this.connect();
-      } catch (err) {
-        // Try to reconnect if initial connection failed
-        if (this.reconnectAttempts < this.maxReconnectAttempts) {
-          this.reconnectAttempts++;
-          console.error(`[MCP Client] Reconnect attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts}`);
-          await new Promise(resolve => setTimeout(resolve, 1000));
-          return this.sendCommand(command, payload);
+      let lastErr: unknown;
+      for (let attempt = 0; attempt <= this.maxReconnectAttempts; attempt++) {
+        try {
+          if (attempt > 0) {
+            console.error(`[MCP Client] Reconnect attempt ${attempt}/${this.maxReconnectAttempts}`);
+            await new Promise(resolve => setTimeout(resolve, 1000 * attempt));
+          }
+          await this.connect();
+          break;
+        } catch (err) {
+          lastErr = err;
+          if (attempt === this.maxReconnectAttempts) {
+            throw lastErr;
+          }
         }
-        throw err;
-      }
-    }
-
-    // Wait if there's already a pending request
-    if (this.pendingRequest) {
-      console.error(`[MCP Client] Waiting for pending request: ${this.pendingRequest.command}`);
-      await new Promise(resolve => setTimeout(resolve, 100));
-      if (this.pendingRequest) {
-        throw new Error("Another request is already in progress");
       }
     }
 
     return new Promise((resolve, reject) => {
-      if (!this.socket) {
+      if (!this.socket || !this.connected) {
         reject(new Error("Socket not connected"));
         return;
       }
 
+      const requestId = this.nextRequestId++;
       const timeout = this.getTimeout(command);
-      
+
       const timeoutId = setTimeout(() => {
-        if (this.pendingRequest) {
+        if (this.pendingRequests.has(requestId)) {
+          this.pendingRequests.delete(requestId);
           console.error(`[MCP Client] Request timeout for command: ${command} (${timeout}ms)`);
-          this.pendingRequest = null;
           reject(new Error(`Request timed out after ${timeout}ms`));
         }
       }, timeout);
 
-      this.pendingRequest = {
+      this.pendingRequests.set(requestId, {
         resolve,
         reject,
         timeoutId,
         command,
-      };
+      });
 
-      const request: SocketRequest = { command, payload };
+      const request: SocketRequest = { id: requestId, command, payload };
       const json = JSON.stringify(request) + "\n";
-      
+
       this.socket.write(json, (err) => {
         if (err) {
-          this.clearPendingRequest(err);
+          if (this.pendingRequests.has(requestId)) {
+            const pending = this.pendingRequests.get(requestId)!;
+            this.pendingRequests.delete(requestId);
+            clearTimeout(pending.timeoutId);
+            pending.reject(err);
+          }
         }
       });
     });
   }
 
   disconnect(): void {
-    if (this.pendingRequest) {
-      this.clearPendingRequest(new Error("Client disconnected"));
-    }
-    
+    this.rejectAllPending(new Error("Client disconnected"));
+
     if (this.socket) {
       this.socket.destroy();
       this.socket = null;
       this.connected = false;
     }
+
+    this.responseBuffer = "";
+    this.connectWaiters = [];
   }
 
   isConnected(): boolean {
-    return this.connected;
+    return this.connected && this.socket !== null;
   }
 }
 
+function parsePort(value: string | undefined, fallback: number): number {
+  if (value == null) return fallback;
+  const parsed = parseInt(value, 10);
+  return Number.isFinite(parsed) && parsed > 0 && parsed <= 65535 ? parsed : fallback;
+}
+
 // Singleton instance
 export const socketClient = new CortexSocketClient(
   process.env.CORTEX_MCP_HOST || "127.0.0.1",
-  parseInt(process.env.CORTEX_MCP_PORT || "4000", 10)
+  parsePort(process.env.CORTEX_MCP_PORT, 4000),
 );