fix: trigger automatic repair for affected conversations - WPB-22447 🍒 (#4063)

Co-authored-by: KaterinaWire <[email protected]>
Co-authored-by: John Nguyen <[email protected]>

Author: 3 people

WireDomain/Sources/WireDomain/Components/ClientSessionComponent.swift

@@ -856,4 +856,11 @@ public final class ClientSessionComponent {
         let brokenGroupIds = journal[.brokenMLSGroupIDs]
         return brokenGroupIds.contains(groupID.description)
     }
+
+    public lazy var repairFaultyMLSRemovalKeysGenerator = RepairFaultyMLSRemovalKeysGenerator(
+        journal: journal,
+        repairUseCase: repairFaultyRemovalKeysUsecase,
+        workAgent: workAgent
+    )
+
 }

WireDomain/Sources/WireDomain/Synchronization/RepairFaultyMLSRemovalKeysGenerator.swift

@@ -0,0 +1,63 @@
+//
+// Wire
+// Copyright (C) 2025 Wire Swiss GmbH
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see http://www.gnu.org/licenses/.
+//
+
+import Foundation
+
+/// Generates and submits a work item to repair faulty MLS removal keys if needed.
+///
+/// This generator checks the journal flag and submits the repair work item
+/// to the work agent if repair is required. This pattern keeps the work item internal
+/// while allowing it to be triggered from the session initialization or after migrations.
+
+public final class RepairFaultyMLSRemovalKeysGenerator {
+
+    private let journal: any JournalProtocol
+    private let repairUseCase: any RepairRemovalKeysUseCaseProtocol
+    private let workAgent: WorkAgent
+
+    public init(
+        journal: any JournalProtocol,
+        repairUseCase: any RepairRemovalKeysUseCaseProtocol,
+        workAgent: WorkAgent
+    ) {
+        self.journal = journal
+        self.repairUseCase = repairUseCase
+        self.workAgent = workAgent
+    }
+
+    /// Checks if repair is needed and submits the work item if so.
+    ///
+    /// This method can be called multiple times safely - it only submits
+    /// the work item if the journal flag is set.
+
+    public func submitWorkItemIfNeeded() {
+        guard journal[.isRepairFaultyMLSRemovalKeysRequired] else {
+            return
+        }
+
+        let workItem = RepairFaultyMLSRemovalKeysWorkItem(
+            journal: journal,
+            repairUseCase: repairUseCase
+        )
+
+        Task { [workAgent] in
+            await workAgent.submitItem(workItem)
+        }
+    }
+
+}

WireDomain/Sources/WireDomain/Utilities/Journal/Journal.swift

