aibroker 0.6.1 → 0.6.2

package/dist/adapters/pailot/gateway.js

@@ -13,7 +13,7 @@
 import { WebSocketServer, WebSocket } from "ws";
 import { join } from "node:path";
 import { writeFileSync, readFileSync, existsSync, unlinkSync, appendFileSync } from "node:fs";
-import { tmpdir } from "node:os";
+import { tmpdir, homedir } from "node:os";
 const DEBUG_LOG = process.env.PAILOT_DEBUG ? "/tmp/pailot-ws-debug.log" : null;
 function dbg(msg) {
     if (DEBUG_LOG)
@@ -24,8 +24,8 @@ import { promisify } from "node:util";
 import { execFile } from "node:child_process";
 import { log } from "../../core/log.js";
 import { WHISPER_BIN, WHISPER_MODEL } from "../kokoro/media.js";
-import { setMessageSource, activeItermSessionId, setActiveItermSessionId, } from "../../core/state.js";
-import { setItermSessionVar, setItermTabName, killSession, createClaudeSession } from "../iterm/sessions.js";
+import { setMessageSource, activeItermSessionId, setActiveItermSessionId, lastRoutedSessionId, setLastRoutedSessionId, getAibpBridge, } from "../../core/state.js";
+import { setItermSessionVar, setItermTabName, setItermBadge, killSession, createClaudeSession } from "../iterm/sessions.js";
 import { listPaiProjects, launchPaiProject } from "../../daemon/pai-projects.js";
 import { runAppleScript, sendKeystrokeToSession, sendEscapeSequenceToSession, pasteTextIntoSession, snapshotAllSessions } from "../iterm/core.js";
 import { hybridManager } from "../../core/hybrid.js";
@@ -33,35 +33,66 @@ const WS_PORT = parseInt(process.env.PAILOT_PORT ?? "8765", 10);
 // --- State ---
 let wss = null;
 const clients = new Set();
-// --- Message outbox for offline clients ---
-// Buffers messages when no client is connected so they can be replayed on reconnect.
-// Only text messages are buffered fully; voice/image are counted but not stored (too large).
-const MAX_OUTBOX = 50;
-const outbox = [];
-let missedVoiceCount = 0;
+/** Maps PAILot session ID → iTerm session ID for explicit reply routing.
+ *  When PAILot sends a message tagged with sessionId, we record it here so
+ *  outbound replies go back tagged with the same sessionId — regardless of
+ *  which iTerm tab happens to be focused on the Mac at response time.
+ */
+const pailotReplyMap = new Map();
+/** Track which session each WebSocket client is currently viewing.
+ *  Messages tagged with a different sessionId are NOT delivered to that client.
+ *  This prevents cross-session content bleed in multi-session PAILot views.
+ */
+const clientActiveSession = new Map();
+/** Track when each client last proved it's alive (pong, message, or connect). */
+const clientLastActive = new Map();
+/** Consider a client "alive" if it responded within this window. */
+const CLIENT_ALIVE_THRESHOLD = 90_000; // 90s (3x the 30s heartbeat)
+// --- Per-session message outbox for offline/backgrounded clients ---
+// Buffers messages when no live client can receive them.
+// Stored per-session so drain delivers to the correct session.
+const MAX_OUTBOX_PER_SESSION = 50;
+const OUTBOX_DIR = join(homedir(), ".aibroker", "outbox");
+const outboxMap = new Map(); // sessionId → entries
 let missedImageCount = 0;
 function addToOutbox(msg) {
     const type = msg.type;
-    // Skip typing indicators — not useful to replay
     if (type === "typing")
         return;
-    // Count but don't buffer screenshots (very large, less important)
     if (type === "image") {
         missedImageCount++;
         return;
     }
-    // Buffer text, voice, sessions, errors, etc.
-    outbox.push({ msg, timestamp: Date.now() });
-    if (outbox.length > MAX_OUTBOX)
-        outbox.shift();
+    const sessionId = msg.sessionId || "_global";
+    let queue = outboxMap.get(sessionId);
+    if (!queue) {
+        queue = [];
+        outboxMap.set(sessionId, queue);
+    }
+    queue.push({ msg, timestamp: Date.now() });
+    if (queue.length > MAX_OUTBOX_PER_SESSION)
+        queue.shift();
+    // Persist to disk (best-effort, async)
+    persistOutbox();
 }
 function drainOutbox(ws) {
-    if (outbox.length === 0 && missedVoiceCount === 0 && missedImageCount === 0)
+    // Collect all entries across all sessions
+    let totalEntries = 0;
+    let textCount = 0;
+    let voiceCount = 0;
+    for (const entries of outboxMap.values()) {
+        totalEntries += entries.length;
+        for (const e of entries) {
+            const t = e.msg.type;
+            if (t === "text")
+                textCount++;
+            else if (t === "voice")
+                voiceCount++;
+        }
+    }
+    if (totalEntries === 0 && missedImageCount === 0)
         return;
-    // Send a summary of what was missed
-    const textCount = outbox.filter(e => e.msg.type === "text").length;
-    const voiceCount = outbox.filter(e => e.msg.type === "voice").length;
-    const otherCount = outbox.length - textCount - voiceCount;
+    const otherCount = totalEntries - textCount - voiceCount;
     const parts = [];
     if (textCount > 0)
         parts.push(`${textCount} text message(s)`);
@@ -75,14 +106,57 @@ function drainOutbox(ws) {
         type: "text",
         content: `📬 While you were away: ${parts.join(", ")}`,
     });
-    // Replay buffered messages
-    for (const entry of outbox) {
+    // Replay all buffered messages (sorted by timestamp across sessions)
+    const allEntries = [];
+    for (const entries of outboxMap.values())
+        allEntries.push(...entries);
+    allEntries.sort((a, b) => a.timestamp - b.timestamp);
+    for (const entry of allEntries) {
         sendTo(ws, entry.msg);
     }
     // Clear outbox
-    outbox.length = 0;
-    missedVoiceCount = 0;
+    outboxMap.clear();
     missedImageCount = 0;
+    persistOutbox();
+}
+/** Persist outbox to disk so daemon restarts don't lose messages. */
+function persistOutbox() {
+    try {
+        const { mkdirSync, writeFileSync: writeSync } = require("node:fs");
+        mkdirSync(OUTBOX_DIR, { recursive: true });
+        const data = {};
+        for (const [k, v] of outboxMap)
+            data[k] = v;
+        writeSync(join(OUTBOX_DIR, "pending.json"), JSON.stringify({ messages: data, missedImageCount }));
+    }
+    catch { /* best-effort */ }
+}
+/** Restore outbox from disk on startup. */
+function restoreOutbox() {
+    try {
+        const path = join(OUTBOX_DIR, "pending.json");
+        if (!existsSync(path))
+            return;
+        const raw = readFileSync(path, "utf-8");
+        const data = JSON.parse(raw);
+        if (data.messages) {
+            for (const [k, v] of Object.entries(data.messages)) {
+                outboxMap.set(k, v);
+            }
+        }
+        if (data.missedImageCount)
+            missedImageCount = data.missedImageCount;
+        log(`[PAILot] Restored ${outboxMap.size} outbox queue(s) from disk`);
+        // Clean up the file after restoring
+        unlinkSync(path);
+    }
+    catch { /* ignore */ }
+}
+function isClientAlive(ws) {
+    if (ws.readyState !== WebSocket.OPEN)
+        return false;
+    const lastActive = clientLastActive.get(ws) ?? 0;
+    return (Date.now() - lastActive) < CLIENT_ALIVE_THRESHOLD;
 }
 // Reference to the screenshot handler — set via setScreenshotHandler()
 // to avoid circular imports (screenshot.ts imports from state.ts which
@@ -141,19 +215,37 @@ function handleSyncCommand(ws, args) {
     }
     // If the client had a session open, try to restore it
     if (clientActiveId) {
-        const sessions = hybridManager.listSessions();
-        const idx = sessions.findIndex(s => s.backendSessionId === clientActiveId);
+        // First pass: check if it was already registered (by auto-discovery above or a prior sync)
+        let sessions = hybridManager.listSessions();
+        let idx = sessions.findIndex(s => s.backendSessionId === clientActiveId);
+        // Second pass: the session may exist in liveSnapshots but wasn't registered because
+        // isClaudeRelated() returned false (e.g. session is at a shell prompt between Claude runs
+        // after a daemon restart). Since the client explicitly asked for this session by ID, register
+        // it unconditionally if it is still alive in iTerm.
+        if (idx < 0 && liveIds.has(clientActiveId)) {
+            const snap = liveSnapshots.find(s => s.id === clientActiveId);
+            if (snap) {
+                const displayName = snap.tabTitle ?? snap.paiName ?? snap.name;
+                hybridManager.registerVisualSession(displayName, "", clientActiveId);
+                sessions = hybridManager.listSessions();
+                idx = sessions.findIndex(s => s.backendSessionId === clientActiveId);
+                log(`[PAILot] sync: force-registered client session "${displayName}" (${clientActiveId.slice(0, 8)}...)`);
+            }
+        }
         if (idx >= 0) {
             hybridManager.switchToIndex(idx + 1);
             setActiveItermSessionId(clientActiveId);
+            setLastRoutedSessionId(clientActiveId);
+            clientActiveSession.set(ws, clientActiveId);
             log(`[PAILot] sync: restored client session "${sessions[idx].name}" (${clientActiveId.slice(0, 8)}...)`);
             handleSessionsCommand(ws);
             drainOutbox(ws);
             return;
         }
-        // Client's session no longer exists — fall through to iTerm focus
+        // Client's session no longer exists in iTerm — fall through to iTerm focus
+        log(`[PAILot] sync: client session ${clientActiveId.slice(0, 8)}... not found in live iTerm — falling back to focused session`);
     }
-    // No client preference — ask iTerm2 which session is focused right now
+    // No client preference (or client's session is gone) — ask iTerm2 which session is focused
     const focusedId = runAppleScript(`tell application "iTerm2"
   try
     return id of current session of current tab of current window
@@ -162,16 +254,40 @@ function handleSyncCommand(ws, args) {
   end try
 end tell`)?.trim() ?? "";
     if (focusedId) {
-        // Find this session in the hybrid manager and activate it
-        const sessions = hybridManager.listSessions();
-        const idx = sessions.findIndex(s => s.backendSessionId === focusedId);
+        // Find this session in the hybrid manager and activate it.
+        // If it wasn't registered by the Claude-related filter, register it now since the
+        // user is actively looking at it.
+        let sessions = hybridManager.listSessions();
+        let idx = sessions.findIndex(s => s.backendSessionId === focusedId);
+        if (idx < 0 && liveIds.has(focusedId)) {
+            const snap = liveSnapshots.find(s => s.id === focusedId);
+            if (snap) {
+                const displayName = snap.tabTitle ?? snap.paiName ?? snap.name;
+                hybridManager.registerVisualSession(displayName, "", focusedId);
+                sessions = hybridManager.listSessions();
+                idx = sessions.findIndex(s => s.backendSessionId === focusedId);
+            }
+        }
         if (idx >= 0) {
             hybridManager.switchToIndex(idx + 1);
             setActiveItermSessionId(focusedId);
+            setLastRoutedSessionId(focusedId);
+            clientActiveSession.set(ws, focusedId);
             log(`[PAILot] sync: activated focused session "${sessions[idx].name}" (${focusedId.slice(0, 8)}...)`);
         }
         else {
-            log(`[PAILot] sync: focused session ${focusedId.slice(0, 8)}... not registered`);
+            log(`[PAILot] sync: focused session ${focusedId.slice(0, 8)}... not found in live iTerm`);
+        }
+    }
+    // Last resort: if nothing is active yet but we have registered sessions, activate the first one
+    if (!hybridManager.activeSession) {
+        const sessions = hybridManager.listSessions();
+        if (sessions.length > 0) {
+            hybridManager.switchToIndex(1);
+            setActiveItermSessionId(sessions[0].backendSessionId);
+            setLastRoutedSessionId(sessions[0].backendSessionId);
+            clientActiveSession.set(ws, sessions[0].backendSessionId);
+            log(`[PAILot] sync: no focused session found — defaulting to first session "${sessions[0].name}"`);
         }
     }
     // Return sessions with updated active state, then drain buffered messages
@@ -272,8 +388,13 @@ end tell`);
         if (session.kind === "visual") {
             setItermSessionVar(session.backendSessionId, newName);
             setItermTabName(session.backendSessionId, newName);
+            setItermBadge(session.backendSessionId, newName);
         }
     }
+    // Record this as the last routed session so outbound replies go here,
+    // regardless of which iTerm tab is focused on the Mac.
+    setLastRoutedSessionId(session.backendSessionId);
+    clientActiveSession.set(ws, session.backendSessionId);
     sendTo(ws, { type: "session_switched", name: session.name, sessionId: session.backendSessionId });
     log(`[PAILot] switched to ${session.kind} session "${session.name}" (${session.id})`);
 }
@@ -292,6 +413,7 @@ function handleRenameCommand(ws, args) {
         if (session.kind === "visual") {
             setItermSessionVar(sessionId, name);
             setItermTabName(sessionId, name);
+            setItermBadge(sessionId, name);
         }
     }
     sendTo(ws, { type: "session_renamed", sessionId, name });
@@ -342,6 +464,7 @@ function handleCreateCommand(ws, args = {}) {
     }
     setItermSessionVar(sessionId, name);
     setItermTabName(sessionId, name);
+    setItermBadge(sessionId, name);
     if (hybridManager) {
         hybridManager.registerVisualSession(name, "", sessionId);
         const sessions = hybridManager.listSessions();
@@ -349,6 +472,7 @@ function handleCreateCommand(ws, args = {}) {
         if (idx >= 0) {
             hybridManager.switchToIndex(idx + 1);
             setActiveItermSessionId(sessionId);
+            setLastRoutedSessionId(sessionId);
         }
     }
     log(`[PAILot] created new session "${name}" (${sessionId.slice(0, 8)}...)`);
@@ -462,16 +586,28 @@ function sendTo(ws, msg) {
     }
 }
 function broadcast(msg) {
-    // Check if any client is actually ready to receive
+    // Only deliver to clients that have proven liveness recently.
+    // iOS can keep a WebSocket "open" while the app is backgrounded —
+    // ws.send() succeeds but the app never processes the data.
+    //
+    // Session filtering: if the message carries a sessionId, only deliver to
+    // clients viewing that session. This prevents cross-session content bleed.
+    const msgSessionId = msg.sessionId;
     let delivered = false;
     const payload = JSON.stringify(msg);
     for (const ws of clients) {
-        if (ws.readyState === WebSocket.OPEN) {
-            ws.send(payload);
-            delivered = true;
+        if (!isClientAlive(ws))
+            continue;
+        // Session gate: if both the message and client have a session, they must match
+        if (msgSessionId) {
+            const clientSession = clientActiveSession.get(ws);
+            if (clientSession && clientSession !== msgSessionId)
+                continue;
         }
+        ws.send(payload);
+        delivered = true;
     }
-    // No connected clients — buffer for later
+    // No live clients — buffer for later
     if (!delivered) {
         addToOutbox(msg);
     }
@@ -483,16 +619,29 @@ const BATCH_WINDOW_MS = 3000;
 let voiceBatchTimer = null;
 let voiceBatchTranscripts = [];
 let voiceBatchOnMessage = null;
+/** iTerm session ID resolved when the first voice chunk of this batch arrived. */
+let voiceBatchSessionId = "";
 function flushVoiceBatch() {
     if (voiceBatchTranscripts.length === 0)
         return;
     const combined = voiceBatchTranscripts.join(" ");
     const handler = voiceBatchOnMessage;
+    const batchSession = voiceBatchSessionId;
     voiceBatchTranscripts = [];
     voiceBatchOnMessage = null;
+    voiceBatchSessionId = "";
     voiceBatchTimer = null;
-    if (handler) {
-        log(`[PAILot] Flushing voice batch (${combined.length} chars)`);
+    log(`[PAILot] Flushing voice batch (${combined.length} chars)`);
+    const routeSession = batchSession || activeItermSessionId;
+    if (routeSession) {
+        pailotReplyMap.set(routeSession, routeSession);
+        setLastRoutedSessionId(routeSession);
+    }
+    const bridge = getAibpBridge();
+    if (bridge && routeSession) {
+        bridge.routeFromMobile(routeSession, `[PAILot:voice] ${combined}`);
+    }
+    else if (handler) {
         setMessageSource("pailot");
         handler(`[PAILot:voice] ${combined}`, Date.now());
         setMessageSource("whatsapp");
@@ -565,6 +714,8 @@ async function transcribeAndRoute(audioBase64, onMessage, messageId) {
  */
 export function startWsGateway(onMessage) {
     wss = new WebSocketServer({ port: WS_PORT });
+    // Restore any buffered messages from a previous daemon run
+    restoreOutbox();
     wss.on("listening", () => {
         log(`WebSocket gateway listening on ws://0.0.0.0:${WS_PORT}`);
         // Pre-populate hybrid manager with live iTerm sessions so messages
@@ -612,12 +763,15 @@ end tell`)?.trim() ?? "";
         const addr = req.socket.remoteAddress ?? "unknown";
         log(`PAILot client connected from ${addr}`);
         clients.add(ws);
+        clientLastActive.set(ws, Date.now());
         ws.on("message", (raw) => {
+            clientLastActive.set(ws, Date.now());
             try {
                 const rawStr = raw.toString();
                 const msg = JSON.parse(rawStr);
                 dbg(`RAW msg (${rawStr.length} chars): type=${msg.type}, hasAudio=${!!msg.audioBase64}, content=${(msg.content ?? "").slice(0, 50)}`);
-                // Heartbeat ping — reply with pong immediately
+                // Heartbeat ping — reply with pong immediately.
+                // The clientLastActive update above already covers liveness.
                 if (msg.type === "ping") {
                     sendTo(ws, { type: "pong" });
                     return;
@@ -672,11 +826,37 @@ end tell`)?.trim() ?? "";
                             break;
                     }
                 }
