Fix: Handle essential tabs and container-specific essentials

Refactor workspace switching logic to correctly handle
essential tabs and introduce container-specific essentials.

This change addresses issues with essential tabs not being
handled correctly during workspace switches.

The `changeWorkspace` function is refactored to improve
clarity and maintainability.  The logic for showing and
hiding tabs is streamlined, and tab selection is handled
more robustly.  A new `_shouldShowTab` function is
introduced to centralize the logic for determining tab
visibility based on workspace and container settings.  The
logic also handles pinned essential tabs in the different workspace types.

The pinned tab manager is updated to support container-
specific essentials and to refresh pinned tabs on workspace
changes.  The `_shouldShowPin` function is introduced to
manage visibility of pinned tabs in different workspaces
considering essential tabs, pinned tabs and containers.

This change also fixes a bug where the selected tab would
sometimes be changed unexpectedly when switching
workspaces.

src/browser/base/content/ZenStartup.mjs

@@ -34,7 +34,6 @@
         // Disable smooth scroll
         gBrowser.tabContainer.arrowScrollbox.smoothScroll = false;
 
-        gZenPinnedTabManager.initTabs();
         ZenWorkspaces.init();
         gZenUIManager.init();
         gZenVerticalTabsManager.init();

src/browser/base/zen-components/ZenPinnedTabManager.mjs

@@ -48,13 +48,16 @@
       this.observer.addPinnedTabListener(this._onPinnedTabEvent.bind(this));
 
       this._zenClickEventListener = this._onTabClick.bind(this);
+      ZenWorkspaces.addChangeListeners(this.onWorkspaceChange.bind(this));
+
     }
 
