Adjust real-time suggestion task cancellation

Author: intitni

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -141,7 +141,7 @@ public actor RealtimeSuggestionController {
         }
     }
 
-    func cancelInFlightTasks() async {
+    func cancelInFlightTasks(excluding: Task<Void, Never>? = nil) async {
         inflightPrefetchTask?.cancel()
 
         // cancel in-flight tasks
@@ -153,17 +153,24 @@ public actor RealtimeSuggestionController {
             }
             group.addTask {
                 await { @ServiceActor in
-                    inflightRealtimeSuggestionsTasks.forEach { $0.cancel() }
+                    inflightRealtimeSuggestionsTasks.forEach {
+                        if $0 == excluding { return }
+                        $0.cancel()
+                    }
                     inflightRealtimeSuggestionsTasks.removeAll()
+                    if let excluded = excluding {
+                        inflightRealtimeSuggestionsTasks.insert(excluded)
+                    }
                 }()
             }
         }
     }
 
+    #warning("TODO: Find a better way to prevent that from happening!")
     /// Prevent prefetch to be triggered by commands. Quick and dirty.
-    func cancelInFlightTasksAndIgnoreTriggerForAWhile() async {
+    func cancelInFlightTasksAndIgnoreTriggerForAWhile(excluding: Task<Void, Never>? = nil) async {
         ignoreUntil = Date(timeIntervalSinceNow: 5)
-        await cancelInFlightTasks()
+        await cancelInFlightTasks(excluding: excluding)
     }
 }
 

Core/Sources/Service/SuggestionCommandHandler/CommentBaseCommandHandler.swift

@@ -143,38 +143,35 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
     }
 
     func generateRealtimeSuggestions(editor: EditorContent) async throws -> UpdatedContent? {
-        // We don't need to wait for this.
-        Task { @ServiceActor in
-            let fileURL = try await Environment.fetchCurrentFileURL()
-            let (workspace, filespace) = try await Workspace
-                .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
-
-            try Task.checkCancellation()
-
-            let snapshot = Filespace.Snapshot(
-                linesHash: editor.lines.hashValue,
-                cursorPosition: editor.cursorPosition
-            )
-
-            // There is no need to regenerate suggestions for the same editor content.
-            guard filespace.suggestionSourceSnapshot != snapshot else { return }
-
-            let suggestions = try await workspace.generateSuggestions(
-                forFileAt: fileURL,
-                content: editor.content,
-                lines: editor.lines,
-                cursorPosition: editor.cursorPosition,
-                tabSize: editor.tabSize,
-                indentSize: editor.indentSize,
-                usesTabsForIndentation: editor.usesTabsForIndentation
-            )
-
-            try Task.checkCancellation()
-
-            // If there is a suggestion available, call another command to present it.
-            guard !suggestions.isEmpty else { return }
-            try await Environment.triggerAction("Real-time Suggestions")
-        }
+        let fileURL = try await Environment.fetchCurrentFileURL()
+        let (workspace, filespace) = try await Workspace
+            .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+
+        try Task.checkCancellation()
+
+        let snapshot = Filespace.Snapshot(
+            linesHash: editor.lines.hashValue,
+            cursorPosition: editor.cursorPosition
+        )
+
+        // There is no need to regenerate suggestions for the same editor content.
+        guard filespace.suggestionSourceSnapshot != snapshot else { return nil }
+
+        let suggestions = try await workspace.generateSuggestions(
+            forFileAt: fileURL,
+            content: editor.content,
+            lines: editor.lines,
+            cursorPosition: editor.cursorPosition,
+            tabSize: editor.tabSize,
+            indentSize: editor.indentSize,
+            usesTabsForIndentation: editor.usesTabsForIndentation
+        )
+
+        try Task.checkCancellation()
+
+        // If there is a suggestion available, call another command to present it.
+        guard !suggestions.isEmpty else { return nil }
+        try await Environment.triggerAction("Real-time Suggestions")
 
         return nil
     }

Core/Sources/Service/XPCService.swift

@@ -12,6 +12,8 @@ import XPCShared
 
 @ServiceActor
 var workspaces = [URL: Workspace]()
+
+#warning("Todo: Find a better place to store it!")
 @ServiceActor
 var inflightRealtimeSuggestionsTasks = Set<Task<Void, Never>>()
 
@@ -95,18 +97,14 @@ public class XPCService: NSObject, XPCServiceProtocol {
         editorContent: Data,
         file: StaticString = #file,
         line: UInt = #line,
-        cancelInFlightTasks: Bool = true,
+        isRealtimeSuggestionRelatedCommand: Bool = false,
         withReply reply: @escaping (Data?, Error?) -> Void,
         getUpdatedContent: @escaping @ServiceActor (
             SuggestionCommandHanlder,
             EditorContent
         ) async throws -> UpdatedContent?
     ) -> Task<Void, Never> {
-        Task {
-            if cancelInFlightTasks {
-                await RealtimeSuggestionController.shared
-                    .cancelInFlightTasksAndIgnoreTriggerForAWhile()
-            }
+        let task = Task {
             do {
                 let editor = try JSONDecoder().decode(EditorContent.self, from: editorContent)
                 let handler = CommentBaseCommandHandler()
@@ -120,6 +118,17 @@ public class XPCService: NSObject, XPCServiceProtocol {
                 reply(nil, NSError.from(error))
             }
         }
+
+        Task {
+            if isRealtimeSuggestionRelatedCommand {
+                await RealtimeSuggestionController.shared
+                    .cancelInFlightTasks(excluding: task)
+            } else {
+                await RealtimeSuggestionController.shared
+                    .cancelInFlightTasksAndIgnoreTriggerForAWhile(excluding: task)
+            }
+        }
+        return task
     }
 
     public func getSuggestedCode(
@@ -173,7 +182,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
     ) {
         let task = replyWithUpdatedContent(
             editorContent: editorContent,
-            cancelInFlightTasks: false,
+            isRealtimeSuggestionRelatedCommand: true,
             withReply: reply
         ) { handler, editor in
             try await handler.presentRealtimeSuggestions(editor: editor)
@@ -186,9 +195,13 @@ public class XPCService: NSObject, XPCServiceProtocol {
         editorContent: Data,
         withReply reply: @escaping () -> Void
     ) {
+        // We don't need to wait for this.
+        reply()
+
         let task = replyWithUpdatedContent(
             editorContent: editorContent,
-            withReply: { _, _ in reply() }
+            isRealtimeSuggestionRelatedCommand: true,
+            withReply: { _, _ in }
         ) { handler, editor in
             try await handler.generateRealtimeSuggestions(editor: editor)
         }