@blackbelt-technology/pi-agent-dashboard 0.2.0

package/packages/server/src/pending-load-manager.ts

@@ -0,0 +1,110 @@
+/**
+ * Tracks in-flight on-demand session load requests from bridge extensions.
+ * Handles deduplication, timeouts, and bridge disconnect cleanup.
+ */
+import type WebSocket from "ws";
+
+export interface PendingLoad {
+  sessionId: string;
+  requestedAt: number;
+  browsers: Set<WebSocket>;
+  bridgeSessionId: string;
+  timer: ReturnType<typeof setTimeout>;
+}
+
+export interface PendingLoadManager {
+  /** Start a pending load. Returns false if already pending (dedup). */
+  start(sessionId: string, browser: WebSocket, bridgeSessionId: string): boolean;
+  /** Add a browser to an existing pending load. Returns false if no pending load exists. */
+  addBrowser(sessionId: string, browser: WebSocket): boolean;
+  /** Check if a session has a pending load. */
+  isPending(sessionId: string): boolean;
+  /** Complete a pending load (success or error). Returns the waiting browsers, or null if not found. */
+  complete(sessionId: string): Set<WebSocket> | null;
+  /** Cancel all pending loads for a specific bridge. Returns Map of sessionId → browsers. */
+  cancelForBridge(bridgeSessionId: string): Map<string, Set<WebSocket>>;
+  /** Cancel a specific pending load. Returns the waiting browsers, or null. */
+  cancel(sessionId: string): Set<WebSocket> | null;
+  /** Clean up all timers. */
+  dispose(): void;
+}
+
+const DEFAULT_TIMEOUT_MS = 10_000;
+
+export function createPendingLoadManager(
+  onTimeout: (sessionId: string, browsers: Set<WebSocket>) => void,
+  timeoutMs: number = DEFAULT_TIMEOUT_MS,
+): PendingLoadManager {
+  const pending = new Map<string, PendingLoad>();
+
+  function createTimer(sessionId: string): ReturnType<typeof setTimeout> {
+    return setTimeout(() => {
+      const load = pending.get(sessionId);
+      if (load) {
+        pending.delete(sessionId);
+        onTimeout(sessionId, load.browsers);
+      }
+    }, timeoutMs);
+  }
+
+  return {
+    start(sessionId: string, browser: WebSocket, bridgeSessionId: string): boolean {
+      if (pending.has(sessionId)) return false;
+      const load: PendingLoad = {
+        sessionId,
+        requestedAt: Date.now(),
+        browsers: new Set([browser]),
+        bridgeSessionId,
+        timer: createTimer(sessionId),
+      };
+      pending.set(sessionId, load);
+      return true;
+    },
+
+    addBrowser(sessionId: string, browser: WebSocket): boolean {
+      const load = pending.get(sessionId);
+      if (!load) return false;
+      load.browsers.add(browser);
+      return true;
+    },
+
+    isPending(sessionId: string): boolean {
+      return pending.has(sessionId);
+    },
+
+    complete(sessionId: string): Set<WebSocket> | null {
+      const load = pending.get(sessionId);
+      if (!load) return null;
+      clearTimeout(load.timer);
+      pending.delete(sessionId);
+      return load.browsers;
+    },
+
+    cancelForBridge(bridgeSessionId: string): Map<string, Set<WebSocket>> {
+      const cancelled = new Map<string, Set<WebSocket>>();
+      for (const [sessionId, load] of pending) {
+        if (load.bridgeSessionId === bridgeSessionId) {
+          clearTimeout(load.timer);
+          pending.delete(sessionId);
+          cancelled.set(sessionId, load.browsers);
+        }
+      }
+      return cancelled;
+    },
+
+    cancel(sessionId: string): Set<WebSocket> | null {
+      const load = pending.get(sessionId);
+      if (!load) return null;
+      clearTimeout(load.timer);
+      pending.delete(sessionId);
+      return load.browsers;
+    },
+
+    dispose(): void {
+      for (const load of pending.values()) {
+        clearTimeout(load.timer);
+      }
+      pending.clear();
+    },
+  };
+}

package/packages/server/src/pending-resume-registry.ts