-    async initTabs() {
-      if (!this.enabled) {
+    async onWorkspaceChange(newWorkspace, onInit) {
+      if (!this.enabled || PrivateBrowsingUtils.isWindowPrivate(window)) {
         return;
       }
-      await ZenPinnedTabsStorage.init();
+
+      await this._refreshPinnedTabs(newWorkspace,{ init: onInit });
     }
 
     get enabled() {
@@ -68,9 +71,12 @@
       return this._enabled;
     }
 
-    async _refreshPinnedTabs({ init = false } = {}) {
+    async _refreshPinnedTabs(currentWorkspace,{ init = false } = {}) {
+      if(init) {
+        await ZenPinnedTabsStorage.init();
+      }
       await this._initializePinsCache();
-      this._initializePinnedTabs(init);
+      await this._initializePinnedTabs(init,currentWorkspace);
     }
 
     async _initializePinsCache() {
@@ -109,12 +115,14 @@
       return this._pinsCache;
     }
 
-    _initializePinnedTabs(init = false) {
+    async _initializePinnedTabs(init = false, currentWorkspace) {
       const pins = this._pinsCache;
       if (!pins?.length) {
         return;
       }
 
+      const workspaces = await ZenWorkspaces._workspaces();
+
       const activeTab = gBrowser.selectedTab;
       const pinnedTabsByUUID = new Map();
       const pinsToCreate = new Set(pins.map(p => p.uuid));
@@ -131,7 +139,7 @@
           pinnedTabsByUUID.set(pinId, tab);
           pinsToCreate.delete(pinId);
 
-          if(lazy.zenPinnedTabRestorePinnedTabsToPinnedUrl && init) {
+          if (lazy.zenPinnedTabRestorePinnedTabsToPinnedUrl && init) {
             this._resetTabToStoredState(tab);
           }
         } else {
@@ -146,6 +154,10 @@
           continue; // Skip pins that already have tabs
         }
 
+        if (!this._shouldShowPin(pin, currentWorkspace, workspaces)) {
+          continue; // Skip pins not relevant to current workspace
+        }
+
         let params = {
           skipAnimation: true,
           allowInheritPrincipal: false,
@@ -196,9 +208,6 @@
 
         gBrowser.pinTab(newTab);
 
-        if(newTab.getAttribute("zen-workspace-id") !== ZenWorkspaces.activeWorkspace && newTab.getAttribute("zen-essential") !== "true") {
-          gBrowser.hideTab(newTab, undefined, true);
-        }
 
         newTab.initialize();
       }
@@ -211,6 +220,41 @@
       gBrowser._updateTabBarForPinnedTabs();
     }
 
+    _shouldShowPin(pin, currentWorkspace, workspaces) {
+      const isEssential = pin.isEssential;
+      const pinWorkspaceUuid = pin.workspaceUuid;
+      const pinContextId = pin.containerTabId ? pin.containerTabId.toString() : "0";
+      const workspaceContextId = currentWorkspace.containerTabId?.toString() || "0";
+      const containerSpecificEssentials = ZenWorkspaces.containerSpecificEssentials;
+
+      // Handle essential pins
+      if (isEssential) {
+        if (!containerSpecificEssentials) {
+          return true; // Show all essential pins when containerSpecificEssentials is false
+        }
+
+        if (workspaceContextId !== "0") {
+          // In workspaces with default container: Show essentials that match the container
+          return pinContextId === workspaceContextId;
+        } else {
+          // In workspaces without a default container: Show essentials that aren't in container-specific workspaces
+          // or have userContextId="0" or no userContextId
+          return !pinContextId || pinContextId === "0" || !workspaces.workspaces.some(
+              workspace => workspace.containerTabId === parseInt(pinContextId, 10)
+          );
+        }
+      }
+
+      // For non-essential pins
+      if (!pinWorkspaceUuid) {
+        // Pins without a workspace belong to all workspaces (if that's your desired behavior)
+        return true;
+      }
+
+      // Show if pin belongs to current workspace
+      return pinWorkspaceUuid === currentWorkspace.uuid;
+    }
+
     _onPinnedTabEvent(action, event) {
       if (!this.enabled) return;
       const tab = event.target;
@@ -275,7 +319,8 @@
       pin.userContextId = userContextId ? parseInt(userContextId, 10) : 0;
 
       await ZenPinnedTabsStorage.savePin(pin);
-      await this._refreshPinnedTabs();
+      const currentWorkspace = await ZenWorkspaces.getActiveWorkspace();
+      await this._refreshPinnedTabs(currentWorkspace);
     }
 
     async _setPinnedAttributes(tab) {
@@ -311,8 +356,8 @@
         tab.removeAttribute("zen-pinned-entry");
         return;
       }
-
-      await this._refreshPinnedTabs();
+      const currentWorkspace = await ZenWorkspaces.getActiveWorkspace();
+      await this._refreshPinnedTabs(currentWorkspace);
     }
 
     async _removePinnedAttributes(tab, isClosing = false) {
@@ -330,8 +375,8 @@
           tab.setAttribute("zen-workspace-id", workspace.uuid);
         }
       }
-
-      await this._refreshPinnedTabs();
+      const currentWorkspace = await ZenWorkspaces.getActiveWorkspace();
+      await this._refreshPinnedTabs(currentWorkspace);
     }
 
     _initClosePinnedTabShortcut() {

src/browser/base/zen-components/ZenPinnedTabsStorage.mjs

@@ -47,8 +47,6 @@ var ZenPinnedTabsStorage = {
       await db.execute(`
         CREATE INDEX IF NOT EXISTS idx_zen_pins_changes_uuid ON zen_pins_changes(uuid)
       `);
-
-      await gZenPinnedTabManager._refreshPinnedTabs({init: true});
     });
   },
 

src/browser/base/zen-components/ZenWorkspaces.mjs

