Update AXNotificationStream to retry automatically

Author: intitni

Core/Sources/AXNotificationStream/AXNotificationStream.swift

@@ -1,6 +1,7 @@
 import AppKit
 import ApplicationServices
 import Foundation
+import Logger
 
 public final class AXNotificationStream: AsyncSequence {
     public typealias Stream = AsyncStream<Element>
@@ -18,7 +19,7 @@ public final class AXNotificationStream: AsyncSequence {
     deinit {
         continuation.finish()
     }
-    
+
     public convenience init(
         app: NSRunningApplication,
         element: AXUIElement? = nil,
@@ -72,24 +73,56 @@ public final class AXNotificationStream: AsyncSequence {
                 .commonModes
             )
         }
-        
-        Task {
-            for name in notificationNames {
-                var error = AXError.cannotComplete
-                var retryCount = 0
-                while error == AXError.cannotComplete, retryCount < 5 {
-                    error = AXObserverAddNotification(observer, observingElement, name as CFString, &continuation)
-                    if error == .cannotComplete {
-                        try await Task.sleep(nanoseconds: 1_000_000_000)
-                    }
-                    retryCount += 1
-                }
-            }
+
+        Task { [weak self] in
             CFRunLoopAddSource(
                 CFRunLoopGetMain(),
                 AXObserverGetRunLoopSource(observer),
                 .commonModes
             )
+            var pendingRegistrationNames = Set(notificationNames)
+            var retry = 0
+            while !pendingRegistrationNames.isEmpty, retry < 100 {
+                guard let self else { return }
+                retry += 1
+                for name in notificationNames {
+                    let e = AXObserverAddNotification(
+                        observer,
+                        observingElement,
+                        name as CFString,
+                        &self.continuation
+                    )
+                    switch e {
+                    case .success:
+                        pendingRegistrationNames.remove(name)
+                    case .actionUnsupported:
+                        Logger.service.error("AXObserver: Action unsupported: \(name)")
+                        pendingRegistrationNames.remove(name)
+                    case .apiDisabled:
+                        Logger.service.error("AXObserver: Accessibility API disabled, will try again later")
+                        retry -= 1
+                    case .invalidUIElement:
+                        Logger.service.error("AXObserver: Invalid UI element")
+                        pendingRegistrationNames.remove(name)
+                    case .invalidUIElementObserver:
+                        Logger.service.error("AXObserver: Invalid UI element observer")
+                        pendingRegistrationNames.remove(name)
+                    case .cannotComplete:
+                        Logger.service
+                            .error("AXObserver: Failed to observe \(name), will try again later")
+                    case .notificationUnsupported:
+                        Logger.service.error("AXObserver: Notification unsupported: \(name)")
+                        pendingRegistrationNames.remove(name)
+                    case .notificationAlreadyRegistered:
+                        pendingRegistrationNames.remove(name)
+                    default:
+                        Logger.service
+                            .error("AXObserver: Unrecognized error \(e) when registering \(name), will try again later")
+                    }
+                }
+                try await Task.sleep(nanoseconds: 1_500_000_000)
+            }
         }
     }
 }
+