From 3d3c7ba372f39021eededb47a6d9371dc683dd34 Mon Sep 17 00:00:00 2001
From: Paul Bakaus <paul.bakaus@gmail.com>
Date: Sun, 12 Apr 2026 18:07:06 -0700
Subject: [PATCH] Fix live session bugs found during E2E testing

Three bugs found and fixed during real end-to-end testing:

1. MutationObserver infinite loop: the observer watched all of
   document.body, so our own bar DOM updates triggered it, which
   rebuilt the bar, which triggered it again, freezing the page.
   Fix: filter mutations to only react when nodes with
   data-impeccable-variant attributes are added inside the variant
   wrapper. Added a re-entrancy guard as a safety net.

2. Premature exit on transient WS disconnect: the server fired an
   exit event the instant the last WebSocket client disconnected.
   HMR page reloads cause brief disconnects that triggered false
   exits. Fix: 8-second debounce before sending exit, cancelled
   if a client reconnects within that window.

3. WS auth_ok clobbering resumed session state: after a page reload,
   resumeSession() correctly set state to CYCLING, but then the
   async WS auth_ok handler overwrote it to PICKING. Fix: only
   transition to PICKING from IDLE, not from an active session state.

Also fixed: highlight tracking during variant cycling (update
selectedElement to the newly visible variant's content element so
the highlight follows the active variant, not the hidden one).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/live/browser.js | 115 ++++++++++++++++++++++++++++++++++++--------
 src/live/server.mjs |  15 +++++-
 2 files changed, 108 insertions(+), 22 deletions(-)

diff --git a/src/live/browser.js b/src/live/browser.js
index a542af543..223bd36a0 100644
--- a/src/live/browser.js
+++ b/src/live/browser.js
@@ -13,16 +13,19 @@
   'use strict';
   if (typeof window === 'undefined') return;
 
+  // Guard against double-init. Bun's HTML loader may process the <script> tag
+  // and create a bundled copy alongside the external load, or HMR may re-execute.
+  // Check BEFORE reading token/port to catch all cases.
+  if (window.__IMPECCABLE_LIVE_INIT__) return;
+  window.__IMPECCABLE_LIVE_INIT__ = true;
+
   const TOKEN = window.__IMPECCABLE_TOKEN__;
   const PORT = window.__IMPECCABLE_PORT__;
   if (!TOKEN || !PORT) {
-    console.warn('[impeccable] Live script loaded without token/port. Aborting.');
+    window.__IMPECCABLE_LIVE_INIT__ = false; // reset so the real load can init
     return;
   }
 
-  if (window.__IMPECCABLE_LIVE_INIT__) return;
-  window.__IMPECCABLE_LIVE_INIT__ = true;
-
   // ---------------------------------------------------------------------------
   // Design tokens
   // ---------------------------------------------------------------------------
@@ -498,6 +501,7 @@
           e.stopPropagation();
           visibleVariant = idx;
           showVariantInDOM(currentSessionId, idx);
+          updateSelectedElement();
           updateBarContent('cycling');
         });
       }
@@ -639,39 +643,71 @@
     if (next < 1 || next > arrivedVariants) return;
     visibleVariant = next;
     showVariantInDOM(currentSessionId, next);
+    // Update selectedElement to the newly visible variant's content
+    updateSelectedElement();
     updateBarContent('cycling');
   }
 
