Add custom commands to widget context menu

intitni · intitni · commit 720a67e9edf6 · 2023-04-21T00:01:01.000+08:00
diff --git a/Core/Sources/Service/GUI/GraphicalUserInterfaceController.swift.swift b/Core/Sources/Service/GUI/GraphicalUserInterfaceController.swift.swift
@@ -25,6 +25,12 @@ public final class GraphicalUserInterfaceController {
                     ))
                 }
             }
+            suggestionWidget.onCustomCommandClicked = { name in
+                Task {
+                    let commandHandler = PseudoCommandHandler()
+                    await commandHandler.handleCustomDomain(name: name)
+                }
+            }
         }
     }
 }
diff --git a/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift b/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
@@ -68,14 +68,36 @@ struct PseudoCommandHandler {
         ))
     }
 
+    func handleCustomDomain(name: String) async {
+        guard let editor = await getEditorContent(sourceEditor: nil)
+        else {
+            do {
+                try await Environment.triggerAction(name)
+            } catch {
+                let presenter = PresentInWindowSuggestionPresenter()
+                presenter.presentError(error)
+            }
+            return
+        }
+
+        let handler = WindowBaseCommandHandler()
+        do {
+            try await handler.handleCustomCommand(name: name, editor: editor)
+        } catch {
+            let presenter = PresentInWindowSuggestionPresenter()
+            presenter.presentError(error)
+        }
+    }
+
     func acceptSuggestion() async {
         if UserDefaults.shared.value(for: \.acceptSuggestionWithAccessibilityAPI) {
-            guard let xcode = ActiveApplicationMonitor.activeXcode ?? ActiveApplicationMonitor.latestXcode else { return }
+            guard let xcode = ActiveApplicationMonitor.activeXcode ?? ActiveApplicationMonitor
+                .latestXcode else { return }
             let application = AXUIElementCreateApplication(xcode.processIdentifier)
             guard let focusElement = application.focusedElement,
                   focusElement.description == "Source Editor"
             else { return }
-            guard let (content, lines, cursorPosition) = await getFileContent(sourceEditor: nil)
+            guard let (content, lines, _, cursorPosition) = await getFileContent(sourceEditor: nil)
             else {
                 PresentInWindowSuggestionPresenter()
                     .presentErrorMessage("Unable to get file content.")
@@ -153,9 +175,14 @@ struct PseudoCommandHandler {
     }
 }
 
-private extension PseudoCommandHandler {
+extension PseudoCommandHandler {
     func getFileContent(sourceEditor: AXUIElement?) async
-        -> (content: String, lines: [String], cursorPosition: CursorPosition)?
+        -> (
+            content: String,
+            lines: [String],
+            selections: [CursorRange],
+            cursorPosition: CursorPosition
+        )?
     {
         guard let xcode = ActiveApplicationMonitor.activeXcode
             ?? ActiveApplicationMonitor.latestXcode else { return nil }
@@ -166,20 +193,8 @@ private extension PseudoCommandHandler {
         guard let selectionRange = focusElement.selectedTextRange else { return nil }
         let content = focusElement.value
         let split = content.breakLines()
-        let selectedPosition = selectionRange.upperBound
-        // find row and col from content at selected position
-        var rowIndex = 0
-        var count = 0
-        var colIndex = 0
-        for (i, row) in split.enumerated() {
-            if count + row.count > selectedPosition {
-                rowIndex = i
-                colIndex = selectedPosition - count
-                break
-            }
-            count += row.count
-        }
-        return (content, split, CursorPosition(line: rowIndex, character: colIndex))
+        let range = convertRangeToCursorRange(selectionRange, in: content)
+        return (content, split, [range], range.end)
     }
 
     func getFileURL() async -> URL? {
@@ -211,7 +226,9 @@ private extension PseudoCommandHandler {
             lines: content.lines,
             uti: uti,
             cursorPosition: content.cursorPosition,
-            selections: [],
+            selections: content.selections.map {
+                .init(start: $0.start, end: $0.end)
+            },
             tabSize: tabSize,
             indentSize: indentSize,
             usesTabsForIndentation: usesTabsForIndentation
@@ -241,6 +258,32 @@ private extension PseudoCommandHandler {
         }
         return range
     }
+
+    func convertRangeToCursorRange(
+        _ range: ClosedRange<Int>,
+        in content: String
+    ) -> CursorRange {
+        let lines = content.breakLines()
+        guard !lines.isEmpty else { return CursorRange(start: .zero, end: .zero) }
+        var countS = 0
+        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 {
+                cursorRange.start = .init(line: i, character: range.lowerBound - countS)
+            }
+            if countE <= range.upperBound && range.upperBound < countE + line.count {
+                cursorRange.end = .init(line: i, character: range.upperBound - countE)
+                break
+            }
+            countS += line.count
+            countE += line.count
+        }
+        if cursorRange.end == .outOfScope {
+            cursorRange.end = .init(line: lines.endIndex - 1, character: lines.last?.count ?? 0)
+        }
+        return cursorRange
+    }
 }
 
 public extension String {
diff --git a/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift b/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift
@@ -39,6 +39,9 @@ public final class SuggestionWidgetController {
                 panelViewModel: suggestionPanelViewModel,
                 onOpenChatClicked: { [weak self] in
                     self?.onOpenChatClicked()
+                },
+                onCustomCommandClicked: { [weak self] name in
+                    self?.onCustomCommandClicked(name)
                 }
             )
         )
@@ -101,6 +104,7 @@ public final class SuggestionWidgetController {
     private var colorScheme: ColorScheme = .light
     
     public var onOpenChatClicked: () -> Void = {}
+    public var onCustomCommandClicked: (String) -> Void = { _ in }
     public var dataSource: SuggestionWidgetDataSource?
 
     public nonisolated init() {
diff --git a/Core/Sources/SuggestionWidget/WidgetView.swift b/Core/Sources/SuggestionWidget/WidgetView.swift
@@ -17,6 +17,7 @@ struct WidgetView: View {
     @State var isHovering: Bool = false
     @State var processingProgress: Double = 0
     var onOpenChatClicked: () -> Void = {}
+    var onCustomCommandClicked: (String) -> Void = { _ in }
 
     var body: some View {
         Circle().fill(isHovering ? .white.opacity(0.8) : .white.opacity(0.3))
@@ -77,7 +78,8 @@ struct WidgetView: View {
                 WidgetContextMenu(
                     widgetViewModel: viewModel,
                     isChatOpen: panelViewModel.isPanelDisplayed && panelViewModel.chat != nil,
-                    onOpenChatClicked: onOpenChatClicked
+                    onOpenChatClicked: onOpenChatClicked,
+                    onCustomCommandClicked: onCustomCommandClicked
                 )
             }
     }
@@ -103,84 +105,92 @@ struct WidgetContextMenu: View {
     @AppStorage(\.forceOrderWidgetToFront) var forceOrderWidgetToFront
     @AppStorage(\.disableSuggestionFeatureGlobally) var disableSuggestionFeatureGlobally
     @AppStorage(\.suggestionFeatureEnabledProjectList) var suggestionFeatureEnabledProjectList
+    @AppStorage(\.customCommands) var customCommands
     @ObservedObject var widgetViewModel: WidgetViewModel
     @State var projectPath: String?
     var isChatOpen: Bool
     var onOpenChatClicked: () -> Void = {}
+    var onCustomCommandClicked: (String) -> Void = { _ in }
 
     var body: some View {
         Group {
-            if !isChatOpen {
-                Button(action: {
-                    onOpenChatClicked()
-                }) {
-                    Text("Open Chat")
+            Group { // Commands
+                if !isChatOpen {
+                    Button(action: {
+                        onOpenChatClicked()
+                    }) {
+                        Text("Open Chat")
+                    }
                 }
+                
+                customCommandMenu()
             }
 
             Divider()
 
-            Button(action: {
-                useGlobalChat.toggle()
-            }) {
-                Text("Use Global Chat")
-                if useGlobalChat {
-                    Image(systemName: "checkmark")
-                }
-            }
-
-            Button(action: {
-                realtimeSuggestionToggle.toggle()
-            }) {
-                Text("Realtime Suggestion")
-                if realtimeSuggestionToggle {
-                    Image(systemName: "checkmark")
-                }
-            }
-
-            Button(action: {
-                acceptSuggestionWithAccessibilityAPI.toggle()
-            }, label: {
-                Text("Accept Suggestion with Accessibility API")
-                if acceptSuggestionWithAccessibilityAPI {
-                    Image(systemName: "checkmark")
-                }
-            })
-
-            Button(action: {
-                hideCommonPrecedingSpacesInSuggestion.toggle()
-            }, label: {
-                Text("Hide Common Preceding Spaces in Suggestion")
-                if hideCommonPrecedingSpacesInSuggestion {
-                    Image(systemName: "checkmark")
-                }
-            })
-
-            Button(action: {
-                forceOrderWidgetToFront.toggle()
-            }, label: {
-                Text("Force Order Widget to Front")
-                if forceOrderWidgetToFront {
-                    Image(systemName: "checkmark")
+            Group { // Settings
+                Button(action: {
+                    useGlobalChat.toggle()
+                }) {
+                    Text("Use Global Chat")
+                    if useGlobalChat {
+                        Image(systemName: "checkmark")
+                    }
                 }
-            })
-
-            if let projectPath, disableSuggestionFeatureGlobally {
-                let matchedPath = suggestionFeatureEnabledProjectList.first { path in
-                    projectPath.hasPrefix(path)
+                
+                Button(action: {
+                    realtimeSuggestionToggle.toggle()
+                }) {
+                    Text("Realtime Suggestion")
+                    if realtimeSuggestionToggle {
+                        Image(systemName: "checkmark")
+                    }
                 }
+                
                 Button(action: {
-                    if matchedPath != nil {
-                        suggestionFeatureEnabledProjectList
-                            .removeAll { path in path == matchedPath }
-                    } else {
-                        suggestionFeatureEnabledProjectList.append(projectPath)
+                    acceptSuggestionWithAccessibilityAPI.toggle()
+                }, label: {
+                    Text("Accept Suggestion with Accessibility API")
+                    if acceptSuggestionWithAccessibilityAPI {
+                        Image(systemName: "checkmark")
                     }
-                }) {
-                    if matchedPath == nil {
-                        Text("Add to Suggestion-Enabled Project List")
-                    } else {
-                        Text("Remove from Suggestion-Enabled Project List")
+                })
+                
+                Button(action: {
+                    hideCommonPrecedingSpacesInSuggestion.toggle()
+                }, label: {
+                    Text("Hide Common Preceding Spaces in Suggestion")
+                    if hideCommonPrecedingSpacesInSuggestion {
+                        Image(systemName: "checkmark")
+                    }
+                })
+                
+                Button(action: {
+                    forceOrderWidgetToFront.toggle()
+                }, label: {
+                    Text("Force Order Widget to Front")
+                    if forceOrderWidgetToFront {
+                        Image(systemName: "checkmark")
+                    }
+                })
+                
+                if let projectPath, disableSuggestionFeatureGlobally {
+                    let matchedPath = suggestionFeatureEnabledProjectList.first { path in
+                        projectPath.hasPrefix(path)
+                    }
+                    Button(action: {
+                        if matchedPath != nil {
+                            suggestionFeatureEnabledProjectList
+                                .removeAll { path in path == matchedPath }
+                        } else {
+                            suggestionFeatureEnabledProjectList.append(projectPath)
+                        }
+                    }) {
+                        if matchedPath == nil {
+                            Text("Add to Suggestion-Enabled Project List")
+                        } else {
+                            Text("Remove from Suggestion-Enabled Project List")
+                        }
                     }
                 }
             }
@@ -211,6 +221,18 @@ struct WidgetContextMenu: View {
             }
         }
     }
+    
+    func customCommandMenu() -> some View {
+        Menu("Custom Commands") {
+            ForEach(customCommands, id: \.name) { command in
+                Button(action: {
+                    onCustomCommandClicked(command.name)
+                }) {
+                    Text(command.name)
+                }
+            }
+        }
+    }
 }
 
 struct WidgetView_Preview: PreviewProvider {
diff --git a/Core/Tests/ServiceTests/PseudoCommandHandlerFileProcessingTests.swift b/Core/Tests/ServiceTests/PseudoCommandHandlerFileProcessingTests.swift
@@ -0,0 +1,43 @@
+import CopilotModel
+import XCTest
+@testable import Service
+
+class PseudoCommandHandlerFileProcessingTests: XCTestCase {
+    func test_convert_range_0_0() {
+        XCTAssertEqual(
+            PseudoCommandHandler().convertRangeToCursorRange(0...0, in: "\n"),
+            CursorRange(start: .zero, end: .init(line: 0, character: 0))
+        )
+    }
+
+    func test_convert_range_same_line() {
+        XCTAssertEqual(
+            PseudoCommandHandler().convertRangeToCursorRange(1...5, in: "123456789\n"),
+            CursorRange(start: .init(line: 0, character: 1), end: .init(line: 0, character: 5))
+        )
+    }
+
+    func test_convert_range_multiple_line() {
+        XCTAssertEqual(
+            PseudoCommandHandler()
+                .convertRangeToCursorRange(5...25, in: "123456789\n123456789\n123456789\n"),
+            CursorRange(start: .init(line: 0, character: 5), end: .init(line: 2, character: 5))
+        )
+    }
+    
+    func test_convert_range_all_line() {
+        XCTAssertEqual(
+            PseudoCommandHandler()
+                .convertRangeToCursorRange(0...29, in: "123456789\n123456789\n123456789\n"),
+            CursorRange(start: .init(line: 0, character: 0), end: .init(line: 2, character: 9))
+        )
+    }
+    
+    func test_convert_range_out_of_range() {
+        XCTAssertEqual(
+            PseudoCommandHandler()
+                .convertRangeToCursorRange(0...70, in: "123456789\n123456789\n123456789\n"),
+            CursorRange(start: .init(line: 0, character: 0), end: .init(line: 3, character: 0))
+        )
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,12 @@ public final class GraphicalUserInterfaceController {`
`25`	`25`	`))`
`26`	`26`	`}`
`27`	`27`	`}`
	`28`	`+ suggestionWidget.onCustomCommandClicked = { name in`
	`29`	`+ Task {`
	`30`	`+ let commandHandler = PseudoCommandHandler()`
	`31`	`+ await commandHandler.handleCustomDomain(name: name)`
	`32`	`+ }`
	`33`	`+ }`
`28`	`34`	`}`
`29`	`35`	`}`
`30`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,9 @@ public final class SuggestionWidgetController {`
`39`	`39`	`panelViewModel: suggestionPanelViewModel,`
`40`	`40`	`onOpenChatClicked: { [weak self] in`
`41`	`41`	`self?.onOpenChatClicked()`
	`42`	`+ },`
	`43`	`+ onCustomCommandClicked: { [weak self] name in`
	`44`	`+ self?.onCustomCommandClicked(name)`
`42`	`45`	`}`
`43`	`46`	`)`
`44`	`47`	`)`
`@@ -101,6 +104,7 @@ public final class SuggestionWidgetController {`
`101`	`104`	`private var colorScheme: ColorScheme = .light`
`102`	`105`
`103`	`106`	`public var onOpenChatClicked: () -> Void = {}`
	`107`	`+ public var onCustomCommandClicked: (String) -> Void = { _ in }`
`104`	`108`	`public var dataSource: SuggestionWidgetDataSource?`
`105`	`109`
`106`	`110`	`public nonisolated init() {`