Fix that realtime suggestion is cancelling itself

Author: intitni

Core/Sources/Service/SuggestionCommandHandler/CommentBaseCommandHandler.swift

@@ -6,10 +6,10 @@ import XPCShared
 @ServiceActor
 struct CommentBaseCommandHandler: SuggestionCommandHanlder {
     nonisolated init() {}
-    
+
     func presentSuggestions(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         try await workspace.generateSuggestions(
             forFileAt: fileURL,
             content: editor.content,
@@ -33,7 +33,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
 
     func presentNextSuggestion(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         workspace.selectNextSuggestion(
             forFileAt: fileURL,
             content: editor.content,
@@ -53,7 +53,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
 
     func presentPreviousSuggestion(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         workspace.selectPreviousSuggestion(
             forFileAt: fileURL,
             content: editor.content,
@@ -73,7 +73,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
 
     func rejectSuggestion(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         workspace.rejectSuggestion(forFileAt: fileURL)
 
         guard let filespace = workspace.filespaces[fileURL] else { return nil }
@@ -89,7 +89,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
 
     func acceptSuggestion(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
 
         guard let acceptedSuggestion = workspace.acceptSuggestion(forFileAt: fileURL)
         else { return nil }
@@ -110,7 +110,6 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
             extraInfo: &extraInfo
         )
 
-        let presenter = PresentInCommentSuggestionPresenter()
         return .init(
             content: String(lines.joined(separator: "")),
             newCursor: cursorPosition,
@@ -120,8 +119,8 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
 
     func presentRealtimeSuggestions(editor: EditorContent) async throws -> UpdatedContent? {
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
-        guard let filespace = workspace.filespaces[fileURL] else { return nil }
+        let (workspace, filespace) = try await Workspace
+            .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
 
         try Task.checkCancellation()
 
@@ -147,8 +146,8 @@ struct CommentBaseCommandHandler: SuggestionCommandHanlder {
         // We don't need to wait for this.
         Task { @ServiceActor in
             let fileURL = try await Environment.fetchCurrentFileURL()
-            let workspace = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
-            guard let filespace = workspace.filespaces[fileURL] else { return }
+            let (workspace, filespace) = try await Workspace
+                .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
 
             try Task.checkCancellation()
 

Core/Sources/Service/Workspace.swift

@@ -72,12 +72,16 @@ final class Workspace {
         return false
     }
 
-    static func fetchOrCreateWorkspaceIfNeeded(fileURL: URL) async throws -> Workspace {
+    static func fetchOrCreateWorkspaceIfNeeded(fileURL: URL) async throws -> (workspace: Workspace, filespace: Filespace) {
         let projectURL = try await Environment.fetchCurrentProjectRootURL(fileURL)
         let workspaceURL = projectURL ?? fileURL
         let workspace = workspaces[workspaceURL] ?? Workspace(projectRootURL: workspaceURL)
+        let filespace = workspace.filespaces[fileURL] ?? .init(fileURL: fileURL)
+        if workspace.filespaces[fileURL] == nil {
+            workspace.filespaces[fileURL] = filespace
+        }
         workspaces[workspaceURL] = workspace
-        return workspace
+        return (workspace, filespace)
     }
 }
 

Core/Sources/Service/XPCService.swift

@@ -93,24 +93,30 @@ public class XPCService: NSObject, XPCServiceProtocol {
     @discardableResult
     private func replyWithUpdatedContent(
         editorContent: Data,
+        file: StaticString = #file,
+        line: UInt = #line,
+        cancelInFlightTasks: Bool = true,
         withReply reply: @escaping (Data?, Error?) -> Void,
         getUpdatedContent: @escaping @ServiceActor (
             SuggestionCommandHanlder,
             EditorContent
         ) async throws -> UpdatedContent?
     ) -> Task<Void, Never> {
         Task {
-            await RealtimeSuggestionController.shared.cancelInFlightTasksAndIgnoreTriggerForAWhile()
+            if cancelInFlightTasks {
+                await RealtimeSuggestionController.shared
+                    .cancelInFlightTasksAndIgnoreTriggerForAWhile()
+            }
             do {
                 let editor = try JSONDecoder().decode(EditorContent.self, from: editorContent)
-                let handler = await CommentBaseCommandHandler()
+                let handler = CommentBaseCommandHandler()
                 guard let updatedContent = try await getUpdatedContent(handler, editor) else {
                     reply(nil, nil)
                     return
                 }
                 reply(try JSONEncoder().encode(updatedContent), nil)
             } catch {
-                os_log(.error, "%@", error.localizedDescription)
+                os_log(.error, "%@", "\(file):\(line) \(error.localizedDescription)")
                 reply(nil, NSError.from(error))
             }
         }
@@ -167,6 +173,7 @@ public class XPCService: NSObject, XPCServiceProtocol {
     ) {
         let task = replyWithUpdatedContent(
             editorContent: editorContent,
+            cancelInFlightTasks: false,
             withReply: reply
         ) { handler, editor in
             try await handler.presentRealtimeSuggestions(editor: editor)
@@ -179,12 +186,9 @@ public class XPCService: NSObject, XPCServiceProtocol {
         editorContent: Data,
         withReply reply: @escaping () -> Void
     ) {
-        // prefetching never wait.
-        reply()
-
         let task = replyWithUpdatedContent(
             editorContent: editorContent,
-            withReply: { _, _ in }
+            withReply: { _, _ in reply() }
         ) { handler, editor in
             try await handler.generateRealtimeSuggestions(editor: editor)
         }