Replace display link with AXObserver

Author: intitni

Core/Sources/Service/GUI/RealtimeSuggestionIndicatorController.swift

@@ -1,16 +1,19 @@
+import ActiveApplicationMonitor
 import AppKit
+import AXNotificationStream
 import DisplayLink
 import Environment
 import QuartzCore
 import SwiftUI
 import XPCShared
 
 /// Present a tiny dot next to mouse cursor if real-time suggestion is enabled.
+@MainActor
 final class RealtimeSuggestionIndicatorController {
     class IndicatorContentViewModel: ObservableObject {
         @Published var isPrefetching = false
         private var prefetchTask: Task<Void, Error>?
-
+        
         @MainActor
         func prefetch() {
             prefetchTask?.cancel()
@@ -63,6 +66,9 @@ final class RealtimeSuggestionIndicatorController {
 
     private let viewModel = IndicatorContentViewModel()
     private var userDefaultsObserver = UserDefaultsObserver()
+    private var windowChangeObservationTask: Task<Void, Error>?
+    private var activeApplicationMonitorTask: Task<Void, Error>?
+    private var xcode: NSRunningApplication?
     var isObserving = false {
         didSet {
             Task {
@@ -71,8 +77,6 @@ final class RealtimeSuggestionIndicatorController {
         }
     }
 
-    private var displayLinkTask: Task<Void, Never>?
-
     @MainActor
     lazy var window = {
         let it = NSWindow(
@@ -92,35 +96,34 @@ final class RealtimeSuggestionIndicatorController {
         return it
     }()
 
-    init() {
-        Task {
-            let sequence = NSWorkspace.shared.notificationCenter
-                .notifications(named: NSWorkspace.didActivateApplicationNotification)
-            for await notification in sequence {
-                guard let app = notification
-                    .userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication
-                else { continue }
-                guard app.bundleIdentifier == "com.apple.dt.Xcode" else { continue }
-                await updateIndicatorVisibility()
-            }
-        }
-
-        Task {
-            let sequence = NSWorkspace.shared.notificationCenter
-                .notifications(named: NSWorkspace.didDeactivateApplicationNotification)
-            for await notification in sequence {
-                guard let app = notification
-                    .userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication
-                else { continue }
-                guard app.bundleIdentifier == "com.apple.dt.Xcode" else { continue }
-                await updateIndicatorVisibility()
+    nonisolated init() {
+        Task { @MainActor in
+            activeApplicationMonitorTask = Task { [weak self] in
+                var previousApp: NSRunningApplication?
+                for await app in ActiveApplicationMonitor.createStream() {
+                    guard let self else { return }
+                    try Task.checkCancellation()
+                    defer { previousApp = app }
+                    if let app = ActiveApplicationMonitor.activeXcode {
+                        if app != previousApp {
+                            windowChangeObservationTask?.cancel()
+                            windowChangeObservationTask = nil
+                            self.observeXcodeWindowChangeIfNeeded(app)
+                        }
+                        await self.updateIndicatorVisibility()
+                        self.updateIndicatorLocation()
+                    } else {
+                        await self.updateIndicatorVisibility()
+                    }
+                }
             }
         }
 
-        Task {
+        Task { @MainActor in
             userDefaultsObserver.onChange = { [weak self] in
                 Task { [weak self] in
                     await self?.updateIndicatorVisibility()
+                    self?.updateIndicatorLocation()
                 }
             }
             UserDefaults.shared.addObserver(
@@ -132,74 +135,77 @@ final class RealtimeSuggestionIndicatorController {
         }
     }
 
+    private func observeXcodeWindowChangeIfNeeded(_ app: NSRunningApplication) {
+        xcode = app
+        guard windowChangeObservationTask == nil else { return }
+        windowChangeObservationTask = Task {
+            let notifications = AXNotificationStream(
+                app: app,
+                notificationNames:
+                kAXMovedNotification,
+                kAXResizedNotification,
+                kAXMainWindowChangedNotification,
+                kAXFocusedWindowChangedNotification,
+                kAXFocusedUIElementChangedNotification,
+                kAXSelectedTextChangedNotification
+            )
+            for await _ in notifications {
+                try Task.checkCancellation()
+                updateIndicatorLocation()
+            }
+        }
+    }
+
     private func updateIndicatorVisibility() async {
         let isVisible = await {
             let isOn = UserDefaults.shared.bool(forKey: SettingsKey.realtimeSuggestionToggle)
             let isXcodeActive = await Environment.isXcodeActive()
-            return isOn && isXcodeActive && isObserving
+            return isOn && isXcodeActive
         }()
 
-        await { @MainActor in
-            guard window.isVisible != isVisible else { return }
-            if isVisible {
-                if displayLinkTask == nil {
-                    displayLinkTask = Task {
-                        for await _ in DisplayLink.createStream() {
-                            self.updateIndicatorLocation()
-                        }
-                    }
-                }
-            } else {
-                displayLinkTask?.cancel()
-                displayLinkTask = nil
-            }
-            window.setIsVisible(isVisible)
-        }()
+        guard window.isVisible != isVisible else { return }
+        window.setIsVisible(isVisible)
     }
 
     private func updateIndicatorLocation() {
-        Task { @MainActor in
-            if !window.isVisible {
-                return
-            }
+        if !window.isVisible {
+            return
+        }
 
-            if let activeXcode = NSRunningApplication
-                .runningApplications(withBundleIdentifier: "com.apple.dt.Xcode")
-                .first(where: \.isActive)
+        if let activeXcode = NSRunningApplication
+            .runningApplications(withBundleIdentifier: "com.apple.dt.Xcode")
+            .first(where: \.isActive)
+        {
+            let application = AXUIElementCreateApplication(activeXcode.processIdentifier)
+            if let focusElement: AXUIElement = try? application
+                .copyValue(key: kAXFocusedUIElementAttribute),
+                let selectedRange: AXValue = try? focusElement
+                .copyValue(key: kAXSelectedTextRangeAttribute),
+                let rect: AXValue = try? focusElement.copyParameterizedValue(
+                    key: kAXBoundsForRangeParameterizedAttribute,
+                    parameters: selectedRange
+                )
             {
-                let application = AXUIElementCreateApplication(activeXcode.processIdentifier)
-                if let focusElement: AXUIElement = try? application
-                    .copyValue(key: kAXFocusedUIElementAttribute),
-                    let selectedRange: AXValue = try? focusElement
-                    .copyValue(key: kAXSelectedTextRangeAttribute),
-                    let rect: AXValue = try? focusElement.copyParameterizedValue(
-                        key: kAXBoundsForRangeParameterizedAttribute,
-                        parameters: selectedRange
+                var frame: CGRect = .zero
+                let found = AXValueGetValue(rect, .cgRect, &frame)
+                let screen = NSScreen.screens.first
+                if found, let screen {
+                    frame.origin = .init(
+                        x: frame.maxX + 2,
+                        y: screen.frame.height - frame.minY - 4
                     )
-                {
-                    var frame: CGRect = .zero
-                    let found = AXValueGetValue(rect, .cgRect, &frame)
-                    let screen = NSScreen.screens.first
-                    if found, let screen {
-                        frame.origin = .init(
-                            x: frame.maxX + 2,
-                            y: screen.frame.height - frame.minY - 4
-                        )
-                        frame.size = .init(width: 10, height: 10)
-                        window.alphaValue = 1
-                        window.setFrame(frame, display: false, animate: true)
-                        return
-                    }
+                    frame.size = .init(width: 10, height: 10)
+                    window.alphaValue = 1
+                    window.setFrame(frame, display: false)
+                    return
                 }
             }
-
-            window.alphaValue = 0
         }
+
+        window.alphaValue = 0
     }
 
     func triggerPrefetchAnimation() {
-        Task { @MainActor in
-            viewModel.prefetch()
-        }
+        viewModel.prefetch()
     }
 }

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -44,9 +44,7 @@ public actor RealtimeSuggestionController {
                 }
             }
         }
-        if eventObserver.activateIfPossible() {
-            realtimeSuggestionIndicatorController.isObserving = true
-        }
+        eventObserver.activateIfPossible()
     }
 
     private func stop(by listener: AnyHashable) {
@@ -55,7 +53,6 @@ public actor RealtimeSuggestionController {
         task?.cancel()
         task = nil
         eventObserver.deactivate()
-        realtimeSuggestionIndicatorController.isObserving = false
     }
 
     func handleKeyboardEvent(event: CGEvent) async {
@@ -107,14 +104,14 @@ public actor RealtimeSuggestionController {
                     .value(forKey: SettingsKey.realtimeSuggestionDebounce) as? Double
                     ?? 0.7
             ) * 1_000_000_000))
-            
+
             guard UserDefaults.shared.bool(forKey: SettingsKey.realtimeSuggestionToggle)
             else { return }
-            
+
             if Task.isCancelled { return }
-            
+
             os_log(.info, "Prefetch suggestions.")
-            
+
             await realtimeSuggestionIndicatorController.triggerPrefetchAnimation()
             do {
                 try await Environment.triggerAction("Prefetch Suggestions")

Core/Sources/Service/SuggestionCommandHandler/CommentBaseCommandHandler.swift

@@ -173,7 +173,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
         // If there is a suggestion available, call another command to present it.
         guard !suggestions.isEmpty else { return nil }
         try await Environment.triggerAction("Real-time Suggestions")
-        GraphicalUserInterfaceController.shared.realtimeSuggestionIndicatorController
+        await GraphicalUserInterfaceController.shared.realtimeSuggestionIndicatorController
             .triggerPrefetchAnimation()
 
         return nil