@@ -1015,7 +1015,8 @@ var ZenWorkspaces = new (class extends ZenMultiWindowFeature {
     document.documentElement.setAttribute('zen-workspace-id', window.uuid);
     let tabCount = 0;
     for (let tab of gBrowser.tabs) {
-      if (!tab.hasAttribute('zen-workspace-id') && !tab.pinned) {
+      const isEssential = tab.getAttribute("zen-essential") === "true";
+      if (!tab.hasAttribute('zen-workspace-id') && !tab.pinned && !isEssential) {
         tab.setAttribute('zen-workspace-id', window.uuid);
         tabCount++;
       }
@@ -1097,82 +1098,168 @@ var ZenWorkspaces = new (class extends ZenMultiWindowFeature {
     }
 
     this._inChangingWorkspace = true;
+    try {
+      await this._performWorkspaceChange(window, onInit);
+    } finally {
+      this._inChangingWorkspace = false;
+    }
+  }
+
+  async _performWorkspaceChange(window, onInit) {
     this.activeWorkspace = window.uuid;
     const containerId = window.containerTabId?.toString();
     const workspaces = await this._workspaces();
 
+    // Refresh tab cache
     this.tabContainer._invalidateCachedTabs();
-    let firstTab = undefined;
-    for (let tab of gBrowser.tabs) {
-      if (tab.getAttribute('zen-workspace-id') === window.uuid || !tab.hasAttribute('zen-workspace-id')
-      ) {
-        if (!firstTab && (onInit || !tab.pinned)) {
-          firstTab = tab;
-        } else if (gBrowser.selectedTab === tab) {
-          // If the selected tab is already in the workspace, we don't want to change it
-          firstTab = null; // note: Do not add "undefined" here, a new tab would be created
-        }
-        gBrowser.showTab(tab);
-        if (!tab.hasAttribute('zen-workspace-id') && tab.getAttribute('zen-essential') !== 'true') {
-          // We add the id to those tabs that got inserted before we initialize the workspaces or those who lost the id for any reason
-          // example use case: opening a link from an external app
-          tab.setAttribute('zen-workspace-id', window.uuid);
-        }
-      }
-    }
-    if (firstTab) {
-      // Don't change the selected tab if it's an essential tab, it becomes annoying
-      if (!gBrowser.selectedTab.hasAttribute('zen-essential')) {
-        gBrowser.selectedTab = this._lastSelectedWorkspaceTabs[window.uuid] ?? firstTab;
-      }
-    }
-    if (typeof firstTab === 'undefined' && !onInit) {
-      this._createNewTabForWorkspace(window);
-    }
-    for (let tab of gBrowser.tabs) {
-      // Skip tabs that are in the current workspace
-      if (tab.getAttribute('zen-workspace-id') === window.uuid) {
+
+    // First pass: Handle tab visibility and workspace ID assignment
+    const visibleTabs = this._processTabVisibility(window.uuid, containerId, workspaces);
+
+    // Second pass: Handle tab selection
+    await this._handleTabSelection(window, onInit, visibleTabs, containerId, workspaces);
+
+    // Update UI and state
+    await this._updateWorkspaceState(window, onInit);
+  }
+
+
+  _processTabVisibility(workspaceUuid, containerId, workspaces) {
+    const visibleTabs = new Set();
+    const lastSelectedTab = this._lastSelectedWorkspaceTabs[workspaceUuid];
+
+    for (const tab of gBrowser.tabs) {
+      const tabWorkspaceId = tab.getAttribute('zen-workspace-id');
+      const isEssential = tab.getAttribute("zen-essential") === "true";
+      const tabContextId = tab.getAttribute("usercontextid");
+
+      // Always hide last selected tabs from other workspaces
+      if (lastSelectedTab === tab && tabWorkspaceId !== workspaceUuid && !isEssential) {
+        gBrowser.hideTab(tab, undefined, true);
         continue;
       }
 
-      // Handle essentials
-      if (tab.getAttribute("zen-essential") === "true") {
-        if(this.containerSpecificEssentials) {
-          if (containerId) {
-            // In workspaces with default container: Hide essentials that don't match the container
-            if (tab.getAttribute("usercontextid") !== containerId) {
-              gBrowser.hideTab(tab, undefined, true);
-            }
-          } else {
-            // In workspaces without a default container: Hide essentials that are opened in a container and some workspace has that container as default
-            if (tab.hasAttribute("usercontextid") && workspaces.workspaces.some((workspace) => workspace.containerTabId === parseInt(tab.getAttribute("usercontextid") || "0" , 10))) {
-              gBrowser.hideTab(tab, undefined, true);
-            }
-          }
+      if (this._shouldShowTab(tab, workspaceUuid, containerId, workspaces)) {
+        gBrowser.showTab(tab);
+        visibleTabs.add(tab);
+
+        // Assign workspace ID if needed
+        if (!tabWorkspaceId && !isEssential) {
+          tab.setAttribute('zen-workspace-id', workspaceUuid);
         }
       } else {
-        // For non-pinned tabs: Hide if they're not in the current workspace
         gBrowser.hideTab(tab, undefined, true);
       }
     }
-    this.tabContainer._invalidateCachedTabs();
-    document.documentElement.setAttribute('zen-workspace-id', window.uuid);
-    await this._updateWorkspacesChangeContextMenu();
 
+    return visibleTabs;
+  }
+
+  _shouldShowTab(tab, workspaceUuid, containerId, workspaces) {
+    const isEssential = tab.getAttribute("zen-essential") === "true";
+    const tabWorkspaceId = tab.getAttribute('zen-workspace-id');
+    const tabContextId = tab.getAttribute("usercontextid");
+
+    // Handle essential tabs
+    if (isEssential) {
+      if (!this.containerSpecificEssentials) {
+        return true; // Show all essential tabs when containerSpecificEssentials is false
+      }
+
+      if (containerId) {
+        // In workspaces with default container: Show essentials that match the container
+        return tabContextId === containerId;
+      } else {
+        // In workspaces without a default container: Show essentials that aren't in container-specific workspaces
+        // or have usercontextid="0" or no usercontextid
+        return !tabContextId || tabContextId === "0" || !workspaces.workspaces.some(
+            workspace => workspace.containerTabId === parseInt(tabContextId, 10)
+        );
+      }
+    }
+
+    // For non-essential tabs (both normal and pinned)
+    if (!tabWorkspaceId) {
+      // Assign workspace ID to tabs without one
+      tab.setAttribute('zen-workspace-id', workspaceUuid);
+      return true;
+    }
+
+    // Show if tab belongs to current workspace
+    return tabWorkspaceId === workspaceUuid;
+  }
+
+  async _handleTabSelection(window, onInit, visibleTabs, containerId, workspaces) {
+    const currentSelectedTab = gBrowser.selectedTab;
+    const oldWorkspaceId = currentSelectedTab.getAttribute('zen-workspace-id');
+    const lastSelectedTab = this._lastSelectedWorkspaceTabs[window.uuid];
+
+    // Save current tab as last selected for old workspace if it shouldn't be visible in new workspace
+    if (oldWorkspaceId && oldWorkspaceId !== window.uuid) {
+      this._lastSelectedWorkspaceTabs[oldWorkspaceId] = currentSelectedTab;
+    }
+
+    let tabToSelect = null;
+
+    // If current tab is visible in new workspace, keep it
+    if (this._shouldShowTab(currentSelectedTab, window.uuid, containerId, workspaces) && visibleTabs.has(currentSelectedTab)) {
+      tabToSelect = currentSelectedTab;
+    }
+    // Try last selected tab if it is visible
+    else if (lastSelectedTab && this._shouldShowTab(lastSelectedTab, window.uuid, containerId, workspaces) && visibleTabs.has(lastSelectedTab)) {
+      tabToSelect = lastSelectedTab;
+    }
+    // Find first suitable tab
+    else {
+      tabToSelect = Array.from(visibleTabs)
+          .find(tab => !tab.pinned);
+    }
+
+    const previousSelectedTab = gBrowser.selectedTab;
+
+    // If we found a tab to select, select it
+    if (tabToSelect) {
+      gBrowser.selectedTab = tabToSelect;
+      this._lastSelectedWorkspaceTabs[window.uuid] = tabToSelect;
+    } else if (!onInit) {
+      // Create new tab if needed and no suitable tab was found
+      const newTab = this._createNewTabForWorkspace(window);
+      gBrowser.selectedTab = newTab;
+      this._lastSelectedWorkspaceTabs[window.uuid] = newTab;
+    }
+
+    // After selecting the new tab, hide the previous selected tab if it shouldn't be visible in the new workspace
+    if (!this._shouldShowTab(previousSelectedTab, window.uuid, containerId, workspaces)) {
+      gBrowser.hideTab(previousSelectedTab, undefined, true);
+    }
+  }
+
+
+  async _updateWorkspaceState(window, onInit) {
+    // Update document state
+    document.documentElement.setAttribute('zen-workspace-id', window.uuid);
+
+    // Update workspace UI
+    await this._updateWorkspacesChangeContextMenu();
     document.getElementById('tabbrowser-tabs')._positionPinnedTabs();
     gZenUIManager.updateTabsToolbar();
+    await this._propagateWorkspaceData({ clearCache: false });
 
-    await this._propagateWorkspaceData({clearCache: onInit});
-    for (let listener of this._changeListeners ?? []) {
-      listener(window);
+    // Notify listeners
+    if (this._changeListeners?.length) {
+      for (const listener of this._changeListeners) {
+        await listener(window, onInit);
+      }
     }
-    // reset bookmark toolbars
+
+    // Reset bookmarks toolbar
     const placesToolbar = document.getElementById("PlacesToolbar");
-    if(placesToolbar?._placesView) {
+    if (placesToolbar?._placesView) {
       placesToolbar._placesView.invalidateContainer(placesToolbar._placesView._resultNode);
     }
+
+    // Update workspace indicator
     await this.updateWorkspaceIndicator();
-    this._inChangingWorkspace = false;
   }
 
   async updateWorkspaceIndicator() {
@@ -1237,7 +1324,8 @@ var ZenWorkspaces = new (class extends ZenMultiWindowFeature {
 
   async onTabBrowserInserted(event) {
     let tab = event.originalTarget;
-    if (tab.getAttribute('zen-workspace-id') || !this.workspaceEnabled) {
+    const isEssential = tab.getAttribute("zen-essential") === "true";
+    if (tab.getAttribute('zen-workspace-id') || !this.workspaceEnabled || isEssential) {
       return;
     }
 
@@ -1252,18 +1340,22 @@ var ZenWorkspaces = new (class extends ZenMultiWindowFeature {
     if (!this.workspaceEnabled || this._inChangingWorkspace) {
       return;
     }
+
     const parent = browser.ownerGlobal;
-    let tab = gBrowser.getTabForBrowser(browser);
-    let workspaceID = tab.getAttribute('zen-workspace-id');
-    if (!workspaceID || tab.pinned) {
-      return;
+    const tab = gBrowser.getTabForBrowser(browser);
+    const workspaceID = tab.getAttribute('zen-workspace-id');
+    const isEssential = tab.getAttribute("zen-essential") === "true";
+    const activeWorkspace = await parent.ZenWorkspaces.getActiveWorkspace();
+
+    // Only update last selected tab for non-essential tabs in their workspace
+    if (!isEssential && workspaceID === activeWorkspace.uuid) {
+      this._lastSelectedWorkspaceTabs[workspaceID] = tab;
     }
-    let activeWorkspace = await parent.ZenWorkspaces.getActiveWorkspace();
-    this._lastSelectedWorkspaceTabs[workspaceID] = tab;
-    if (workspaceID === activeWorkspace.uuid) {
-      return;
+
+    // Switch workspace if needed
+    if (workspaceID && workspaceID !== activeWorkspace.uuid) {
+      await parent.ZenWorkspaces.changeWorkspace({ uuid: workspaceID });
     }
-    await parent.ZenWorkspaces.changeWorkspace({ uuid: workspaceID });
   }
 
   // Context menu management