agent-desk-mcp 1.0.4 → 1.1.0

package/dist/index.js

@@ -2,6 +2,7 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { CallToolRequestSchema, ListToolsRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
 import axios from "axios";
+import { WebSocket } from "ws";
 class APIClient {
     client;
     constructor(baseUrl) {
@@ -44,6 +45,8 @@ class AgentDeskServer {
     server;
     apiClient = null;
     apiUrl;
+    ws = null;
+    pendingRequests = new Map();
     constructor(apiUrl) {
         this.apiUrl = apiUrl;
         this.server = new Server({
@@ -157,41 +160,84 @@ Returns: User's answers concatenated with newlines, or timeout message.`,
             id: q.id || `Q${i + 1}`,
         }));
         await this.apiClient.addAgent(sessionId, agent_summary, context, processedQuestions, timeout_seconds);
-        // Poll for answer
-        const pollInterval = 1000; // 1 second
-        const startTime = Date.now();
-        while (true) {
-            // Check timeout
-            if (timeout_seconds && (Date.now() - startTime) >= timeout_seconds * 1000) {
+        // Wait for answer via WebSocket push
+        const answerPromise = new Promise((resolve) => {
+            this.pendingRequests.set(sessionId, resolve);
+        });
+        // Set up timeout if specified
+        let timeoutHandle = null;
+        if (timeout_seconds) {
+            timeoutHandle = setTimeout(() => {
+                this.pendingRequests.delete(sessionId);
                 console.error(`[${new Date().toISOString()}] TIMEOUT handleAskUser: ${sessionId}`);
-                return {
-                    content: [
-                        {
-                            type: "text",
-                            text: "User did not respond within the timeout period.",
-                        },
-                    ],
-                };
+            }, timeout_seconds * 1000);
+        }
+        try {
+            const answer = await answerPromise;
+            if (timeoutHandle)
+                clearTimeout(timeoutHandle);
+            console.error(`[${new Date().toISOString()}] ANSWERED handleAskUser: ${sessionId}`);
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: answer,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            if (timeoutHandle)
+                clearTimeout(timeoutHandle);
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: "User did not respond within the timeout period.",
+                    },
+                ],
+            };
+        }
+    }
+    setupWebSocket() {
+        const wsUrl = this.apiUrl.replace(/^http/, "ws") + "/ws";
+        console.error(`Connecting to WebSocket: ${wsUrl}`);
+        this.ws = new WebSocket(wsUrl);
+        this.ws.on("open", () => {
+            console.error("WebSocket connected");
+        });
+        this.ws.on("message", (data) => {
+            try {
+                const message = JSON.parse(data.toString());
+                console.error(`WebSocket message received: ${JSON.stringify(message)}`);
+                if (message.type === "agent_answered") {
+                    const sessionId = message.session_id;
+                    const resolver = this.pendingRequests.get(sessionId);
+                    if (resolver) {
+                        // Fetch the full answer from API
+                        this.apiClient.getAgent(sessionId).then((agent) => {
+                            if (agent && agent.status === "answered") {
+                                const answers = agent.questions
+                                    .filter((q) => q.answer)
+                                    .map((q) => `${q.id}: ${q.answer}`);
+                                resolver(answers.join("\n"));
+                                this.pendingRequests.delete(sessionId);
+                            }
+                        });
+                    }
+                }
             }
-            // Check for answer
-            const agent = await this.apiClient.getAgent(sessionId);
-            if (agent && agent.status === "answered") {
-                console.error(`[${new Date().toISOString()}] ANSWERED handleAskUser: ${sessionId}`);
-                const answers = agent.questions
-                    .filter((q) => q.answer)
-                    .map((q) => `${q.id}: ${q.answer}`);
-                return {
-                    content: [
-                        {
-                            type: "text",
-                            text: answers.join("\n"),
-                        },
-                    ],
-                };
+            catch (error) {
+                console.error("Error processing WebSocket message:", error);
             }
-            console.error(`[${new Date().toISOString()}] POLLING handleAskUser: ${sessionId}`);
-            await new Promise((resolve) => setTimeout(resolve, pollInterval));
-        }
+        });
+        this.ws.on("error", (error) => {
+            console.error("WebSocket error:", error);
+        });
+        this.ws.on("close", () => {
+            console.error("WebSocket closed, reconnecting in 5s...");
+            setTimeout(() => this.setupWebSocket(), 5000);
+        });
     }
     async run() {
         this.apiClient = new APIClient(this.apiUrl);
@@ -202,6 +248,8 @@ Returns: User's answers concatenated with newlines, or timeout message.`,
             console.error("Please ensure the API Server is running.");
             process.exit(1);
         }
+        // Setup WebSocket connection
+        this.setupWebSocket();
         const transport = new StdioServerTransport();
         await this.server.connect(transport);
         console.error("Agent Desk MCP Server started");

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-desk-mcp",
-  "version": "1.0.4",
+  "version": "1.1.0",
   "description": "MCP server for multi-agent interaction with users",
   "main": "dist/index.js",
   "type": "module",

package/src/index.ts

@@ -6,6 +6,7 @@ import {
   ServerCapabilities,
 } from "@modelcontextprotocol/sdk/types.js";
 import axios, { AxiosInstance } from "axios";
+import { WebSocket } from "ws";
 
 interface Question {
   id: string;
@@ -73,6 +74,8 @@ class AgentDeskServer {
   private server: Server;
   private apiClient: APIClient | null = null;
   private apiUrl: string;
+  private ws: WebSocket | null = null;
+  private pendingRequests: Map<string, (answer: string) => void> = new Map();
 
   constructor(apiUrl: string) {
     this.apiUrl = apiUrl;
@@ -215,44 +218,90 @@ Returns: User's answers concatenated with newlines, or timeout message.`,
       timeout_seconds
     );
 
-    // Poll for answer
-    const pollInterval = 1000; // 1 second
-    const startTime = Date.now();
+    // Wait for answer via WebSocket push
+    const answerPromise = new Promise<string>((resolve) => {
+      this.pendingRequests.set(sessionId, resolve);
+    });
 
-    while (true) {
-      // Check timeout
-      if (timeout_seconds && (Date.now() - startTime) >= timeout_seconds * 1000) {
+    // Set up timeout if specified
+    let timeoutHandle: NodeJS.Timeout | null = null;
+    if (timeout_seconds) {
+      timeoutHandle = setTimeout(() => {
+        this.pendingRequests.delete(sessionId);
         console.error(`[${new Date().toISOString()}] TIMEOUT handleAskUser: ${sessionId}`);
-        return {
-          content: [
-            {
-              type: "text",
-              text: "User did not respond within the timeout period.",
-            },
-          ],
-        };
-      }
+      }, timeout_seconds * 1000);
+    }
 
-      // Check for answer
-      const agent = await this.apiClient!.getAgent(sessionId);
-      if (agent && agent.status === "answered") {
-        console.error(`[${new Date().toISOString()}] ANSWERED handleAskUser: ${sessionId}`);
-        const answers = agent.questions
-          .filter((q) => q.answer)
-          .map((q) => `${q.id}: ${q.answer}`);
-        return {
-          content: [
-            {
-              type: "text",
-              text: answers.join("\n"),
-            },
-          ],
-        };
+    try {
+      const answer = await answerPromise;
+      if (timeoutHandle) clearTimeout(timeoutHandle);
+      console.error(`[${new Date().toISOString()}] ANSWERED handleAskUser: ${sessionId}`);
+      return {
+        content: [
+          {
+            type: "text",
+            text: answer,
+          },
+        ],
+      };
+    } catch (error) {
+      if (timeoutHandle) clearTimeout(timeoutHandle);
+      return {
+        content: [
+          {
+            type: "text",
+            text: "User did not respond within the timeout period.",
+          },
+        ],
+      };
+    }
+  }
+
+  private setupWebSocket(): void {
+    const wsUrl = this.apiUrl.replace(/^http/, "ws") + "/ws";
+    console.error(`Connecting to WebSocket: ${wsUrl}`);
+
+    this.ws = new WebSocket(wsUrl);
+
+    this.ws.on("open", () => {
+      console.error("WebSocket connected");
+    });
+
+    this.ws.on("message", (data: Buffer) => {
+      try {
+        const message = JSON.parse(data.toString());
+        console.error(`WebSocket message received: ${JSON.stringify(message)}`);
+
+        if (message.type === "agent_answered") {
+          const sessionId = message.session_id;
+          const resolver = this.pendingRequests.get(sessionId);
+
+          if (resolver) {
+            // Fetch the full answer from API
+            this.apiClient!.getAgent(sessionId).then((agent) => {
+              if (agent && agent.status === "answered") {
+                const answers = agent.questions
+                  .filter((q) => q.answer)
+                  .map((q) => `${q.id}: ${q.answer}`);
+                resolver(answers.join("\n"));
+                this.pendingRequests.delete(sessionId);
+              }
+            });
+          }
+        }
+      } catch (error) {
+        console.error("Error processing WebSocket message:", error);
       }
+    });
 
-      console.error(`[${new Date().toISOString()}] POLLING handleAskUser: ${sessionId}`);
-      await new Promise((resolve) => setTimeout(resolve, pollInterval));
-    }
+    this.ws.on("error", (error) => {
+      console.error("WebSocket error:", error);
+    });
+
+    this.ws.on("close", () => {
+      console.error("WebSocket closed, reconnecting in 5s...");
+      setTimeout(() => this.setupWebSocket(), 5000);
+    });
   }
 
   async run(): Promise<void> {
@@ -266,6 +315,9 @@ Returns: User's answers concatenated with newlines, or timeout message.`,
       process.exit(1);
     }
 
+    // Setup WebSocket connection
+    this.setupWebSocket();
+
     const transport = new StdioServerTransport();
     await this.server.connect(transport);