Introduce a local mechanism for managing the app's badge count

ElementX/Sources/Application/Settings/AppSettings.swift

@@ -19,6 +19,7 @@ import SwiftUI
 nonisolated protocol CommonSettingsProtocol: AnyObject, Sendable {
     var lastNotificationBootTime: TimeInterval? { get set }
     var selectedNotificationTone: NotificationTone? { get set }
+    var lastKnownBadgeCount: Int { get set }
     
     var logLevel: LogLevel { get }
     var traceLogPacks: Set<TraceLogPack> { get }
@@ -30,6 +31,7 @@ nonisolated protocol CommonSettingsProtocol: AnyObject, Sendable {
     var enableOnlySignedDeviceIsolationMode: Bool { get }
     var threadsEnabled: Bool { get }
     var hideQuietNotificationAlerts: Bool { get }
+    var roomListNotificationCountEnabled: Bool { get }
 }
 
 nonisolated enum AppBuildType {
@@ -261,6 +263,10 @@ final nonisolated class AppSettings: @unchecked Sendable {
     @UserPreference
     var lastNotificationBootTime: TimeInterval?
     
+    /// The app icon badge value the app last computed from the SDK's unread notification counts.
+    @UserPreference(defaultValue: 0)
+    var lastKnownBadgeCount: Int
+    
     /// The sound played when delivering noisy notifications. If nil, use the ElementX default
     @UserPreference
     var selectedNotificationTone: NotificationTone?

ElementX/Sources/FlowCoordinators/ChatsTabFlowCoordinator.swift

@@ -341,6 +341,7 @@ class ChatsTabFlowCoordinator: FlowCoordinatorProtocol {
                     Task {
                         let roomSummaries = self.userSession.clientProxy.staticRoomSummaryProvider.roomListPublisher.value
                         await self.flowParameters.notificationManager.removeDeliveredNotificationsForFullyReadRooms(roomSummaries)
+                        await self.flowParameters.notificationManager.updateAppBadgeCount()
                     }
                 default:
                     break

ElementX/Sources/Mocks/ClientProxyMock.swift

@@ -36,6 +36,8 @@ struct ClientProxyMockConfiguration {
     
     var maxMediaUploadSize: UInt = 100 * 1024 * 1024
     
+    var totalUnreadNotifications: UInt64 = 0
+    
     class Overrides {
         var joinedRoomIDs: Set<String> = []
     }
@@ -56,6 +58,8 @@ extension ClientProxyMock {
         homeserver = configuration.homeserver
         userIDServerName = configuration.userIDServerName
         
+        totalUnreadNotifications = configuration.totalUnreadNotifications
+        
         roomSummaryProvider = configuration.roomSummaryProvider
         alternateRoomSummaryProvider = RoomSummaryProviderMock(.init())
         staticRoomSummaryProvider = RoomSummaryProviderMock(.init())

ElementX/Sources/Mocks/Generated/GeneratedMocks.swift

@@ -2322,6 +2322,11 @@ nonisolated class ClientProxyMock: ClientProxyProtocol, @unchecked Sendable {
     }
     nonisolated(unsafe) var underlyingHideInviteAvatarsPublisher: CurrentValuePublisher<Bool, Never>!
     nonisolated(unsafe) var pusherNotificationClientIdentifier: String?
+    var totalUnreadNotifications: UInt64 {
+        get { return underlyingTotalUnreadNotifications }
+        set(value) { underlyingTotalUnreadNotifications = value }
+    }
+    nonisolated(unsafe) var underlyingTotalUnreadNotifications: UInt64!
     var mediaLoader: MediaLoaderProtocol {
         get { return underlyingMediaLoader }
         set(value) { underlyingMediaLoader = value }
@@ -8663,6 +8668,23 @@ nonisolated class NotificationManagerMock: NotificationManagerProtocol, @uncheck
         removeDeliveredNotificationsForFullyReadRoomsReceivedInvocationsLock.withLock { removeDeliveredNotificationsForFullyReadRoomsUnderlyingReceivedInvocations.append(rooms) }
         await removeDeliveredNotificationsForFullyReadRoomsClosure?(rooms)
     }
+    //MARK: - updateAppBadgeCount
+
+    private let updateAppBadgeCountCallsCountLock = NSLock()
+    private nonisolated(unsafe) var updateAppBadgeCountUnderlyingCallsCount = 0
+    var updateAppBadgeCountCallsCount: Int {
+        get { updateAppBadgeCountCallsCountLock.withLock { updateAppBadgeCountUnderlyingCallsCount } }
+        set { updateAppBadgeCountCallsCountLock.withLock { updateAppBadgeCountUnderlyingCallsCount = newValue } }
+    }
+    var updateAppBadgeCountCalled: Bool {
+        return updateAppBadgeCountCallsCount > 0
+    }
+    nonisolated(unsafe) var updateAppBadgeCountClosure: (() async -> Void)?
+
+    @concurrent func updateAppBadgeCount() async {
+        updateAppBadgeCountCallsCountLock.withLock { updateAppBadgeCountUnderlyingCallsCount += 1 }
+        await updateAppBadgeCountClosure?()
+    }
 }
 nonisolated class NotificationSettingsProxyMock: NotificationSettingsProxyProtocol, @unchecked Sendable {
     var callbacks: PassthroughSubject<NotificationSettingsProxyCallback, Never> {
@@ -14103,6 +14125,41 @@ nonisolated class UserNotificationCenterMock: UserNotificationCenterProtocol, @u
         removeDeliveredNotificationsWithIdentifiersReceivedInvocationsLock.withLock { removeDeliveredNotificationsWithIdentifiersUnderlyingReceivedInvocations.append(identifiers) }
         removeDeliveredNotificationsWithIdentifiersClosure?(identifiers)
     }
+    //MARK: - setBadgeCount
+
+    nonisolated(unsafe) var setBadgeCountThrowableError: Error?
+    private let setBadgeCountCallsCountLock = NSLock()
+    private nonisolated(unsafe) var setBadgeCountUnderlyingCallsCount = 0
+    var setBadgeCountCallsCount: Int {
+        get { setBadgeCountCallsCountLock.withLock { setBadgeCountUnderlyingCallsCount } }
+        set { setBadgeCountCallsCountLock.withLock { setBadgeCountUnderlyingCallsCount = newValue } }
+    }
+    var setBadgeCountCalled: Bool {
+        return setBadgeCountCallsCount > 0
+    }
+    private let setBadgeCountReceivedNewBadgeCountLock = NSLock()
+    private nonisolated(unsafe) var setBadgeCountUnderlyingReceivedNewBadgeCount: Int?
+    var setBadgeCountReceivedNewBadgeCount: Int? {
+        get { setBadgeCountReceivedNewBadgeCountLock.withLock { setBadgeCountUnderlyingReceivedNewBadgeCount } }
+        set { setBadgeCountReceivedNewBadgeCountLock.withLock { setBadgeCountUnderlyingReceivedNewBadgeCount = newValue } }
+    }
+    private let setBadgeCountReceivedInvocationsLock = NSLock()
+    private nonisolated(unsafe) var setBadgeCountUnderlyingReceivedInvocations: [Int] = []
+    var setBadgeCountReceivedInvocations: [Int] {
+        get { setBadgeCountReceivedInvocationsLock.withLock { setBadgeCountUnderlyingReceivedInvocations } }
+        set { setBadgeCountReceivedInvocationsLock.withLock { setBadgeCountUnderlyingReceivedInvocations = newValue } }
+    }
+    nonisolated(unsafe) var setBadgeCountClosure: ((Int) async throws -> Void)?
+
+    @concurrent func setBadgeCount(_ newBadgeCount: Int) async throws {
+        if let error = setBadgeCountThrowableError {
+            throw error
+        }
+        setBadgeCountCallsCountLock.withLock { setBadgeCountUnderlyingCallsCount += 1 }
+        setBadgeCountReceivedNewBadgeCount = newBadgeCount
+        setBadgeCountReceivedInvocationsLock.withLock { setBadgeCountUnderlyingReceivedInvocations.append(newBadgeCount) }
+        try await setBadgeCountClosure?(newBadgeCount)
+    }
     //MARK: - setNotificationCategories
 
     private let setNotificationCategoriesCallsCountLock = NSLock()

ElementX/Sources/Screens/Settings/DeveloperOptionsScreen/View/DeveloperOptionsScreen.swift

@@ -69,6 +69,7 @@ struct DeveloperOptionsScreen: View {
                 
                 Toggle(isOn: $context.roomListNotificationCountEnabled) {
                     Text("Show unread notification count")
+                    Text("Also makes the app icon badge use the SDK's own unread notification count")
                 }
                 
                 Toggle(isOn: $context.fuzzyRoomListSearchEnabled) {

ElementX/Sources/Services/Client/ClientProxy.swift

@@ -339,6 +339,10 @@ class ClientProxy: ClientProxyProtocol {
         client.canDeactivateAccount()
     }
     
+    var totalUnreadNotifications: UInt64 {
+        client.totalUnreadNotifications()
+    }
+    
     var userIDServerName: String? {
         do {
             return try client.userIdServerName()

ElementX/Sources/Services/Client/ClientProxyProtocol.swift

@@ -137,6 +137,9 @@ protocol ClientProxyProtocol: AnyObject {
     
     var pusherNotificationClientIdentifier: String? { get }
     
+    /// The total number of unread notifications across all joined, non-muted rooms, as computed by the SDK.
+    var totalUnreadNotifications: UInt64 { get }
+    
     var mediaLoader: MediaLoaderProtocol { get }
     
     var contentScanner: ContentScannerProxyProtocol? { get }

ElementX/Sources/Services/Notification/Manager/NotificationManager.swift

@@ -158,6 +158,22 @@ final class NotificationManager: NSObject, NotificationManagerProtocol {
         notificationCenter.removeDeliveredNotifications(withIdentifiers: notificationsIdentifiers)
     }
     
+    func updateAppBadgeCount() async {
+        guard let userSession, appSettings.roomListNotificationCountEnabled else { return }
+        
+        let badgeCount = Int(userSession.clientProxy.totalUnreadNotifications)
+        
+        appSettings.lastKnownBadgeCount = badgeCount
+        
+        MXLog.debug("Updating app badge count to \(badgeCount)")
+        
+        do {
+            try await notificationCenter.setBadgeCount(badgeCount)
+        } catch {
+            MXLog.error("Failed updating the app badge count with error: \(error)")
+        }
+    }
+    
     private func removeReceivedWhileOfflineNotification() {
         notificationCenter.removeDeliveredNotifications(withIdentifiers: [NotificationServiceExtensionActor.receivedWhileOfflineNotificationID])
     }

ElementX/Sources/Services/Notification/Manager/NotificationManagerProtocol.swift

@@ -36,4 +36,6 @@ protocol NotificationManagerProtocol: AnyObject {
     func removeDeliveredMessageNotifications(for roomID: String) async
     
     func removeDeliveredNotificationsForFullyReadRooms(_ rooms: [RoomSummary]) async
+    
+    func updateAppBadgeCount() async
 }

ElementX/Sources/Services/Notification/Manager/UserNotificationCenterProtocol.swift

@@ -15,6 +15,7 @@ protocol UserNotificationCenterProtocol: AnyObject {
     func requestAuthorization(options: UNAuthorizationOptions) async throws -> Bool
     func deliveredNotifications() async -> [UNNotification]
     func removeDeliveredNotifications(withIdentifiers identifiers: [String])
+    func setBadgeCount(_ newBadgeCount: Int) async throws
     func setNotificationCategories(_ categories: Set<UNNotificationCategory>)
     func authorizationStatus() async -> UNAuthorizationStatus
     func notificationSettings() async -> UNNotificationSettings

NSE/Sources/NotificationHandler.swift

@@ -20,6 +20,9 @@ nonisolated class NotificationHandler {
     
     private let notificationContentBuilder: NotificationContentBuilder
     
+    /// Whether this handler has already counted its notification towards the app icon badge.
+    private var hasCountedTowardsBadge = false
+    
     init(userSession: NSEUserSession,
          settings: CommonSettingsProtocol,
          contentHandler: @escaping (UNNotificationContent) -> Void,
@@ -42,9 +45,11 @@ nonisolated class NotificationHandler {
     func processEvent(_ eventID: String, roomID: String) async {
         MXLog.info("\(tag) Processing event: \(eventID) in room: \(roomID)")
         
-        // Copy over the unread information to the notification badge
-        notificationContent.badge = notificationContent.unreadCount as NSNumber?
-        MXLog.info("\(tag) New badge value: \(notificationContent.badge?.stringValue ?? "nil")")
+        if !settings.roomListNotificationCountEnabled {
+            // Copy over the unread information provided by the push payload to the notification badge.
+            notificationContent.badge = notificationContent.unreadCount as NSNumber?
+            MXLog.info("\(tag) New badge value: \(notificationContent.badge?.stringValue ?? "nil")")
+        }
         
         guard let notificationItemProxy = await userSession.notificationItemProxy(roomID: roomID, eventID: eventID) else {
             MXLog.error("\(tag) Failed retrieving notification item")
@@ -79,15 +84,37 @@ nonisolated class NotificationHandler {
     // MARK: - Private
     
     private func deliverNotification() {
-        MXLog.info("\(tag) Delivering notification")
+        if settings.roomListNotificationCountEnabled {
+            notificationContent.badge = NSNumber(value: incrementBadgeCount())
+            MXLog.info("\(tag) Delivering notification, new badge value: \(settings.lastKnownBadgeCount)")
+        } else {
+            MXLog.info("\(tag) Delivering notification")
+        }
         contentHandler(notificationContent)
     }
     
+    private func incrementBadgeCount() -> Int {
+        // `handleTimeExpiration` can deliver a notification this handler already delivered.
+        guard !hasCountedTowardsBadge else {
+            return settings.lastKnownBadgeCount
+        }
+        
+        hasCountedTowardsBadge = true
+        settings.lastKnownBadgeCount += 1
+        
+        return settings.lastKnownBadgeCount
+    }
+    
     private func discardNotification() {
         MXLog.info("\(tag) Discarding notification")
         
         let content = UNMutableNotificationContent()
-        content.badge = notificationContent.unreadCount as NSNumber?
+        if settings.roomListNotificationCountEnabled {
+            // Nothing new is shown to the user, so leave the badge where the app last put it.
+            content.badge = NSNumber(value: settings.lastKnownBadgeCount)
+        } else {
+            content.badge = notificationContent.unreadCount as NSNumber?
+        }
         MXLog.info("\(tag) New badge value: \(content.badge?.stringValue ?? "nil")")
         
         contentHandler(content)
@@ -122,6 +149,9 @@ nonisolated class NotificationHandler {
                 
                 if let targetNotification = deliveredNotifications.first(where: { $0.request.content.eventID == redactedEventID }) {
                     UNUserNotificationCenter.current().removeDeliveredNotifications(withIdentifiers: [targetNotification.request.identifier])
+                    if settings.roomListNotificationCountEnabled {
+                        settings.lastKnownBadgeCount = max(0, settings.lastKnownBadgeCount - 1)
+                    }
                 }
                 
                 return .processedShouldDiscard

NSE/Sources/NotificationServiceExtension.swift

@@ -258,7 +258,9 @@ actor NotificationServiceExtensionActor {
         
         let content = UNMutableNotificationContent()
         content.body = L10n.notificationReceivedWhileOfflineIos
-        content.badge = originalRequest.content.unreadCount as NSNumber?
+        if !settings.roomListNotificationCountEnabled {
+            content.badge = originalRequest.content.unreadCount as NSNumber?
+        }
         content.sound = settings.notificationSound
         
         let request = UNNotificationRequest(identifier: Self.receivedWhileOfflineNotificationID, content: content, trigger: nil)

UnitTests/Sources/NotificationManager/NotificationManagerTests.swift

@@ -237,6 +237,37 @@ final class NotificationManagerTests {
         await notificationManager.userNotificationCenter(UNUserNotificationCenter.current(), didReceive: response)
         #expect(notificationTappedDelegateCalled)
     }
+    
+    @Test
+    func updatingAppBadgeCountUsesTheClientSideCount() async {
+        appSettings.roomListNotificationCountEnabled = true
+        clientProxy.totalUnreadNotifications = 7
+        
+        await notificationManager.updateAppBadgeCount()
+        
+        #expect(notificationCenter.setBadgeCountReceivedNewBadgeCount == 7)
+        #expect(appSettings.lastKnownBadgeCount == 7)
+    }
+    
+    @Test
+    func updatingAppBadgeCountWithoutASessionDoesNothing() async {
+        appSettings.roomListNotificationCountEnabled = true
+        notificationManager.setUserSession(nil)
+        
+        await notificationManager.updateAppBadgeCount()
+        
+        #expect(!notificationCenter.setBadgeCountCalled)
+    }
+    
+    @Test
+    func updatingAppBadgeCountWithFeatureDisabledDoesNothing() async {
+        appSettings.roomListNotificationCountEnabled = false
+        clientProxy.totalUnreadNotifications = 7
+        
+        await notificationManager.updateAppBadgeCount()
+        
+        #expect(!notificationCenter.setBadgeCountCalled)
+    }
 }
 
 extension NotificationManagerTests: @MainActor NotificationManagerDelegate {