Restore tree structure more correctly around tabs moved across windows #3918

Author: piroor

webextensions/background/api-tabs-listener.js

@@ -1115,6 +1115,7 @@ async function onAttached(tabId, attachInfo) {
   try {
     log('tabs.onAttached, id: ', tabId, attachInfo);
     let tab = Tab.get(tabId);
+    await ensureSessionTabValueStored(tab);
     let attachedTab = await browser.tabs.get(tabId).catch(ApiTabs.createErrorHandler(ApiTabs.handleMissingTabError));
     if (!attachedTab) {
       // We sometimes fail to get window and tab via API if it is opened
@@ -1212,6 +1213,26 @@ async function onAttached(tabId, attachInfo) {
   }
 }
 
+// Workaround for https://github.com/piroor/treestyletab/issues/3918
+async function ensureSessionTabValueStored(tab) {
+  const [oldId, states] = await Promise.all([
+    browser.sessions.getTabValue(tab.id, Constants.kPERSISTENT_ID).catch(ApiTabs.createErrorHandler()),
+    browser.sessions.getTabValue(tab.id, Constants.kPERSISTENT_STATES).catch(ApiTabs.createErrorHandler()),
+  ]);
+  if ((!tab.$TST.uniqueId || oldId !== undefined) &&
+      (!tab.$TST.$lastPermanentStates || states !== undefined))
+    return;
+
+  log(`ensureSessionTabValueStored: The tab ${tab.id} has been moved from another window, and lost its session data unexpectedly, so now we try to save them again for safety.`);
+  await Promise.all([
+    tab.$TST.uniqueId && browser.sessions.setTabValue(tab.id, Constants.kPERSISTENT_ID, {
+      id:    tab.$TST.uniqueId.id,
+      tabId: tab.id,
+    }).catch(ApiTabs.createErrorSuppressor()),
+    tab.$TST.$lastPermanentStates && browser.sessions.setTabValue(tab.id, Constants.kPERSISTENT_STATES, tab.$TST.$lastPermanentStates).catch(ApiTabs.createErrorSuppressor()),
+  ]);
+}
+
 async function onDetached(tabId, detachInfo) {
   if (mPromisedStarted)
     await mPromisedStarted;

webextensions/common/TreeItem.js

@@ -1325,6 +1325,11 @@ export class Tab extends TreeItem {
     // We should initialize private properties with blank value for better performance with a fixed shape.
     this.delayedInheritSoundStateFromChildren = null;
     this.delayedInheritSharingStateFromChildren = null;
+
+    this.$lastPermanentStates = null;
+    browser.sessions.getTabValue(raw.id, Constants.kPERSISTENT_STATES).then(states => {
+      this.$lastPermanentStates = states;
+    }).catch(ApiTabs.handleMissingTabError);
   }
 
   destroy() {
@@ -2812,6 +2817,7 @@ export class Tab extends TreeItem {
     const states = this.raw && await browser.sessions.getTabValue(this.id, Constants.kPERSISTENT_STATES).catch(ApiTabs.handleMissingTabError);
     // We need to cleanup invalid values stored accidentally.
     // See also: https://github.com/piroor/treestyletab/issues/2882
+    this.$lastPermanentStates = states;
     return states && mapAndFilterUniq(states, state => state && String(state) || undefined) || [];
   }
 

webextensions/common/unique-id.js

@@ -198,11 +198,13 @@ export async function request(tabOrId, options = {}) {
 
     let oldId;
     if (mRestoredPersistentIdMap.has(tab.id)) {
+      log('tab', tab.id, ': in mRestoredPersistentIdMap');
       oldId = mRestoredPersistentIdMap.get(tab.id);
       mRestoredPersistentIdMap.delete(tab.id);
     }
     else {
       oldId = await browser.sessions.getTabValue(tab.id, Constants.kPERSISTENT_ID).catch(ApiTabs.createErrorHandler());
+      log('tab', tab.id, ': oldId from session tabValue: ', oldId);
     }
     if (oldId && !oldId.tabId) // ignore broken information!
       oldId = null;