+  function updateSelectedElement() {
+    if (!currentSessionId) return;
+    const wrapper = document.querySelector('[data-impeccable-variants="' + currentSessionId + '"]');
+    if (!wrapper) return;
+    const visEl = wrapper.querySelector('[data-impeccable-variant="' + visibleVariant + '"] > :first-child');
+    if (visEl) selectedElement = visEl;
+  }
+
   // ---------------------------------------------------------------------------
   // MutationObserver for progressive variant reveal
   // ---------------------------------------------------------------------------
 
   function startVariantObserver(sessionId) {
-    const obs = new MutationObserver(() => {
+    let updating = false; // re-entrancy guard
+
+    const obs = new MutationObserver((mutations) => {
+      if (updating) return;
+
+      // Only react to mutations that add nodes with data-impeccable-variant,
+      // or mutations inside the variant wrapper. Ignore our own bar/UI changes.
+      let dominated = false;
+      for (const m of mutations) {
+        if (m.target.closest?.('[data-impeccable-variants]')) { dominated = true; break; }
+        for (const n of m.addedNodes) {
+          if (n.nodeType === 1 && (n.dataset?.impeccableVariants || n.dataset?.impeccableVariant)) {
+            dominated = true; break;
+          }
+        }
+        if (dominated) break;
+      }
+      if (!dominated) return;
+
       const wrapper = document.querySelector('[data-impeccable-variants="' + sessionId + '"]');
       if (!wrapper) return;
+
       const variants = wrapper.querySelectorAll('[data-impeccable-variant]:not([data-impeccable-variant="original"])');
       const count = variants.length;
-      if (count > arrivedVariants) {
-        arrivedVariants = count;
-        if (visibleVariant === 0 && arrivedVariants > 0) {
-          visibleVariant = 1;
-          showVariantInDOM(sessionId, 1);
-        }
-        // Update bar to reflect new dots
-        if (state === 'GENERATING') updateBarContent('generating');
-        else if (state === 'CYCLING') updateBarContent('cycling');
+
+      // Nothing new
+      if (count <= arrivedVariants) return;
+
+      updating = true;
+      arrivedVariants = count;
+      if (visibleVariant === 0 && arrivedVariants > 0) {
+        visibleVariant = 1;
+        showVariantInDOM(sessionId, 1);
       }
+
       const expected = parseInt(wrapper.dataset.impeccableVariantCount || '0');
-      if (expected > 0 && expected !== expectedVariants) {
-        expectedVariants = expected;
-        if (state === 'GENERATING') updateBarContent('generating');
-      }
+      if (expected > 0) expectedVariants = expected;
+
       if (arrivedVariants >= expectedVariants && expectedVariants > 0) {
         state = 'CYCLING';
         updateBarContent('cycling');
+      } else if (state === 'GENERATING') {
+        updateBarContent('generating');
       }
+      updating = false;
     });
+
     obs.observe(document.body, { childList: true, subtree: true });
     return obs;
   }
@@ -709,7 +745,8 @@
           hasProjectContext = !!msg.hasProjectContext;
           if (!hasProjectContext) showToast('No .impeccable.md found. Variants will be brand-agnostic.', 6000);
           console.log('[impeccable] Live mode connected.');
-          state = 'PICKING';
+          // Only go to PICKING if we're not already in a resumed session
+          if (state === 'IDLE') state = 'PICKING';
           break;
         case 'auth_fail':
           console.error('[impeccable] Auth failed:', msg.reason);
@@ -894,6 +931,38 @@
   // Init
   // ---------------------------------------------------------------------------
 
+  // Resume an active variant session after HMR/page reload.
+  // If a [data-impeccable-variants] wrapper exists in the DOM, the agent wrote
+  // variants before HMR fired. Pick up where we left off.
+  function resumeSession() {
+    const wrapper = document.querySelector('[data-impeccable-variants]');
+    if (!wrapper) return false;
+
+    currentSessionId = wrapper.dataset.impeccableVariants;
+    expectedVariants = parseInt(wrapper.dataset.impeccableVariantCount || '0');
+    const variants = wrapper.querySelectorAll('[data-impeccable-variant]:not([data-impeccable-variant="original"])');
+    arrivedVariants = variants.length;
+    visibleVariant = arrivedVariants > 0 ? 1 : 0;
+
+    // Find the visible variant's content element for highlight positioning.
+    // The wrapper has display:contents (no box), so we target the actual
+    // content element inside the currently visible variant.
+    const visEl = wrapper.querySelector('[data-impeccable-variant]:not([style*="display: none"]):not([data-impeccable-variant="original"])');
+    selectedElement = (visEl && visEl.firstElementChild) || visEl || wrapper.parentElement;
+
+    // Set display state BEFORE starting observer (avoid triggering it)
+    if (arrivedVariants > 0) showVariantInDOM(currentSessionId, 1);
+
+    state = arrivedVariants >= expectedVariants ? 'CYCLING' : 'GENERATING';
+    showBar(state === 'CYCLING' ? 'cycling' : 'generating');
+    startScrollTracking();
+
+    // Start observing for more variants AFTER initial setup
+    if (variantObserver) variantObserver.disconnect();
+    variantObserver = startVariantObserver(currentSessionId);
+    return true;
+  }
+
   function init() {
     initHighlight();
     initBar();
@@ -902,7 +971,13 @@
     document.addEventListener('click', handleClick, true);
     document.addEventListener('keydown', handleKeyDown, true);
     connectWS();
-    console.log('[impeccable] Live variant mode ready. Hover over elements to pick one.');
+
+    // Check for an active session to resume (variant wrapper already in DOM after HMR)
+    if (!resumeSession()) {
+      console.log('[impeccable] Live variant mode ready. Hover over elements to pick one.');
+    } else {
+      console.log('[impeccable] Resumed active variant session ' + currentSessionId + ' (' + arrivedVariants + '/' + expectedVariants + ' variants).');
+    }
   }
 
   if (document.readyState === 'loading') {
diff --git a/src/live/server.mjs b/src/live/server.mjs
index 95f3390c6..85b06ba9c 100644
--- a/src/live/server.mjs
+++ b/src/live/server.mjs
@@ -51,6 +51,8 @@ const state = {
   pendingEvents: [],
   // Queue: agent poll response callbacks waiting for browser events
   pendingPolls: [],
+  // Debounce timer for exit event (avoids false exits on transient disconnects)
+  exitTimer: null,
 };
 
 /** Push an event from the browser into the queue or resolve a waiting poll. */
@@ -293,6 +295,8 @@ function setupWebSocket(server) {
         if (msg.type === 'auth' && msg.token === state.token) {
           authenticated = true;
           state.wsClients.add(ws);
+          // Cancel any pending exit timer (client reconnected)
+          clearTimeout(state.exitTimer);
           ws.send(JSON.stringify({
             type: 'auth_ok',
             hasProjectContext: hasProjectContext(),
@@ -316,9 +320,16 @@ function setupWebSocket(server) {
 
     ws.on('close', () => {
       state.wsClients.delete(ws);
-      // If all browser clients disconnected, signal exit to agent
+      // If all browser clients disconnected, debounce before signaling exit.
+      // The browser script reconnects within 3s, and HMR page reloads cause
+      // brief disconnects. Wait 8s to avoid false exits.
       if (authenticated && state.wsClients.size === 0) {
-        enqueueEvent({ type: 'exit' });
+        clearTimeout(state.exitTimer);
+        state.exitTimer = setTimeout(() => {
+          if (state.wsClients.size === 0) {
+            enqueueEvent({ type: 'exit' });
+          }
+        }, 8000);
       }
     });