Fix undo all tabs

OrlaM · Mar 1, 2024 · e64fc13 · e64fc13
1 parent 24a5186
commit e64fc13
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/firefox-ios/Client/TabManagement/Legacy/LegacyTabManager.swift b/firefox-ios/Client/TabManagement/Legacy/LegacyTabManager.swift
@@ -647,17 +647,26 @@ class LegacyTabManager: NSObject, FeatureFlaggable, TabManager, TabEventHandler
     @MainActor
     func removeAllTabs(isPrivateMode: Bool) async {
         let currentModeTabs = tabs.filter { $0.isPrivate == isPrivateMode }
-        backupCloseTabs = currentModeTabs
+        if let tab = selectedTab, tab.isPrivate == isPrivateMode {
+            backupCloseTab = BackupCloseTab(tab: tab,
+                                            restorePosition: tabs.firstIndex(of: tab))
+        }
+        backupCloseTabs = tabs
         for tab in currentModeTabs {
             await self.removeTab(tab.tabUUID)
         }
         storeChanges()
     }
 
     func undoCloseAllTabs() {
+        guard !backupCloseTabs.isEmpty else { return }
         tabs = backupCloseTabs
         storeChanges()
         backupCloseTabs = [Tab]()
+        if backupCloseTab != nil {
+            selectTab(backupCloseTab?.tab)
+            backupCloseTab = nil
+        }
     }
 
     @MainActor