Implement cancellation in SuggestionService

Author: intitni

Core/Sources/CodeiumService/CodeiumService.swift

@@ -18,6 +18,7 @@ public protocol CodeiumSuggestionServiceType {
     func notifyOpenTextDocument(fileURL: URL, content: String) async throws
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
+    func cancelRequest() async
 }
 
 enum CodeiumError: Error, LocalizedError {
@@ -43,6 +44,7 @@ public class CodeiumSuggestionService {
     var server: CodeiumLSP?
     var heartbeatTask: Task<Void, Error>?
     var requestCounter: UInt64 = 0
+    var cancellationCounter: UInt64 = 0
     let openedDocumentPool = OpenedDocumentPool()
     let onServiceLaunched: () -> Void
 
@@ -118,6 +120,7 @@ public class CodeiumSuggestionService {
             self?.server = nil
             self?.heartbeatTask?.cancel()
             self?.requestCounter = 0
+            self?.cancellationCounter = 0
             Logger.codeium.info("Language server is terminated, will be restarted when needed.")
         }
 
@@ -253,8 +256,16 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
                     )
                 }
         ))
+        
+        if request.requestBody.metadata.request_id <= cancellationCounter {
+            throw CancellationError()
+        }
 
         let result = try await (try await setupServerIfNeeded()).sendRequest(request)
+        
+        if request.requestBody.metadata.request_id <= cancellationCounter {
+            throw CancellationError()
+        }
 
         return result.completionItems?.filter { item in
             if ignoreSpaceOnlySuggestions {
@@ -280,6 +291,10 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
             )
         } ?? []
     }
+    
+    public func cancelRequest() async {
+        cancellationCounter = requestCounter
+    }
 
     public func notifyAccepted(_ suggestion: CodeSuggestion) async {
         _ = try? await (try setupServerIfNeeded())

Core/Sources/GitHubCopilotService/GitHubCopilotService.swift

@@ -31,6 +31,7 @@ public protocol GitHubCopilotSuggestionServiceType {
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
+    func cancelRequest() async
 }
 
 protocol GitHubCopilotLSP {
@@ -311,6 +312,10 @@ public final class GitHubCopilotSuggestionService: GitHubCopilotBaseService,
 
         return try await task.value
     }
+    
+    public func cancelRequest() async {
+        await localProcessServer?.cancelOngoingTasks()
+    }
 
     public func notifyAccepted(_ completion: CodeSuggestion) async {
         _ = try? await server.sendRequest(

Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift

@@ -70,5 +70,10 @@ extension CodeiumSuggestionProvider {
     }
 
     func notifySaveTextDocument(fileURL: URL) async throws {}
+    
+    func cancelRequest() async {
+        await (try? createCodeiumServiceIfNeeded())?
+            .cancelRequest()
+    }
 }
 

Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift

@@ -73,5 +73,10 @@ extension GitHubCopilotSuggestionProvider {
         try await (try? createGitHubCopilotServiceIfNeeded())?
             .notifySaveTextDocument(fileURL: fileURL)
     }
+    
+    func cancelRequest() async {
+        await (try? createGitHubCopilotServiceIfNeeded())?
+            .cancelRequest()
+    }
 }
 

Core/Sources/SuggestionService/SuggestionService.swift

@@ -20,6 +20,7 @@ public protocol SuggestionServiceType {
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
+    func cancelRequest() async
 }
 
 protocol SuggestionServiceProvider: SuggestionServiceType {}
@@ -116,5 +117,9 @@ public extension SuggestionService {
     func notifySaveTextDocument(fileURL: URL) async throws {
         try await suggestionProvider.notifySaveTextDocument(fileURL: fileURL)
     }
+    
+    func cancelRequest() async {
+        await suggestionProvider.cancelRequest()
+    }
 }