+                // Extract the session ID that PAILot says this message belongs to.
+                // This is the iTerm session ID the user was viewing when they typed/spoke.
+                // We use it for routing instead of guessing from activeItermSessionId.
+                const pailotSessionId = typeof msg.sessionId === "string" ? msg.sessionId : undefined;
+                // Helper: resolve the routing target for this message.
+                // Prefers the explicit PAILot session ID; falls back to active session.
+                const routeTarget = pailotSessionId || activeItermSessionId;
+                // If PAILot told us the session, switch iTerm to it (for stdin routing)
+                // and record it in the reply map so outbound replies go to the same session.
+                if (pailotSessionId) {
+                    pailotReplyMap.set(pailotSessionId, pailotSessionId);
+                    setLastRoutedSessionId(pailotSessionId);
+                    // Switch iTerm to the correct session if it differs from the active one
+                    if (pailotSessionId !== activeItermSessionId) {
+                        setActiveItermSessionId(pailotSessionId);
+                        if (hybridManager) {
+                            const sessions = hybridManager.listSessions();
+                            const idx = sessions.findIndex(s => s.backendSessionId === pailotSessionId);
+                            if (idx >= 0)
+                                hybridManager.switchToIndex(idx + 1);
+                        }
+                    }
+                }
                 // Voice message — transcribe with Whisper then route
                 if (msg.type === "voice" && msg.audioBase64) {
                     dbg(`Voice message received, audioBase64 length: ${msg.audioBase64.length}`);
-                    broadcast({ type: "typing", typing: true });
+                    broadcast({ type: "typing", typing: true, ...(routeTarget && { sessionId: routeTarget }) });
                     const voiceMsgId = typeof msg.messageId === "string" ? msg.messageId : undefined;
+                    // Capture the routing session for this batch (first chunk wins)
+                    if (!voiceBatchSessionId && routeTarget)
+                        voiceBatchSessionId = routeTarget;
                     transcribeAndRoute(msg.audioBase64, onMessage, voiceMsgId).catch((err) => {
                         log(`[PAILot] voice transcription error: ${err}`);
                     });
@@ -698,20 +878,40 @@ end tell`)?.trim() ?? "";
                     const routeText = caption
                         ? `${caption} (image at ${imgPath})`
                         : `(image at ${imgPath})`;
-                    setMessageSource("pailot");
-                    onMessage(routeText, Date.now());
-                    setMessageSource("whatsapp");
+                    if (!pailotSessionId)
+                        setLastRoutedSessionId(activeItermSessionId);
+                    const imgBridge = getAibpBridge();
+                    if (imgBridge && routeTarget) {
+                        imgBridge.routeFromMobile(routeTarget, routeText, "IMAGE", {
+                            imageBase64: msg.imageBase64,
+                            mimeType: msg.mimeType ?? "image/jpeg",
+                        });
+                    }
+                    else {
+                        setMessageSource("pailot");
+                        onMessage(routeText, Date.now());
+                        setMessageSource("whatsapp");
+                    }
                     return;
                 }
-                // Plain text message — route through handleMessage
+                // Plain text message — route through AIBP bridge (or legacy fallback)
                 const text = msg.content ?? "";
                 if (!text.trim())
                     return;
                 log(`[PAILot] ← ${text.slice(0, 80)}${text.length > 80 ? "..." : ""}`);
-                broadcast({ type: "typing", typing: true });
-                setMessageSource("pailot");
-                onMessage(text, Date.now());
-                setMessageSource("whatsapp");
+                broadcast({ type: "typing", typing: true, ...(routeTarget && { sessionId: routeTarget }) });
+                if (!pailotSessionId)
+                    setLastRoutedSessionId(activeItermSessionId);
+                const bridge = getAibpBridge();
+                if (bridge && routeTarget) {
+                    bridge.routeFromMobile(routeTarget, text);
+                }
+                else {
+                    // Fallback to legacy routing
+                    setMessageSource("pailot");
+                    onMessage(text, Date.now());
+                    setMessageSource("whatsapp");
+                }
             }
             catch {
                 log(`[PAILot] Invalid message from ${addr}`);
@@ -720,10 +920,14 @@ end tell`)?.trim() ?? "";
         ws.on("close", () => {
             log(`PAILot client disconnected from ${addr}`);
             clients.delete(ws);
+            clientLastActive.delete(ws);
+            clientActiveSession.delete(ws);
         });
         ws.on("error", (err) => {
             log(`[PAILot] WebSocket error: ${err.message}`);
             clients.delete(ws);
+            clientLastActive.delete(ws);
+            clientActiveSession.delete(ws);
         });
         // Welcome — outbox drains after client sends "sync" command
         // (so activeSessionId is set before messages arrive)