@@ -0,0 +1,69 @@
+/**
+ * Tracks pending auto-resume operations: prompts queued for ended sessions
+ * being resumed. Entries expire after 30 seconds if not consumed.
+ */
+
+import type { ImageContent } from "@blackbelt-technology/pi-dashboard-shared/types.js";
+
+const EXPIRY_MS = 30_000;
+
+export interface PendingResumeEntry {
+  text: string;
+  images?: ImageContent[];
+  oldSessionId: string;
+  sessionFile: string;
+}
+
+interface InternalEntry extends PendingResumeEntry {
+  timer: ReturnType<typeof setTimeout>;
+}
+
+export interface PendingResumeRegistryOptions {
+  /** Called when a pending resume expires without being consumed. */
+  onTimeout?: (oldSessionId: string) => void;
+}
+
+export interface PendingResumeRegistry {
+  /** Record a pending resume for a cwd. Overwrites any previous entry for the same cwd. */
+  record(cwd: string, entry: PendingResumeEntry): void;
+  /** Consume and return the pending resume for a cwd, or undefined if none pending. */
+  consume(cwd: string): PendingResumeEntry | undefined;
+  /** Clear all pending entries and timers. */
+  dispose(): void;
+}
+
+export function createPendingResumeRegistry(
+  options?: PendingResumeRegistryOptions,
+): PendingResumeRegistry {
+  const pending = new Map<string, InternalEntry>();
+
+  return {
+    record(cwd: string, entry: PendingResumeEntry): void {
+      const existing = pending.get(cwd);
+      if (existing) {
+        clearTimeout(existing.timer);
+      }
+      const timer = setTimeout(() => {
+        pending.delete(cwd);
+        options?.onTimeout?.(entry.oldSessionId);
+      }, EXPIRY_MS);
+      pending.set(cwd, { ...entry, timer });
+    },
+
+    consume(cwd: string): PendingResumeEntry | undefined {
+      const entry = pending.get(cwd);
+      if (!entry) return undefined;
+      clearTimeout(entry.timer);
+      pending.delete(cwd);
+      const { timer: _, ...result } = entry;
+      return result;
+    },
+
+    dispose(): void {
+      for (const entry of pending.values()) {
+        clearTimeout(entry.timer);
+      }
+      pending.clear();
+    },
+  };
+}

package/packages/server/src/pi-gateway.ts

