Merge branch 'feature/real-time-suggestion-tweak' into develop

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.")
         }
 
@@ -227,7 +230,7 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
 
         let relativePath = getRelativePath(of: fileURL)
 
-        let request = CodeiumRequest.GetCompletion(requestBody: .init(
+        let request = await CodeiumRequest.GetCompletion(requestBody: .init(
             metadata: try getMetadata(),
             document: .init(
                 absolute_path: fileURL.path,
@@ -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())
@@ -291,7 +306,7 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
 
     public func notifyOpenTextDocument(fileURL: URL, content: String) async throws {
         let relativePath = getRelativePath(of: fileURL)
-        openedDocumentPool.openDocument(
+        await openedDocumentPool.openDocument(
             url: fileURL,
             relativePath: relativePath,
             content: content
@@ -300,15 +315,15 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
 
     public func notifyChangeTextDocument(fileURL: URL, content: String) async throws {
         let relativePath = getRelativePath(of: fileURL)
-        openedDocumentPool.updateDocument(
+        await openedDocumentPool.updateDocument(
             url: fileURL,
             relativePath: relativePath,
             content: content
         )
     }
 
     public func notifyCloseTextDocument(fileURL: URL) async throws {
-        openedDocumentPool.closeDocument(url: fileURL)
+        await openedDocumentPool.closeDocument(url: fileURL)
     }
 }
 

Core/Sources/CodeiumService/OpendDocumentPool.swift

@@ -2,7 +2,7 @@ import Foundation
 
 private let maxSize: Int = 1_000_000 // Byte
 
-final class OpenedDocumentPool {
+actor OpenedDocumentPool {
     var openedDocuments = [URL: OpenedDocument]()
 
     func getOtherDocuments(exceptURL: URL) -> [OpenedDocument] {

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/Service/RealtimeSuggestionController.swift

@@ -9,6 +9,7 @@ import Foundation
 import Logger
 import Preferences
 import QuartzCore
+import XcodeInspector
 
 @ServiceActor
 public class RealtimeSuggestionController {
@@ -25,7 +26,7 @@ public class RealtimeSuggestionController {
     private var activeApplicationMonitorTask: Task<Void, Error>?
     private var editorObservationTask: Task<Void, Error>?
     private var focusedUIElement: AXUIElement?
-    private var sourceEditor: AXUIElement?
+    private var sourceEditor: SourceEditor?
 
     var isCommentMode: Bool {
         UserDefaults.shared.value(for: \.suggestionPresentationMode) == .comment
@@ -48,9 +49,6 @@ public class RealtimeSuggestionController {
                     await self.handleXcodeChanged(app)
                 }
 
-                #warning(
-                    "TODO: Is it possible to get rid of hid event observation with only AXObserver?"
-                )
                 if ActiveApplicationMonitor.activeXcode != nil {
                     await startHIDObservation(by: 1)
                 } else {
@@ -118,7 +116,7 @@ public class RealtimeSuggestionController {
         }
 
         guard focusElementType == "Source Editor" else { return }
-        sourceEditor = focusElement
+        sourceEditor = SourceEditor(runningApplication: activeXcode, element: focusElement)
 
         editorObservationTask?.cancel()
         editorObservationTask = nil
@@ -127,7 +125,7 @@ public class RealtimeSuggestionController {
             let notificationsFromEditor = AXNotificationStream(
                 app: activeXcode,
                 element: focusElement,
-                notificationNames: kAXValueChangedNotification
+                notificationNames: kAXValueChangedNotification, kAXSelectedTextChangedNotification
             )
 
             for await notification in notificationsFromEditor {
@@ -139,6 +137,10 @@ public class RealtimeSuggestionController {
                 case kAXValueChangedNotification:
                     self.triggerPrefetchDebounced()
                     await self.notifyEditingFileChange(editor: focusElement)
+                case kAXSelectedTextChangedNotification:
+                    guard let sourceEditor else { continue }
+                    await PseudoCommandHandler()
+                        .invalidateRealtimeSuggestionsIfNeeded(sourceEditor: sourceEditor)
                 default:
                     continue
                 }
@@ -169,29 +171,16 @@ public class RealtimeSuggestionController {
     func handleHIDEvent(event: CGEvent) async {
         guard await Environment.isXcodeActive() else { return }
 
-        // Mouse clicks should cancel in-flight tasks.
-        if [CGEventType.rightMouseDown, .leftMouseDown].contains(event.type) {
-            await cancelInFlightTasks()
-            return
-        }
-
         let keycode = Int(event.getIntegerValueField(.keyboardEventKeycode))
         let escape = 0x35
-        let arrowKeys = [0x7B, 0x7C, 0x7D, 0x7E]
-
-        // Arrow keys should cancel in-flight tasks.
-        if arrowKeys.contains(keycode) {
-            await cancelInFlightTasks()
-            return
-        }
 
         // Escape should cancel in-flight tasks.
         // Except that when the completion panel is presented, it should trigger prefetch instead.
         if keycode == escape {
             if event.type == .keyDown {
                 await cancelInFlightTasks()
             } else {
-                let task = Task {
+                Task {
                     #warning(
                         "TODO: Any method to avoid using AppleScript to check that completion panel is presented?"
                     )
@@ -200,7 +189,6 @@ public class RealtimeSuggestionController {
                         self.triggerPrefetchDebounced(force: true)
                     }
                 }
-                inflightRealtimeSuggestionsTasks.insert(task)
             }
         }
     }
@@ -247,18 +235,6 @@ public class RealtimeSuggestionController {
                     await workspace.cancelInFlightRealtimeSuggestionRequests()
                 }
             }
-            group.addTask {
-                await { @ServiceActor in
-                    inflightRealtimeSuggestionsTasks.forEach {
-                        if $0 == excluding { return }
-                        $0.cancel()
-                    }
-                    inflightRealtimeSuggestionsTasks.removeAll()
-                    if let excluded = excluding {
-                        inflightRealtimeSuggestionsTasks.insert(excluded)
-                    }
-                }()
-            }
         }
     }
 

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -1,9 +1,10 @@
 import ActiveApplicationMonitor
 import AppKit
-import SuggestionModel
 import Environment
 import Preferences
 import SuggestionInjector
+import SuggestionModel
+import XcodeInspector
 import XPCShared
 
 /// It's used to run some commands without really triggering the menu bar item.
@@ -38,9 +39,21 @@ struct PseudoCommandHandler {
         ))
     }
 
-    func generateRealtimeSuggestions(sourceEditor: AXUIElement?) async {
+    func generateRealtimeSuggestions(sourceEditor: SourceEditor?) async {
         // Can't use handler if content is not available.
-        guard let editor = await getEditorContent(sourceEditor: sourceEditor) else { return }
+        guard
+            let editor = await getEditorContent(sourceEditor: sourceEditor),
+            let filespace = await getFilespace()
+        else { return }
+
+        if await filespace.validateSuggestions(
+            lines: editor.lines,
+            cursorPosition: editor.cursorPosition
+        ) {
+            return
+        } else {
+            PresentInWindowSuggestionPresenter().discardSuggestion(fileURL: filespace.fileURL)
+        }
 
         // Otherwise, get it from pseudo handler directly.
         let mode = UserDefaults.shared.value(for: \.suggestionPresentationMode)
@@ -54,6 +67,19 @@ struct PseudoCommandHandler {
         }
     }
 
+    func invalidateRealtimeSuggestionsIfNeeded(sourceEditor: SourceEditor) async {
+        guard let fileURL = try? await Environment.fetchCurrentFileURL(),
+              let (_, filespace) = try? await Workspace
+              .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL) else { return }
+
+        if await !filespace.validateSuggestions(
+            lines: sourceEditor.content.lines,
+            cursorPosition: sourceEditor.content.cursorPosition
+        ) {
+            PresentInWindowSuggestionPresenter().discardSuggestion(fileURL: fileURL)
+        }
+    }
+
     func rejectSuggestions() async {
         let handler = WindowBaseCommandHandler()
         _ = try? await handler.rejectSuggestion(editor: .init(
@@ -214,7 +240,7 @@ extension PseudoCommandHandler {
         let content = focusElement.value
         let split = content.breakLines()
         let range = convertRangeToCursorRange(selectionRange, in: content)
-        return (content, split, [range], range.end)
+        return (content, split, [range], range.start)
     }
 
     func getFileURL() async -> URL? {
@@ -232,11 +258,9 @@ extension PseudoCommandHandler {
     }
 
     @ServiceActor
-    func getEditorContent(sourceEditor: AXUIElement?) async -> EditorContent? {
-        guard
-            let filespace = await getFilespace(),
-            let content = await getFileContent(sourceEditor: sourceEditor)
-        else { return nil }
+    func getEditorContent(sourceEditor: SourceEditor?) async -> EditorContent? {
+        guard let filespace = await getFilespace(), let sourceEditor else { return nil }
+        let content = sourceEditor.content
         let uti = filespace.uti ?? ""
         let tabSize = filespace.tabSize ?? 4
         let indentSize = filespace.indentSize ?? 4
@@ -289,10 +313,10 @@ extension PseudoCommandHandler {
         var countE = 0
         var cursorRange = CursorRange(start: .zero, end: .outOfScope)
         for (i, line) in lines.enumerated() {
-            if countS <= range.lowerBound && range.lowerBound < countS + line.count {
+            if countS <= range.lowerBound, range.lowerBound < countS + line.count {
                 cursorRange.start = .init(line: i, character: range.lowerBound - countS)
             }
-            if countE <= range.upperBound && range.upperBound < countE + line.count {
+            if countE <= range.upperBound, range.upperBound < countE + line.count {
                 cursorRange.end = .init(line: i, character: range.upperBound - countE)
                 break
             }
@@ -321,3 +345,4 @@ public extension String {
         return all
     }
 }
+