Support terminating suggestion services on normal exits

Author: intitni

Core/Sources/CodeiumService/CodeiumLanguageServer.swift

@@ -7,6 +7,7 @@ import Preferences
 
 protocol CodeiumLSP {
     func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
+    func terminate()
 }
 
 final class CodeiumLanguageServer {
@@ -119,6 +120,14 @@ final class CodeiumLanguageServer {
         self.port = port
         launchHandler?()
     }
+    
+    func terminate() {
+        process.terminationHandler = nil
+        if process.isRunning {
+            process.terminate()
+        }
+        transport.close()
+    }
 }
 
 extension CodeiumLanguageServer: CodeiumLSP {

Core/Sources/CodeiumService/CodeiumService.swift

@@ -19,6 +19,7 @@ public protocol CodeiumSuggestionServiceType {
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate()
 }
 
 enum CodeiumError: Error, LocalizedError {
@@ -325,6 +326,11 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
     public func notifyCloseTextDocument(fileURL: URL) async throws {
         await openedDocumentPool.closeDocument(url: fileURL)
     }
+    
+    public func terminate() {
+        server?.terminate()
+        server = nil
+    }
 }
 
 func getXcodeVersion() async throws -> String {

Core/Sources/GitHubCopilotService/GitHubCopilotService.swift

@@ -32,6 +32,7 @@ public protocol GitHubCopilotSuggestionServiceType {
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate() async
 }
 
 protocol GitHubCopilotLSP {
@@ -379,6 +380,10 @@ public final class GitHubCopilotSuggestionService: GitHubCopilotBaseService,
 //        Logger.service.debug("Close \(uri)")
         try await server.sendNotification(.didCloseTextDocument(.init(uri: uri)))
     }
+    
+    public func terminate() async {
+        // automatically handled
+    }
 }
 
 extension InitializingServer: GitHubCopilotLSP {

Core/Sources/Service/ScheduledCleaner.swift

@@ -70,5 +70,12 @@ public final class ScheduledCleaner {
             }
         }
     }
+    
+    @ServiceActor
+    public func closeAllChildProcesses() async {
+        for (_, workspace) in workspaces {
+            await workspace.terminateSuggestionService()
+        }
+    }
 }
 

Core/Sources/Service/Workspace.swift

@@ -428,4 +428,8 @@ extension Workspace {
         guard let suggestionService else { return }
         await suggestionService.cancelRequest()
     }
+    
+    func terminateSuggestionService() async {
+        await _suggestionService?.terminate()
+    }
 }

Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift

@@ -75,5 +75,9 @@ extension CodeiumSuggestionProvider {
         await (try? createCodeiumServiceIfNeeded())?
             .cancelRequest()
     }
+    
+    func terminate() async {
+        (try? createCodeiumServiceIfNeeded())?.terminate()
+    }
 }
 

Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift

@@ -78,5 +78,9 @@ extension GitHubCopilotSuggestionProvider {
         await (try? createGitHubCopilotServiceIfNeeded())?
             .cancelRequest()
     }
+    
+    func terminate() async {
+        await (try? createGitHubCopilotServiceIfNeeded())?.terminate()
+    }
 }
 

Core/Sources/SuggestionService/SuggestionService.swift

@@ -1,3 +1,4 @@
+import AppKit
 import Foundation
 import Preferences
 import SuggestionModel
@@ -21,6 +22,7 @@ public protocol SuggestionServiceType {
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate() async
 }
 
 protocol SuggestionServiceProvider: SuggestionServiceType {}
@@ -117,9 +119,13 @@ public extension SuggestionService {
     func notifySaveTextDocument(fileURL: URL) async throws {
         try await suggestionProvider.notifySaveTextDocument(fileURL: fileURL)
     }
-    
+
     func cancelRequest() async {
         await suggestionProvider.cancelRequest()
     }
+
+    func terminate() async {
+        await suggestionProvider.terminate()
+    }
 }
 

ExtensionService/AppDelegate.swift

@@ -34,7 +34,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
         _ = RealtimeSuggestionController.shared
         _ = XcodeInspector.shared
         AXIsProcessTrustedWithOptions([
-            kAXTrustedCheckOptionPrompt.takeRetainedValue() as NSString: true
+            kAXTrustedCheckOptionPrompt.takeRetainedValue() as NSString: true,
         ] as CFDictionary)
         setupQuitOnUpdate()
         setupQuitOnUserTerminated()
@@ -105,7 +105,10 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
     }
 
     @objc func quit() {
-        exit(0)
+        Task { @MainActor in
+            await scheduledCleaner.closeAllChildProcesses()
+            exit(0)
+        }
     }
 
     @objc func openCopilotForXcode() {
@@ -148,7 +151,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
                 Logger.service.info("Extension Service will quit.")
                 #if DEBUG
                 #else
-                exit(0)
+                quit()
                 #endif
             }
         }
@@ -172,7 +175,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
                 if NSWorkspace.shared.runningApplications.contains(where: \.isUserOfService) {
                     continue
                 }
-                exit(0)
+                quit()
             }
         }
     }