agenthub-multiagent-mcp 1.1.5 → 1.3.0

package/src/websocket.ts

@@ -0,0 +1,237 @@
+/**
+ * WebSocket client for real-time push notifications from AgentHub server.
+ *
+ * Receives pushed tasks and messages and stores them in a queue
+ * for inclusion in tool responses.
+ */
+
+import WebSocket from "ws";
+import type { Message, PendingTask } from "./client.js";
+
+export interface WSMessage {
+  type: "task" | "message" | "ping" | "pong" | "ack";
+  timestamp: string;
+  task?: PendingTask;
+  message?: Message;
+  data?: unknown;
+}
+
+export interface PushedItems {
+  tasks: PendingTask[];
+  messages: Message[];
+}
+
+export class WebSocketClient {
+  private ws: WebSocket | null = null;
+  private baseUrl: string;
+  private apiKey: string;
+  private agentId: string | null = null;
+  private reconnectAttempts = 0;
+  private maxReconnectAttempts = 5;
+  private reconnectDelay = 1000; // Start with 1 second
+  private reconnectTimer: NodeJS.Timeout | null = null;
+  private isConnecting = false;
+  private isClosing = false;
+
+  // Queue for pushed items
+  private pushedTasks: PendingTask[] = [];
+  private pushedMessages: Message[] = [];
+
+  constructor(baseUrl: string, apiKey: string) {
+    // Convert HTTP URL to WebSocket URL
+    this.baseUrl = baseUrl
+      .replace(/^https:/, "wss:")
+      .replace(/^http:/, "ws:")
+      .replace(/\/$/, "");
+    this.apiKey = apiKey;
+  }
+
+  /**
+   * Connect to WebSocket for a specific agent
+   */
+  connect(agentId: string): void {
+    if (this.isConnecting || (this.ws && this.ws.readyState === WebSocket.OPEN)) {
+      return;
+    }
+
+    this.agentId = agentId;
+    this.isConnecting = true;
+    this.isClosing = false;
+
+    const url = `${this.baseUrl}/ws/agents/${agentId}?api_key=${encodeURIComponent(this.apiKey)}`;
+
+    try {
+      this.ws = new WebSocket(url);
+
+      this.ws.on("open", () => {
+        this.isConnecting = false;
+        this.reconnectAttempts = 0;
+        this.reconnectDelay = 1000;
+        console.error(`[AgentHub] WebSocket connected for agent: ${agentId}`);
+      });
+
+      this.ws.on("message", (data: WebSocket.RawData) => {
+        try {
+          const msg = JSON.parse(data.toString()) as WSMessage;
+          this.handleMessage(msg);
+        } catch (error) {
+          console.error("[AgentHub] Failed to parse WebSocket message:", error);
+        }
+      });
+
+      this.ws.on("close", () => {
+        this.isConnecting = false;
+        if (!this.isClosing) {
+          console.error("[AgentHub] WebSocket disconnected, will reconnect...");
+          this.scheduleReconnect();
+        }
+      });
+
+      this.ws.on("error", (error) => {
+        this.isConnecting = false;
+        console.error("[AgentHub] WebSocket error:", error.message);
+      });
+    } catch (error) {
+      this.isConnecting = false;
+      console.error("[AgentHub] Failed to create WebSocket:", error);
+      this.scheduleReconnect();
+    }
+  }
+
+  /**
+   * Handle incoming WebSocket message
+   */
+  private handleMessage(msg: WSMessage): void {
+    switch (msg.type) {
+      case "task":
+        if (msg.task) {
+          console.error(`[AgentHub] Received pushed task: ${msg.task.id}`);
+          this.pushedTasks.push(msg.task);
+          // Send ack
+          this.sendAck(msg.task.id);
+        }
+        break;
+
+      case "message":
+        if (msg.message) {
+          console.error(`[AgentHub] Received pushed message: ${msg.message.id}`);
+          this.pushedMessages.push(msg.message);
+          // Send ack
+          this.sendAck(msg.message.id);
+        }
+        break;
+
+      case "ping":
+        // Respond with pong
+        this.send({ type: "pong", timestamp: new Date().toISOString() });
+        break;
+
+      default:
+        // Ignore other message types
+        break;
+    }
+  }
+
+  /**
+   * Send acknowledgment for received item
+   */
+  private sendAck(itemId: string): void {
+    this.send({
+      type: "ack",
+      timestamp: new Date().toISOString(),
+      data: { id: itemId },
+    });
+  }
+
+  /**
+   * Send a message over WebSocket
+   */
+  private send(msg: Partial<WSMessage>): void {
+    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
+      try {
+        this.ws.send(JSON.stringify(msg));
+      } catch (error) {
+        console.error("[AgentHub] Failed to send WebSocket message:", error);
+      }
+    }
+  }
+
+  /**
+   * Schedule reconnection with exponential backoff
+   */
+  private scheduleReconnect(): void {
+    if (this.isClosing || !this.agentId) {
+      return;
+    }
+
+    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
+      console.error("[AgentHub] Max reconnect attempts reached, giving up");
+      return;
+    }
+
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+    }
+
+    this.reconnectTimer = setTimeout(() => {
+      this.reconnectAttempts++;
+      console.error(`[AgentHub] Reconnect attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts}`);
+      if (this.agentId) {
+        this.connect(this.agentId);
+      }
+    }, this.reconnectDelay);
+
+    // Exponential backoff with max of 30 seconds
+    this.reconnectDelay = Math.min(this.reconnectDelay * 2, 30000);
+  }
+
+  /**
+   * Get and clear pushed items
+   */
+  getPushedItems(): PushedItems {
+    const items = {
+      tasks: [...this.pushedTasks],
+      messages: [...this.pushedMessages],
+    };
+
+    // Clear the queues
+    this.pushedTasks = [];
+    this.pushedMessages = [];
+
+    return items;
+  }
+
+  /**
+   * Check if connected
+   */
+  isConnected(): boolean {
+    return this.ws !== null && this.ws.readyState === WebSocket.OPEN;
+  }
+
+  /**
+   * Close the WebSocket connection
+   */
+  close(): void {
+    this.isClosing = true;
+
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+
+    if (this.ws) {
+      try {
+        this.ws.close();
+      } catch (error) {
+        // Ignore close errors
+      }
+      this.ws = null;
+    }
+
+    this.agentId = null;
+    this.pushedTasks = [];
+    this.pushedMessages = [];
+
+    console.error("[AgentHub] WebSocket closed");
+  }
+}