chainlesschain 0.45.12 → 0.45.19

package/src/lib/ws-server.js

@@ -11,6 +11,38 @@ import { spawn } from "node:child_process";
 import { fileURLToPath } from "node:url";
 import { dirname, join } from "node:path";
 import { WebSocketServer } from "ws";
+import { createTaskRecord } from "../runtime/contracts/task-record.js";
+import {
+  RUNTIME_EVENTS,
+  createRuntimeEvent,
+} from "../runtime/runtime-events.js";
+import { createWsMessageDispatcher } from "../gateways/ws/message-dispatcher.js";
+import {
+  handleTaskDetail,
+  handleTaskHistory,
+} from "../gateways/ws/task-protocol.js";
+import {
+  handleSessionCreate,
+  handleSessionResume,
+  handleSessionMessage,
+  handleSessionPolicyUpdate,
+  handleSessionList,
+  handleSessionClose,
+  handleSessionAnswer,
+  handleHostToolResult,
+} from "../gateways/ws/session-protocol.js";
+import {
+  handleSlashCommand,
+  handleOrchestrate,
+} from "../gateways/ws/action-protocol.js";
+import {
+  handleWorktreeDiff,
+  handleWorktreeMerge,
+  handleWorktreeMergePreview,
+  handleWorktreeAutomationApply,
+  handleWorktreeList,
+  handleCompressionStats,
+} from "../gateways/ws/worktree-protocol.js";
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = dirname(__filename);
@@ -99,6 +131,7 @@ export class ChainlessChainWSServer extends EventEmitter {
 
     /** Session handlers: sessionId → WSAgentHandler | WSChatHandler */
     this.sessionHandlers = new Map();
+    this._dispatcher = createWsMessageDispatcher(this);
 
     this._heartbeatTimer = null;
     this._clientCounter = 0;
@@ -227,83 +260,7 @@ export class ChainlessChainWSServer extends EventEmitter {
 
   /** @private */
   async _handleMessage(clientId, ws, message) {
-    const { id, type } = message;
-
-    if (!id) {
-      this._send(ws, {
-        type: "error",
-        code: "MISSING_ID",
-        message: 'Message must include an "id" field',
-      });
-      return;
-    }
-
-    // Check authentication
-    const client = this.clients.get(clientId);
-    if (this.token && !client.authenticated && type !== "auth") {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "AUTH_REQUIRED",
-        message: "Authentication required. Send an auth message first.",
-      });
-      return;
-    }
-
-    switch (type) {
-      case "auth":
-        this._handleAuth(clientId, ws, message);
-        break;
-      case "ping":
-        this._send(ws, { id, type: "pong", serverTime: Date.now() });
-        break;
-      case "execute":
-        this._executeCommand(id, ws, message.command, false);
-        break;
-      case "stream":
-        this._executeCommand(id, ws, message.command, true);
-        break;
-      case "cancel":
-        this._cancelRequest(id, ws);
-        break;
-      case "session-create":
-        await this._handleSessionCreate(id, ws, message);
-        break;
-      case "session-resume":
-        await this._handleSessionResume(id, ws, message);
-        break;
-      case "session-message":
-        this._handleSessionMessage(id, ws, message);
-        break;
-      case "session-list":
-        this._handleSessionList(id, ws);
-        break;
-      case "session-close":
-        this._handleSessionClose(id, ws, message);
-        break;
-      case "slash-command":
-        this._handleSlashCommand(id, ws, message);
-        break;
-      case "session-answer":
-        this._handleSessionAnswer(id, ws, message);
-        break;
-      case "orchestrate":
-        this._handleOrchestrate(id, ws, message);
-        break;
-      case "tasks-list":
-        this._handleTasksList(id, ws);
-        break;
-      case "tasks-stop":
-        this._handleTasksStop(id, ws, message);
-        break;
-      default:
-        this._send(ws, {
-          id,
-          type: "error",
-          code: "UNKNOWN_TYPE",
-          message: `Unknown message type: ${type}`,
-        });
-    }
+    return this._dispatcher.dispatch(clientId, ws, message);
   }
 
   /**
@@ -320,84 +277,13 @@ export class ChainlessChainWSServer extends EventEmitter {
    *   { type: "error", code: "ORCHESTRATE_FAILED", ... }
    */
   async _handleOrchestrate(id, ws, message) {
-    const {
-      task,
-      cwd,
-      agents = 3,
-      ci = "npm test",
-      noCi = false,
-      strategy,
-    } = message;
-
-    if (!task || typeof task !== "string") {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "INVALID_TASK",
-        message: "task field required",
-      });
-      return;
-    }
-
-    try {
-      const { Orchestrator, TASK_SOURCE } = await import("./orchestrator.js");
-
-      const orch = new Orchestrator({
-        cwd: cwd || this.projectRoot || process.cwd(),
-        maxParallel: Math.min(parseInt(agents, 10) || 3, 10),
-        ciCommand: ci,
-        agents: strategy ? { strategy } : undefined,
-        verbose: false,
-      });
-
-      // Add WebSocket as a notification channel — all events go back to this client
-      const wsNotifier = orch.notifier.addWebSocketChannel({
-        send: (data) => this._send(ws, data),
-        requestId: id,
-      });
-
-      // Forward real-time agent output
-      orch.on("agent:output", (ev) => wsNotifier.sendAgentOutput(ev));
-
-      // Forward task status changes
-      orch.on("task:added", (t) => wsNotifier.sendStatus(t));
-      orch.on("task:decomposing", (t) => wsNotifier.sendStatus(t));
-      orch.on("ci:checking", ({ task: t }) => wsNotifier.sendStatus(t));
-      orch.on("task:retrying", ({ task: t }) => wsNotifier.sendStatus(t));
-
-      const result = await orch.addTask(task, {
-        source: TASK_SOURCE.CLI,
-        cwd: cwd || this.projectRoot || process.cwd(),
-        runCI: !noCi,
-        notify: true,
-      });
-
-      this._send(ws, {
-        id,
-        type: "orchestrate:done",
-        taskId: result.id,
-        status: result.status,
-        retries: result.retries,
-        subtasks: result.subtasks?.length || 0,
-      });
-    } catch (err) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "ORCHESTRATE_FAILED",
-        message: err.message,
-      });
-    }
+    return handleOrchestrate(this, id, ws, message);
   }
 
