Make generateRealtimeSuggestions almost always work

Author: intitni

Core/Sources/AXExtension/AXUIElement.swift

@@ -19,6 +19,10 @@ public extension AXUIElement {
     var description: String {
         (try? copyValue(key: kAXDescriptionAttribute)) ?? ""
     }
+    
+    var isSourceEditor: Bool {
+        description == "Source Editor"
+    }
 
     var selectedTextRange: Range<Int>? {
         guard let value: AXValue = try? copyValue(key: kAXSelectedTextRangeAttribute)

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -25,6 +25,7 @@ public class RealtimeSuggestionController {
     private var activeApplicationMonitorTask: Task<Void, Error>?
     private var editorObservationTask: Task<Void, Error>?
     private var focusedUIElement: AXUIElement?
+    private var sourceEditor: AXUIElement?
 
     var isCommentMode: Bool {
         UserDefaults.shared.value(for: \.suggestionPresentationMode) == .comment
@@ -108,8 +109,9 @@ public class RealtimeSuggestionController {
         let application = AXUIElementCreateApplication(activeXcode.processIdentifier)
         guard let focusElement = application.focusedElement else { return }
         let focusElementType = focusElement.description
-        guard focusElementType == "Source Editor" else { return }
         focusedUIElement = focusElement
+        guard focusElementType == "Source Editor" else { return }
+        sourceEditor = focusElement
 
         editorObservationTask?.cancel()
         editorObservationTask = nil
@@ -216,7 +218,7 @@ public class RealtimeSuggestionController {
             }
 
             // So the editor won't be blocked (after information are cached)!
-            await PseudoCommandHandler().generateRealtimeSuggestions()
+            await PseudoCommandHandler().generateRealtimeSuggestions(sourceEditor: sourceEditor)
         }
     }
 

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -38,13 +38,10 @@ struct PseudoCommandHandler {
         ))
     }
 
-    func generateRealtimeSuggestions() async {
-        // Can't use handler directly if content is not available.
-        guard let editor = await getEditorContent() else {
-            try? await Environment.triggerAction("Prefetch Suggestions")
-            return
-        }
-        
+    func generateRealtimeSuggestions(sourceEditor: AXUIElement?) async {
+        // Can't use handler if content is not available.
+        guard let editor = await getEditorContent(sourceEditor: sourceEditor) else { return }
+
         // Otherwise, get it from pseudo handler directly.
         let mode = UserDefaults.shared.value(for: \.suggestionPresentationMode)
         switch mode {
@@ -78,7 +75,8 @@ struct PseudoCommandHandler {
             guard let focusElement = application.focusedElement,
                   focusElement.description == "Source Editor"
             else { return }
-            guard let (content, lines, cursorPosition) = await getFileContent() else {
+            guard let (content, lines, cursorPosition) = await getFileContent(sourceEditor: nil)
+            else {
                 PresentInWindowSuggestionPresenter()
                     .presentErrorMessage("Unable to get file content.")
                 return
@@ -147,13 +145,13 @@ struct PseudoCommandHandler {
 }
 
 private extension PseudoCommandHandler {
-    func getFileContent() async
+    func getFileContent(sourceEditor: AXUIElement?) async
         -> (content: String, lines: [String], cursorPosition: CursorPosition)?
     {
         guard let xcode = ActiveApplicationMonitor.activeXcode
             ?? ActiveApplicationMonitor.latestXcode else { return nil }
         let application = AXUIElementCreateApplication(xcode.processIdentifier)
-        guard let focusElement = application.focusedElement,
+        guard let focusElement = sourceEditor ?? application.focusedElement,
               focusElement.description == "Source Editor"
         else { return nil }
         guard let selectionRange = focusElement.selectedTextRange else { return nil }
@@ -190,10 +188,10 @@ private extension PseudoCommandHandler {
     }
 
     @ServiceActor
-    func getEditorContent() async -> EditorContent? {
+    func getEditorContent(sourceEditor: AXUIElement?) async -> EditorContent? {
         guard
             let filespace = await getFilespace(),
-            let content = await getFileContent()
+            let content = await getFileContent(sourceEditor: sourceEditor)
         else { return nil }
         let uti = filespace.uti ?? ""
         let tabSize = filespace.tabSize ?? 4