@@ -137,7 +137,8 @@ public extension Journal {
             JournalKey.isInitialSyncRequired,
             JournalKey.isSyncV2Enabled,
             JournalKey.isBackendMLSEnabled,
-            JournalKey.isFederationMigrationRequired
+            JournalKey.isFederationMigrationRequired,
+            JournalKey.isRepairFaultyMLSRemovalKeysRequired
 
         ].forEach {
             result[$0.name] = "\(self[$0] == true ? "Yes" : "No")"

WireDomain/Sources/WireDomain/Utilities/Journal/JournalKey.swift

@@ -103,6 +103,13 @@ public extension JournalKey where Value == Bool {
         defaultValue: false
     )
 
+    /// Whether faulty MLS removal keys need to be repaired.
+
+    static let isRepairFaultyMLSRemovalKeysRequired = Self(
+        "isRepairFaultyMLSRemovalKeysRequired",
+        defaultValue: false
+    )
+
 }
 
 public extension JournalKey where Value == Set<String> {

WireDomain/Sources/WireDomain/WorkAgent/WorkItem/RepairFaultyMLSRemovalKeysWorkItem.swift

@@ -0,0 +1,82 @@
+//
+// Wire
+// Copyright (C) 2025 Wire Swiss GmbH
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see http://www.gnu.org/licenses/.
+//
+
+import Foundation
+import WireLogging
+
+/// Repairs conversations with faulty MLS removal keys.
+///
+/// This work item checks if the journal flag `isRepairFaultyMLSRemovalKeysRequired`
+/// is set, and if so, invokes the repair use case. On successful completion,
+/// the flag is cleared so the repair only runs once per migration.
+
+struct RepairFaultyMLSRemovalKeysWorkItem: WorkItem {
+
+    let id = UUID()
+    var priority: WorkItemPriority {
+        .low
+    }
+
+    private let journal: any JournalProtocol
+    private let repairUseCase: any RepairRemovalKeysUseCaseProtocol
+
+    init(
+        journal: any JournalProtocol,
+        repairUseCase: any RepairRemovalKeysUseCaseProtocol
+    ) {
+        self.journal = journal
+        self.repairUseCase = repairUseCase
+    }
+
+    func start() async throws {
+        // Check if repair is needed
+        guard journal[.isRepairFaultyMLSRemovalKeysRequired] else {
+            WireLogger.mls.debug(
+                "faulty MLS removal keys repair not needed, skipping",
+                attributes: .safePublic
+            )
+            return
+        }
+
+        WireLogger.mls.info(
+            "starting faulty MLS removal keys repair",
+            attributes: .safePublic
+        )
+
+        do {
+            let result = try await repairUseCase.invoke()
+
+            WireLogger.mls.info(
+                "faulty MLS removal keys repair completed: found \(result.faultyConversationsFound), repaired \(result.conversationsRepaired)",
+                attributes: .safePublic
+            )
+
+            // Clear the flag on success
+            journal[.isRepairFaultyMLSRemovalKeysRequired] = false
+
+        } catch {
+            WireLogger.mls.error(
+                "faulty MLS removal keys repair failed: \(String(describing: error))",
+                attributes: .safePublic
+            )
+            // Don't clear the flag so it will be retried
+            throw error
+        }
+    }
+
+}

WireDomain/Tests/WireDomainTests/Utilities/JournalTests.swift

@@ -70,7 +70,7 @@ class JournalStoreTests {
     @Test("Values contain all declared values")
     func values() {
         // Given
-        let exhaustiveKeysCount = 10
+        let exhaustiveKeysCount = 11
         sut[.isSyncV2Enabled] = true
 
         // When

wire-ios-sync-engine/Source/UserSession/ZMUserSession/AppVersionMigrations/AppVersionMigration_4_12_0.swift

@@ -0,0 +1,37 @@
+//
+// Wire
+// Copyright (C) 2025 Wire Swiss GmbH
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see http://www.gnu.org/licenses/.
+//
+
+import Foundation
+import WireDomain
+
+/// **Issue:**: Faulty MLS removal keys need repair - [WPB-22447]
+struct AppVersionMigration_4_12_0: AppVersionMigration {
+
+    let version: SemanticVersion = "4.12.0"
+    let journal: any JournalProtocol
+    let repairGenerator: RepairFaultyMLSRemovalKeysGenerator?
+
+    func perform() async throws {
+        // Mark that faulty MLS removal keys need to be repaired
+        journal[.isRepairFaultyMLSRemovalKeysRequired] = true
+
+        // Submit the work item now that the flag is set
+        repairGenerator?.submitWorkItemIfNeeded()
+    }
+
+}

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession.swift

@@ -669,6 +669,9 @@ public final class ZMUserSession: NSObject {
             await clientSessionComponent.workAgent.setAutoStartEnabled(true)
             await clientSessionComponent.workAgent.start()
             clientSessionComponent.generatorsDirectory.observeSyncState()
+
+            // Initialize the generator to enqueue repair work item if needed
+            clientSessionComponent.repairFaultyMLSRemovalKeysGenerator.submitWorkItemIfNeeded()
         }
     }
 
@@ -1711,7 +1714,11 @@ extension ZMUserSession {
                 sessionManager: sessionManager
             ),
             AppVersionMigration_4_3_0(coreCryptoProvider: coreCryptoProvider),
-            AppVersionMigration_4_10_0(journal: journal)
+            AppVersionMigration_4_10_0(journal: journal),
+            AppVersionMigration_4_12_0(
+                journal: journal,
+                repairGenerator: clientSessionComponent?.repairFaultyMLSRemovalKeysGenerator
+            )
         ]
     }