@duckmind/dm-darwin-arm64 0.13.5 → 0.13.7

package/extensions/dm-phone/src/session-pool/session-pool.ts

@@ -0,0 +1,470 @@
+import type { WebSocket } from "ws";
+import type {
+  ClientState,
+  PhoneSessionPoolOptions,
+  SessionController,
+  SessionSnapshot,
+  SessionSummary,
+} from "./types";
+
+export class PhoneSessionPool {
+  private readonly options: PhoneSessionPoolOptions;
+  private readonly workers = new Map<string, SessionController>();
+  private readonly clients = new Map<WebSocket, ClientState>();
+  private readonly statusSignatures = new Map<WebSocket, string>();
+  private readonly catalogSignatures = new Map<WebSocket, string>();
+  private defaultWorkerId: string | null = null;
+  private defaultWorkerPromise: Promise<SessionController> | null = null;
+
+  constructor(options: PhoneSessionPoolOptions) {
+    this.options = options;
+  }
+
+  setCwd(cwd: string) {
+    this.options.cwd = cwd;
+    this.broadcastStatus();
+  }
+
+  get clientCount() {
+    return this.clients.size;
+  }
+
+  getClients() {
+    return [...this.clients.keys()];
+  }
+
+  getSelectedSessionId() {
+    return this.defaultWorkerId;
+  }
+
+  getSession(sessionId: string | null | undefined) {
+    if (!sessionId) return null;
+    return this.workers.get(sessionId) || null;
+  }
+
+  getSessions() {
+    return [...this.workers.values()];
+  }
+
+  private attachSession(session: SessionController) {
+    this.workers.set(session.id, session);
+    return session;
+  }
+
+  private createDefaultSession() {
+    const session = this.options.createDefaultSession();
+    return this.attachSession(session);
+  }
+
+  private createParallelSession(sessionFile: string | null = null) {
+    const session = this.options.createParallelSession(sessionFile);
+    return this.attachSession(session);
+  }
+
+  private sortedWorkers() {
+    return [...this.workers.values()].sort((left, right) => right.lastActivityAt - left.lastActivityAt);
+  }
+
+  private serializeSessions() {
+    return this.sortedWorkers().map((worker) => worker.getSummary());
+  }
+
+  async ensureDefaultWorker() {
+    const existing = this.defaultWorkerId ? this.workers.get(this.defaultWorkerId) : this.sortedWorkers()[0];
+    if (existing) {
+      this.defaultWorkerId = existing.id;
+      return existing;
+    }
+
+    if (this.defaultWorkerPromise) {
+      return this.defaultWorkerPromise;
+    }
+
+    this.defaultWorkerPromise = (async () => {
+      const worker = this.createDefaultSession();
+      this.defaultWorkerId = worker.id;
+
+      try {
+        await worker.ensureStarted();
+        await worker.refreshCachedSnapshot(5000).catch(() => {});
+        this.broadcastCatalog();
+        this.broadcastStatus();
+        return worker;
+      } catch (error) {
+        this.workers.delete(worker.id);
+        if (this.defaultWorkerId === worker.id) {
+          this.defaultWorkerId = null;
+        }
+        throw error;
+      }
+    })().finally(() => {
+      this.defaultWorkerPromise = null;
+    });
+
+    return this.defaultWorkerPromise;
+  }
+
+  private async getWorkerForClient(ws: WebSocket) {
+    const client = this.clients.get(ws);
+    if (!client) {
+      const worker = await this.ensureDefaultWorker();
+      this.clients.set(ws, { activeSessionId: worker.id });
+      return worker;
+    }
+
+    const activeWorker = client.activeSessionId ? this.workers.get(client.activeSessionId) : null;
+    if (activeWorker) {
+      return activeWorker;
+    }
+
+    const fallback = await this.ensureDefaultWorker();
+    client.activeSessionId = fallback.id;
+    return fallback;
+  }
+
+  async getActiveWorker(ws: WebSocket) {
+    return this.getWorkerForClient(ws);
+  }
+
+  private buildBaseStatus() {
+    const meta = this.options.buildStatusMeta();
+    return {
+      ...meta,
+      connectedClients: this.clients.size,
+      sessionCount: this.workers.size,
+      isRunning: Boolean((meta as any).serverRunning),
+    };
+  }
+
+  private normalizeStatusSignature(status: Record<string, unknown>) {
+    const { lastActivityAt: _ignored, ...rest } = status;
+    return JSON.stringify(rest);
+  }
+
+  private normalizeCatalogSignature(data: { activeSessionId: string | null; sessions: SessionSummary[] }) {
+    return JSON.stringify({
+      activeSessionId: data.activeSessionId,
+      sessions: data.sessions.map(({ lastActivityAt: _ignored, ...session }) => session),
+    });
+  }
+
+  private handleWorkerStateChange(_worker: SessionController) {
+    this.broadcastCatalog();
+    this.broadcastStatus();
+  }
+
+  buildOverallStatus() {
+    const worker = this.defaultWorkerId ? this.workers.get(this.defaultWorkerId) : this.sortedWorkers()[0] || null;
+    return {
+      ...this.buildBaseStatus(),
+      ...(worker ? worker.getStatus() : {
+        childRunning: false,
+        isStreaming: false,
+        isCompacting: false,
+        lastError: "",
+        childPid: null,
+        sessionWorkerId: null,
+        sessionKind: "parallel",
+      }),
+    };
+  }
+
+  private buildClientStatus(ws: WebSocket) {
+    const client = this.clients.get(ws);
+    const worker = client?.activeSessionId ? this.workers.get(client.activeSessionId) : null;
+    return {
+      ...this.buildBaseStatus(),
+      ...(worker ? worker.getStatus() : {
+        childRunning: false,
+        isStreaming: false,
+        isCompacting: false,
+        lastError: "",
+        childPid: null,
+        sessionWorkerId: null,
+        sessionKind: "parallel",
+      }),
+      activeSessionId: client?.activeSessionId || null,
+    };
+  }
+
+  private sendStatus(ws: WebSocket, options: { force?: boolean } = {}) {
+    const data = this.buildClientStatus(ws);
+    const signature = this.normalizeStatusSignature(data);
+    if (!options.force && this.statusSignatures.get(ws) === signature) {
+      return;
+    }
+
+    this.statusSignatures.set(ws, signature);
+    this.options.send(ws, { channel: "server", event: "status", data });
+  }
+
+  private sendSnapshot(ws: WebSocket, worker: SessionController, snapshot: SessionSnapshot) {
+    this.options.send(ws, {
+      channel: "snapshot",
+      sessionWorkerId: worker.id,
+      state: snapshot.state,
+      messages: snapshot.messages || [],
+      commands: snapshot.commands || [],
+      liveAssistantMessage: snapshot.liveAssistantMessage || null,
+      liveTools: snapshot.liveTools || [],
+    });
+  }
+
+  broadcastStatus() {
+    for (const ws of this.clients.keys()) {
+      this.sendStatus(ws);
+    }
+  }
+
+  sendCatalog(ws: WebSocket, options: { force?: boolean } = {}) {
+    const client = this.clients.get(ws);
+    const data = {
+      activeSessionId: client?.activeSessionId || null,
+      sessions: this.serializeSessions(),
+    };
+    const signature = this.normalizeCatalogSignature(data);
+    if (!options.force && this.catalogSignatures.get(ws) === signature) {
+      return;
+    }
+
+    this.catalogSignatures.set(ws, signature);
+    this.options.send(ws, {
+      channel: "sessions",
+      event: "catalog",
+      data,
+    });
+  }
+
+  broadcastCatalog() {
+    for (const ws of this.clients.keys()) {
+      this.sendCatalog(ws);
+    }
+  }
+
+  private forwardEnvelope(worker: SessionController, envelope: any) {
+    for (const [ws, client] of this.clients.entries()) {
+      if (client.activeSessionId === worker.id) {
+        this.options.send(ws, envelope);
+      }
+    }
+  }
+
+  async addClient(ws: WebSocket) {
+    const worker = await this.ensureDefaultWorker();
+    this.clients.set(ws, { activeSessionId: this.defaultWorkerId || worker.id });
+    this.sendStatus(ws, { force: true });
+    this.sendCatalog(ws, { force: true });
+    await this.refreshActiveSnapshot(ws);
+  }
+
+  removeClient(ws: WebSocket) {
+    this.clients.delete(ws);
+    this.statusSignatures.delete(ws);
+    this.catalogSignatures.delete(ws);
+    this.broadcastStatus();
+  }
+
+  async refreshActiveSnapshot(ws: WebSocket) {
+    const worker = await this.getWorkerForClient(ws);
+    const requestedWorkerId = worker.id;
+
+    try {
+      const snapshot = await worker.getSnapshot();
+      const client = this.clients.get(ws);
+      if (client?.activeSessionId !== requestedWorkerId) {
+        return;
+      }
+
+      this.sendSnapshot(ws, worker, snapshot);
+      this.sendStatus(ws);
+      if (worker.pendingUiRequest) {
+        this.options.send(ws, { channel: "rpc", payload: worker.pendingUiRequest });
+      }
+    } catch (error) {
+      const client = this.clients.get(ws);
+      if (client?.activeSessionId !== requestedWorkerId) {
+        return;
+      }
+
+      this.options.send(ws, {
+        channel: "server",
+        event: "snapshot-error",
+        data: { message: error instanceof Error ? error.message : String(error) },
+      });
+    }
+  }
+
+  async broadcastSnapshots() {
+    await Promise.all(this.getClients().map(async (ws) => this.refreshActiveSnapshot(ws)));
+  }
+
+  async selectSession(ws: WebSocket, sessionId: string) {
+    const worker = this.workers.get(sessionId);
+    if (!worker) {
+      this.options.send(ws, { channel: "server", event: "client-error", data: { message: "That session no longer exists." } });
+      return;
+    }
+
+    const client = this.clients.get(ws);
+    if (!client) {
+      this.clients.set(ws, { activeSessionId: sessionId });
+    } else {
+      client.activeSessionId = sessionId;
+    }
+
+    this.defaultWorkerId = worker.id;
+
+    this.sendCatalog(ws, { force: true });
+    this.sendStatus(ws, { force: true });
+    this.sendSnapshot(ws, worker, worker.getCachedSnapshot());
+    await this.refreshActiveSnapshot(ws);
+  }
+
+  async spawnSession(ws: WebSocket, sessionFile: string | null = null) {
+    const worker = this.createParallelSession(sessionFile);
+    let added = false;
+    const existingClient = this.clients.get(ws);
+    const previousActiveSessionId = existingClient?.activeSessionId || null;
+
+    try {
+      added = true;
+      this.defaultWorkerId = worker.id;
+
+      if (existingClient) {
+        existingClient.activeSessionId = worker.id;
+      } else {
+        this.clients.set(ws, { activeSessionId: worker.id });
+      }
+
+      this.sendCatalog(ws, { force: true });
+      this.sendStatus(ws, { force: true });
+      this.sendSnapshot(ws, worker, worker.getCachedSnapshot());
+
+      await worker.ensureStarted();
+      await worker.refreshCachedSnapshot(5000).catch(() => {});
+      this.broadcastCatalog();
+      this.broadcastStatus();
+      await this.refreshActiveSnapshot(ws);
+      return worker;
+    } catch (error) {
+      if (added) {
+        this.workers.delete(worker.id);
+      }
+      const fallbackWorker = previousActiveSessionId ? this.workers.get(previousActiveSessionId) : this.sortedWorkers()[0] || null;
+      if (this.defaultWorkerId === worker.id) {
+        this.defaultWorkerId = fallbackWorker?.id || null;
+      }
+
+      const client = this.clients.get(ws);
+      if (client) {
+        client.activeSessionId = fallbackWorker?.id || null;
+      }
+
+      await worker.dispose().catch(() => {});
+      this.sendCatalog(ws, { force: true });
+      this.sendStatus(ws, { force: true });
+      if (fallbackWorker) {
+        this.sendSnapshot(ws, fallbackWorker, fallbackWorker.getCachedSnapshot());
+        await this.refreshActiveSnapshot(ws).catch(() => {});
+      }
+      this.broadcastCatalog();
+      this.broadcastStatus();
+      throw error;
+    }
+  }
+
+  async removeSession(sessionId: string, options: { dispose?: boolean; fallbackSessionId?: string | null } = {}) {
+    const worker = this.workers.get(sessionId);
+    if (!worker) return false;
+
+    const { dispose = true } = options;
+    this.workers.delete(sessionId);
+
+    let fallbackWorker = options.fallbackSessionId ? this.workers.get(options.fallbackSessionId) || null : null;
+    if (!fallbackWorker) {
+      fallbackWorker = this.defaultWorkerId && this.defaultWorkerId !== sessionId
+        ? this.workers.get(this.defaultWorkerId) || null
+        : null;
+    }
+    if (!fallbackWorker) {
+      fallbackWorker = this.sortedWorkers()[0] || null;
+    }
+
+    if (this.defaultWorkerId === sessionId) {
+      this.defaultWorkerId = fallbackWorker?.id || null;
+    }
+
+    for (const client of this.clients.values()) {
+      if (client.activeSessionId === sessionId) {
+        client.activeSessionId = fallbackWorker?.id || null;
+      }
+    }
+
+    if (dispose) {
+      await worker.dispose().catch(() => {});
+    }
+
+    this.broadcastCatalog();
+    this.broadcastStatus();
+    await Promise.all(this.getClients().map(async (ws) => {
+      const client = this.clients.get(ws);
+      if (!client?.activeSessionId) return;
+      await this.refreshActiveSnapshot(ws).catch(() => {});
+    }));
+
+    return true;
+  }
+
+  setDefaultWorker(sessionId: string | null) {
+    this.defaultWorkerId = sessionId && this.workers.has(sessionId) ? sessionId : this.sortedWorkers()[0]?.id || null;
+    for (const client of this.clients.values()) {
+      if (!client.activeSessionId || !this.workers.has(client.activeSessionId)) {
+        client.activeSessionId = this.defaultWorkerId;
+      }
+    }
+    this.broadcastCatalog();
+    this.broadcastStatus();
+  }
+
+  bindExternalSession(session: SessionController) {
+    this.attachSession(session);
+    if (!this.defaultWorkerId) {
+      this.defaultWorkerId = session.id;
+    }
+  }
+
+  notifySessionStateChanged(session: SessionController) {
+    this.handleWorkerStateChange(session);
+  }
+
+  forwardSessionEnvelope(session: SessionController, envelope: any) {
+    this.forwardEnvelope(session, envelope);
+  }
+
+  async closeAllClients(options: { payload?: unknown; code?: number; reason?: string } = {}) {
+    const { payload, code = 1000, reason = "closing" } = options;
+    const sockets = this.getClients();
+    this.clients.clear();
+    this.statusSignatures.clear();
+    this.catalogSignatures.clear();
+
+    for (const ws of sockets) {
+      if (payload) {
+        this.options.send(ws, payload);
+      }
+      try {
+        ws.close(code, reason);
+      } catch {
+        // ignore
+      }
+    }
+  }
+
+  async dispose() {
+    await this.closeAllClients();
+    await Promise.all([...this.workers.values()].map(async (worker) => worker.dispose().catch(() => {})));
+    this.workers.clear();
+    this.defaultWorkerId = null;
+    this.defaultWorkerPromise = null;
+  }
+}