Fix transcript handoff barrier race

Author: esokullu

src/chrome/src/background.js

@@ -165,7 +165,21 @@ function getContextMenuPromptStore() {
 }
 
 const contextMenuStorage = createContextMenuStorage(getContextMenuPromptStore);
-const tabChatHandoff = createTabChatHandoffCoordinator(chrome.storage.session);
+const tabChatHandoff = createTabChatHandoffCoordinator(chrome.storage.session, {
+  requestHandoff: async (tabId, { ownerId, generation }) => {
+    try {
+      return await chrome.runtime.sendMessage({
+        target: 'sidepanel',
+        action: 'tab_chat_handoff_request',
+        tabId,
+        ownerId,
+        generation,
+      });
+    } catch {
+      return null;
+    }
+  },
+});
 
 function createContextMenus() {
   if (!chrome.contextMenus?.create) return;
@@ -2658,11 +2672,15 @@ async function handleMessage(msg, sender) {
     }
 
     case 'persist_tab_chat':
-      return await tabChatHandoff.save(msg.tabId || sender.tab?.id, msg.html);
+      return await tabChatHandoff.save(msg.tabId || sender.tab?.id, msg.html, {
+        ownerId: msg.handoffOwnerId,
+        handoffGeneration: msg.handoffGeneration,
+      });
 
     case 'load_tab_chat':
       return await tabChatHandoff.load(msg.tabId || sender.tab?.id, {
         waitForHandoff: msg.waitForHandoff === true,
+        claimantId: msg.handoffOwnerId,
       });
 
     case 'clear_tab_chat':

src/chrome/src/ui/sidepanel.js

@@ -1538,6 +1538,10 @@ function isSuccessfulAskCompletion(mode, response) {
 // conversation survives the side panel being closed and reopened.
 const tabChats = new Map();
 const tabChatOperations = new Map();
+const tabChatHandoffGenerations = new Map();
+const tabChatHandoffOwnerId = `sidepanel-chat-${
+  globalThis.crypto?.randomUUID?.() || `${Date.now()}-${Math.random().toString(36).slice(2)}`
+}`;
 const tabInputDrafts = new Map();
 const permissionSkipCommandContextsByTab = new Map();
 const queuedComposerMessagesByTab = new Map();
@@ -1568,7 +1572,12 @@ async function loadTabChat(tabId, { waitForHandoff = false } = {}) {
       const stored = await sendToBackground('load_tab_chat', {
         tabId: queuedTabId,
         waitForHandoff,
+        handoffOwnerId: tabChatHandoffOwnerId,
       });
+      if (stored?.handoffOwnerId === tabChatHandoffOwnerId
+          && Number.isFinite(Number(stored?.handoffGeneration))) {
+        tabChatHandoffGenerations.set(queuedTabId, Number(stored.handoffGeneration));
+      }
       const html = stored?.found ? stored.html : null;
       if (typeof html === 'string') {
         tabChats.set(queuedTabId, html);
@@ -1587,18 +1596,25 @@ function persistTabChat(tabId, html, { allowHidden = false } = {}) {
   }
   const numericTabId = Number(tabId);
   if (!Number.isFinite(numericTabId)) return Promise.resolve({ ok: false, error: 'No tab ID' });
+  const handoffGeneration = tabChatHandoffGenerations.get(numericTabId);
+  const payload = {
+    tabId: numericTabId,
+    html,
+    handoffOwnerId: tabChatHandoffOwnerId,
+  };
+  if (Number.isFinite(handoffGeneration)) payload.handoffGeneration = handoffGeneration;
   if (document.visibilityState === 'hidden' && allowHidden) {
     // Do not wait behind this document's local queue. The shared background
     // coordinator orders this authoritative handoff behind any earlier write
     // and ahead of the newly visible document's restore.
     tabChats.set(numericTabId, html);
-    return sendToBackground('persist_tab_chat', { tabId: numericTabId, html });
+    return sendToBackground('persist_tab_chat', payload);
   }
   return enqueueTabChatOperation(tabId, async (numericTabId) => {
     // Keep the live transcript lossless. The background may compact only the
     // storage.session copy when the shared quota requires it.
     tabChats.set(numericTabId, html);
-    return sendToBackground('persist_tab_chat', { tabId: numericTabId, html });
+    return sendToBackground('persist_tab_chat', payload);
   });
 }
 
@@ -1626,6 +1642,7 @@ function clearCachedTabChat(tabId) {
     persistTimerTabId = null;
   }
   tabChats.delete(tabId);
+  tabChatHandoffGenerations.delete(Number(tabId));
   return enqueueTabChatOperation(tabId, async (numericTabId) => {
     tabChats.delete(numericTabId);
     try {
@@ -7468,6 +7485,21 @@ chrome.runtime.onMessage.addListener((msg) => {
   acceptContextMenuPrompt(msg.prompt || msg);
 });
 
+chrome.runtime.onMessage.addListener((msg, _sender, sendResponse) => {
+  if (msg?.target !== 'sidepanel'
+      || msg.action !== 'tab_chat_handoff_request'
+      || String(msg.ownerId || '') !== tabChatHandoffOwnerId) return;
+  const snapshot = lastVisibleTabChatSnapshot;
+  if (!snapshot || !sameTabId(snapshot.tabId, msg.tabId)) return;
+  sendResponse({
+    ok: true,
+    tabId: Number(snapshot.tabId),
+    ownerId: tabChatHandoffOwnerId,
+    generation: Number(msg.generation),
+    html: snapshot.html,
+  });
+});
+
 document.addEventListener('visibilitychange', () => {
   if (document.visibilityState === 'hidden') {
     visibleStateRefreshPending = false;

src/chrome/src/ui/tab-chat-persistence.js

@@ -192,25 +192,25 @@ export async function persistTabChatToSession(storageArea, key, html, warn = con
   }
 }
 
-export const TAB_CHAT_HANDOFF_SETTLE_MS = 25;
+export const TAB_CHAT_HANDOFF_PREFIX = 'tabChatHandoff:';
 
 /**
  * Serialize tab-chat reads and writes in the background's shared JavaScript
- * realm. The short settle window lets the outgoing panel enqueue its final
- * visibility snapshot even if the newly visible panel's message is delivered
- * first.
+ * realm. Coordinated loads explicitly request and acknowledge the outgoing
+ * document's final snapshot before assigning a new handoff generation.
  *
  * @param {chrome.storage.StorageArea | browser.storage.StorageArea} storageArea
  * @param {{
  *   persist?: typeof persistTabChatToSession,
- *   settleHandoff?: () => Promise<void>,
+ *   requestHandoff?: (tabId: number, handoff: { ownerId: string, generation: number }) => Promise<object | null>,
  * }} options
  */
 export function createTabChatHandoffCoordinator(storageArea, {
   persist = persistTabChatToSession,
-  settleHandoff = () => new Promise(resolve => setTimeout(resolve, TAB_CHAT_HANDOFF_SETTLE_MS)),
+  requestHandoff = async () => null,
 } = {}) {
   const operations = new Map();
+  const handoffOperations = new Map();
   const latestHtml = new Map();
 
   function normalizeTabId(tabId) {
@@ -232,34 +232,110 @@ export function createTabChatHandoffCoordinator(storageArea, {
     return operation;
   }
 
-  function save(tabId, html) {
+  function enqueueHandoff(tabId, fn) {
+    const previous = handoffOperations.get(tabId) || Promise.resolve();
+    const operation = previous.catch(() => {}).then(fn);
+    handoffOperations.set(tabId, operation);
+    operation.finally(() => {
+      if (handoffOperations.get(tabId) === operation) handoffOperations.delete(tabId);
+    }).catch(() => {});
+    return operation;
+  }
+
+  async function readHandoffState(tabId) {
+    const key = TAB_CHAT_HANDOFF_PREFIX + tabId;
+    const stored = await storageArea.get(key);
+    const state = stored?.[key];
+    const ownerId = String(state?.ownerId || '');
+    const generation = Number(state?.generation);
+    return ownerId && Number.isFinite(generation) && generation > 0
+      ? { ownerId, generation }
+      : null;
+  }
+
+  async function readLatest(tabId) {
+    if (latestHtml.has(tabId)) {
+      return { ok: true, found: true, html: latestHtml.get(tabId) };
+    }
+    const key = TAB_CHAT_PREFIX + tabId;
+    const stored = await storageArea.get(key);
+    const html = stored?.[key];
+    if (typeof html === 'string') {
+      latestHtml.set(tabId, html);
+      return { ok: true, found: true, html };
+    }
+    return { ok: true, found: false, html: null };
+  }
+
+  function save(tabId, html, { ownerId = '', handoffGeneration = null } = {}) {
     const numericTabId = normalizeTabId(tabId);
     if (numericTabId == null) return Promise.resolve({ ok: false, error: 'No tab ID' });
     const source = String(html || '');
     return enqueue(numericTabId, async (queuedTabId) => {
+      const normalizedOwnerId = String(ownerId || '');
+      const normalizedGeneration = Number(handoffGeneration);
+      if (normalizedOwnerId && Number.isFinite(normalizedGeneration) && normalizedGeneration > 0) {
+        const state = await readHandoffState(queuedTabId);
+        if (!state
+            || state.ownerId !== normalizedOwnerId
+            || state.generation !== normalizedGeneration) {
+          return { ok: true, skipped: true, reason: 'stale-handoff' };
+        }
+      }
       // Retain the lossless copy even if persistence has to compact the
       // storage.session value for quota recovery.
       latestHtml.set(queuedTabId, source);
       return persist(storageArea, TAB_CHAT_PREFIX + queuedTabId, source);
     });
   }
 
-  async function load(tabId, { waitForHandoff = false } = {}) {
+  async function load(tabId, { waitForHandoff = false, claimantId = '' } = {}) {
     const numericTabId = normalizeTabId(tabId);
     if (numericTabId == null) return { ok: false, found: false, error: 'No tab ID' };
-    if (waitForHandoff) await settleHandoff();
-    return enqueue(numericTabId, async (queuedTabId) => {
-      if (latestHtml.has(queuedTabId)) {
-        return { ok: true, found: true, html: latestHtml.get(queuedTabId) };
-      }
-      const key = TAB_CHAT_PREFIX + queuedTabId;
-      const stored = await storageArea.get(key);
-      const html = stored?.[key];
-      if (typeof html === 'string') {
-        latestHtml.set(queuedTabId, html);
-        return { ok: true, found: true, html };
+    if (!waitForHandoff) {
+      return enqueue(numericTabId, queuedTabId => readLatest(queuedTabId));
+    }
+
+    return enqueueHandoff(numericTabId, async () => {
+      const outgoing = await enqueue(numericTabId, queuedTabId => readHandoffState(queuedTabId));
+      if (outgoing) {
+        let acknowledgement = null;
+        try {
+          acknowledgement = await requestHandoff(numericTabId, outgoing);
+        } catch { /* an unavailable outgoing document has no snapshot to acknowledge */ }
+        if (acknowledgement?.ok
+            && String(acknowledgement.ownerId || '') === outgoing.ownerId
+            && Number(acknowledgement.generation) === outgoing.generation
+            && typeof acknowledgement.html === 'string') {
+          await save(numericTabId, acknowledgement.html, {
+            ownerId: outgoing.ownerId,
+            handoffGeneration: outgoing.generation,
+          });
+        }
       }
-      return { ok: true, found: false, html: null };
+
+      return enqueue(numericTabId, async (queuedTabId) => {
+        const current = await readHandoffState(queuedTabId);
+        const generation = Math.max(
+          Number(outgoing?.generation || 0),
+          Number(current?.generation || 0),
+        ) + 1;
+        const normalizedClaimantId = String(claimantId || '');
+        if (normalizedClaimantId) {
+          await storageArea.set({
+            [TAB_CHAT_HANDOFF_PREFIX + queuedTabId]: {
+              ownerId: normalizedClaimantId,
+              generation,
+            },
+          });
+        }
+        const result = await readLatest(queuedTabId);
+        return {
+          ...result,
+          handoffOwnerId: normalizedClaimantId || null,
+          handoffGeneration: normalizedClaimantId ? generation : null,
+        };
+      });
     });
   }
 
@@ -268,7 +344,10 @@ export function createTabChatHandoffCoordinator(storageArea, {
     if (numericTabId == null) return Promise.resolve({ ok: false, error: 'No tab ID' });
     return enqueue(numericTabId, async (queuedTabId) => {
       latestHtml.delete(queuedTabId);
-      await storageArea.remove(TAB_CHAT_PREFIX + queuedTabId);
+      await storageArea.remove([
+        TAB_CHAT_PREFIX + queuedTabId,
+        TAB_CHAT_HANDOFF_PREFIX + queuedTabId,
+      ]);
       return { ok: true };
     });
   }

src/firefox/src/background.js

@@ -126,7 +126,21 @@ function getContextMenuPromptStore() {
 }
 
 const contextMenuStorage = createContextMenuStorage(getContextMenuPromptStore);
-const tabChatHandoff = createTabChatHandoffCoordinator(browser.storage.session);
+const tabChatHandoff = createTabChatHandoffCoordinator(browser.storage.session, {
+  requestHandoff: async (tabId, { ownerId, generation }) => {
+    try {
+      return await browser.runtime.sendMessage({
+        target: 'sidepanel',
+        action: 'tab_chat_handoff_request',
+        tabId,
+        ownerId,
+        generation,
+      });
+    } catch {
+      return null;
+    }
+  },
+});
 
 function createContextMenus() {
   const api = getContextMenuApi();
@@ -2345,11 +2359,15 @@ async function handleMessage(msg, sender) {
     }
 
     case 'persist_tab_chat':
-      return await tabChatHandoff.save(msg.tabId || sender.tab?.id, msg.html);
+      return await tabChatHandoff.save(msg.tabId || sender.tab?.id, msg.html, {
+        ownerId: msg.handoffOwnerId,
+        handoffGeneration: msg.handoffGeneration,
+      });
 
     case 'load_tab_chat':
       return await tabChatHandoff.load(msg.tabId || sender.tab?.id, {
         waitForHandoff: msg.waitForHandoff === true,
+        claimantId: msg.handoffOwnerId,
       });
 
     case 'clear_tab_chat':