@@ -0,0 +1,419 @@
+/**
+ * Pi Gateway - WebSocket server for bridge extension connections.
+ */
+import { WebSocketServer, WebSocket } from "ws";
+import type { ExtensionToServerMessage, ServerToExtensionMessage } from "@blackbelt-technology/pi-dashboard-shared/protocol.js";
+import type { DashboardSession } from "@blackbelt-technology/pi-dashboard-shared/types.js";
+import type { SessionManager } from "./memory-session-manager.js";
+
+export const HEARTBEAT_TIMEOUT = 180_000;
+export const WS_PING_INTERVAL = 60_000;
+
+export interface PiGatewayOptions {
+  heartbeatTimeout?: number;
+  pingInterval?: number;
+}
+
+export interface PiGateway {
+  start(port: number): void;
+  stop(): void;
+  sendToSession(sessionId: string, msg: ServerToExtensionMessage): boolean;
+  broadcast(msg: ServerToExtensionMessage): void;
+  connectionCount(): number;
+  findSessionByCwd(cwd: string): string | undefined;
+  getConnectedSessionIds(): string[];
+  isSessionConnected(sessionId: string): boolean;
+  /** Force-close the WebSocket connection for a session */
+  closeSession(sessionId: string): boolean;
+  onEvent?: (sessionId: string, msg: ExtensionToServerMessage) => void;
+  onEmpty?: () => void;
+  onConnection?: () => void;
+  onDisconnect?: (sessionId: string) => void;
+  onSessionCreated?: (sessionId: string) => void;
+}
+
+export function createPiGateway(
+  sessionManager: SessionManager,
+  options?: PiGatewayOptions,
+): PiGateway {
+  const hbTimeout = options?.heartbeatTimeout ?? HEARTBEAT_TIMEOUT;
+  const pingMs = options?.pingInterval ?? WS_PING_INTERVAL;
+  let wss: WebSocketServer | null = null;
+  let pingTimer: ReturnType<typeof setInterval> | null = null;
+
+  // Map sessionId → WebSocket
+  const connections = new Map<string, WebSocket>();
+  // Track connection liveness for WS ping/pong (miss counter: kill after 2 consecutive misses)
+  const aliveMisses = new Map<WebSocket, number>();
+  // Map sessionId → heartbeat timeout
+  const heartbeatTimers = new Map<string, ReturnType<typeof setTimeout>>();
+  // Map sessionId → { setAt: timestamp, sleepRetried: boolean } for sleep detection
+  const heartbeatMeta = new Map<string, { setAt: number; sleepRetried: boolean }>();
+
+  let onEvent: ((sessionId: string, msg: ExtensionToServerMessage) => void) | undefined;
+  let onEmpty: (() => void) | undefined;
+  let onConnection: (() => void) | undefined;
+  let onDisconnect: ((sessionId: string) => void) | undefined;
+  let onSessionCreated: ((sessionId: string) => void) | undefined;
+
+  function checkEmpty() {
+    if (connections.size === 0) {
+      onEmpty?.();
+    }
+  }
+
+  function resetHeartbeat(sessionId: string) {
+    const existing = heartbeatTimers.get(sessionId);
+    if (existing) clearTimeout(existing);
+
+    const now = Date.now();
+    heartbeatMeta.set(sessionId, { setAt: now, sleepRetried: false });
+
+    heartbeatTimers.set(
+      sessionId,
+      setTimeout(() => {
+        // If the WebSocket TCP connection is still open, don't kill the session.
+        // The bridge is just busy (e.g. running a long tool execution) and can't
+        // send heartbeats, but the connection itself is alive. Reschedule.
+        const ws = connections.get(sessionId);
+        if (ws && ws.readyState === WebSocket.OPEN) {
+          console.error(`[gateway] heartbeat timeout but WS still OPEN for ${sessionId}, rescheduling`);
+          resetHeartbeat(sessionId);
+          return;
+        }
+        // Session status check: if the session is still streaming/active
+        // (not manually ended), give it more time to reconnect.
+        // Forked child processes (vitest) can kill the WS connection by
+        // inheriting and closing the FD, but the bridge will reconnect
+        // once the event loop is free.
+        const session = sessionManager.get(sessionId);
+        const meta = heartbeatMeta.get(sessionId);
+        if (session && session.status !== "ended" && !meta?.sleepRetried) {
+          console.error(`[gateway] heartbeat timeout but session ${sessionId} still active, giving reconnect grace period`);
+          if (meta) {
+            meta.sleepRetried = true;
+            meta.setAt = Date.now();
+          }
+          heartbeatTimers.set(
+            sessionId,
+            setTimeout(() => {
+              const ws2 = connections.get(sessionId);
+              if (ws2 && ws2.readyState === WebSocket.OPEN) {
+                resetHeartbeat(sessionId);
+                return;
+              }
+              console.error(`[gateway] session timed out: ${sessionId} (reconnect grace period expired)`);
+              sessionManager.unregister(sessionId);
+              connections.delete(sessionId);
+              heartbeatTimers.delete(sessionId);
+              heartbeatMeta.delete(sessionId);
+              checkEmpty();
+            }, hbTimeout),
+          );
+          return;
+        }
+        console.error(`[gateway] heartbeat timeout, WS state=${ws?.readyState} for ${sessionId}`);
+
+        const meta2 = heartbeatMeta.get(sessionId);
+        const elapsed = Date.now() - (meta2?.setAt ?? now);
+
+        // Detect sleep: elapsed >> expected means system was suspended
+        if (meta2 && !meta2.sleepRetried && elapsed > hbTimeout * 2) {
+          // Give one more cycle for the extension to reconnect
+          meta2.sleepRetried = true;
+          meta2.setAt = Date.now();
+          heartbeatTimers.set(
+            sessionId,
+            setTimeout(() => {
+              const ws2 = connections.get(sessionId);
+              if (ws2 && ws2.readyState === WebSocket.OPEN) {
+                resetHeartbeat(sessionId);
+                return;
+              }
+              console.error(`[gateway] session timed out: ${sessionId} (sleep recovery failed)`);
+              sessionManager.unregister(sessionId);
+              connections.delete(sessionId);
+              heartbeatTimers.delete(sessionId);
+              heartbeatMeta.delete(sessionId);
+              checkEmpty();
+            }, hbTimeout),
+          );
+          return;
+        }
+
+        console.error(`[gateway] session timed out: ${sessionId} (no heartbeat for ${hbTimeout}ms)`);
+        sessionManager.unregister(sessionId);
+        connections.delete(sessionId);
+        heartbeatTimers.delete(sessionId);
+        heartbeatMeta.delete(sessionId);
+        checkEmpty();
+      }, hbTimeout)
+    );
+  }
+
+  return {
+    set onEvent(handler: ((sessionId: string, msg: ExtensionToServerMessage) => void) | undefined) {
+      onEvent = handler;
+    },
+
+    set onEmpty(handler: (() => void) | undefined) {
+      onEmpty = handler;
+    },
+
+    set onConnection(handler: (() => void) | undefined) {
+      onConnection = handler;
+    },
+
+    set onDisconnect(handler: ((sessionId: string) => void) | undefined) {
+      onDisconnect = handler;
+    },
+
+    set onSessionCreated(handler: ((sessionId: string) => void) | undefined) {
+      onSessionCreated = handler;
+    },
+
+    start(port: number) {
+      wss = new WebSocketServer({ port });
+
+      // WS-level ping/pong: detect truly dead connections.
+      // Pong responses are processed in the event loop, so a busy bridge
+      // won't respond to pings. We check the underlying TCP socket's
+      // writable state as a fallback — if TCP is alive, the bridge is just
+      // busy, not dead.
+      const PING_MISS_THRESHOLD = 3;
+      if (pingMs > 0) pingTimer = setInterval(() => {
+        if (!wss) return;
+        for (const client of wss.clients) {
+          const misses = aliveMisses.get(client) ?? 0;
+          if (misses >= PING_MISS_THRESHOLD) {
+            // Check if the underlying TCP socket is still alive.
+            // If the socket is writable, the connection is physically intact —
+            // the bridge is just too busy to process pong frames.
+            const socket = (client as any)._socket;
+            const socketAlive = socket && !socket.destroyed && socket.writable;
+            if (socketAlive) {
+              // TCP alive but no pong — bridge is busy. Reset counter, keep alive.
+              console.error(`[gateway] ping: ${misses} misses but TCP alive, keeping session (socket.destroyed=${socket?.destroyed} writable=${socket?.writable})`);
+              aliveMisses.set(client, 0);
+              client.ping();
+              continue;
+            }
+            // TCP is dead — clean up
+            console.error(`[gateway] ping: TCP dead (socket=${!!socket} destroyed=${socket?.destroyed} writable=${socket?.writable})`);
+            
+            for (const [sid, ws] of connections) {
+              if (ws === client) {
+                console.error(`[gateway] connection dead (ping timeout, ${misses} misses): ${sid}`);
+                sessionManager.unregister(sid);
+                connections.delete(sid);
+                const timer = heartbeatTimers.get(sid);
+                if (timer) clearTimeout(timer);
+                heartbeatTimers.delete(sid);
+                heartbeatMeta.delete(sid);
+                break;
+              }
+            }
+            client.terminate();
+            aliveMisses.delete(client);
+            checkEmpty();
+            continue;
+          }
+          aliveMisses.set(client, misses + 1);
+          client.ping();
+        }
+      }, pingMs);
+
+      wss.on("connection", (ws) => {
+        let currentSessionId: string | null = null;
+        aliveMisses.set(ws, 0);
+        ws.on("pong", () => { aliveMisses.set(ws, 0); });
+
+        ws.on("message", (raw) => {
+          // Any received message proves the connection is alive
+          aliveMisses.set(ws, 0);
+          try {
+            const msg = JSON.parse(raw.toString()) as ExtensionToServerMessage;
+
+            // Track session identity from any message with a sessionId
+            if (!currentSessionId && "sessionId" in msg && (msg as any).sessionId) {
+              const sid: string = (msg as any).sessionId;
+              currentSessionId = sid;
+              connections.set(sid, ws);
+              // Auto-create a placeholder session so events aren't lost
+              if (!sessionManager.get(sid)) {
+                sessionManager.register({
+                  id: sid,
+                  cwd: "",
+                  source: "unknown",
+                });
+                onSessionCreated?.(sid);
+              }
+              resetHeartbeat(sid);
+              onConnection?.();
+            }
+
+            if (msg.type === "session_register") {
+              // If session ID changed (e.g., after /reload), clean up the old placeholder
+              if (currentSessionId && currentSessionId !== msg.sessionId) {
+                const oldSession = sessionManager.get(currentSessionId);
+                // Clean up if it's an auto-created placeholder (source unknown)
+                // or a ghost session (no sessionFile, created by duplicate bridge)
+                if (oldSession && (oldSession.source === "unknown" || !oldSession.sessionFile)) {
+                  sessionManager.unregister(currentSessionId);
+                  connections.delete(currentSessionId);
+                }
+              }
+              currentSessionId = msg.sessionId;
+              connections.set(msg.sessionId, ws);
+
+              sessionManager.register({
+                id: msg.sessionId,
+                cwd: msg.cwd,
+                name: msg.name,
+                source: msg.source,
+                model: msg.model,
+                thinkingLevel: msg.thinkingLevel,
+                sessionFile: msg.sessionFile,
+                sessionDir: msg.sessionDir,
+                firstMessage: msg.firstMessage,
+                pid: msg.pid,
+              });
+              console.error(`[gateway] session registered: ${msg.sessionId} cwd=${msg.cwd}`);
+
+              resetHeartbeat(msg.sessionId);
+              onConnection?.();
+            }
+
+            if (msg.type === "session_heartbeat" && msg.sessionId) {
+              resetHeartbeat(msg.sessionId);
+              // Store process metrics on the session if provided
+              if (msg.metrics) {
+                sessionManager.update(msg.sessionId, {
+                  processMetrics: { ...msg.metrics, updatedAt: Date.now() },
+                });
+              }
+              // Respond with ack so the bridge can track server liveness
+              if (ws.readyState === WebSocket.OPEN) {
+                ws.send(JSON.stringify({ type: "heartbeat_ack" }));
+              }
+            }
+
+            if (msg.type === "session_unregister" && msg.sessionId) {
+              console.error(`[gateway] session unregistered: ${msg.sessionId} (explicit)`);
+              sessionManager.unregister(msg.sessionId);
+              connections.delete(msg.sessionId);
+              const timer = heartbeatTimers.get(msg.sessionId);
+              if (timer) {
+                clearTimeout(timer);
+                heartbeatTimers.delete(msg.sessionId);
+              }
+              heartbeatMeta.delete(msg.sessionId);
+              checkEmpty();
+            }
+
+            if (msg.type === "model_update") {
+              const session = sessionManager.get(msg.sessionId);
+              if (session) {
+                const updates: Partial<typeof session> = { model: msg.model };
+                if (msg.thinkingLevel !== undefined) {
+                  updates.thinkingLevel = msg.thinkingLevel;
+                }
+                sessionManager.update(msg.sessionId, updates);
+              }
+            }
+
+            // Notify listeners
+            const eventSessionId = "sessionId" in msg ? (msg as any).sessionId : undefined;
+            onEvent?.(eventSessionId ?? currentSessionId ?? "", msg);
+          } catch {
+            // Ignore malformed messages
+          }
+        });
+
+        ws.on("close", () => {
+          if (currentSessionId) {
+            console.error(`[gateway] connection closed: ${currentSessionId}`);
+            // Don't immediately unregister - wait for heartbeat timeout
+            // This handles temporary disconnects
+            onDisconnect?.(currentSessionId);
+          }
+          aliveMisses.delete(ws);
+        });
+      });
+    },
+
+    stop() {
+      if (pingTimer) {
+        clearInterval(pingTimer);
+        pingTimer = null;
+      }
+      for (const timer of heartbeatTimers.values()) {
+        clearTimeout(timer);
+      }
+      heartbeatTimers.clear();
+      heartbeatMeta.clear();
+      aliveMisses.clear();
+      // Forcibly terminate all extension connections
+      for (const ws of connections.values()) {
+        ws.terminate();
+      }
+      connections.clear();
+      wss?.close();
+      wss = null;
+    },
+
+    sendToSession(sessionId: string, msg: ServerToExtensionMessage): boolean {
+      const ws = connections.get(sessionId);
+      if (ws && ws.readyState === WebSocket.OPEN) {
+        ws.send(JSON.stringify(msg));
+        return true;
+      }
+      return false;
+    },
+
+    broadcast(msg: ServerToExtensionMessage): void {
+      const payload = JSON.stringify(msg);
+      for (const ws of connections.values()) {
+        if (ws.readyState === WebSocket.OPEN) {
+          ws.send(payload);
+        }
+      }
+    },
+
+    connectionCount(): number {
+      return connections.size;
+    },
+
+    isSessionConnected(sessionId: string): boolean {
+      const ws = connections.get(sessionId);
+      return ws !== undefined && ws.readyState === WebSocket.OPEN;
+    },
+
+    findSessionByCwd(cwd: string): string | undefined {
+      // Find a connected session whose cwd matches or is a prefix
+      for (const sid of connections.keys()) {
+        const session = sessionManager.get(sid);
+        if (session && (session.cwd === cwd || session.cwd.startsWith(cwd + "/") || cwd.startsWith(session.cwd + "/"))) {
+          return sid;
+        }
+      }
+      return undefined;
+    },
+
+    getConnectedSessionIds(): string[] {
+      return [...connections.keys()].filter(
+        (sid) => connections.get(sid)?.readyState === WebSocket.OPEN,
+      );
+    },
+
+    closeSession(sessionId: string): boolean {
+      const ws = connections.get(sessionId);
+      if (ws) {
+        ws.close();
+        connections.delete(sessionId);
+        return true;
+      }
+      return false;
+    },
+  };
+}