Support request cancellation

Author: intitni

Core/Sources/GitHubCopilotService/CopilotLocalProcessServer.swift

@@ -9,9 +9,11 @@ import ProcessEnv
 /// We need it because the original one does not allow us to handle custom notifications.
 class CopilotLocalProcessServer {
     private let transport: StdioDataTransport
+    private let customTransport: CustomDataTransport
     private let process: Process
     private var wrappedServer: CustomJSONRPCLanguageServer?
     var terminationHandler: (() -> Void)?
+    @MainActor var ongoingCompletionRequestIDs: [JSONId] = []
 
     public convenience init(
         path: String,
@@ -29,10 +31,27 @@ class CopilotLocalProcessServer {
 
     init(executionParameters parameters: Process.ExecutionParameters) {
         transport = StdioDataTransport()
-        wrappedServer = CustomJSONRPCLanguageServer(dataTransport: transport)
+        let framing = SeperatedHTTPHeaderMessageFraming()
+        let messageTransport = MessageTransport(
+            dataTransport: transport,
+            messageProtocol: framing
+        )
+        customTransport = CustomDataTransport(nextTransport: messageTransport)
+        wrappedServer = CustomJSONRPCLanguageServer(dataTransport: customTransport)
 
         process = Process()
 
+        // Because the implementation of LanguageClient is so closed,
+        // we need to get the request IDs from a custom transport before the data
+        // is written to the language server.
+        customTransport.onWriteRequest = { [weak self] request in
+            if request.method == "getCompletionsCycling" {
+                Task { @MainActor [weak self] in
+                    self?.ongoingCompletionRequestIDs.append(request.id)
+                }
+            }
+        }
+
         process.standardInput = transport.stdinPipe
         process.standardOutput = transport.stdoutPipe
         process.standardError = transport.stderrPipe
@@ -89,6 +108,27 @@ extension CopilotLocalProcessServer: LanguageServerProtocol.Server {
 
         server.sendNotification(notif, completionHandler: completionHandler)
     }
+    
+    /// Cancel ongoing completion requests.
+    public func cancelOngoingTasks() async {
+        guard let server = wrappedServer, process.isRunning else {
+            return
+        }
+        
+        let task = Task { @MainActor in
+            for id in self.ongoingCompletionRequestIDs {
+                switch id {
+                case let .numericId(id):
+                    try? await server.sendNotification(.protocolCancelRequest(.init(id: id)))
+                case let .stringId(id):
+                    try? await server.sendNotification(.protocolCancelRequest(.init(id: id)))
+                }
+            }
+            self.ongoingCompletionRequestIDs = []
+        }
+        
+        await task.value
+    }
 
     public func sendRequest<Response: Codable>(
         _ request: ClientRequest,
@@ -139,13 +179,7 @@ final class CustomJSONRPCLanguageServer: Server {
     }
 
     convenience init(dataTransport: DataTransport) {
-        let framing = SeperatedHTTPHeaderMessageFraming()
-        let messageTransport = MessageTransport(
-            dataTransport: dataTransport,
-            messageProtocol: framing
-        )
-
-        self.init(protocolTransport: ProtocolTransport(dataTransport: messageTransport))
+        self.init(protocolTransport: ProtocolTransport(dataTransport: dataTransport))
     }
 
     deinit {
@@ -219,3 +253,4 @@ extension CustomJSONRPCLanguageServer {
         internalServer.sendRequest(request, completionHandler: completionHandler)
     }
 }
+

Core/Sources/GitHubCopilotService/CustomStdioTransport.swift

@@ -0,0 +1,30 @@
+import Foundation
+import JSONRPC
+import os.log
+
+public class CustomDataTransport: DataTransport {
+    let nextTransport: DataTransport
+    
+    var onWriteRequest: (JSONRPCRequest<JSONValue>) -> Void = { _ in }
+    
+    init(nextTransport: DataTransport) {
+        self.nextTransport = nextTransport
+    }
+    
+    public func write(_ data: Data) {
+        if let request = try? JSONDecoder().decode(JSONRPCRequest<JSONValue>.self, from: data) {
+            onWriteRequest(request)
+        }
+        
+        nextTransport.write(data)
+    }
+    
+    public func setReaderHandler(_ handler: @escaping ReadHandler) {
+        nextTransport.setReaderHandler(handler)
+    }
+    
+    public func close() {
+        nextTransport.close()
+    }
+}
+

Core/Sources/GitHubCopilotService/GitHubCopilotService.swift

@@ -52,6 +52,7 @@ enum GitHubCopilotError: Error, LocalizedError {
 public class GitHubCopilotBaseService {
     let projectRootURL: URL
     var server: GitHubCopilotLSP
+    var localProcessServer: CopilotLocalProcessServer?
 
     init(designatedServer: GitHubCopilotLSP) {
         projectRootURL = URL(fileURLWithPath: "/")
@@ -60,7 +61,7 @@ public class GitHubCopilotBaseService {
 
     init(projectRootURL: URL) throws {
         self.projectRootURL = projectRootURL
-        server = try {
+        let (server, localServer) = try {
             let urls = try GitHubCopilotBaseService.createFoldersIfNeeded()
             var userEnvPath = ProcessInfo.processInfo.userEnvironment["PATH"] ?? ""
             if userEnvPath.isEmpty {
@@ -124,6 +125,7 @@ public class GitHubCopilotBaseService {
                 }()
             }
             let localServer = CopilotLocalProcessServer(executionParameters: executionParams)
+            
             localServer.logMessages = UserDefaults.shared.value(for: \.gitHubCopilotVerboseLog)
             localServer.notificationHandler = { _, respond in
                 respond(.timeout)
@@ -156,8 +158,11 @@ public class GitHubCopilotBaseService {
                 )
             }
 
-            return server
+            return (server, localServer)
         }()
+        
+        self.server = server
+        self.localProcessServer = localServer
     }
 
     public static func createFoldersIfNeeded() throws -> (
@@ -238,6 +243,8 @@ public final class GitHubCopilotAuthService: GitHubCopilotBaseService,
 public final class GitHubCopilotSuggestionService: GitHubCopilotBaseService,
     GitHubCopilotSuggestionServiceType
 {
+    private var ongoingTasks = Set<Task<[CodeSuggestion], Error>>()
+
     override public init(projectRootURL: URL = URL(fileURLWithPath: "/")) throws {
         try super.init(projectRootURL: projectRootURL)
     }
@@ -272,27 +279,37 @@ public final class GitHubCopilotSuggestionService: GitHubCopilotBaseService,
             return filePath
         }()
 
-        let completions = try await server
-            .sendRequest(GitHubCopilotRequest.GetCompletionsCycling(doc: .init(
-                source: content,
-                tabSize: tabSize,
-                indentSize: indentSize,
-                insertSpaces: !usesTabsForIndentation,
-                path: fileURL.path,
-                uri: fileURL.path,
-                relativePath: relativePath,
-                languageId: languageId,
-                position: cursorPosition
-            )))
-            .completions
-            .filter { completion in
-                if ignoreSpaceOnlySuggestions {
-                    return !completion.text.allSatisfy { $0.isWhitespace || $0.isNewline }
+        ongoingTasks.forEach { $0.cancel() }
+        ongoingTasks.removeAll()
+        await localProcessServer?.cancelOngoingTasks()
+
+        let task = Task {
+            let completions = try await server
+                .sendRequest(GitHubCopilotRequest.GetCompletionsCycling(doc: .init(
+                    source: content,
+                    tabSize: tabSize,
+                    indentSize: indentSize,
+                    insertSpaces: !usesTabsForIndentation,
+                    path: fileURL.path,
+                    uri: fileURL.path,
+                    relativePath: relativePath,
+                    languageId: languageId,
+                    position: cursorPosition
+                )))
+                .completions
+                .filter { completion in
+                    if ignoreSpaceOnlySuggestions {
+                        return !completion.text.allSatisfy { $0.isWhitespace || $0.isNewline }
+                    }
+                    return true
                 }
-                return true
-            }
+            try Task.checkCancellation()
+            return completions
+        }
+
+        ongoingTasks.insert(task)
 
-        return completions
+        return try await task.value
     }
 
     public func notifyAccepted(_ completion: CodeSuggestion) async {