kibbutz-mcp 1.0.1 → 1.0.3

package/build/index.js

@@ -8,13 +8,15 @@ import { z } from 'zod';
 import { spawn } from 'child_process';
 import fs from 'fs';
 import path from 'path';
-let wss;
-let extensionSocket = null;
+let mcpSocket = null;
+let pingSocket = null;
 let token = null;
+let port = null;
+let wsPromise = null;
 const map = new Map();
 const server = new McpServer({
     name: 'kibbutz-mcp',
-    version: '1.0.1',
+    version: '1.0.3',
 });
 function findChromeExecutable() {
     const platform = process.platform;
@@ -56,20 +58,32 @@ function openChromeWithExtension(wsPort) {
     });
     child.unref();
 }
-const sendAndWaitForReply = (message, args) => {
-    if (!extensionSocket || extensionSocket.readyState !== WebSocket.OPEN) {
+const sendAndWaitForReply = async (message, args) => {
+    if (!token) {
+        try {
+            await new Promise((resolve, reject) => {
+                const timeout = setTimeout(() => {
+                    wsPromise = null;
+                    reject(new Error('Timeout waiting for MCP extension WebSocket connection'));
+                }, 2000);
+                wsPromise = () => {
+                    clearTimeout(timeout);
+                    resolve();
+                };
+                openChromeWithExtension(port);
+            });
+        }
+        catch (error) { }
+    }
+    if (!mcpSocket ||
+        !pingSocket ||
+        mcpSocket.readyState !== WebSocket.OPEN ||
+        pingSocket.readyState !== WebSocket.OPEN) {
         return Promise.resolve({
             content: [
                 {
                     type: 'text',
-                    text: `Connection failed: The MCP server cannot reach the Chrome extension.
-
-Please check the following:
-1. Is the browser open?
-2. Is the MCP switch inside the extension popup turned ON?
-
-If the issue persists, copy and paste this URL into your browser address bar:
-chrome-extension://bpfjmggaaiigpfahhmpmacfhlemnhhip/kibbutz-mcp.html`,
+                    text: `Connection failed: The MCP server cannot reach the Chrome extension. Open chrome-extension://bpfjmggaaiigpfahhmpmacfhlemnhhip/KIBBUTZ-MCP.html for troubleshooting.`,
                 },
             ],
             isError: true,
@@ -86,7 +100,7 @@ chrome-extension://bpfjmggaaiigpfahhmpmacfhlemnhhip/kibbutz-mcp.html`,
             },
             reject,
         });
-        extensionSocket?.send(JSON.stringify({
+        mcpSocket?.send(JSON.stringify({
             id,
             message,
             args,
@@ -220,30 +234,36 @@ server.registerTool('UPDATE_GROUP_MCP', {
     },
 }, (args) => sendAndWaitForReply('UPDATE_GROUP_MCP', args));
 async function main() {
-    const httpServer = createServer((req, res) => {
-        try {
-            const url = new URL(req.url || '', `http://${req.headers.host}`);
-            const retToken = url.searchParams.get('token');
-            const verify = retToken === token;
-            if (verify) {
-                token = uuidv4();
+    const httpServer = createServer();
+    const wss = new WebSocketServer({ noServer: true });
+    const wssMcp = new WebSocketServer({ noServer: true });
+    wss.on('connection', (ws) => {
+        pingSocket = ws;
+        ws.on('error', console.error);
+        ws.on('close', () => {
+            if (pingSocket === ws) {
+                pingSocket = null;
+                if (mcpSocket) {
+                    wssMcp.clients.forEach((ws) => ws.terminate());
+                    mcpSocket = null;
+                }
+                for (const { reject } of map.values()) {
+                    reject(new Error('Extension WebSocket closed'));
+                }
+                map.clear();
             }
-            res.writeHead(404);
-            res.end(JSON.stringify({ verify }));
-        }
-        catch (e) {
-            console.error('Failed to parse URL:', req.url, e);
-            res.writeHead(404);
-            res.end(JSON.stringify({ verify: false }));
-        }
+        });
     });
-    wss = new WebSocketServer({ noServer: true });
-    wss.on('connection', (ws, request) => {
-        extensionSocket = ws;
+    wssMcp.on('connection', (ws) => {
+        mcpSocket = ws;
+        if (wsPromise) {
+            wsPromise();
+            wsPromise = null;
+        }
         ws.on('error', console.error);
         ws.on('close', () => {
-            if (extensionSocket === ws) {
-                extensionSocket = null;
+            if (mcpSocket === ws) {
+                mcpSocket = null;
             }
             for (const { reject } of map.values()) {
                 reject(new Error('Extension WebSocket closed'));
@@ -251,49 +271,57 @@ async function main() {
             map.clear();
         });
         ws.on('message', (data) => {
-            try {
-                const text = data.toString();
-                if (text === 'ping')
-                    return;
-                const response = JSON.parse(text);
-                if (response.id && map.has(response.id)) {
-                    const pending = map.get(response.id);
-                    pending.resolve(response.result);
-                    map.delete(response.id);
-                }
-            }
-            catch (_e) {
-                console.error('Failed to parse message from extension:', data);
+            const text = data.toString();
+            const response = JSON.parse(text);
+            if (response.id && map.has(response.id)) {
+                const pending = map.get(response.id);
+                pending.resolve(response.result);
+                map.delete(response.id);
             }
         });
     });
     httpServer.on('upgrade', (request, socket, head) => {
-        if (request.headers.origin !== 'chrome-extension://bpfjmggaaiigpfahhmpmacfhlemnhhip') {
-            socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        const { pathname } = new URL(request.url || '', `wss:127.0.0.1:${port}`);
+        if (pathname === '/ping') {
+            if (request.headers['sec-websocket-protocol'] !== token) {
+                socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+                socket.destroy();
+                return;
+            }
+            wss.clients.forEach((ws) => ws.terminate());
+            wss.handleUpgrade(request, socket, head, (ws) => {
+                wss.emit('connection', ws, request);
+            });
+        }
+        else if (pathname === '/mcp') {
+            if (pingSocket?.readyState !== WebSocket.OPEN) {
+                socket.write('HTTP/1.1 403 Forbidden\r\n\r\n');
+                socket.destroy();
+                return;
+            }
+            wssMcp.clients.forEach((ws) => ws.terminate());
+            wssMcp.handleUpgrade(request, socket, head, (ws) => {
+                wssMcp.emit('connection', ws, request);
+            });
+        }
+        else {
             socket.destroy();
-            return;
         }
-        wss.clients.forEach((ws) => ws.terminate());
-        wss.handleUpgrade(request, socket, head, (ws) => {
-            wss.emit('connection', ws, request);
-        });
     });
     httpServer.listen(0, () => {
-        const port = httpServer.address().port;
-        openChromeWithExtension(port);
+        port = httpServer.address().port;
     });
     const transport = new StdioServerTransport();
     await server.connect(transport);
     process.stdin.on('end', () => {
-        try {
-            openChromeWithExtension(0);
-        }
-        catch (err) {
-            console.error('Failed to open Chrome with extension:', err);
-        }
+        wss.clients.forEach((ws) => ws.terminate());
+        wssMcp.clients.forEach((ws) => ws.terminate());
+    });
+    process.on('SIGTERM', () => {
+        wss.clients.forEach((ws) => ws.terminate());
+        wssMcp.clients.forEach((ws) => ws.terminate());
     });
 }
-main().catch((error) => {
-    console.error('Server error:', error);
+main().catch(() => {
     process.exit(1);
 });

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "kibbutz-mcp",
-    "version": "1.0.1",
+    "version": "1.0.3",
     "description": "",
     "main": "build/index.js",
     "scripts": {