[codex] Fix migration recovery and switch races

TypeSwitch/Sources/App/AppFeature.swift

@@ -108,7 +108,10 @@ struct AppFeature {
         case frontmostApplicationLoaded(AppInfo?)
         case inputMethods([InputMethod])
         case launchAtLoginLoaded(LaunchAtLoginStatus)
-        case legacyRulesLoaded([String: AppRuleRecord])
+        case legacyRulesLoaded(
+            [String: AppRuleRecord],
+            didCompleteLegacyMigration: Bool
+        )
         case programmaticSwitchFinished(bundleId: String, inputMethodId: String, didSwitch: Bool)
         case runningApps([AppInfo])
     }
@@ -131,6 +134,7 @@ struct AppFeature {
     private enum CancelID {
         case inputMethodAvailability
         case inputMethodSelection
+        case programmaticSwitch
         case workspaceEvents
     }
 
@@ -234,10 +238,11 @@ struct AppFeature {
                 state.launchAtLoginStatus = status
                 return .none
 
-            case .response(.legacyRulesLoaded(let legacyRules)):
+            case let .response(.legacyRulesLoaded(legacyRules, didCompleteLegacyMigration)):
                 let mergedRules = AppRulesStoreMigration.merge(
                     currentRules: state.appRules,
-                    legacyRules: legacyRules
+                    legacyRules: legacyRules,
+                    didCompleteLegacyMigration: didCompleteLegacyMigration
                 )
                 guard mergedRules != state.appRules else {
                     return markLegacyMigrationCompletedEffect()
@@ -266,6 +271,8 @@ struct AppFeature {
                 return .none
 
             case .view(.ignoreAppTapped(let appInfo)):
+                let shouldCancelProgrammaticSwitch = state.currentFrontmostBundleId == appInfo.bundleId
+                    || state.pendingProgrammaticSwitch?.bundleId == appInfo.bundleId
                 let updateDate = now
                 state.$appRulesStore.withLock { store in
                     let currentRule = store.rules[appInfo.bundleId] ?? AppRuleRecord(
@@ -285,7 +292,9 @@ struct AppFeature {
                     updatedRule.updatedAt = updateDate
                     store.rules[appInfo.bundleId] = updatedRule
                 }
-                return .none
+                guard shouldCancelProgrammaticSwitch else { return .none }
+                state.pendingProgrammaticSwitch = nil
+                return .cancel(id: CancelID.programmaticSwitch)
 
             case .view(.removeMissingInputMethodRulesTapped):
                 let updateDate = now
@@ -389,18 +398,28 @@ struct AppFeature {
                 return refreshRunningAppsEffect()
 
             case .system(.workspaceEvent(.terminated(let bundleId))):
-                if state.currentFrontmostBundleId == bundleId {
+                let wasCurrentApp = state.currentFrontmostBundleId == bundleId
+                if wasCurrentApp {
                     state.currentFrontmostBundleId = nil
                 }
-                return refreshRunningAppsEffect()
+                let shouldCancelProgrammaticSwitch = wasCurrentApp
+                    || state.pendingProgrammaticSwitch?.bundleId == bundleId
+                guard shouldCancelProgrammaticSwitch else {
+                    return refreshRunningAppsEffect()
+                }
+                state.pendingProgrammaticSwitch = nil
+                return .merge(
+                    .cancel(id: CancelID.programmaticSwitch),
+                    refreshRunningAppsEffect()
+                )
 
             case .system(.workspaceEvent(.activated(let appInfo))):
                 state.currentFrontmostBundleId = appInfo.bundleId
                 upsertRecord(for: appInfo, in: &state)
 
                 guard let inputMethodId = targetInputMethodId(for: appInfo.bundleId, state: state) else {
                     state.pendingProgrammaticSwitch = nil
-                    return .none
+                    return .cancel(id: CancelID.programmaticSwitch)
                 }
 
                 state.pendingProgrammaticSwitch = .init(
@@ -411,19 +430,23 @@ struct AppFeature {
                 return .run { send in
                     var didSwitch = false
                     if (try? await inputMethodClient.currentInputMethodId()) != inputMethodId {
+                        guard !Task.isCancelled else { return }
                         do {
                             try await inputMethodClient.switchToInputMethod(inputMethodId)
                             didSwitch = true
                         } catch {
+                            guard !Task.isCancelled else { return }
                             didSwitch = false
                         }
                     }
+                    guard !Task.isCancelled else { return }
                     await send(.response(.programmaticSwitchFinished(
                         bundleId: appInfo.bundleId,
                         inputMethodId: inputMethodId,
                         didSwitch: didSwitch
                     )))
                 }
+                .cancellable(id: CancelID.programmaticSwitch, cancelInFlight: true)
             }
         }
     }
@@ -435,7 +458,11 @@ struct AppFeature {
             }
 
             let legacyRules = await legacyDefaultsMigrationClient.loadRules(now)
-            await send(.response(.legacyRulesLoaded(legacyRules)))
+            let didCompleteLegacyMigration = await legacyDefaultsMigrationClient.didCompleteLegacyMigration()
+            await send(.response(.legacyRulesLoaded(
+                legacyRules,
+                didCompleteLegacyMigration: didCompleteLegacyMigration
+            )))
         }
     }
 

TypeSwitch/Sources/Migrations/AppRulesStoreMigration.swift

@@ -4,21 +4,37 @@ import Sharing
 enum AppRulesStoreMigration {
     static func merge(
         currentRules: [String: AppRuleRecord],
-        legacyRules: [String: AppRuleRecord]
+        legacyRules: [String: AppRuleRecord],
+        didCompleteLegacyMigration: Bool
     ) -> [String: AppRuleRecord] {
-        currentRules.merging(legacyRules) { currentRule, legacyRule in
-            guard currentRule.strategy == .none else {
-                return currentRule
+        guard !didCompleteLegacyMigration || currentRules.isEmpty else {
+            return currentRules
+        }
+
+        return legacyRules.reduce(into: currentRules) { rules, entry in
+            guard let currentRule = rules[entry.key] else {
+                rules[entry.key] = entry.value
+                return
+            }
+
+            guard currentRule.strategy == .none,
+                  currentRule.createdAt == currentRule.updatedAt
+            else {
+                return
             }
 
             var recoveredRule = currentRule
-            if legacyRule.lastKnownPath != nil {
-                recoveredRule.lastKnownPath = legacyRule.lastKnownPath
-                recoveredRule.lastKnownName = legacyRule.lastKnownName
+            recoveredRule.strategy = entry.value.strategy
+            recoveredRule.updatedAt = entry.value.updatedAt
+            if recoveredRule.lastKnownPath == nil {
+                recoveredRule.lastKnownPath = entry.value.lastKnownPath
+            }
+            if recoveredRule.lastKnownName == recoveredRule.bundleId,
+               entry.value.lastKnownName != entry.value.bundleId
+            {
+                recoveredRule.lastKnownName = entry.value.lastKnownName
             }
-            recoveredRule.strategy = legacyRule.strategy
-            recoveredRule.updatedAt = legacyRule.updatedAt
-            return recoveredRule
+            rules[entry.key] = recoveredRule
         }
     }
 }

TypeSwitch/Sources/Migrations/LegacyDefaultsMigration.swift

@@ -4,12 +4,17 @@ import Foundation
 // Migration: legacy defaults -> app rules store
 enum LegacyDefaultsMigration {
     static let currentVersion = 2
+    static let legacyCompletionKey = "didMigrateLegacyAppRules"
     static let versionKey = "legacyAppRulesMigrationVersion"
 
     static func completedVersion(in defaults: UserDefaults) -> Int {
         defaults.integer(forKey: versionKey)
     }
 
+    static func didCompleteLegacyMigration(in defaults: UserDefaults) -> Bool {
+        defaults.bool(forKey: legacyCompletionKey)
+    }
+
     static func makeRules(
         legacyMappings: [String: String],
         matchedApplications: [String: AppInfo],
@@ -32,6 +37,7 @@ enum LegacyDefaultsMigration {
 
 struct LegacyDefaultsMigrationClient {
     var completedVersion: @Sendable () async -> Int
+    var didCompleteLegacyMigration: @Sendable () async -> Bool
     var loadRules: @Sendable (_ migrationDate: Date) async -> [String: AppRuleRecord]
     var markCompleted: @Sendable (_ version: Int) async -> Void
 }
@@ -41,6 +47,9 @@ extension LegacyDefaultsMigrationClient: DependencyKey {
         completedVersion: {
             LegacyDefaultsMigration.completedVersion(in: .standard)
         },
+        didCompleteLegacyMigration: {
+            LegacyDefaultsMigration.didCompleteLegacyMigration(in: .standard)
+        },
         loadRules: { migrationDate in
             let defaults = UserDefaults(suiteName: "group.top.ygsgdbd.TypeSwitch") ?? .standard
             let legacyMappings = defaults.dictionary(forKey: "appInputMethodSettings")?
@@ -64,6 +73,7 @@ extension LegacyDefaultsMigrationClient: DependencyKey {
 
     static let testValue = Self(
         completedVersion: { LegacyDefaultsMigration.currentVersion },
+        didCompleteLegacyMigration: { false },
         loadRules: { _ in [:] },
         markCompleted: { _ in }
     )