@@ -738,8 +942,19 @@ end tell`)?.trim() ?? "";
  * @param sessionId — iTerm session ID of the originating Claude session
  */
 function resolveSessionId(sessionId) {
-    if (sessionId)
+    // If an explicit sessionId was passed (from MCP detectSessionId), check the reply map
+    // first — the reply map records which session the user was actually talking to,
+    // which may differ from whatever iTerm tab is currently focused on the Mac.
+    if (sessionId) {
+        const mapped = pailotReplyMap.get(sessionId);
+        if (mapped)
+            return mapped;
         return sessionId;
+    }
+    // Prefer the session that last received user input from PAILot —
+    // this survives session switches that happen while Claude is thinking.
+    if (lastRoutedSessionId)
+        return lastRoutedSessionId;
     if (activeItermSessionId)
         return activeItermSessionId;
     // Last resort: ask hybrid manager for the active session's backend ID
@@ -747,7 +962,7 @@ function resolveSessionId(sessionId) {
 }
 export function broadcastText(text, sessionId) {
     const resolvedSession = resolveSessionId(sessionId);
-    broadcast({ type: "typing", typing: false });
+    broadcast({ type: "typing", typing: false, ...(resolvedSession && { sessionId: resolvedSession }) });
     broadcast({ type: "text", content: text, ...(resolvedSession && { sessionId: resolvedSession }) });
 }
 /**
@@ -757,7 +972,7 @@ export function broadcastText(text, sessionId) {
  */
 export async function broadcastVoice(audioBuffer, transcript, sessionId) {
     const resolvedSession = resolveSessionId(sessionId);
-    broadcast({ type: "typing", typing: false });
+    broadcast({ type: "typing", typing: false, ...(resolvedSession && { sessionId: resolvedSession }) });
     let sendBuffer = audioBuffer;
     // Convert OGG Opus → M4A for iOS compatibility
     try {