Correctly handle custom command when it's not in source editor

Author: intitni

Core/Sources/Environment/Environment.swift

@@ -105,6 +105,24 @@ public enum Environment {
         }
         throw FailedToFetchFileURLError()
     }
+    
+    public static var fetchFocusedElementURI: () async throws -> URL = {
+        guard let xcode = ActiveApplicationMonitor.activeXcode
+            ?? ActiveApplicationMonitor.latestXcode
+        else {
+            throw FailedToFetchFileURLError()
+        }
+        
+        let application = AXUIElementCreateApplication(xcode.processIdentifier)
+        let focusedElement = application.focusedElement
+        if focusedElement?.description != "Source Editor" {
+            let window = application.focusedWindow
+            let id = window?.identifier.hashValue
+            return URL(fileURLWithPath: "/xcode-focused-element/\(id ?? 0)")
+        }
+        
+        return try await fetchCurrentFileURL()
+    }
 
     public static var createAuthService: () -> CopilotAuthServiceType = {
         CopilotAuthService()
@@ -148,6 +166,15 @@ public enum Environment {
                         .error("Trigger action \(name) failed: \(error.localizedDescription)")
                     throw error
                 }
+            } else {
+                struct CantRunCommand: Error, LocalizedError {
+                    let name: String
+                    var errorDescription: String? {
+                        "Can't run command \(name)."
+                    }
+                }
+                
+                throw CantRunCommand(name: name)
             }
         } else {
             /// check if menu is open, if not, click the menu item.

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -69,8 +69,26 @@ struct PseudoCommandHandler {
     }
 
     func handleCustomCommand(_ command: CustomCommand) async {
-        guard let editor = await getEditorContent(sourceEditor: nil)
-        else {
+        guard let editor = await {
+            if let it = await getEditorContent(sourceEditor: nil) {
+                return it
+            }
+            switch command.feature {
+            case .customChat:
+                return .init(
+                    content: "",
+                    lines: [],
+                    uti: "",
+                    cursorPosition: .outOfScope,
+                    selections: [],
+                    tabSize: 0,
+                    indentSize: 0,
+                    usesTabsForIndentation: false
+                )
+            case .chatWithSelection, .promptToCode:
+                return nil
+            }
+        }() else {
             do {
                 try await Environment.triggerAction(command.name)
             } catch {

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -297,8 +297,7 @@ extension WindowBaseCommandHandler {
                 name: command.name
             )
         case let .customChat(systemPrompt, prompt):
-            try await startChatWithSelection(
-                editor: editor,
+            try await startChat(
                 specifiedSystemPrompt: systemPrompt,
                 extraSystemPrompt: nil,
                 sendingMessageImmediately: prompt,
@@ -435,7 +434,7 @@ extension WindowBaseCommandHandler {
 
         Task {
             let customCommandPrefix = {
-                if let name { return "[\(name)]" }
+                if let name { return "[\(name)] " }
                 return ""
             }()
 
@@ -444,23 +443,23 @@ extension WindowBaseCommandHandler {
                     history.append(.init(
                         role: .assistant,
                         content: "",
-                        summary: "\(customCommandPrefix) System prompt is updated."
+                        summary: "\(customCommandPrefix)System prompt is updated."
                     ))
                 }
             } else if !code.isEmpty, let selection = editor.selections.last {
                 await chat.chatGPTService.mutateHistory { history in
                     history.append(.init(
                         role: .assistant,
                         content: "",
-                        summary: "\(customCommandPrefix) Chatting about selected code in `\(fileURL.lastPathComponent)` from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character)`.\nThe code will persist in the conversation."
+                        summary: "\(customCommandPrefix)Chatting about selected code in `\(fileURL.lastPathComponent)` from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character)`.\nThe code will persist in the conversation."
                     ))
                 }
             } else if !customCommandPrefix.isEmpty {
                 await chat.chatGPTService.mutateHistory { history in
                     history.append(.init(
                         role: .assistant,
                         content: "",
-                        summary: "\(customCommandPrefix) System prompt is updated."
+                        summary: "\(customCommandPrefix)System prompt is updated."
                     ))
                 }
             }
@@ -472,4 +471,63 @@ extension WindowBaseCommandHandler {
 
         presenter.presentChatRoom(fileURL: fileURL)
     }
+    
+    private func startChat(
+        specifiedSystemPrompt: String?,
+        extraSystemPrompt: String?,
+        sendingMessageImmediately: String?,
+        name: String?
+    ) async throws {
+        presenter.markAsProcessing(true)
+        defer { presenter.markAsProcessing(false) }
+
+        let focusedElementURI = try await Environment.fetchFocusedElementURI()
+        let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
+
+        var systemPrompt = specifiedSystemPrompt ?? """
+        \(language.isEmpty ? "" : "You must always reply in \(language)")
+        You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
+        
+        You don't have any code in advance, ask me to provide it when needed.
+        """
+
+        if let extraSystemPrompt {
+            systemPrompt += "\n\(extraSystemPrompt)"
+        }
+
+        let chat = WidgetDataSource.shared.createChatIfNeeded(for: focusedElementURI)
+
+        await chat.mutateSystemPrompt(systemPrompt)
+
+        Task {
+            let customCommandPrefix = {
+                if let name { return "[\(name)] " }
+                return ""
+            }()
+            
+            if specifiedSystemPrompt != nil {
+                await chat.chatGPTService.mutateHistory { history in
+                    history.append(.init(
+                        role: .assistant,
+                        content: "",
+                        summary: "\(customCommandPrefix)System prompt is updated."
+                    ))
+                }
+            } else if !customCommandPrefix.isEmpty {
+                await chat.chatGPTService.mutateHistory { history in
+                    history.append(.init(
+                        role: .assistant,
+                        content: "",
+                        summary: "\(customCommandPrefix)System prompt is updated."
+                    ))
+                }
+            }
+
+            if let sendingMessageImmediately, !sendingMessageImmediately.isEmpty {
+                try await chat.send(content: sendingMessageImmediately)
+            }
+        }
+
+        presenter.presentChatRoom(fileURL: focusedElementURI)
+    }
 }

Core/Sources/SuggestionWidget/SuggestionWidgetController.swift

@@ -374,20 +374,7 @@ extension SuggestionWidgetController {
                     observeEditorChangeIfNeeded(app)
 
                     guard let fileURL = await {
-                        // if it's not editor, use the window id as url.
-                        if let activeNonEditorWindow = {
-                            let application = AXUIElementCreateApplication(app.processIdentifier)
-                            let focusedElement = application.focusedElement
-                            if focusedElement?.description != "Source Editor" {
-                                return application.focusedWindow
-                            }
-                            return nil
-                        }() {
-                            let id = activeNonEditorWindow.identifier.hashValue
-                            return URL(fileURLWithPath: "/xcode-focused-element/\(id)")
-                        }
-
-                        return try? await Environment.fetchCurrentFileURL()
+                        return try? await Environment.fetchFocusedElementURI()
                     }() else {
                         continue
                     }