Merge Wave 2 refactor: cross-window state, persistence integrity, lifecycle leak fixes

datlechin · datlechin · commit e2cc21a7475c · 2026-05-01T15:40:09.000+07:00
diff --git a/TablePro/Core/Database/DatabaseManager+Health.swift b/TablePro/Core/Database/DatabaseManager+Health.swift
@@ -51,6 +51,7 @@ extension DatabaseManager {
             reconnectHandler: { [weak self] in
                 guard let self else { return false }
                 guard let session = await self.activeSessions[connectionId] else { return false }
+                await SchemaService.shared.invalidate(connectionId: connectionId)
                 do {
                     let result = try await self.trackOperation(sessionId: connectionId) {
                         try await self.reconnectDriver(for: session)
@@ -207,6 +208,8 @@ extension DatabaseManager {
             session.status = .connecting
         }
 
+        await SchemaService.shared.invalidate(connectionId: sessionId)
+
         // Stop existing health monitor
         await stopHealthMonitor(for: sessionId)
 
diff --git a/TablePro/Core/Database/DatabaseManager+Sessions.swift b/TablePro/Core/Database/DatabaseManager+Sessions.swift
@@ -247,6 +247,7 @@ extension DatabaseManager {
                 session.currentSchema = nil
             }
             appSettingsStorage.saveLastSchema(nil, for: connectionId)
+            await SchemaService.shared.invalidate(connectionId: connectionId)
             await reconnectSession(connectionId)
         } else if pm?.capabilities.supportsSchemaSwitching == true,
                   let schemaDriver = driver as? SchemaSwitchable {
diff --git a/TablePro/Core/MCP/MCPServer.swift b/TablePro/Core/MCP/MCPServer.swift
@@ -132,12 +132,21 @@ actor MCPServer {
         cleanupTask?.cancel()
         cleanupTask = nil
 
+        let sessionIds = Array(sessions.keys)
         for (_, session) in sessions {
             await session.cancelAllTasks()
             await session.cancelSSEConnection()
         }
         sessions.removeAll()
 
+        // Run the per-session cleanup handler so policy-side state
+        // (e.g. MCPAuthPolicy.sessionApprovals) doesn't leak across server restarts.
+        if let cleanupHandler = sessionCleanupHandler {
+            for id in sessionIds {
+                await cleanupHandler(id)
+            }
+        }
+
         if let currentListener = listener {
             listener = nil
             await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
diff --git a/TablePro/Core/Services/Infrastructure/SessionStateFactory.swift b/TablePro/Core/Services/Infrastructure/SessionStateFactory.swift
@@ -30,8 +30,28 @@ enum SessionStateFactory {
     /// only one coordinator exists per window — no duplicate-tab side effects.
     private static var pendingSessionStates: [UUID: SessionState] = [:]
 
+    /// Pending entries that aren't claimed by a `ContentView.init` within this TTL
+    /// are dropped to prevent leaks if the window-open path errors out before the
+    /// hand-off completes.
+    private static let pendingEntryTTL: Duration = .seconds(5)
+
     static func registerPending(_ state: SessionState, for payloadId: UUID) {
         pendingSessionStates[payloadId] = state
+        Task { [payloadId] in
+            try? await Task.sleep(for: pendingEntryTTL)
+            await MainActor.run {
+                guard let abandoned = pendingSessionStates.removeValue(forKey: payloadId) else {
+                    return
+                }
+                // Coordinator was eagerly registered with `activeCoordinators`
+                // by `create(...)`. If no `ContentView.init` consumed the pending
+                // entry within the TTL, the window-open path never completed —
+                // unregister so the coordinator can deallocate.
+                MainContentCoordinator.activeCoordinators.removeValue(
+                    forKey: abandoned.coordinator.instanceId
+                )
+            }
+        }
     }
 
     static func consumePending(for payloadId: UUID) -> SessionState? {
@@ -158,6 +178,13 @@ enum SessionStateFactory {
             toolbarState: toolbarSt
         )
 
+        // Eagerly publish to the active-coordinator registry so concurrent
+        // window opens for the same connection both observe each other when
+        // computing globals like nextQueryTitle. Without this, two windows
+        // opened back-to-back can both compute "Query 1" before either has
+        // run onAppear.
+        coord.registerEagerly()
+
         return SessionState(
             tabManager: tabMgr,
             changeManager: changeMgr,
diff --git a/TablePro/Core/Services/Infrastructure/TabPersistenceCoordinator+AggregatedSave.swift b/TablePro/Core/Services/Infrastructure/TabPersistenceCoordinator+AggregatedSave.swift
@@ -0,0 +1,33 @@
+//
+//  TabPersistenceCoordinator+AggregatedSave.swift
+//  TablePro
+//
+
+import Foundation
+
+extension TabPersistenceCoordinator {
+    /// Save or clear persisted state based on tabs aggregated across all windows
+    /// for the connection. Prevents the per-window close path from clobbering
+    /// state when sibling windows still have open tabs.
+    func saveOrClearAggregated() {
+        let aggregatedTabs = MainContentCoordinator.aggregatedTabs(for: connectionId)
+        if aggregatedTabs.isEmpty {
+            clearSavedState()
+        } else {
+            let selectedId = MainContentCoordinator.aggregatedSelectedTabId(for: connectionId)
+            saveNow(tabs: aggregatedTabs, selectedTabId: selectedId)
+        }
+    }
+
+    /// Synchronous variant for the window-close path, where the run loop may
+    /// not be available to service Tasks before the window tears down.
+    func saveOrClearAggregatedSync() {
+        let aggregatedTabs = MainContentCoordinator.aggregatedTabs(for: connectionId)
+        if aggregatedTabs.isEmpty {
+            saveNowSync(tabs: [], selectedTabId: nil)
+        } else {
+            let selectedId = MainContentCoordinator.aggregatedSelectedTabId(for: connectionId)
+            saveNowSync(tabs: aggregatedTabs, selectedTabId: selectedId)
+        }
+    }
+}
diff --git a/TablePro/Core/Services/Infrastructure/TabPersistenceCoordinator.swift b/TablePro/Core/Services/Infrastructure/TabPersistenceCoordinator.swift
@@ -30,6 +30,8 @@ internal final class TabPersistenceCoordinator {
     private static let logger = Logger(subsystem: "com.TablePro", category: "NativeTabLifecycle")
     let connectionId: UUID
 
+    @ObservationIgnored private var saveTask: Task<Void, Never>?
+
     init(connectionId: UUID) {
         self.connectionId = connectionId
     }
@@ -45,43 +47,25 @@ internal final class TabPersistenceCoordinator {
             return
         }
         let persisted = nonPreviewTabs.map { convertToPersistedTab($0) }
-        let connId = connectionId
         let normalizedSelectedId = nonPreviewTabs.contains(where: { $0.id == selectedTabId })
             ? selectedTabId : nonPreviewTabs.first?.id
-        Self.logger.debug("[persist] saveNow queued tabCount=\(nonPreviewTabs.count) connId=\(connId, privacy: .public)")
-
-        Task {
-            let t0 = Date()
-            do {
-                try await TabDiskActor.shared.save(connectionId: connId, tabs: persisted, selectedTabId: normalizedSelectedId)
-                Self.logger.debug("[persist] saveNow written tabCount=\(persisted.count) connId=\(connId, privacy: .public) ms=\(Int(Date().timeIntervalSince(t0) * 1_000))")
-            } catch {
-                TabDiskActor.logSaveError(connectionId: connId, error: error)
-            }
-        }
+        scheduleSave(tabs: persisted, selectedTabId: normalizedSelectedId)
     }
 
     /// Save pre-aggregated tabs for the quit path, where the caller has already
     /// collected and converted tabs from all windows for this connection.
     internal func saveNow(persistedTabs: [PersistedTab], selectedTabId: UUID?) {
-        let connId = connectionId
-        let selectedId = selectedTabId
-
-        Task {
-            do {
-                try await TabDiskActor.shared.save(connectionId: connId, tabs: persistedTabs, selectedTabId: selectedId)
-            } catch {
-                TabDiskActor.logSaveError(connectionId: connId, error: error)
-            }
-        }
+        scheduleSave(tabs: persistedTabs, selectedTabId: selectedTabId)
     }
 
     /// Synchronous save for `applicationWillTerminate` where no run loop
     /// remains to service async Tasks. Bypasses the actor and writes directly.
     internal func saveNowSync(tabs: [QueryTab], selectedTabId: UUID?) {
         let nonPreviewTabs = tabs.filter { !$0.isPreview }
         guard !nonPreviewTabs.isEmpty else {
-            TabDiskActor.saveSync(connectionId: connectionId, tabs: [], selectedTabId: nil)
+            saveTask?.cancel()
+            saveTask = nil
+            TabDiskActor.clearSync(connectionId: connectionId)
             return
         }
         let persisted = nonPreviewTabs.map { convertToPersistedTab($0) }
@@ -94,12 +78,37 @@ internal final class TabPersistenceCoordinator {
 
     /// Clear all saved state for this connection (user closed all tabs).
     internal func clearSavedState() {
+        saveTask?.cancel()
+        saveTask = nil
         let connId = connectionId
         Task {
             await TabDiskActor.shared.clear(connectionId: connId)
         }
     }
 
+    // MARK: - Private save scheduling
+
+    private func scheduleSave(tabs: [PersistedTab], selectedTabId: UUID?) {
+        saveTask?.cancel()
+        let connId = connectionId
+        let tabsCopy = tabs
+        let selectedId = selectedTabId
+        Self.logger.debug("[persist] saveNow queued tabCount=\(tabsCopy.count) connId=\(connId, privacy: .public)")
+
+        saveTask = Task {
+            guard !Task.isCancelled else { return }
+            let t0 = Date()
+            do {
+                try await TabDiskActor.shared.save(connectionId: connId, tabs: tabsCopy, selectedTabId: selectedId)
+                Self.logger.debug("[persist] saveNow written tabCount=\(tabsCopy.count) connId=\(connId, privacy: .public) ms=\(Int(Date().timeIntervalSince(t0) * 1_000))")
+            } catch is CancellationError {
+                return
+            } catch {
+                Self.logger.fault("Failed to save tab state for connection \(connId, privacy: .public): \(error.localizedDescription, privacy: .public)")
+            }
+        }
+    }
+
     // MARK: - Restore
 
     /// Restore tabs from disk. Called once at window creation.
diff --git a/TablePro/Core/Storage/TabDiskActor.swift b/TablePro/Core/Storage/TabDiskActor.swift
@@ -60,11 +60,6 @@ internal actor TabDiskActor {
         try data.write(to: fileURL, options: .atomic)
     }
 
-    /// Log a save error from callers that handle errors externally.
-    nonisolated static func logSaveError(connectionId: UUID, error: Error) {
-        logger.error("Failed to save tab state for \(connectionId): \(error.localizedDescription)")
-    }
-
     /// Load tab state for a connection. Returns nil if the file is missing or corrupt.
     internal func load(connectionId: UUID) -> TabDiskState? {
         let fileURL = tabStateFileURL(for: connectionId)
@@ -96,6 +91,8 @@ internal actor TabDiskActor {
     }
 
     /// List all connection IDs that have saved tab state on disk.
+    /// Self-cleans legacy empty-payload files: if a file decodes with no tabs,
+    /// it is deleted and its connection ID is excluded from the result.
     internal func connectionIdsWithSavedState() -> [UUID] {
         let fm = FileManager.default
         guard let files = try? fm.contentsOfDirectory(
@@ -104,10 +101,18 @@ internal actor TabDiskActor {
         ) else {
             return []
         }
-        return files.compactMap { url -> UUID? in
-            guard url.pathExtension == "json" else { return nil }
-            return UUID(uuidString: url.deletingPathExtension().lastPathComponent)
+        var validIds: [UUID] = []
+        for url in files where url.pathExtension == "json" {
+            guard let id = UUID(uuidString: url.deletingPathExtension().lastPathComponent) else { continue }
+            if let data = try? Data(contentsOf: url),
+               let state = try? decoder.decode(TabDiskState.self, from: data),
+               !state.tabs.isEmpty {
+                validIds.append(id)
+            } else {
+                try? fm.removeItem(at: url)
+            }
         }
+        return validIds
     }
 
     // MARK: - Static Path Helpers
@@ -145,7 +150,20 @@ internal actor TabDiskActor {
             let fileURL = tabStateFileURL(for: connectionId)
             try data.write(to: fileURL, options: .atomic)
         } catch {
-            logger.error("saveSync failed for \(connectionId): \(error.localizedDescription)")
+            logger.fault("saveSync failed for \(connectionId): \(error.localizedDescription)")
+        }
+    }
+
+    /// Synchronous clear for `applicationWillTerminate`, where no run loop
+    /// remains to execute an async Task. Mirrors `saveSync` — deletes the
+    /// connection's tab state file directly.
+    nonisolated internal static func clearSync(connectionId: UUID) {
+        let fileURL = tabStateFileURL(for: connectionId)
+        guard FileManager.default.fileExists(atPath: fileURL.path) else { return }
+        do {
+            try FileManager.default.removeItem(at: fileURL)
+        } catch {
+            logger.fault("clearSync failed for \(connectionId): \(error.localizedDescription)")
         }
     }
 
diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+WindowLifecycle.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+WindowLifecycle.swift
@@ -102,19 +102,11 @@ extension MainContentCoordinator {
             "[close] coordinator.handleWindowWillClose connId=\(self.connectionId, privacy: .public) tabs=\(self.tabManager.tabs.count)"
         )
 
-        // Persist remaining non-preview tabs synchronously. saveNowSync writes
-        // directly without spawning a Task — required here because the window
-        // is closing and we cannot rely on async tasks being serviced.
-        let persistableTabs = tabManager.tabs.filter { !$0.isPreview }
-        if persistableTabs.isEmpty {
-            // Empty → clear saved state so next open shows a default empty window.
-            persistence.saveNowSync(tabs: [], selectedTabId: nil)
-        } else {
-            let normalizedSelectedId =
-                persistableTabs.contains(where: { $0.id == tabManager.selectedTabId })
-                ? tabManager.selectedTabId : persistableTabs.first?.id
-            persistence.saveNowSync(tabs: persistableTabs, selectedTabId: normalizedSelectedId)
-        }
+        // Persist tabs aggregated across all windows for this connection.
+        // Writing this window's tabs in isolation can clobber sibling windows'
+        // state on disk — for example, closing an empty window would erase the
+        // saved tabs of an open sibling window.
+        persistence.saveOrClearAggregatedSync()
 
         // Cancel the pending eviction task before teardown drops it.
         evictionTask?.cancel()
diff --git a/TablePro/Views/Main/Extensions/MainContentView+EventHandlers.swift b/TablePro/Views/Main/Extensions/MainContentView+EventHandlers.swift
@@ -58,20 +58,9 @@ extension MainContentView {
             coordinator.promotePreviewTab()
         }
 
-        let persistableTabs = tabManager.tabs.filter { !$0.isPreview }
-        if persistableTabs.isEmpty {
-            coordinator.persistence.clearSavedState()
-        } else {
-            let normalizedSelectedId =
-                persistableTabs.contains(where: { $0.id == tabManager.selectedTabId })
-                ? tabManager.selectedTabId : persistableTabs.first?.id
-            coordinator.persistence.saveNow(
-                tabs: persistableTabs,
-                selectedTabId: normalizedSelectedId
-            )
-        }
+        coordinator.persistence.saveOrClearAggregated()
         MainContentView.lifecycleLogger.debug(
-            "[switch] handleStructureChange tabCount=\(tabManager.tabs.count) persistableCount=\(persistableTabs.count) ms=\(Int(Date().timeIntervalSince(t0) * 1_000))"
+            "[switch] handleStructureChange tabCount=\(tabManager.tabs.count) ms=\(Int(Date().timeIntervalSince(t0) * 1_000))"
         )
     }
 
diff --git a/TablePro/Views/Main/Extensions/MainContentView+Setup.swift b/TablePro/Views/Main/Extensions/MainContentView+Setup.swift
@@ -137,17 +137,13 @@ extension MainContentView {
 
             if !remainingTabs.isEmpty {
                 let selectedWasFirst = firstTab.id == selectedId
-                Task { @MainActor in
-                    for tab in remainingTabs {
-                        let restorePayload = EditorTabPayload(
-                            from: tab, connectionId: connection.id, skipAutoExecute: true)
-                        WindowManager.shared.openTab(payload: restorePayload)
-                    }
-                    // Bring the first window to front only if it had the selected tab.
-                    // Otherwise let the last restored window stay focused.
-                    if selectedWasFirst {
-                        viewWindow?.makeKeyAndOrderFront(nil)
-                    }
+                for tab in remainingTabs {
+                    let restorePayload = EditorTabPayload(
+                        from: tab, connectionId: connection.id, skipAutoExecute: true)
+                    WindowManager.shared.openTab(payload: restorePayload)
+                }
+                if selectedWasFirst {
+                    viewWindow?.makeKeyAndOrderFront(nil)
                 }
             }
 
diff --git a/TablePro/Views/Main/MainContentCoordinator.swift b/TablePro/Views/Main/MainContentCoordinator.swift

Original file line number	Diff line number	Diff line change
`@@ -247,6 +247,7 @@ extension DatabaseManager {`
`247`	`247`	`session.currentSchema = nil`
`248`	`248`	`}`
`249`	`249`	`appSettingsStorage.saveLastSchema(nil, for: connectionId)`
	`250`	`+ await SchemaService.shared.invalidate(connectionId: connectionId)`
`250`	`251`	`await reconnectSession(connectionId)`
`251`	`252`	`} else if pm?.capabilities.supportsSchemaSwitching == true,`
`252`	`253`	`let schemaDriver = driver as? SchemaSwitchable {`