-  /** @private — list background tasks */
-  _handleTasksList(id, ws) {
+  /** @private – list background tasks */
+  async _handleTasksList(id, ws) {
     try {
-      const { BackgroundTaskManager } = require("./background-task-manager.js");
-      // Reuse singleton or create lightweight instance for listing
-      if (!this._taskManager) {
-        this._taskManager = new BackgroundTaskManager();
-      }
+      await this._ensureTaskManager();
       const tasks = this._taskManager.list();
       this._send(ws, { id, type: "tasks-list", tasks });
     } catch (err) {
@@ -405,9 +291,35 @@ export class ChainlessChainWSServer extends EventEmitter {
     }
   }
 
-  /** @private — stop a background task */
-  _handleTasksStop(id, ws, message) {
+  /** @private — subscribe to task completion events and broadcast to all clients */
+  _subscribeTaskNotifications() {
+    if (!this._taskManager || this._taskNotificationsSubscribed) return;
+    this._taskNotificationsSubscribed = true;
+
+    this._taskManager.on("task:complete", (task) => {
+      const record = createTaskRecord(task, {
+        source: "background-task-manager",
+      });
+      this.emit(
+        RUNTIME_EVENTS.TASK_NOTIFICATION,
+        createRuntimeEvent(
+          RUNTIME_EVENTS.TASK_NOTIFICATION,
+          { task: record },
+          { kind: "server" },
+        ),
+      );
+      this._broadcast({
+        type: "task:notification",
+        task: record,
+      });
+    });
+  }
+
+  /** @private – stop a background task */
+  async _handleTasksStop(id, ws, message) {
     try {
+      await this._ensureTaskManager();
+
       if (this._taskManager && message.taskId) {
         this._taskManager.stop(message.taskId);
         this._send(ws, { id, type: "tasks-stopped", taskId: message.taskId });
@@ -429,6 +341,46 @@ export class ChainlessChainWSServer extends EventEmitter {
     }
   }
 
+  /** @private */
+  async _handleTaskDetail(id, ws, message) {
+    return handleTaskDetail(this, id, ws, message);
+  }
+
+  /** @private */
+  async _handleTaskHistory(id, ws, message) {
+    return handleTaskHistory(this, id, ws, message);
+  }
+
+  /** @private — diff preview for agent worktree branch */
+  async _handleWorktreeDiff(id, ws, message) {
+    return handleWorktreeDiff(this, id, ws, message);
+  }
+
+  /** @private — one-click merge of agent worktree branch */
+  async _handleWorktreeMerge(id, ws, message) {
+    return handleWorktreeMerge(this, id, ws, message);
+  }
+
+  /** @private - dry-run merge preview for an agent worktree branch */
+  async _handleWorktreeMergePreview(id, ws, message) {
+    return handleWorktreeMergePreview(this, id, ws, message);
+  }
+
+  /** @private - apply a safe automation candidate inside an agent worktree */
+  async _handleWorktreeAutomationApply(id, ws, message) {
+    return handleWorktreeAutomationApply(this, id, ws, message);
+  }
+
+  /** @private - list agent worktrees */
+  async _handleWorktreeList(id, ws) {
+    return handleWorktreeList(this, id, ws);
+  }
+
+  /** @private */
+  async _handleCompressionStats(id, ws, message) {
+    return handleCompressionStats(this, id, ws, message);
+  }
+
   /** @private */
   _handleAuth(clientId, ws, message) {
     const { id, token } = message;
@@ -619,268 +571,58 @@ export class ChainlessChainWSServer extends EventEmitter {
 
   /** @private */
   async _handleSessionCreate(id, ws, message) {
-    if (!this.sessionManager) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "NO_SESSION_SUPPORT",
-        message: "Session support not configured on this server",
-      });
-      return;
-    }
-
-    const { sessionType, provider, model, apiKey, baseUrl, projectRoot } =
-      message;
-
-    try {
-      const { sessionId } = this.sessionManager.createSession({
-        type: sessionType || "agent",
-        provider,
-        model,
-        apiKey,
-        baseUrl,
-        projectRoot,
-      });
-
-      const session = this.sessionManager.getSession(sessionId);
-
-      // Lazy-load handler modules to avoid circular deps
-      try {
-        const { WebSocketInteractionAdapter } =
-          await import("./interaction-adapter.js");
-        session.interaction = new WebSocketInteractionAdapter(ws, sessionId);
-
-        let handler;
-        if ((sessionType || "agent") === "chat") {
-          const { WSChatHandler } = await import("./ws-chat-handler.js");
-          handler = new WSChatHandler({
-            session,
-            interaction: session.interaction,
-          });
-        } else {
-          const { WSAgentHandler } = await import("./ws-agent-handler.js");
-          handler = new WSAgentHandler({
-            session,
-            interaction: session.interaction,
-            db: this.sessionManager.db,
-          });
-        }
-        this.sessionHandlers.set(sessionId, handler);
-      } catch (_err) {
-        // Handler creation failed — session still created, handler can be set later
-      }
-
-      this.emit("session:create", { sessionId, type: sessionType || "agent" });
-
-      this._send(ws, {
-        id,
-        type: "session-created",
-        sessionId,
-        sessionType: sessionType || "agent",
-      });
-    } catch (err) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "SESSION_CREATE_FAILED",
-        message: err.message,
-      });
-    }
+    return handleSessionCreate(this, id, ws, message);
   }
 
   /** @private */
   async _handleSessionResume(id, ws, message) {
-    if (!this.sessionManager) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "NO_SESSION_SUPPORT",
-        message: "Session support not configured",
-      });
-      return;
-    }
-
-    const { sessionId } = message;
-    const session = this.sessionManager.resumeSession(sessionId);
-
-    if (!session) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "SESSION_NOT_FOUND",
-        message: `Session not found: ${sessionId}`,
-      });
-      return;
-    }
-
-    // Rebuild interaction adapter and handler for the resumed session
-    if (!this.sessionHandlers.has(sessionId)) {
-      try {
-        const { WebSocketInteractionAdapter } =
-          await import("./interaction-adapter.js");
-        session.interaction = new WebSocketInteractionAdapter(ws, sessionId);
-
-        let handler;
-        if (session.type === "chat") {
-          const { WSChatHandler } = await import("./ws-chat-handler.js");
-          handler = new WSChatHandler({
-            session,
-            interaction: session.interaction,
-          });
-        } else {
-          const { WSAgentHandler } = await import("./ws-agent-handler.js");
-          handler = new WSAgentHandler({
-            session,
-            interaction: session.interaction,
-            db: this.sessionManager.db,
-          });
-        }
-        this.sessionHandlers.set(sessionId, handler);
-      } catch (_err) {
-        // Handler creation failed — session resumed without handler
-      }
-    }
-
-    // Filter out system messages for history
-    const history = (session.messages || []).filter((m) => m.role !== "system");
-
-    this._send(ws, {
-      id,
-      type: "session-resumed",
-      sessionId: session.id,
-      history,
-    });
+    return handleSessionResume(this, id, ws, message);
   }
 
   /** @private */
   _handleSessionMessage(id, ws, message) {
-    const { sessionId, content } = message;
-    const handler = this.sessionHandlers.get(sessionId);
-
-    if (!handler) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "SESSION_NOT_FOUND",
-        message: `No active session handler for: ${sessionId}`,
-      });
-      return;
-    }
+    return handleSessionMessage(this, id, ws, message);
+  }
 
