Add package DisplayLink

Author: intitni

Core/Package.swift

@@ -51,7 +51,8 @@ let package = Package(
         ),
         .target(
             name: "Service",
-            dependencies: ["CopilotModel", "CopilotService", "XPCShared", "CGEventObserver"]
+            dependencies: ["CopilotModel", "CopilotService", "XPCShared", "CGEventObserver",
+                           "DisplayLink"]
         ),
         .target(
             name: "XPCShared",
@@ -63,5 +64,6 @@ let package = Package(
         ),
         .target(name: "FileChangeChecker"),
         .target(name: "LaunchAgentManager"),
+        .target(name: "DisplayLink"),
     ]
 )

Core/Sources/DisplayLink/DisplayLink.swift

@@ -0,0 +1,58 @@
+import Foundation
+import QuartzCore
+
+public actor DisplayLink {
+    private var displayLink: CVDisplayLink!
+    static let shared = DisplayLink()
+    private var continuations: [UUID: AsyncStream<Void>.Continuation] = [:]
+
+    public static func createStream() -> AsyncStream<Void> {
+        .init { continuation in
+            Task {
+                let id = UUID()
+                await DisplayLink.shared?.addContinuation(continuation, id: id)
+                continuation.onTermination = { _ in
+                    Task {
+                        await DisplayLink.shared?.removeContinuation(id: id)
+                    }
+                }
+            }
+        }
+    }
+
+    private init?() {
+        _ = CVDisplayLinkCreateWithCGDisplay(CGMainDisplayID(), &displayLink)
+        guard displayLink != nil else { return nil }
+        CVDisplayLinkSetOutputHandler(displayLink) { [weak self] _, _, _, _, _ in
+            guard let self else { return kCVReturnSuccess }
+            Task { await self.notifyContinuations() }
+            return kCVReturnSuccess
+        }
+    }
+
+    deinit {
+        for continuation in continuations {
+            continuation.value.finish()
+        }
+    }
+
+    func addContinuation(_ continuation: AsyncStream<Void>.Continuation, id: UUID) {
+        continuations[id] = continuation
+        if !continuations.isEmpty {
+            CVDisplayLinkStart(displayLink)
+        }
+    }
+
+    func removeContinuation(id: UUID) {
+        continuations[id] = nil
+        if continuations.isEmpty {
+            CVDisplayLinkStop(displayLink)
+        }
+    }
+
+    private func notifyContinuations() {
+        for continuation in continuations {
+            continuation.value.yield()
+        }
+    }
+}

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -1,5 +1,6 @@
 import AppKit
 import CGEventObserver
+import DisplayLink
 import Foundation
 import os.log
 import QuartzCore
@@ -66,7 +67,7 @@ public actor RealtimeSuggestionController {
             }
         }
         if eventObserver.activateIfPossible() {
-            realtimeSuggestionIndicatorController?.isObserving = true
+            realtimeSuggestionIndicatorController.isObserving = true
         }
     }
 
@@ -78,7 +79,7 @@ public actor RealtimeSuggestionController {
         task?.cancel()
         task = nil
         eventObserver.deactivate()
-        realtimeSuggestionIndicatorController?.isObserving = false
+        realtimeSuggestionIndicatorController.isObserving = false
     }
 
     func handleKeyboardEvent(event: CGEvent) async {
@@ -132,7 +133,7 @@ public actor RealtimeSuggestionController {
             else { return }
             if Task.isCancelled { return }
             os_log(.info, "Prefetch suggestions.")
-            realtimeSuggestionIndicatorController?.triggerPrefetchAnimation()
+            realtimeSuggestionIndicatorController.triggerPrefetchAnimation()
             do {
                 try await Environment.triggerAction("Prefetch Suggestions")
             } catch {
@@ -239,8 +240,6 @@ final class RealtimeSuggestionIndicatorController {
     }
 
     private let viewModel = IndicatorContentViewModel()
-    private var displayLink: CVDisplayLink!
-    private var isDisplayLinkStarted: Bool = false
     private var userDefaultsObserver = UserDefaultsObserver()
     var isObserving = false {
         didSet {
@@ -250,6 +249,8 @@ final class RealtimeSuggestionIndicatorController {
         }
     }
 
+    private var displayLinkTask: Task<Void, Never>?
+
     @MainActor
     lazy var window = {
         let it = NSWindow(
@@ -269,15 +270,7 @@ final class RealtimeSuggestionIndicatorController {
         return it
     }()
 
-    init?() {
-        _ = CVDisplayLinkCreateWithCGDisplay(CGMainDisplayID(), &displayLink)
-        guard displayLink != nil else { return nil }
-        CVDisplayLinkSetOutputHandler(displayLink) { [weak self] _, _, _, _, _ in
-            guard let self else { return kCVReturnSuccess }
-            self.updateIndicatorLocation()
-            return kCVReturnSuccess
-        }
-
+    init() {
         Task {
             let sequence = NSWorkspace.shared.notificationCenter
                 .notifications(named: NSWorkspace.didActivateApplicationNotification)
@@ -327,9 +320,16 @@ final class RealtimeSuggestionIndicatorController {
         await { @MainActor in
             guard window.isVisible != isVisible else { return }
             if isVisible {
-                CVDisplayLinkStart(self.displayLink)
+                if displayLinkTask == nil {
+                    displayLinkTask = Task {
+                        for await _ in DisplayLink.createStream() {
+                            self.updateIndicatorLocation()
+                        }
+                    }
+                }
             } else {
-                CVDisplayLinkStop(self.displayLink)
+                displayLinkTask?.cancel()
+                displayLinkTask = nil
             }
             window.setIsVisible(isVisible)
         }()