Tweak update process

Author: intitni

Copilot for Xcode/App.swift

@@ -6,10 +6,22 @@ import UpdateChecker
 import XPCShared
 
 struct VisualEffect: NSViewRepresentable {
-  func makeNSView(context: Self.Context) -> NSView { return NSVisualEffectView() }
-  func updateNSView(_ nsView: NSView, context: Context) { }
+    func makeNSView(context: Self.Context) -> NSView { return NSVisualEffectView() }
+    func updateNSView(_ nsView: NSView, context: Context) {}
 }
 
+class TheUpdateCheckerDelegate: UpdateCheckerDelegate {
+    func prepareForRelaunch(finish: @escaping () -> Void) {
+        Task {
+            let service = try? getService()
+            try? await service?.quitService()
+            finish()
+        }
+    }
+}
+
+let updateCheckerDelegate = TheUpdateCheckerDelegate()
+
 @main
 struct CopilotForXcodeApp: App {
     var body: some Scene {
@@ -20,7 +32,17 @@ struct CopilotForXcodeApp: App {
                 .onAppear {
                     UserDefaults.setupDefaultSettings()
                 }
-                .environment(\.updateChecker, UpdateChecker(hostBundle: Bundle.main))
+                .environment(
+                    \.updateChecker,
+                    {
+                        let checker = UpdateChecker(
+                            hostBundle: Bundle.main,
+                            shouldAutomaticallyCheckForUpdate: false
+                        )
+                        checker.updateCheckerDelegate = updateCheckerDelegate
+                        return checker
+                    }()
+                )
         }
     }
 }

Core/Sources/HostApp/TabContainer.swift