-    // Fire and forget — handler emits events via interaction adapter
-    handler
-      .handleMessage(content, id)
-      .then(() => {
-        // Persist messages after each turn
-        if (this.sessionManager) {
-          try {
-            this.sessionManager.persistMessages(sessionId);
-          } catch (_err) {
-            // Non-critical
-          }
-        }
-      })
-      .catch((err) => {
-        this._send(ws, {
-          id,
-          type: "error",
-          code: "MESSAGE_FAILED",
-          message: err.message,
-        });
-      });
+  /** @private */
+  _handleSessionPolicyUpdate(id, ws, message) {
+    return handleSessionPolicyUpdate(this, id, ws, message);
   }
 
   /** @private */
   _handleSessionList(id, ws) {
-    if (!this.sessionManager) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "NO_SESSION_SUPPORT",
-        message: "Session support not configured",
-      });
-      return;
-    }
-
-    const sessions = this.sessionManager.listSessions();
-    this._send(ws, {
-      id,
-      type: "session-list-result",
-      sessions,
-    });
+    return handleSessionList(this, id, ws);
   }
 
   /** @private */
   _handleSessionClose(id, ws, message) {
-    const { sessionId } = message;
-
-    // Remove handler
-    const handler = this.sessionHandlers.get(sessionId);
-    if (handler && handler.destroy) {
-      handler.destroy();
-    }
-    this.sessionHandlers.delete(sessionId);
-
-    // Close session in manager
-    if (this.sessionManager) {
-      try {
-        this.sessionManager.closeSession(sessionId);
-      } catch (_err) {
-        // Non-critical
-      }
-    }
-
-    this.emit("session:close", { sessionId });
-
-    this._send(ws, {
-      id,
-      type: "result",
-      success: true,
-      sessionId,
-    });
+    return handleSessionClose(this, id, ws, message);
   }
 
   /** @private */
   _handleSlashCommand(id, ws, message) {
-    const { sessionId, command } = message;
-    const handler = this.sessionHandlers.get(sessionId);
-
-    if (!handler) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "SESSION_NOT_FOUND",
-        message: `No active session handler for: ${sessionId}`,
-      });
-      return;
-    }
-
-    handler.handleSlashCommand(command, id);
+    return handleSlashCommand(this, id, ws, message);
   }
 
   /** @private */
   _handleSessionAnswer(id, ws, message) {
-    const { sessionId, requestId, answer } = message;
-
-    if (!this.sessionManager) {
-      this._send(ws, {
-        id,
-        type: "error",
-        code: "NO_SESSION_SUPPORT",
-        message: "Session support not configured",
-      });
-      return;
-    }
-
-    const session = this.sessionManager.getSession(sessionId);
-    if (session && session.interaction && session.interaction.resolveAnswer) {
-      session.interaction.resolveAnswer(requestId, answer);
-    }
+    return handleSessionAnswer(this, id, ws, message);
+  }
 
-    this._send(ws, { id, type: "result", success: true });
+  _handleHostToolResult(id, ws, message) {
+    return handleHostToolResult(this, id, ws, message);
   }
 
   /** @private — ping/pong heartbeat to detect dead connections */
+  async _ensureTaskManager() {
+    if (this._taskManager) return this._taskManager;
+    const { BackgroundTaskManager } =
+      await import("./background-task-manager.js");
+    this._taskManager = new BackgroundTaskManager({ recoverOnStart: true });
+    this._subscribeTaskNotifications();
+    return this._taskManager;
+  }
+
   _startHeartbeat() {
     this._heartbeatTimer = setInterval(() => {
       for (const [clientId, client] of this.clients) {
@@ -910,4 +652,13 @@ export class ChainlessChainWSServer extends EventEmitter {
       }
     }
   }
+
+  /** @private — broadcast a message to all connected, authenticated clients */
+  _broadcast(data) {
+    for (const [, client] of this.clients) {
+      if (client.authenticated || !this.token) {
+        this._send(client.ws, data);
+      }
+    }
+  }
 }