fix: keep wallet backups consistent

Bitkit/Services/PrivatePaykitService+Backup.swift

@@ -3,13 +3,15 @@ import Foundation
 // MARK: - Backup
 
 extension PrivatePaykitService {
-    func backupSnapshot() async throws -> String? {
+    func backupSnapshot() async throws -> (state: String?, sdkVersion: UInt64) {
+        let initialSdkVersion = await PaykitSdkService.shared.backupStateVersion
         guard try await PaykitSdkService.shared.currentPublicKey() != nil else {
-            return nil
+            return (nil, initialSdkVersion)
         }
         // Wallet backups may be required before broadcast; they cannot wait for payment to finish.
-        let backup = try await Backup(
-            sdkState: PaykitSdkService.shared.exportBackupState(priority: .interactive),
+        let exported = try await PaykitSdkService.shared.exportBackupSnapshot(priority: .interactive)
+        let backup = Backup(
+            sdkState: exported.state,
             consumedPrivatePaymentListVersions: state.contacts.compactMapValues { contactState in
                 contactState.consumedPrivatePaymentListVersion
             }
@@ -18,7 +20,7 @@ extension PrivatePaykitService {
         guard let encoded = String(data: data, encoding: .utf8) else {
             throw CocoaError(.fileWriteInapplicableStringEncoding)
         }
-        return encoded
+        return (encoded, exported.version)
     }
 
     func restoreBackup(_ backup: String?) async throws {

Bitkit/Services/PubkyService.swift

@@ -422,6 +422,16 @@ actor PaykitSdkService {
         let backupRevision: String
     }
 
+    struct BackupExport {
+        let state: String
+        let version: UInt64
+    }
+
+    struct BackupIdentity: Equatable {
+        let runtimeGeneration: UInt64
+        let walletGeneration: Int
+    }
+
     static let shared = PaykitSdkService()
     /// Maximum plaintext size accepted by Paykit's pubky-noise transport.
     static let maximumMessageBytes = 1000
@@ -449,6 +459,8 @@ actor PaykitSdkService {
     private var sdk: PaykitSdk?
     private var cachedPaykitKey: (publicKey: String, generation: UInt64)?
     private var cachedBackupState: BackupStateSnapshot?
+    private(set) var backupStateVersion: UInt64 = 0
+    private var runtimeGeneration: UInt64 = 0
 
     init(
         sdkFactory: (() throws -> PaykitSdk)? = nil,
@@ -1248,17 +1260,33 @@ actor PaykitSdkService {
     }
 
     func exportBackupState(priority: PaykitSdkOperationLock.Priority = .background) async throws -> String {
+        try await exportBackupSnapshot(priority: priority).state
+    }
+
+    func backupIdentity() throws -> BackupIdentity {
+        try BackupIdentity(runtimeGeneration: runtimeGeneration, walletGeneration: operationLock.walletGeneration())
+    }
+
+    func validateBackupSnapshot(version: UInt64, identity: BackupIdentity) async throws {
+        try await operationLock.withLock(priority: .interactive) {
+            guard try backupIdentity() == identity else { throw PubkyServiceError.identityChanged }
+            guard backupStateVersion == version else { throw WalletBackupSnapshotError.changed }
+        }
+    }
+
+    func exportBackupSnapshot(priority: PaykitSdkOperationLock.Priority = .background) async throws -> BackupExport {
         let generation = try operationLock.walletGeneration()
         let instance = try handle()
         while true {
             if priority == .background { try await PaykitPaymentActivity.shared.waitUntilIdle() }
-            let exported: String? = try await withSdk(priority: priority) { sdk in
+            let exported: BackupExport? = try await withSdk(priority: priority) { sdk in
                 guard sdk === instance, try operationLock.walletGeneration() == generation else {
                     throw PubkyServiceError.identityChanged
                 }
                 // Payment may have started while this export waited for the SDK lock.
                 if priority == .background, await PaykitPaymentActivity.shared.isActive { return nil }
-                return try await sdk.exportBackupString()
+                let state = try await sdk.exportBackupString()
+                return BackupExport(state: state, version: backupStateVersion)
             }
             if let exported { return exported }
         }
@@ -1477,10 +1505,12 @@ actor PaykitSdkService {
     }
 
     private func markWalletBackupDataChanged() {
+        backupStateVersion &+= 1
         Self.walletBackupDataChangedSubject.send()
     }
 
     private func resetRuntime() {
+        runtimeGeneration &+= 1
         sdk = nil
         cachedPaykitKey = nil
         cachedBackupState = nil

BitkitTests/BackupServiceTests.swift

@@ -5,6 +5,291 @@ import XCTest
 
 @MainActor
 final class BackupServiceTests: XCTestCase {
+    func testChangesDuringCaptureAndUploadRemainRequired() async throws {
+        for changeDuringCapture in [true, false] {
+            let suite = "BackupServiceTests.\(UUID().uuidString)"
+            let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+            defer { defaults.removePersistentDomain(forName: suite) }
+            var service: BackupService!
+            var shouldChange = true
+            service = BackupService(
+                defaults: defaults,
+                backupData: { _ in
+                    if changeDuringCapture, shouldChange { service.markBackupRequired(category: .wallet) }
+                    return Data("snapshot".utf8)
+                },
+                uploadBackup: { _, _ in
+                    if !changeDuringCapture, shouldChange { service.markBackupRequired(category: .wallet) }
+                }
+            )
+
+            await service.triggerBackup(category: .wallet)
+            XCTAssertTrue(service.getBackupStatus(category: .wallet).isRequired)
+            shouldChange = false
+            await service.triggerBackup(category: .wallet)
+            XCTAssertFalse(service.getBackupStatus(category: .wallet).isRequired)
+            service.setWiping(true)
+            await service.stopObservingBackups()
+        }
+    }
+
+    func testWalletCaptureDiscardsMixedStateAndRecapturesLatestValues() async throws {
+        let suite = "BackupServiceTests.\(UUID().uuidString)"
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+        defer { defaults.removePersistentDomain(forName: suite) }
+        let tracker = WalletBackupChangeTracker()
+        var state = 0
+        var attempts = 0
+        let snapshot = try await tracker.capture {
+            attempts += 1
+            let earlyRead = state
+            if attempts == 1 {
+                state = 1
+                try Bitkit.WatchOnlyAccountStore.save([], defaults: defaults)
+            }
+            await Task.yield()
+            return [earlyRead, state]
+        }
+        XCTAssertEqual(snapshot, [1, 1])
+        XCTAssertEqual(attempts, 2)
+    }
+
+    func testScheduledBackupUploadsChangesFromEveryCategory() async throws {
+        for category in BackupCategory.allCases where category != .lightningConnections {
+            let suite = "BackupServiceTests.\(UUID().uuidString)"
+            let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+            defer { defaults.removePersistentDomain(forName: suite) }
+            let completed = expectation(description: "Latest \(category.rawValue) state backed up")
+            var state = Data("initial".utf8)
+            var uploads: [Data] = []
+            var service: BackupService!
+            service = BackupService(defaults: defaults, backupData: { _ in state }, uploadBackup: { key, data in
+                XCTAssertEqual(key, category.rawValue)
+                uploads.append(data)
+                if uploads.count == 1 {
+                    state = Data("updated".utf8)
+                    service.markBackupRequired(category: category, schedule: false)
+                }
+            }, waitForBackupDelay: { await Task.yield() })
+            let subscription = service.backupStatusesPublisher
+                .filter { $0[category].map { !$0.isRequired && !$0.running } == true }
+                .first()
+                .sink { _ in completed.fulfill() }
+            defer { subscription.cancel() }
+
+            service.markBackupRequired(category: category)
+            await fulfillment(of: [completed], timeout: 2)
+            XCTAssertEqual(uploads, [Data("initial".utf8), Data("updated".utf8)], category.rawValue)
+            XCTAssertFalse(service.getBackupStatus(category: category).isRequired, category.rawValue)
+            service.setWiping(true)
+            await service.stopObservingBackups()
+        }
+    }
+
+    func testAutomaticBackupRetriesContentionAndChangesDuringUpload() async throws {
+        let suite = "BackupServiceTests.\(UUID().uuidString)"
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+        defer { defaults.removePersistentDomain(forName: suite) }
+        let changes = PassthroughSubject<Void, Never>()
+        let tracker = WalletBackupChangeTracker(changes: changes.eraseToAnyPublisher())
+        let completed = expectation(description: "All wallet changes backed up")
+        var captures = 0
+        var uploads: [Data] = []
+        var service: BackupService!
+        service = BackupService(defaults: defaults, backupData: { _ in
+            try await tracker.capture {
+                captures += 1
+                if captures <= 3 { changes.send() }
+                return Data("snapshot-\(captures)".utf8)
+            }
+        }, uploadBackup: { _, data in
+            uploads.append(data)
+            if uploads.count == 1 { service.markBackupRequired(category: .wallet) }
+        })
+        let subscription = service.backupStatusesPublisher
+            .filter { $0[.wallet].map { !$0.isRequired && !$0.running } == true }
+            .first()
+            .sink { _ in completed.fulfill() }
+        defer { subscription.cancel() }
+
+        await service.triggerBackup(category: .wallet)
+        XCTAssertEqual(captures, 3)
+        XCTAssertTrue(uploads.isEmpty)
+        XCTAssertTrue(service.getBackupStatus(category: .wallet).isRequired)
+        await fulfillment(of: [completed], timeout: 15)
+        XCTAssertEqual(uploads, [Data("snapshot-4".utf8), Data("snapshot-5".utf8)])
+        service.setWiping(true)
+        await service.stopObservingBackups()
+    }
+
+    func testWalletCaptureBoundsContentionAndDoesNotRetryOtherFailures() async {
+        let changes = PassthroughSubject<Void, Never>()
+        let tracker = WalletBackupChangeTracker(changes: changes.eraseToAnyPublisher())
+        var attempts = 0
+        do {
+            try await tracker.capture {
+                attempts += 1
+                changes.send()
+            }
+            XCTFail("Expected unstable snapshot to be rejected")
+        } catch WalletBackupSnapshotError.changed {
+        } catch {
+            XCTFail("Unexpected error: \(error)")
+        }
+        XCTAssertEqual(attempts, 3)
+
+        for identityChanged in [true, false] {
+            attempts = 0
+            do {
+                try await tracker.capture {
+                    attempts += 1
+                    changes.send()
+                    if identityChanged { throw PubkyServiceError.identityChanged }
+                    throw OnchainSendAttemptError.unresolved
+                }
+                XCTFail("Expected capture failure")
+            } catch PubkyServiceError.identityChanged {
+                XCTAssertTrue(identityChanged)
+            } catch OnchainSendAttemptError.unresolved {
+                XCTAssertFalse(identityChanged)
+            } catch {
+                XCTFail("Unexpected error: \(error)")
+            }
+            XCTAssertEqual(attempts, 1)
+        }
+    }
+
+    func testAutomaticBackupDoesNotRepeatUnrelatedFailures() async throws {
+        let firstAttempts = expectation(description: "Scheduled backups attempted")
+        firstAttempts.expectedFulfillmentCount = 6
+        let repeatedAttempt = expectation(description: "Unrelated failure retried")
+        repeatedAttempt.isInverted = true
+        var services: [(BackupService, BackupCategory)] = []
+        var suites: [(UserDefaults, String)] = []
+        defer { suites.forEach { $0.0.removePersistentDomain(forName: $0.1) } }
+        for category in [BackupCategory.wallet, .settings] {
+            for failure in 0 ..< 3 {
+                let suite = "BackupServiceTests.\(UUID().uuidString)"
+                let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+                suites.append((defaults, suite))
+                var attempts = 0
+                let service = BackupService(defaults: defaults, backupData: { _ in
+                    attempts += 1
+                    if attempts == 1 { firstAttempts.fulfill() } else { repeatedAttempt.fulfill() }
+                    if failure == 0 { throw PubkyServiceError.identityChanged }
+                    if failure == 1 { throw OnchainSendAttemptError.unresolved }
+                    return Data("snapshot".utf8)
+                }, uploadBackup: { _, _ in throw URLError(.notConnectedToInternet) })
+                services.append((service, category))
+                service.markBackupRequired(category: category)
+            }
+        }
+
+        await fulfillment(of: [firstAttempts], timeout: 8)
+        await fulfillment(of: [repeatedAttempt], timeout: 6)
+        for (service, category) in services {
+            XCTAssertTrue(service.getBackupStatus(category: category).isRequired)
+            XCTAssertFalse(service.getBackupStatus(category: category).running)
+            service.setWiping(true)
+            await service.stopObservingBackups()
+        }
+    }
+
+    func testScheduledBackupRetainsFollowupWhileManualUploadIsBusy() async throws {
+        for category in [BackupCategory.wallet, .settings] {
+            let suite = "BackupServiceTests.\(UUID().uuidString)"
+            let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+            defer { defaults.removePersistentDomain(forName: suite) }
+            let manualUploadStarted = expectation(description: "Manual upload held")
+            let firstDelayStarted = expectation(description: "Scheduled backup held")
+            let secondDelayStarted = expectation(description: "Busy scheduled backup retained")
+            let followupUploaded = expectation(description: "Newer state uploaded")
+            let (manualUpload, releaseManualUpload) = AsyncStream<Void>.makeStream()
+            let (firstDelay, releaseFirstDelay) = AsyncStream<Void>.makeStream()
+            let (secondDelay, releaseSecondDelay) = AsyncStream<Void>.makeStream()
+            defer {
+                releaseManualUpload.finish()
+                releaseFirstDelay.finish()
+                releaseSecondDelay.finish()
+            }
+            var uploads = 0
+            var delays = 0
+            let service = BackupService(defaults: defaults, backupData: { _ in Data("snapshot".utf8) }, uploadBackup: { _, _ in
+                uploads += 1
+                if uploads == 1 {
+                    manualUploadStarted.fulfill()
+                    for await _ in manualUpload {}
+                } else {
+                    followupUploaded.fulfill()
+                }
+            }, waitForBackupDelay: {
+                delays += 1
+                if delays == 1 {
+                    firstDelayStarted.fulfill()
+                    for await _ in firstDelay {}
+                } else {
+                    secondDelayStarted.fulfill()
+                    for await _ in secondDelay {}
+                }
+            })
+            let manual = Task { await service.triggerBackup(category: category) }
+            await fulfillment(of: [manualUploadStarted], timeout: 2)
+            service.markBackupRequired(category: category)
+            await fulfillment(of: [firstDelayStarted], timeout: 2)
+            releaseFirstDelay.finish()
+            await fulfillment(of: [secondDelayStarted], timeout: 2)
+            releaseManualUpload.finish()
+            _ = await manual.value
+            XCTAssertTrue(service.getBackupStatus(category: category).isRequired)
+            releaseSecondDelay.finish()
+            await fulfillment(of: [followupUploaded], timeout: 2)
+            XCTAssertEqual(uploads, 2)
+            service.setWiping(true)
+            await service.stopObservingBackups()
+        }
+    }
+
+    func testCancellationWinsOverFinalSnapshotContention() async {
+        let tracker = WalletBackupChangeTracker(changes: Empty<Void, Never>().eraseToAnyPublisher())
+        let capture = Task { () in
+            var attempts = 0
+            try await tracker.capture {
+                attempts += 1
+                if attempts == 3 { withUnsafeCurrentTask { $0?.cancel() } }
+                throw WalletBackupSnapshotError.changed
+            }
+        }
+        do {
+            try await capture.value
+            XCTFail("Expected cancellation")
+        } catch is CancellationError {
+        } catch {
+            XCTFail("Unexpected error: \(error)")
+        }
+    }
+
+    func testCancelledWalletCaptureDoesNotUploadOrRetry() async throws {
+        let suite = "BackupServiceTests.\(UUID().uuidString)"
+        let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
+        defer { defaults.removePersistentDomain(forName: suite) }
+        let started = expectation(description: "Wallet capture started")
+        let (gate, release) = AsyncStream<Void>.makeStream()
+        defer { release.finish() }
+        let service = BackupService(defaults: defaults, backupData: { _ in
+            started.fulfill()
+            for await _ in gate {}
+            return Data("cancelled".utf8)
+        }, uploadBackup: { _, _ in XCTFail("Cancelled capture must not upload") })
+        let backup = Task { await service.triggerBackup(category: .wallet) }
+        await fulfillment(of: [started], timeout: 2)
+        backup.cancel()
+        _ = await backup.value
+        XCTAssertTrue(service.getBackupStatus(category: .wallet).isRequired)
+        XCTAssertFalse(service.getBackupStatus(category: .wallet).running)
+        service.setWiping(true)
+        await service.stopObservingBackups()
+    }
+
     func testCancelledDeferredWalletBackupRemainsRequiredAndRetryCompletes() async throws {
         let suite = "BackupServiceTests.\(UUID().uuidString)"
         let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))