@@ -215,7 +215,10 @@ private extension EnvironmentValues {
 }
 
 struct UpdateCheckerKey: EnvironmentKey {
-    static var defaultValue: UpdateChecker = .init(hostBundle: nil)
+    static var defaultValue: UpdateChecker = .init(
+        hostBundle: nil,
+        shouldAutomaticallyCheckForUpdate: false
+    )
 }
 
 public extension EnvironmentValues {

Core/Sources/Service/XPCService.swift

@@ -195,6 +195,13 @@ public class XPCService: NSObject, XPCServiceProtocol {
         reply()
         NotificationCenter.default.post(name: .init(name), object: nil)
     }
+    
+    public func quit(reply: @escaping () -> Void) {
+        Task {
+            await Service.shared.prepareForExit()
+            reply()
+        }
+    }
 
     // MARK: - Requests
 

Core/Sources/UpdateChecker/UpdateChecker.swift

@@ -5,15 +5,25 @@ import Sparkle
 public final class UpdateChecker {
     let updater: SPUUpdater
     let hostBundleFound: Bool
-    let delegate = UpdaterDelegate()
+    let delegate: UpdaterDelegate
+    public weak var updateCheckerDelegate: UpdateCheckerDelegate? {
+        get { delegate.updateCheckerDelegate }
+        set { delegate.updateCheckerDelegate = newValue }
+    }
 
-    public init(hostBundle: Bundle?) {
+    public init(
+        hostBundle: Bundle?,
+        shouldAutomaticallyCheckForUpdate: Bool
+    ) {
         if hostBundle == nil {
             hostBundleFound = false
             Logger.updateChecker.error("Host bundle not found")
         } else {
             hostBundleFound = true
         }
+        delegate = .init(
+            shouldAutomaticallyCheckForUpdate: shouldAutomaticallyCheckForUpdate
+        )
         updater = SPUUpdater(
             hostBundle: hostBundle ?? Bundle.main,
             applicationBundle: Bundle.main,
@@ -37,7 +47,32 @@ public final class UpdateChecker {
     }
 }
 
+public protocol UpdateCheckerDelegate: AnyObject {
+    func prepareForRelaunch(finish: @escaping () -> Void)
+}
+
 class UpdaterDelegate: NSObject, SPUUpdaterDelegate {
+    let shouldAutomaticallyCheckForUpdate: Bool
+    weak var updateCheckerDelegate: UpdateCheckerDelegate?
+
+    init(shouldAutomaticallyCheckForUpdate: Bool) {
+        self.shouldAutomaticallyCheckForUpdate = shouldAutomaticallyCheckForUpdate
+    }
+
+    func updater(_ updater: SPUUpdater, mayPerform updateCheck: SPUUpdateCheck) throws {
+        if !shouldAutomaticallyCheckForUpdate, updateCheck == .updatesInBackground {
+            throw CancellationError()
+        }
+    }
+    
+    func updater(_ updater: SPUUpdater, shouldPostponeRelaunchForUpdate item: SUAppcastItem, untilInvokingBlock installHandler: @escaping () -> Void) -> Bool {
+        if let updateCheckerDelegate {
+            updateCheckerDelegate.prepareForRelaunch(finish: installHandler)
+            return true
+        }
+        return false
+    }
+
     func allowedChannels(for updater: SPUUpdater) -> Set<String> {
         if UserDefaults.shared.value(for: \.installBetaBuilds) {
             Set(["beta"])

ExtensionService/AppDelegate.swift

@@ -21,15 +21,16 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
     let service = Service.shared
     var statusBarItem: NSStatusItem!
     var xpcController: XPCController?
-    let updateChecker =
-        UpdateChecker(
-            hostBundle: locateHostBundleURL(url: Bundle.main.bundleURL)
-                .flatMap(Bundle.init(url:))
-        )
+    let updateChecker = UpdateChecker(
+        hostBundle: locateHostBundleURL(url: Bundle.main.bundleURL)
+            .flatMap(Bundle.init(url:)),
+        shouldAutomaticallyCheckForUpdate: true
+    )
 
     func applicationDidFinishLaunching(_: Notification) {
         if ProcessInfo.processInfo.environment["IS_UNIT_TEST"] == "YES" { return }
         _ = XcodeInspector.shared
+        updateChecker.updateCheckerDelegate = self
         service.start()
         AXIsProcessTrustedWithOptions([
             kAXTrustedCheckOptionPrompt.takeRetainedValue() as NSString: true,
@@ -138,6 +139,15 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
     }
 }
 
+extension AppDelegate: UpdateCheckerDelegate {
+    func prepareForRelaunch(finish: @escaping () -> Void) {
+        Task {
+            await service.prepareForExit()
+            finish()
+        }
+    }
+}
+
 extension NSRunningApplication {
     var isUserOfService: Bool {
         [

Pro

@@ -161,6 +161,15 @@ public class XPCExtensionService {
             { service in { service.customCommand(id: id, editorContent: $0, withReply: $1) } }
         )
     }
+    
+    public func quitService() async throws {
+        try await withXPCServiceConnectedWithoutLaunching {
+            service, continuation in
+            service.quit {
+                continuation.resume(())
+            }
+        }
+    }
 
     public func postNotification(name: String) async throws {
         try await withXPCServiceConnected {
@@ -254,6 +263,20 @@ extension XPCExtensionService {
             }
         }
     }
+    
+    @XPCServiceActor
+    private func withXPCServiceConnectedWithoutLaunching<T>(
+        _ fn: @escaping (XPCServiceProtocol, AutoFinishContinuation<T>) -> Void
+    ) async throws -> T {
+        if let service, let connection = service.connection {
+            do {
+                return try await XPCShared.withXPCServiceConnected(connection: connection, fn)
+            } catch {
+                throw XPCExtensionServiceError.xpcServiceError(error)
+            }
+        }
+        throw XPCExtensionServiceError.failedToCreateXPCConnection
+    }
 
     @XPCServiceActor
     private func suggestionRequest(

Tool/Sources/XPCShared/XPCExtensionService.swift

@@ -56,6 +56,7 @@ public protocol XPCServiceProtocol {
     func getXPCServiceAccessibilityPermission(withReply reply: @escaping (Bool) -> Void)
     func postNotification(name: String, withReply reply: @escaping () -> Void)
     func send(endpoint: String, requestBody: Data, reply: @escaping (Data?, Error?) -> Void)
+    func quit(reply: @escaping () -> Void)
 }
 
 public struct NoResponse: Codable {