fix: preserve workspace tab index on workspace close -> approach change

Author: 7ito

src/browser/components/sessionstore/SessionStore-sys-mjs.patch

@@ -1,5 +1,5 @@
 diff --git a/browser/components/sessionstore/SessionStore.sys.mjs b/browser/components/sessionstore/SessionStore.sys.mjs
-index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c2045745f 100644
+index 2a055f0c5f34f0a2667f659185120c07d38f4e41..20cdd6721eeb8d8e2f85785f425f0edc3cb305bf 100644
 --- a/browser/components/sessionstore/SessionStore.sys.mjs
 +++ b/browser/components/sessionstore/SessionStore.sys.mjs
 @@ -127,6 +127,9 @@ const TAB_EVENTS = [
@@ -170,7 +170,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
        // If sidebar command is truthy, i.e. sidebar is open, store sidebar settings
        let shouldStore = hasSaveableTabs || isLastWindow;
 
-@@ -3408,18 +3429,34 @@ var SessionStoreInternal = {
+@@ -3408,12 +3429,17 @@ var SessionStoreInternal = {
      if (!isPrivateWindow && tabState.isPrivate) {
        return;
      }
@@ -181,32 +181,15 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
 
      let permanentKey = aTab.linkedBrowser.permanentKey;
 
-+    let zenWorkspaceTabIndex;
-+    let workspaceId = aTab.getAttribute("zen-workspace-id");
-+    if (workspaceId && aWindow.gZenWorkspaces?.workspaceEnabled) {
-+      let workspaceContainer = aWindow.gZenWorkspaces.workspaceElement(workspaceId);
-+      let wsTabContainer = aTab.pinned
-+        ? workspaceContainer?.pinnedTabsContainer
-+        : workspaceContainer?.tabsContainer;
-+      if (wsTabContainer) {
-+        let wsTabs = [...wsTabContainer.children].filter(
-+          child => aWindow.gBrowser.isTab(child) && !child.hasAttribute("zen-empty-tab")
-+        );
-+        zenWorkspaceTabIndex = wsTabs.indexOf(aTab);
-+      }
++    const zenWorkspaceIndex = getZenWorkspaceIndex(aTab);
++    if (Number.isInteger(zenWorkspaceIndex)) {
++      tabState.zenWorkspaceIndex = zenWorkspaceIndex;
 +    }
 +
      let tabData = {
        permanentKey,
        state: tabState,
-       title: aTab.label,
-       image: aWindow.gBrowser.getIcon(aTab),
-       pos: aTab._tPos,
-+      zenWorkspaceTabIndex,
-       closedAt: Date.now(),
-       closedInGroup: aTab._closedInMultiselection,
-       closedInTabGroupId: closedInTabGroup ? tabState.groupId : null,
-@@ -4129,6 +4166,12 @@ var SessionStoreInternal = {
+@@ -4129,6 +4155,12 @@ var SessionStoreInternal = {
          Math.min(tabState.index, tabState.entries.length)
        );
        tabState.pinned = false;
@@ -219,26 +202,17 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
 
        if (inBackground === false) {
          aWindow.gBrowser.selectedTab = newTab;
-@@ -4534,7 +4577,7 @@ var SessionStoreInternal = {
-       this._getClosedTabStateFromUnifiedIndex(sourceWinData, closedTabState);
- 
-     // fetch the data of closed tab, while removing it from the array
--    let { state, pos } = this.removeClosedTabData(
-+    let { state, pos, zenWorkspaceTabIndex } = this.removeClosedTabData(
-       sourceWinData,
-       closedTabSet,
-       closedTabIndex
-@@ -4565,6 +4608,9 @@ var SessionStoreInternal = {
+@@ -4565,6 +4597,9 @@ var SessionStoreInternal = {
        // Append the tab if we're opening into a different window,
        tabIndex: aSource == aTargetWindow ? pos : Infinity,
        pinned: state.pinned,
 +      essential: state.zenEssential,
 +      zenWorkspaceId: state.zenWorkspace,
-+      zenWorkspaceTabIndex,
++      zenWorkspaceIndex: state.zenWorkspaceIndex,
        userContextId: state.userContextId,
        skipLoad: true,
        preferredRemoteType,
-@@ -5414,7 +5460,7 @@ var SessionStoreInternal = {
+@@ -5414,7 +5449,7 @@ var SessionStoreInternal = {
 
      for (let i = tabbrowser.pinnedTabCount; i < tabbrowser.tabs.length; i++) {
        let tab = tabbrowser.tabs[i];
@@ -247,7 +221,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
          removableTabs.push(tab);
        }
      }
-@@ -5525,7 +5571,7 @@ var SessionStoreInternal = {
+@@ -5525,7 +5560,7 @@ var SessionStoreInternal = {
 
      // collect the data for all windows
      for (ix in this._windows) {
@@ -256,7 +230,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
          // window data is still in _statesToRestore
          continue;
        }
-@@ -5668,11 +5714,12 @@ var SessionStoreInternal = {
+@@ -5668,11 +5703,12 @@ var SessionStoreInternal = {
      }
 
      let tabbrowser = aWindow.gBrowser;
@@ -270,7 +244,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
      // update the internal state data for this window
      for (let tab of tabs) {
        if (tab == aWindow.FirefoxViewHandler.tab) {
-@@ -5683,6 +5730,9 @@ var SessionStoreInternal = {
+@@ -5683,6 +5719,9 @@ var SessionStoreInternal = {
        tabsData.push(tabData);
      }
 
@@ -280,7 +254,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
      // update tab group state for this window
      winData.groups = [];
      for (let tabGroup of aWindow.gBrowser.tabGroups) {
-@@ -5695,7 +5745,7 @@ var SessionStoreInternal = {
+@@ -5695,7 +5734,7 @@ var SessionStoreInternal = {
      // a window is closed, point to the first item in the tab strip instead (it will never be the Firefox View tab,
      // since it's only inserted into the tab strip after it's selected).
      if (aWindow.FirefoxViewHandler.tab?.selected) {
@@ -289,7 +263,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
        winData.title = tabbrowser.tabs[0].label;
      }
      winData.selected = selectedIndex;
-@@ -5810,8 +5860,8 @@ var SessionStoreInternal = {
+@@ -5810,8 +5849,8 @@ var SessionStoreInternal = {
      // selectTab represents.
      let selectTab = 0;
      if (overwriteTabs) {
@@ -300,15 +274,15 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
        selectTab = Math.min(selectTab, winData.tabs.length);
      }
 
-@@ -5833,6 +5883,7 @@ var SessionStoreInternal = {
+@@ -5833,6 +5872,7 @@ var SessionStoreInternal = {
      if (overwriteTabs) {
        for (let i = tabbrowser.browsers.length - 1; i >= 0; i--) {
          if (!tabbrowser.tabs[i].selected) {
 +          aWindow.gZenWorkspaces._shouldOverrideTabs = true;
            tabbrowser.removeTab(tabbrowser.tabs[i]);
          }
        }
-@@ -5866,6 +5917,12 @@ var SessionStoreInternal = {
+@@ -5866,6 +5906,12 @@ var SessionStoreInternal = {
          savedTabGroup => !openTabGroupIdsInWindow.has(savedTabGroup.id)
        );
      }
@@ -321,7 +295,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
 
      // Move the originally open tabs to the end.
      if (initialTabs) {
-@@ -6419,6 +6476,25 @@ var SessionStoreInternal = {
+@@ -6419,6 +6465,25 @@ var SessionStoreInternal = {
 
      // Most of tabData has been restored, now continue with restoring
      // attributes that may trigger external events.
@@ -347,7 +321,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
 
      if (tabData.pinned) {
        tabbrowser.pinTab(tab);
-@@ -6567,6 +6643,9 @@ var SessionStoreInternal = {
+@@ -6567,6 +6632,9 @@ var SessionStoreInternal = {
          aWindow.gURLBar.readOnly = false;
        }
      }
@@ -357,7 +331,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
 
      let promiseParts = Promise.withResolvers();
      aWindow.setTimeout(() => {
-@@ -7343,7 +7422,7 @@ var SessionStoreInternal = {
+@@ -7343,7 +7411,7 @@ var SessionStoreInternal = {
 
        let groupsToSave = new Map();
        for (let tIndex = 0; tIndex < window.tabs.length; ) {
@@ -366,7 +340,7 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
            // Adjust window.selected
            if (tIndex + 1 < window.selected) {
              window.selected -= 1;
-@@ -7358,7 +7437,7 @@ var SessionStoreInternal = {
+@@ -7358,7 +7426,7 @@ var SessionStoreInternal = {
            );
            // We don't want to increment tIndex here.
            continue;
@@ -375,3 +349,63 @@ index 2a055f0c5f34f0a2667f659185120c07d38f4e41..e0a1f2414d752f019c0a826ef71bc14c
            // Convert any open groups into saved groups.
            let groupStateToSave = window.groups.find(
              groupState => groupState.id == window.tabs[tIndex].groupId
+@@ -8786,5 +8854,59 @@ function removeWhere(array, predicate) {
+   }
+ }
+ 
++function getZenWorkspaceIndex(tab) {
++  const win = tab.ownerGlobal;
++  if (!win?.gZenWorkspaces?.workspaceEnabled) {
++    return null;
++  }
++
++  if (
++    tab.pinned ||
++    tab.hasAttribute("zen-essential") ||
++    tab.hasAttribute("zen-empty-tab") ||
++    tab.hasAttribute("zen-glance-tab") ||
++    tab.group?.isZenFolder
++  ) {
++    return null;
++  }
++
++  const workspaceId = tab.getAttribute("zen-workspace-id");
++  if (!workspaceId) {
++    return null;
++  }
++
++  const workspaceElement = win.gZenWorkspaces.workspaceElement(workspaceId);
++  const container = workspaceElement?.tabsContainer;
++  if (!container) {
++    return null;
++  }
++
++  const tabs = [];
++  for (const child of container.children) {
++    if (win.gBrowser.isTab(child)) {
++      tabs.push(child);
++      continue;
++    }
++    if (win.gBrowser.isTabGroup(child)) {
++      for (const groupTab of child.tabs) {
++        tabs.push(groupTab);
++      }
++    }
++  }
++
++  const eligibleTabs = tabs.filter(
++    t =>
++      t.getAttribute("zen-workspace-id") === workspaceId &&
++      !t.pinned &&
++      !t.hasAttribute("zen-essential") &&
++      !t.hasAttribute("zen-empty-tab") &&
++      !t.hasAttribute("zen-glance-tab") &&
++      !t.group?.isZenFolder
++  );
++
++  const index = eligibleTabs.indexOf(tab);
++  return index >= 0 ? index : null;
++}
++
+ // Exposed for tests
+ export const _LastSession = LastSession;

src/browser/components/tabbrowser/content/tabbrowser-js.patch

@@ -1,5 +1,5 @@
 diff --git a/browser/components/tabbrowser/content/tabbrowser.js b/browser/components/tabbrowser/content/tabbrowser.js
-index 0eaca7a58e0026237b71b2ad515efe84d9e8c779..28bbab3086d6d2447fd354d2eeca3673c255af72 100644
+index 0eaca7a58e0026237b71b2ad515efe84d9e8c779..9eba7866a119a4d683564ccfb9f50f3bff389d7d 100644
 --- a/browser/components/tabbrowser/content/tabbrowser.js
 +++ b/browser/components/tabbrowser/content/tabbrowser.js
 @@ -398,6 +398,7 @@
@@ -260,7 +260,7 @@ index 0eaca7a58e0026237b71b2ad515efe84d9e8c779..28bbab3086d6d2447fd354d2eeca3673
 +        _forZenEmptyTab,
 +        essential,
 +        zenWorkspaceId,
-+        zenWorkspaceTabIndex,
++        zenWorkspaceIndex,
        } = {}
      ) {
        // all callers of addTab that pass a params object need to pass
@@ -304,8 +304,8 @@ index 0eaca7a58e0026237b71b2ad515efe84d9e8c779..28bbab3086d6d2447fd354d2eeca3673
 +          t.setAttribute("zen-workspace-id", zenForcedWorkspaceId);
 +          t.setAttribute("change-workspace", "")
 +        }
-+        if (typeof zenWorkspaceTabIndex === "number") {
-+          t.setAttribute("zen-workspace-tab-index", String(zenWorkspaceTabIndex));
++        if (Number.isInteger(zenWorkspaceIndex)) {
++          t.setAttribute("zen-workspace-index", String(zenWorkspaceIndex));
 +        }
 +        if (_forZenEmptyTab) {
 +          t.setAttribute("zen-empty-tab", "true");

src/zen/workspaces/ZenWorkspaces.mjs

@@ -1568,11 +1568,37 @@ class nsZenWorkspaces {
     if (typeof tabIndex !== "number" || tabIndex < 0) {
       return container.lastChild;
     }
-    const tabs = [...container.children].filter(
-      (child) => gBrowser.isTab(child) && !child.hasAttribute("zen-empty-tab")
+    const workspaceId = container.getAttribute("zen-workspace-id");
+    if (!workspaceId) {
+      return container.lastChild;
+    }
+    const tabs = [];
+    for (const child of container.children) {
+      if (gBrowser.isTab(child)) {
+        tabs.push(child);
+        continue;
+      }
+      if (gBrowser.isTabGroup(child)) {
+        for (const groupTab of child.tabs) {
+          tabs.push(groupTab);
+        }
+      }
+    }
+    const eligibleTabs = tabs.filter(
+      (tab) =>
+        tab.getAttribute("zen-workspace-id") === workspaceId &&
+        !tab.pinned &&
+        !tab.hasAttribute("zen-essential") &&
+        !tab.hasAttribute("zen-empty-tab") &&
+        !tab.hasAttribute("zen-glance-tab") &&
+        !tab.group?.isZenFolder
     );
-    if (tabIndex < tabs.length) {
-      return tabs[tabIndex];
+    const targetTab = eligibleTabs[tabIndex];
+    if (targetTab) {
+      const targetNode = targetTab.group ?? targetTab;
+      if (targetNode.parentElement === container) {
+        return targetNode;
+      }
     }
     return container.lastChild;
   }
@@ -2611,9 +2637,9 @@ class nsZenWorkspaces {
     }
 
     if (workspaceID) {
-      const rawTabIndex = tab.getAttribute("zen-workspace-tab-index");
+      const rawTabIndex = tab.getAttribute("zen-workspace-index");
       const tabIndex = rawTabIndex !== null ? parseInt(rawTabIndex, 10) : undefined;
-      tab.removeAttribute("zen-workspace-tab-index");
+      tab.removeAttribute("zen-workspace-index");
       if (
         tab.hasAttribute("change-workspace") &&
         this.moveTabToWorkspace(tab, workspaceID, tabIndex)