fix(snapshot): finish queued writes without blocking quit (CC-196)

Sources/CodeCaps/AppDelegate.swift

@@ -363,6 +363,17 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate, NSMe
         DispatchQueue.main.async { [weak self] in self?.navigatingPlanWarning = false }
     }
 
+    private let snapshotTerminationDrain = SnapshotTerminationDrain()
+
+    func applicationShouldTerminate(_ sender: NSApplication) -> NSApplication.TerminateReply {
+        // This app has no cancelled-quit/resume path: every request replies true.
+        snapshotTerminationDrain.request(
+            stop: { model.stop() },
+            flush: { [self] in await model.flushSnapshotWrites() },
+            reply: { sender.reply(toApplicationShouldTerminate: true) })
+        return .terminateLater
+    }
+
     func applicationWillTerminate(_ notification: Notification) { model.stop() }
     func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool { false }
     func applicationShouldHandleReopen(_ sender: NSApplication, hasVisibleWindows flag: Bool) -> Bool {

Sources/CodeCaps/MonitorModel.swift

@@ -420,6 +420,7 @@ final class MonitorModel: ObservableObject {
     /// actor and hang the UI (BOTFLEET-5B); each write now runs off-main, in
     /// order, after the one before it.
     private var snapshotWriteChain: Task<Void, Never>?
+    private(set) var snapshotWritesClosedForTermination = false
     @Published private(set) var now = Date()
 
     // Local & Remote Reading
@@ -1353,6 +1354,10 @@ final class MonitorModel: ObservableObject {
     }
 
     func start() {
+        // Termination is terminal for this process, even after a drain timeout.
+        // A future cancelled-quit feature must define recovery without letting
+        // old filesystem writes race newer snapshots. Do not reopen admission.
+        guard !snapshotWritesClosedForTermination else { return }
         hasStarted = true
         refresh()
         scheduleSourceTimers()
@@ -1391,6 +1396,14 @@ final class MonitorModel: ObservableObject {
         clockTimer?.invalidate()
     }
 
+    /// Suspend instead of blocking the main actor, which queued completions need.
+    /// The caller may terminate after the deadline even if filesystem I/O is stuck.
+    func flushSnapshotWrites(timeout: Duration = .seconds(2)) async -> Bool {
+        snapshotWritesClosedForTermination = true
+        let chain = snapshotWriteChain
+        return await SnapshotWriteDrain.wait(for: chain, timeout: timeout)
+    }
+
     /// Fires only while something is at risk, so the forecast keeps measuring a
     /// rate from fresh samples instead of a stale gap.
     ///
@@ -1781,7 +1794,7 @@ final class MonitorModel: ObservableObject {
             localWindows = []
             updateActiveAnomalies([])
             activePlanChanges = []
-            if !skipsSnapshotIOForTesting { try? LocalQuotaSnapshot.remove() }
+            if !skipsSnapshotIOForTesting { enqueueHandoffWrite(windows: nil) }
             rebuildLocalState(recordSamples: false)
         }
         refresh()
@@ -2811,14 +2824,13 @@ final class MonitorModel: ObservableObject {
         guard !skipsSnapshotIOForTesting else { return }
         let marks = exportedCustomMarks()
         let stamp = now
-        let order = platformOrder
         enqueueSnapshotWrite({
             try LocalQuotaSnapshot.writeWidgetSnapshot(windows: windows, customMarks: marks, now: stamp)
         }, completion: { [weak self] succeeded in
             guard let self else { return }
             if succeeded {
                 self.widgetSharingError = nil
-                UserDefaults(suiteName: LocalQuotaSnapshot.appGroupId)?.set(order, forKey: "platformOrder")
+                UserDefaults(suiteName: LocalQuotaSnapshot.appGroupId)?.set(self.platformOrder, forKey: "platformOrder")
                 #if canImport(WidgetKit)
                 WidgetCenter.shared.reloadAllTimelines()
                 #endif
@@ -2869,8 +2881,9 @@ final class MonitorModel: ObservableObject {
 
     /// Runs `work` on a background executor after every previously enqueued
     /// write, then reports success back on the main actor.
-    private func enqueueSnapshotWrite(_ work: @escaping @Sendable () throws -> Void,
+    func enqueueSnapshotWrite(_ work: @escaping @Sendable () throws -> Void,
                                       completion: @escaping @MainActor @Sendable (Bool) -> Void) {
+        guard !snapshotWritesClosedForTermination else { return }
         let previous = snapshotWriteChain
         snapshotWriteChain = Task.detached(priority: .utility) {
             await previous?.value

Sources/CodeCaps/SnapshotWriteDrain.swift

@@ -0,0 +1,58 @@
+import Foundation
+
+/// A deadline bounds how long termination waits, not the underlying disk write.
+/// Unstructured waiters intentionally do not join the possibly stuck task on timeout.
+@MainActor
+enum SnapshotWriteDrain {
+    private final class ResultGate {
+        var continuation: CheckedContinuation<Bool, Never>?
+        var timeoutTask: Task<Void, Never>?
+
+        func finish(_ drained: Bool) {
+            guard let continuation else { return }
+            self.continuation = nil
+            timeoutTask?.cancel()
+            timeoutTask = nil
+            continuation.resume(returning: drained)
+        }
+    }
+
+    static func wait(for chain: Task<Void, Never>?, timeout: Duration) async -> Bool {
+        guard let chain else { return true }
+        return await withCheckedContinuation { continuation in
+            let gate = ResultGate()
+            gate.continuation = continuation
+            Task { @MainActor in
+                await chain.value
+                gate.finish(true)
+            }
+            gate.timeoutTask = Task { @MainActor in
+                do { try await Task.sleep(for: timeout) }
+                catch { return }
+                gate.finish(false)
+            }
+        }
+    }
+}
+
+/// One terminal lifecycle per process, with no cancellation or resume path.
+/// Retaining the completed task is intentional: no stale callback can clear a
+/// newer drain because a newer drain is never created after termination starts.
+@MainActor
+final class SnapshotTerminationDrain {
+    private var task: Task<Bool, Never>?
+
+    func request(stop: () -> Void,
+                 flush: @escaping @MainActor () async -> Bool,
+                 reply: @escaping @MainActor () -> Void) {
+        if task == nil {
+            stop()
+            task = Task { @MainActor in await flush() }
+        }
+        guard let task else { return }
+        Task { @MainActor in
+            _ = await task.value
+            reply()
+        }
+    }
+}

Tests/CodeCapsTests/SnapshotWriteDrainTests.swift

@@ -0,0 +1,164 @@
+import Foundation
+import QuotaCore
+import XCTest
+@testable import CodeCaps
+
+@MainActor
+final class SnapshotWriteDrainTests: XCTestCase {
+    func testEmptyQueueDoesNotWait() async {
+        let drained = await SnapshotWriteDrain.wait(for: nil, timeout: .seconds(2))
+        XCTAssertTrue(drained)
+    }
+
+    func testMainActorCompletionRunsBeforeDrainReturns() async {
+        var completed = false
+        let chain = Task { @MainActor in
+            await Task.yield()
+            completed = true
+        }
+        let drained = await SnapshotWriteDrain.wait(for: chain, timeout: .seconds(2))
+        XCTAssertTrue(drained)
+        XCTAssertTrue(completed)
+    }
+
+    func testStopAndFlushWaitForRealQueuedFileAndMainActorCompletion() async throws {
+        let name = "cc196-snapshot-tests-" + UUID().uuidString
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: name))
+        defaults.set(false, forKey: "localEnabled")
+        defer { defaults.removePersistentDomain(forName: name) }
+        let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: directory) }
+        let file = directory.appendingPathComponent("snapshot.txt")
+        let model = MonitorModel(defaults: defaults)
+        model.skipsSnapshotIOForTesting = true
+        var completion = false
+        model.enqueueSnapshotWrite({ try Data("ordered snapshot".utf8).write(to: file) },
+                                   completion: { succeeded in completion = succeeded })
+        model.stop()
+        let drained = await model.flushSnapshotWrites()
+        XCTAssertTrue(drained)
+        XCTAssertTrue(completion)
+        XCTAssertEqual(try String(contentsOf: file), "ordered snapshot")
+    }
+
+    func testFlushClosesQueueBeforeSuspendingAndPreservesPriorOrdering() async throws {
+        let name = "cc196-closed-queue-" + UUID().uuidString
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: name))
+        defaults.set(false, forKey: "localEnabled")
+        defer { defaults.removePersistentDomain(forName: name) }
+        let model = MonitorModel(defaults: defaults)
+        model.skipsSnapshotIOForTesting = true
+        var firstCompleted = false
+        var lateCompleted = false
+        model.enqueueSnapshotWrite({}, completion: { _ in firstCompleted = true })
+        let drain = Task { @MainActor in await model.flushSnapshotWrites() }
+        while !model.snapshotWritesClosedForTermination { await Task.yield() }
+        model.enqueueSnapshotWrite({}, completion: { _ in lateCompleted = true })
+        let drained = await drain.value
+        XCTAssertTrue(drained)
+        XCTAssertTrue(firstCompleted)
+        XCTAssertFalse(lateCompleted)
+    }
+
+    func testTerminationCannotRestartOrAdmitMoreWrites() async throws {
+        let name = "cc196-terminal-queue-" + UUID().uuidString
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: name))
+        defaults.set(false, forKey: "localEnabled")
+        defaults.set(false, forKey: SourceRefreshPreference.providerEnabled)
+        defaults.set(false, forKey: SourceRefreshPreference.sessionEnabled)
+        defer { defaults.removePersistentDomain(forName: name) }
+        let model = MonitorModel(defaults: defaults)
+        model.skipsSnapshotIOForTesting = true
+        model.stop()
+        let drained = await model.flushSnapshotWrites()
+        XCTAssertTrue(drained)
+        model.start()
+        XCTAssertTrue(model.snapshotWritesClosedForTermination)
+        var ran = false
+        model.enqueueSnapshotWrite({}, completion: { _ in ran = true })
+        let secondDrained = await model.flushSnapshotWrites()
+        XCTAssertTrue(secondDrained)
+        XCTAssertFalse(ran)
+    }
+
+    func testRepeatedTerminationCallbacksShareOneDrainAndEachReply() async {
+        let termination = SnapshotTerminationDrain()
+        var stops = 0
+        var flushes = 0
+        var replies = 0
+        var release: CheckedContinuation<Bool, Never>?
+        let flush: @MainActor () async -> Bool = {
+            flushes += 1
+            return await withCheckedContinuation { release = $0 }
+        }
+        termination.request(stop: { stops += 1 }, flush: flush, reply: { replies += 1 })
+        termination.request(stop: { stops += 1 }, flush: flush, reply: { replies += 1 })
+        while release == nil { await Task.yield() }
+        XCTAssertEqual(stops, 1)
+        XCTAssertEqual(flushes, 1)
+        XCTAssertEqual(replies, 0)
+        release?.resume(returning: true)
+        while replies < 2 { await Task.yield() }
+        XCTAssertEqual(replies, 2)
+        // A callback arriving after completion cannot create a newer generation
+        // for a stale waiter to clear, or invoke a new flush/stop.
+        termination.request(stop: { stops += 1 }, flush: flush, reply: { replies += 1 })
+        while replies < 3 { await Task.yield() }
+        XCTAssertEqual(stops, 1)
+        XCTAssertEqual(flushes, 1)
+        XCTAssertEqual(replies, 3)
+    }
+
+    func testDeadlineDoesNotJoinStuckWriteAndLateCompletionIsSafe() async {
+        var release: CheckedContinuation<Void, Never>?
+        var completed = false
+        let chain = Task { @MainActor in
+            await withCheckedContinuation { release = $0 }
+            completed = true
+        }
+        while release == nil { await Task.yield() }
+        let clock = ContinuousClock()
+        let start = clock.now
+        let drained = await SnapshotWriteDrain.wait(for: chain, timeout: .milliseconds(20))
+        XCTAssertFalse(drained)
+        XCTAssertLessThan(start.duration(to: clock.now), .seconds(1))
+        XCTAssertFalse(completed)
+        release?.resume()
+        await chain.value
+        XCTAssertTrue(completed)
+    }
+}
+
+final class SnapshotTerminationIntegrationTests: XCTestCase {
+    private func source(_ file: String) throws -> String {
+        let root = URL(fileURLWithPath: #filePath)
+            .deletingLastPathComponent().deletingLastPathComponent().deletingLastPathComponent()
+        return try String(contentsOf: root.appendingPathComponent("Sources/CodeCaps/" + file))
+    }
+
+    func testTerminationDefersReplyUntilAsyncFlushWithoutBlockingMainActor() throws {
+        let delegate = try source("AppDelegate.swift")
+        XCTAssertTrue(delegate.contains("func applicationShouldTerminate("))
+        XCTAssertTrue(delegate.contains("private let snapshotTerminationDrain = SnapshotTerminationDrain()"))
+        XCTAssertTrue(delegate.contains("snapshotTerminationDrain.request("))
+        XCTAssertFalse(delegate.contains("snapshotFlushTask"))
+        XCTAssertTrue(delegate.contains("sender.reply(toApplicationShouldTerminate: true)"))
+        XCTAssertTrue(delegate.contains("return .terminateLater"))
+        let model = try source("MonitorModel.swift")
+        XCTAssertTrue(model.contains("timeout: Duration = .seconds(2)"))
+        XCTAssertTrue(model.contains("await SnapshotWriteDrain.wait(for: chain, timeout: timeout)"))
+        XCTAssertFalse(model.contains("group.wait("))
+        XCTAssertFalse(model.contains("sem.wait("))
+    }
+
+    func testDisableQueuesRemovalAndWidgetUsesCompletionTimeOrder() throws {
+        let model = try source("MonitorModel.swift")
+        let disable = try XCTUnwrap(model.range(of: "func setLocalEnabled("))
+        let next = try XCTUnwrap(model.range(of: "func setProviderChecksEnabled("))
+        let block = String(model[disable.lowerBound..<next.lowerBound])
+        XCTAssertTrue(block.contains("enqueueHandoffWrite(windows: nil)"))
+        XCTAssertTrue(model.contains(".set(self.platformOrder, forKey: \"platformOrder\")"))
+        XCTAssertFalse(model.contains("let order = platformOrder"))
+    }
+}