Persist real-time suggestions state

Core/Sources/Client/XPCService.swift

@@ -1,5 +1,6 @@
 import Foundation
 import XPCShared
+import os.log
 
 var asyncService: AsyncXPCService?
 var shared = XPCService()
@@ -30,11 +31,11 @@ class XPCService {
         connection.remoteObjectInterface =
             NSXPCInterface(with: XPCServiceProtocol.self)
         connection.invalidationHandler = { [weak self] in
-            print("XPCService Invalidated")
+            os_log(.info, "XPCService Invalidated")
             self?.isInvalidated = true
         }
         connection.interruptionHandler = { [weak self] in
-            print("XPCService interrupted")
+            os_log(.info, "XPCService interrupted")
         }
         connection.resume()
         return connection

Core/Sources/Service/AutoTrigger.swift

@@ -1,14 +1,17 @@
+import AppKit
 import Foundation
 import XPCShared
+import os.log
 
-actor AutoTrigger {
-    static let shared = AutoTrigger()
+public actor AutoTrigger {
+    public static let shared = AutoTrigger()
 
-    private var listeners = Set<ObjectIdentifier>()
+    private var listeners = Set<AnyHashable>()
     var eventObserver: CGEventObserverType = CGEventObserver()
     var task: Task<Void, Error>?
 
     private init() {
+        // Occasionally cleanup workspaces.
         Task { @ServiceActor in
             while !Task.isCancelled {
                 try await Task.sleep(nanoseconds: 8 * 60 * 60 * 1_000_000_000)
@@ -21,15 +24,42 @@ actor AutoTrigger {
                 }
             }
         }
+
+        // Start the auto trigger if Xcode is running.
+        Task {
+            for xcode in await Environment.runningXcodes() {
+                await start(by: xcode.processIdentifier)
+            }
+            let sequence = NSWorkspace.shared.notificationCenter
+                .notifications(named: NSWorkspace.didLaunchApplicationNotification)
+            for await notification in sequence {
+                try Task.checkCancellation()
+                guard let app =  notification.userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication else { continue }
+                guard app.bundleIdentifier == "com.apple.dt.Xcode" else { continue }
+                await start(by: app.processIdentifier)
+            }
+        }
+
+        // Remove listener if Xcode is terminated.
+        Task {
+            let sequence = NSWorkspace.shared.notificationCenter
+                .notifications(named: NSWorkspace.didTerminateApplicationNotification)
+            for await notification in sequence {
+                try Task.checkCancellation()
+                guard let app = notification.userInfo?[NSWorkspace.applicationUserInfoKey] as? NSRunningApplication else { continue }
+                guard app.bundleIdentifier == "com.apple.dt.Xcode" else { continue }
+                await stop(by: app.processIdentifier)
+            }
+        }
     }
 
-    func start(by listener: ObjectIdentifier) {
+    func start(by listener: AnyHashable) {
+        os_log(.info, "Add auto trigger listener: %@.", listener as CVarArg)
         listeners.insert(listener)
 
         if task == nil {
             task = Task { [stream = eventObserver.stream] in
                 var triggerTask: Task<Void, Error>?
-                try? await Environment.triggerAction("Prefetch Suggestions")
                 for await _ in stream {
                     triggerTask?.cancel()
                     if Task.isCancelled { break }
@@ -47,10 +77,11 @@ actor AutoTrigger {
                         try? await Task.sleep(nanoseconds: 2_000_000_000)
                         if Task.isCancelled { return }
                         let fileURL = try? await Environment.fetchCurrentFileURL()
-                        guard let folderURL = try? await Environment.fetchCurrentProjectRootURL(fileURL),
-                              let workspace = workspaces[folderURL],
-                              workspace.isRealtimeSuggestionEnabled
+                        guard let folderURL = try? await Environment.fetchCurrentProjectRootURL(fileURL)
                         else { return }
+                        let workspace = workspaces[folderURL] ?? Workspace(projectRootURL: folderURL)
+                        workspaces[folderURL] = workspace
+                        guard workspace.isRealtimeSuggestionEnabled else { return }
                         if Task.isCancelled { return }
                         try? await Environment.triggerAction("Prefetch Suggestions")
                     }
@@ -60,9 +91,11 @@ actor AutoTrigger {
         eventObserver.activateIfPossible()
     }
 
-    func stop(by listener: ObjectIdentifier) {
+    func stop(by listener: AnyHashable) {
+        os_log(.info, "Remove auto trigger listener: %@.", listener as CVarArg)
         listeners.remove(listener)
         guard listeners.isEmpty else { return }
+        os_log(.info, "Auto trigger is stopped.")
         task?.cancel()
         task = nil
         eventObserver.deactivate()

Core/Sources/Service/CGEventObserver.swift

@@ -1,5 +1,6 @@
 import Cocoa
 import Foundation
+import os.log
 
 public protocol CGEventObserverType {
     @discardableResult
@@ -38,7 +39,7 @@ final class CGEventObserver: CGEventObserverType {
         retryTask?.cancel()
         retryTask = nil
         guard let port = port else { return }
-        print("Deactivate")
+        os_log(.info, "CGEventObserver deactivated.")
         CFMachPortInvalidate(port)
         self.port = nil
     }
@@ -95,6 +96,7 @@ final class CGEventObserver: CGEventObserverType {
         self.port = port
         let runLoopSource = CFMachPortCreateRunLoopSource(kCFAllocatorDefault, port, 0)
         CFRunLoopAddSource(RunLoop.main.getCFRunLoop(), runLoopSource, .commonModes)
+        os_log(.info, "CGEventObserver activated.")
         return true
     }
 }

Core/Sources/Service/Environment.swift

@@ -17,18 +17,28 @@ private struct FailedToFetchFileURLError: Error, LocalizedError {
 enum Environment {
     static var now = { Date() }
 
+    static var runningXcodes: () async -> [NSRunningApplication] = {
+        var xcodes = [NSRunningApplication]()
+        var retryCount = 0
+        // Sometimes runningApplications returns 0 items.
+        while xcodes.isEmpty, retryCount < 3 {
+            xcodes = NSRunningApplication
+                .runningApplications(withBundleIdentifier: "com.apple.dt.Xcode")
+            try? await Task.sleep(nanoseconds: 1_000_000)
+            retryCount += 1
+        }
+        return xcodes
+    }
+
     static var isXcodeActive: () async ->  Bool = {
         var activeXcodes = [NSRunningApplication]()
         var retryCount = 0
         // Sometimes runningApplications returns 0 items.
         while activeXcodes.isEmpty, retryCount < 3 {
             activeXcodes = NSRunningApplication
                 .runningApplications(withBundleIdentifier: "com.apple.dt.Xcode")
-                .sorted { lhs, _ in
-                    if lhs.isActive { return true }
-                    return false
-                }
-            if retryCount > 0 { try? await Task.sleep(nanoseconds: 1_000_000) }
+                .filter { $0.isActive }
+            try? await Task.sleep(nanoseconds: 1_000_000)
             retryCount += 1
         }
         return !activeXcodes.isEmpty

Core/Sources/Service/Workspace.swift

@@ -45,7 +45,12 @@ final class Workspace {
     }
 
     var filespaces = [URL: Filespace]()
-    var isRealtimeSuggestionEnabled = false
+    var isRealtimeSuggestionEnabled: Bool {
+        (UserDefaults.shared.dictionary(
+            forKey: SettingsKey.realtimeSuggestionState
+        )?[projectRootURL.absoluteString]) as? Bool ?? false
+    }
+
     var realtimeSuggestionRequests = Set<Task<Void, Error>>()
 
     private lazy var service: CopilotSuggestionServiceType = Environment

Core/Sources/Service/XPCService.swift

@@ -2,6 +2,7 @@ import AppKit
 import CopilotService
 import Foundation
 import LanguageServerProtocol
+import os.log
 import XPCShared
 
 @globalActor enum ServiceActor {
@@ -16,21 +17,6 @@ public class XPCService: NSObject, XPCServiceProtocol {
     @ServiceActor
     lazy var authService: CopilotAuthServiceType = Environment.createAuthService()
 
-    override public init() {
-        super.init()
-        let identifier = ObjectIdentifier(self)
-        Task {
-            await AutoTrigger.shared.start(by: identifier)
-        }
-    }
-
-    deinit {
-        let identifier = ObjectIdentifier(self)
-        Task {
-            await AutoTrigger.shared.stop(by: identifier)
-        }
-    }
-
     public func checkStatus(withReply reply: @escaping (String?, Error?) -> Void) {
         Task { @ServiceActor in
             do {
@@ -110,7 +96,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -137,7 +123,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -164,7 +150,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -188,7 +174,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 )
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -215,7 +201,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -244,7 +230,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply(nil, NSError.from(error))
             }
         }
@@ -254,11 +240,19 @@ public class XPCService: NSObject, XPCServiceProtocol {
         Task { @ServiceActor in
             let fileURL = try await Environment.fetchCurrentFileURL()
             let workspace = try await fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
-            workspace.isRealtimeSuggestionEnabled = enabled
+            if var state = UserDefaults.shared.dictionary(forKey: SettingsKey.realtimeSuggestionState) {
+                state[workspace.projectRootURL.absoluteString] = enabled
+                UserDefaults.shared.set(state, forKey: SettingsKey.realtimeSuggestionState)
+            } else {
+                UserDefaults.shared.set(
+                    [workspace.projectRootURL.absoluteString: enabled],
+                    forKey: SettingsKey.realtimeSuggestionState
+                )
+            }
             reply(nil)
         }
     }
-    
+
     public func prefetchRealtimeSuggestions(
         editorContent: Data,
         withReply reply: @escaping () -> Void
@@ -279,7 +273,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 )
                 reply()
             } catch {
-                print(error)
+                os_log(.error, "%@", error.localizedDescription)
                 reply()
             }
         }

Core/Sources/XPCShared/UserDefaults.swift

@@ -6,4 +6,5 @@ public extension UserDefaults {
 
 public enum SettingsKey {
     public static let nodePath = "NodePath"
+    public static let realtimeSuggestionState = "RealtimeSuggestionState"
 }

README.md

@@ -38,22 +38,23 @@ The first time the commands run, the extension will ask for 2 types of permissio
 - Previous Suggestion: If there is more than 1 suggestion, switch to the previous one.
 - Accept Suggestion: Add the suggestion to the code.
 - Reject Suggestion: Remove the suggestion comments.
-- Turn On Real-time Suggestions: When turn on, Copilot will auto-insert suggestion comments to your code while editing. You have to manually turn it on for every open window of Xcode.
+- Turn On Real-time Suggestions: When turn on, Copilot will auto-insert suggestion comments to your code while editing.
 - Turn Off Real-time Suggestions: Turns the real-time suggestions off.
 - Real-time Suggestions: It is an entry point only for Copilot for Xcode. When suggestions are successfully fetched, Copilot for Xcode will run this command to present the suggestions.
 - Prefetch Suggestions: It is an entry point only for Copilot for Xcode. In the background, Copilot for Xcode will occasionally run this command to prefetch real-time suggestions. 
 
 **About real-time suggestions**
 
-The implementation won't feel as smooth as that of VSCode.
+- The on/off state is persisted, make sure you turn it off manually if you no longer want it.
+- The implementation won't feel as smooth as that of VSCode.
+
+    The magic behind it is that it will keep calling the command from the menu when you are not typing or clicking the mouse. So it will have to listen to those events, I am not sure if people like it.
 
-The magic behind it is that it will keep calling the command from the menu when you are not typing, or clicking mouse. So it will have to listen to those events, I am not sure if people like it.
-
-Hope that next year, Apple can spend some time on Xcode Extensions.  
+    Hope that next year, Apple can spend some time on Xcode Extensions.  
 
 ## Prevent Suggestions Being Committed
 
-Since the suggestions are presented as comments, they are in your code. If you are not careful enough, they can be committed to your git repo. To avoid that, I would recommend adding a pre-commit git hook to prevent this from happening. Maybe later I will add an command for that.
+Since the suggestions are presented as comments, they are in your code. If you are not careful enough, they can be committed to your git repo. To avoid that, I would recommend adding a pre-commit git hook to prevent this from happening.
 
 ```sh
 #!/bin/sh

XPCService/main.swift

@@ -1,4 +1,5 @@
 import Foundation
+import Service
 
 let listener = NSXPCListener(
     machServiceName: Bundle.main.object(forInfoDictionaryKey: "BUNDLE_IDENTIFIER_BASE") as! String
@@ -7,4 +8,5 @@ let listener = NSXPCListener(
 let delegate = ServiceDelegate()
 listener.delegate = delegate
 listener.resume()
+_ = AutoTrigger.shared
 RunLoop.main.run()