From 7e6e07b0211bdf383ce00e8122e40ac193110e14 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 16:03:13 +0800
Subject: [PATCH 01/43] Implement cancellation in SuggestionService

---
 Core/Sources/CodeiumService/CodeiumService.swift  | 15 +++++++++++++++
 .../GitHubCopilotService.swift                    |  5 +++++
 .../CodeiumSuggestionProvider.swift               |  5 +++++
 .../GitHubCopilotSuggestionProvider.swift         |  5 +++++
 .../SuggestionService/SuggestionService.swift     |  5 +++++
 5 files changed, 35 insertions(+)

diff --git a/Core/Sources/CodeiumService/CodeiumService.swift b/Core/Sources/CodeiumService/CodeiumService.swift
index ba4def38..d0acfe7e 100644
--- a/Core/Sources/CodeiumService/CodeiumService.swift
+++ b/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.")
         }
 
@@ -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())
diff --git a/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift b/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift
index 80848db0..24515808 100644
--- a/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift
+++ b/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(
diff --git a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
index bb3b63f9..52df919d 100644
--- a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
@@ -70,5 +70,10 @@ extension CodeiumSuggestionProvider {
     }
 
     func notifySaveTextDocument(fileURL: URL) async throws {}
+    
+    func cancelRequest() async {
+        await (try? createCodeiumServiceIfNeeded())?
+            .cancelRequest()
+    }
 }
 
diff --git a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
index 083da4fd..5e074c93 100644
--- a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
@@ -73,5 +73,10 @@ extension GitHubCopilotSuggestionProvider {
         try await (try? createGitHubCopilotServiceIfNeeded())?
             .notifySaveTextDocument(fileURL: fileURL)
     }
+    
+    func cancelRequest() async {
+        await (try? createGitHubCopilotServiceIfNeeded())?
+            .cancelRequest()
+    }
 }
 
diff --git a/Core/Sources/SuggestionService/SuggestionService.swift b/Core/Sources/SuggestionService/SuggestionService.swift
index 05e762a3..d74ca905 100644
--- a/Core/Sources/SuggestionService/SuggestionService.swift
+++ b/Core/Sources/SuggestionService/SuggestionService.swift
@@ -20,6 +20,7 @@ public protocol SuggestionServiceType {
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
+    func cancelRequest() async
 }
 
 protocol SuggestionServiceProvider: SuggestionServiceType {}
@@ -116,5 +117,9 @@ public extension SuggestionService {
     func notifySaveTextDocument(fileURL: URL) async throws {
         try await suggestionProvider.notifySaveTextDocument(fileURL: fileURL)
     }
+    
+    func cancelRequest() async {
+        await suggestionProvider.cancelRequest()
+    }
 }
 

From 1b751fe267988cb51cff563f40288f6015b03f4d Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 16:14:19 +0800
Subject: [PATCH 02/43] Make request cancellation simpler

---
 .../RealtimeSuggestionController.swift        | 15 +--------------
 Core/Sources/Service/Workspace.swift          | 19 ++++---------------
 Core/Sources/Service/XPCService.swift         | 12 ++----------
 3 files changed, 7 insertions(+), 39 deletions(-)

diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
index 853e7e91..e15ca81b 100644
--- a/Core/Sources/Service/RealtimeSuggestionController.swift
+++ b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -191,7 +191,7 @@ public class RealtimeSuggestionController {
             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 +200,6 @@ public class RealtimeSuggestionController {
                         self.triggerPrefetchDebounced(force: true)
                     }
                 }
-                inflightRealtimeSuggestionsTasks.insert(task)
             }
         }
     }
@@ -247,18 +246,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)
-                    }
-                }()
-            }
         }
     }
 
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index 7056a054..14a7d634 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -91,7 +91,6 @@ final class Workspace {
         UserDefaults.shared.value(for: \.realtimeSuggestionToggle)
     }
 
-    var realtimeSuggestionRequests = Set<Task<Void, Error>>()
     let userDefaultsObserver = UserDefaultsObserver(
         object: UserDefaults.shared, forKeyPaths: [
             UserDefaultPreferenceKeys().suggestionFeatureEnabledProjectList.key,
@@ -245,12 +244,8 @@ extension Workspace {
     @discardableResult
     func generateSuggestions(
         forFileAt fileURL: URL,
-        editor: EditorContent,
-        shouldcancelInFlightRealtimeSuggestionRequests: Bool = true
+        editor: EditorContent
     ) async throws -> [CodeSuggestion] {
-        if shouldcancelInFlightRealtimeSuggestionRequests {
-            cancelInFlightRealtimeSuggestionRequests()
-        }
         refreshUpdateTime()
 
         let filespace = createFilespaceIfNeeded(fileURL: fileURL)
@@ -291,7 +286,6 @@ extension Workspace {
     }
 
     func selectNextSuggestion(forFileAt fileURL: URL) {
-        cancelInFlightRealtimeSuggestionRequests()
         refreshUpdateTime()
         guard let filespace = filespaces[fileURL],
               filespace.suggestions.count > 1
@@ -303,7 +297,6 @@ extension Workspace {
     }
 
     func selectPreviousSuggestion(forFileAt fileURL: URL) {
-        cancelInFlightRealtimeSuggestionRequests()
         refreshUpdateTime()
         guard let filespace = filespaces[fileURL],
               filespace.suggestions.count > 1
@@ -315,7 +308,6 @@ extension Workspace {
     }
 
     func rejectSuggestion(forFileAt fileURL: URL, editor: EditorContent?) {
-        cancelInFlightRealtimeSuggestionRequests()
         refreshUpdateTime()
 
         if let editor, !editor.uti.isEmpty {
@@ -331,7 +323,6 @@ extension Workspace {
     }
 
     func acceptSuggestion(forFileAt fileURL: URL, editor: EditorContent?) -> CodeSuggestion? {
-        cancelInFlightRealtimeSuggestionRequests()
         refreshUpdateTime()
         guard let filespace = filespaces[fileURL],
               !filespace.suggestions.isEmpty,
@@ -412,10 +403,8 @@ extension Workspace {
         return filespace.isExpired
     }
 
-    func cancelInFlightRealtimeSuggestionRequests() {
-        for task in realtimeSuggestionRequests {
-            task.cancel()
-        }
-        realtimeSuggestionRequests = []
+    func cancelInFlightRealtimeSuggestionRequests() async {
+        guard let suggestionService else { return }
+        await suggestionService.cancelRequest()
     }
 }
diff --git a/Core/Sources/Service/XPCService.swift b/Core/Sources/Service/XPCService.swift
index 091b13e4..036afbb0 100644
--- a/Core/Sources/Service/XPCService.swift
+++ b/Core/Sources/Service/XPCService.swift
@@ -15,10 +15,6 @@ import XPCShared
 @ServiceActor
 var workspaces = [URL: Workspace]()
 
-#warning("TODO: Find a better place to store it!")
-@ServiceActor
-var inflightRealtimeSuggestionsTasks = Set<Task<Void, Never>>()
-
 public class XPCService: NSObject, XPCServiceProtocol {
     // MARK: - Service
 
@@ -127,15 +123,13 @@ public class XPCService: NSObject, XPCServiceProtocol {
         editorContent: Data,
         withReply reply: @escaping (Data?, Error?) -> Void
     ) {
-        let task = replyWithUpdatedContent(
+        replyWithUpdatedContent(
             editorContent: editorContent,
             isRealtimeSuggestionRelatedCommand: true,
             withReply: reply
         ) { handler, editor in
             try await handler.presentRealtimeSuggestions(editor: editor)
         }
-
-        Task { @ServiceActor in inflightRealtimeSuggestionsTasks.insert(task) }
     }
 
     public func prefetchRealtimeSuggestions(
@@ -145,15 +139,13 @@ public class XPCService: NSObject, XPCServiceProtocol {
         // We don't need to wait for this.
         reply()
 
-        let task = replyWithUpdatedContent(
+        replyWithUpdatedContent(
             editorContent: editorContent,
             isRealtimeSuggestionRelatedCommand: true,
             withReply: { _, _ in }
         ) { handler, editor in
             try await handler.generateRealtimeSuggestions(editor: editor)
         }
-
-        Task { @ServiceActor in inflightRealtimeSuggestionsTasks.insert(task) }
     }
 
     public func chatWithSelection(

From c439bc7a0506b2c394c6fcaefc2ff1a3bac8c2c2 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 16:54:26 +0800
Subject: [PATCH 03/43] Discard suggestion when cursor moved to another line

---
 .../RealtimeSuggestionController.swift        | 21 ++++++++---------
 .../PseudoCommandHandler.swift                | 23 +++++++++++++++----
 Core/Sources/Service/Workspace.swift          | 14 +++++++++++
 .../Sources/XcodeInspector/SourceEditor.swift |  2 +-
 4 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
index e15ca81b..c2f56945 100644
--- a/Core/Sources/Service/RealtimeSuggestionController.swift
+++ b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -9,6 +9,7 @@ import Foundation
 import Logger
 import Preferences
 import QuartzCore
+import XcodeInspector
 
 @ServiceActor
 public class RealtimeSuggestionController {
@@ -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 {
@@ -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,14 @@ public class RealtimeSuggestionController {
                 case kAXValueChangedNotification:
                     self.triggerPrefetchDebounced()
                     await self.notifyEditingFileChange(editor: focusElement)
+                case kAXSelectedTextChangedNotification:
+                    guard let editor = sourceEditor else { continue }
+                    let sourceEditor = SourceEditor(
+                        runningApplication: activeXcode,
+                        element: editor
+                    )
+                    await PseudoCommandHandler()
+                        .invalidateRealtimeSuggestionsIfNeeded(sourceEditor: sourceEditor)
                 default:
                     continue
                 }
@@ -177,13 +183,6 @@ public class RealtimeSuggestionController {
 
         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.
diff --git a/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift b/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
index ea7de188..0c9d321e 100644
--- a/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
+++ b/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.
@@ -54,6 +55,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 +228,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? {
@@ -289,10 +303,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 +335,4 @@ public extension String {
         return all
     }
 }
+
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index 14a7d634..7c604c17 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -67,6 +67,20 @@ final class Filespace {
     func refreshUpdateTime() {
         lastSuggestionUpdateTime = Environment.now()
     }
+    
+    func validateSuggestions(lines: [String], cursorPosition: CursorPosition) -> Bool {
+        if cursorPosition.line != suggestionSourceSnapshot.cursorPosition.line {
+            reset()
+            return false
+        }
+        
+        guard cursorPosition.line >= 0, cursorPosition.line < lines.count else {
+            reset()
+            return false
+        }
+        
+        return true
+    }
 }
 
 // MARK: - Workspace
diff --git a/Core/Sources/XcodeInspector/SourceEditor.swift b/Core/Sources/XcodeInspector/SourceEditor.swift
index 1d830a2c..d5b3e772 100644
--- a/Core/Sources/XcodeInspector/SourceEditor.swift
+++ b/Core/Sources/XcodeInspector/SourceEditor.swift
@@ -34,7 +34,7 @@ public class SourceEditor {
                 content: content,
                 lines: split,
                 selections: [range],
-                cursorPosition: range.end,
+                cursorPosition: range.start,
                 lineAnnotations: lineAnnotations
             )
         }

From 5287a5052a1ee0b71324f82f8bc795b7feb8a569 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 17:41:46 +0800
Subject: [PATCH 04/43] Fix tests

---
 Core/Tests/ServiceTests/Environment.swift | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/Core/Tests/ServiceTests/Environment.swift b/Core/Tests/ServiceTests/Environment.swift
index c2e99b06..fd202379 100644
--- a/Core/Tests/ServiceTests/Environment.swift
+++ b/Core/Tests/ServiceTests/Environment.swift
@@ -34,6 +34,10 @@ func completion(text: String, range: CursorRange, uuid: String = "") -> CodeSugg
 }
 
 class MockSuggestionService: GitHubCopilotSuggestionServiceType {
+    func cancelRequest() async {
+        fatalError()
+    }
+    
     func notifyOpenTextDocument(fileURL: URL, content: String) async throws {
         fatalError()
     }

From 5f80ae966f63e83f840bc904a4406352cb112506 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 17:45:42 +0800
Subject: [PATCH 05/43] Make OpenedDocumentPool an actor

---
 Core/Sources/CodeiumService/CodeiumService.swift    | 8 ++++----
 Core/Sources/CodeiumService/OpendDocumentPool.swift | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumService.swift b/Core/Sources/CodeiumService/CodeiumService.swift
index d0acfe7e..93f55f93 100644
--- a/Core/Sources/CodeiumService/CodeiumService.swift
+++ b/Core/Sources/CodeiumService/CodeiumService.swift
@@ -230,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,
@@ -306,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
@@ -315,7 +315,7 @@ 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
@@ -323,7 +323,7 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
     }
 
     public func notifyCloseTextDocument(fileURL: URL) async throws {
-        openedDocumentPool.closeDocument(url: fileURL)
+        await openedDocumentPool.closeDocument(url: fileURL)
     }
 }
 
diff --git a/Core/Sources/CodeiumService/OpendDocumentPool.swift b/Core/Sources/CodeiumService/OpendDocumentPool.swift
index e0f55d87..ca8a7d50 100644
--- a/Core/Sources/CodeiumService/OpendDocumentPool.swift
+++ b/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] {

From 0e4f7f87efc7de2b064f82b597d9445855580743 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 21:32:44 +0800
Subject: [PATCH 06/43] Reuse suggestion if user is typing according to it

---
 .../RealtimeSuggestionController.swift        | 10 ++---
 .../PseudoCommandHandler.swift                | 24 ++++++++---
 Core/Sources/Service/Workspace.swift          |  7 +++
 .../SuggestionInjector.swift                  | 10 +++++
 .../AcceptSuggestionTests.swift               | 43 +++++++++++++++++++
 5 files changed, 80 insertions(+), 14 deletions(-)

diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
index c2f56945..cdc0903f 100644
--- a/Core/Sources/Service/RealtimeSuggestionController.swift
+++ b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -26,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
@@ -116,7 +116,7 @@ public class RealtimeSuggestionController {
         }
 
         guard focusElementType == "Source Editor" else { return }
-        sourceEditor = focusElement
+        sourceEditor = SourceEditor(runningApplication: activeXcode, element: focusElement)
 
         editorObservationTask?.cancel()
         editorObservationTask = nil
@@ -138,11 +138,7 @@ public class RealtimeSuggestionController {
                     self.triggerPrefetchDebounced()
                     await self.notifyEditingFileChange(editor: focusElement)
                 case kAXSelectedTextChangedNotification:
-                    guard let editor = sourceEditor else { continue }
-                    let sourceEditor = SourceEditor(
-                        runningApplication: activeXcode,
-                        element: editor
-                    )
+                    guard let sourceEditor else { continue }
                     await PseudoCommandHandler()
                         .invalidateRealtimeSuggestionsIfNeeded(sourceEditor: sourceEditor)
                 default:
diff --git a/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift b/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
index 0c9d321e..ee4e22cb 100644
--- a/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
+++ b/Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift
@@ -39,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)
@@ -246,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
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index 7c604c17..7bad9812 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -79,6 +79,13 @@ final class Filespace {
             return false
         }
         
+        let editingLine = lines[cursorPosition.line].dropLast(1) // dropping \n
+        let suggestionFirstLine = presentingSuggestion?.text.split(separator: "\n").first ?? ""
+        if !suggestionFirstLine.hasPrefix(editingLine) {
+            reset()
+            return false
+        }
+        
         return true
     }
 }
diff --git a/Core/Sources/SuggestionInjector/SuggestionInjector.swift b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
index d94064aa..797da21b 100644
--- a/Core/Sources/SuggestionInjector/SuggestionInjector.swift
+++ b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
@@ -154,6 +154,7 @@ public struct SuggestionInjector {
 
         var toBeInserted = suggestionContent.breakLines(appendLineBreakToLastLine: true)
 
+        // prepending prefix text not in range if needed.
         if let firstRemovedLine,
            !firstRemovedLine.isEmptyOrNewLine,
            start.character > 0,
@@ -175,8 +176,16 @@ public struct SuggestionInjector {
             )
         }
 
+        // appending suffix text not in range if needed.
         let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1
+        let skipAppendingDueToContinueTyping = {
+            guard let first = toBeInserted.first?.dropLast(1), !first.isEmpty else { return false }
+            let droppedLast = lastRemovedLine?.dropLast(1)
+            guard let droppedLast, !droppedLast.isEmpty else { return false }
+            return first.hasPrefix(droppedLast)
+        }()
         if let lastRemovedLine,
+           !skipAppendingDueToContinueTyping,
            !lastRemovedLine.isEmptyOrNewLine,
            end.character >= 0,
            end.character - 1 < lastRemovedLine.count,
@@ -191,6 +200,7 @@ public struct SuggestionInjector {
                 toBeInserted[toBeInserted.endIndex - 1].removeLast(1)
             }
             let leftover = lastRemovedLine[leftoverRange]
+            
             toBeInserted[toBeInserted.endIndex - 1]
                 .append(contentsOf: leftover)
         }
diff --git a/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift b/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
index 41238f0e..d9d5a4a8 100644
--- a/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
+++ b/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
@@ -131,6 +131,49 @@ final class AcceptSuggestionTests: XCTestCase {
         }
         """)
     }
+    
+    func test_accept_suggestion_overlap_continue_typing() async throws {
+        let content = """
+        struct Cat {
+            var name: Str
+        }
+        """
+        let text = """
+            var name: String
+            var age: String
+        """
+        let suggestion = CodeSuggestion(
+            text: text,
+            position: .init(line: 1, character: 12),
+            uuid: "",
+            range: .init(
+                start: .init(line: 1, character: 0),
+                end: .init(line: 1, character: 12)
+            ),
+            displayText: ""
+        )
+
+        var extraInfo = SuggestionInjector.ExtraInfo()
+        var lines = content.breakLines()
+        var cursor = CursorPosition(line: 0, character: 0)
+        SuggestionInjector().acceptSuggestion(
+            intoContentWithoutSuggestion: &lines,
+            cursorPosition: &cursor,
+            completion: suggestion,
+            extraInfo: &extraInfo
+        )
+        XCTAssertTrue(extraInfo.didChangeContent)
+        XCTAssertTrue(extraInfo.didChangeCursorPosition)
+        XCTAssertNil(extraInfo.suggestionRange)
+        XCTAssertEqual(lines, content.breakLines().applying(extraInfo.modifications))
+        XCTAssertEqual(cursor, .init(line: 2, character: 19))
+        XCTAssertEqual(lines.joined(separator: ""), """
+        struct Cat {
+            var name: String
+            var age: String
+        }
+        """)
+    }
 
     func test_propose_suggestion_partial_overlap() async throws {
         let content = "func quickSort() {}}\n"

From 256032a7f7532807648afffd296e1e83e3a70bcd Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 21:35:58 +0800
Subject: [PATCH 07/43] Disable request cancellation on mouse clicks

---
 Core/Sources/Service/RealtimeSuggestionController.swift | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
index cdc0903f..41722db2 100644
--- a/Core/Sources/Service/RealtimeSuggestionController.swift
+++ b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -171,12 +171,6 @@ 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
 

From d684e32afe0d05cf32d5cb115f9a29b08a750d45 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Wed, 24 May 2023 21:37:18 +0800
Subject: [PATCH 08/43] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 31662599..fc804d2d 100644
--- a/README.md
+++ b/README.md
@@ -148,7 +148,7 @@ The app can provide real-time code suggestions based on the files you have opene
 
 If you're working on a company project and don't want the suggestion feature to be triggered, you can globally disable it and choose to enable it only for specific projects. 
 
-Whenever you stop typing for a few milliseconds, the app will automatically fetch suggestions for you, you can cancel this by clicking the mouse, or pressing **Escape** or the **arrow keys**.
+Whenever your code is updated, the app will automatically fetch suggestions for you, you can cancel this by pressing **Escape**.
 
 *: If a file is already open before the helper app launches, you will need to switch to those files in order to send the open file notification. 
 

From d7d7e3ef8c30a654a806c19b7ed21742a5860b4f Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 13:55:41 +0800
Subject: [PATCH 09/43] Fix typos

---
 Core/Sources/SuggestionInjector/SuggestionInjector.swift | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Core/Sources/SuggestionInjector/SuggestionInjector.swift b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
index 797da21b..ac57de17 100644
--- a/Core/Sources/SuggestionInjector/SuggestionInjector.swift
+++ b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
@@ -104,11 +104,11 @@ public struct SuggestionInjector {
             )
         }
 
-        // if the suggestion is only appeding new lines and spaces, return without modification
+        // if the suggestion is only appending new lines and spaces, return without modification
         if completion.text.dropFirst(commonPrefix.count)
             .allSatisfy({ $0.isWhitespace || $0.isNewline }) { return }
 
-        // determin if it's inserted to the current line or the next line
+        // determine if it's inserted to the current line or the next line
         let lineIndex = start.line + {
             guard let existedLine else { return 0 }
             if existedLine.isEmptyOrNewLine { return 1 }

From f60b59594213d0dd57a62353e1af5c24e878ea87 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 15:20:00 +0800
Subject: [PATCH 10/43] Tweak workspace cleanup

---
 Core/Sources/Service/ScheduledCleaner.swift   | 50 +++++------
 Core/Sources/Service/Workspace.swift          |  2 +-
 Core/Sources/XcodeInspector/Helpers.swift     |  2 +-
 .../XcodeInspector/XcodeInspector.swift       | 90 +++++++++++++------
 4 files changed, 88 insertions(+), 56 deletions(-)

diff --git a/Core/Sources/Service/ScheduledCleaner.swift b/Core/Sources/Service/ScheduledCleaner.swift
index 0a3a8302..c9c229ad 100644
--- a/Core/Sources/Service/ScheduledCleaner.swift
+++ b/Core/Sources/Service/ScheduledCleaner.swift
@@ -3,6 +3,7 @@ import AppKit
 import AXExtension
 import Foundation
 import Logger
+import XcodeInspector
 
 public final class ScheduledCleaner {
     public init() {
@@ -19,7 +20,7 @@ public final class ScheduledCleaner {
             for await app in ActiveApplicationMonitor.createStream() {
                 try Task.checkCancellation()
                 if let app, !app.isXcode {
-                    cleanUp() 
+                    cleanUp()
                 }
             }
         }
@@ -27,51 +28,48 @@ public final class ScheduledCleaner {
 
     @ServiceActor
     func cleanUp() {
-        let availableTabs = findAvailableOpenedTabs()
+        let workspaceInfos = XcodeInspector.shared.xcodes.reduce(
+            into: [
+                XcodeAppInstanceInspector.WorkspaceIdentifier:
+                    XcodeAppInstanceInspector.WorkspaceInfo
+            ]()
+        ) { result, xcode in
+            let infos = xcode.workspaces
+            for (id, info) in infos {
+                if let existed = result[id] {
+                    result[id] = existed.combined(with: info)
+                } else {
+                    result[id] = info
+                }
+            }
+        }
+        dump(workspaceInfos)
         for (url, workspace) in workspaces {
-            if workspace.isExpired {
+            if workspace.isExpired, workspaceInfos[.url(url)] == nil {
                 Logger.service.info("Remove idle workspace")
                 for url in workspace.filespaces.keys {
                     WidgetDataSource.shared.cleanup(for: url)
                 }
-                workspace.cleanUp(availableTabs: availableTabs)
+                workspace.cleanUp(availableTabs: [])
                 workspaces[url] = nil
             } else {
+                let tabs = (workspaceInfos[.url(url)]?.tabs ?? [])
+                    .union(workspaceInfos[.unknown]?.tabs ?? [])
                 // cleanup chats for unused files
                 let filespaces = workspace.filespaces
                 for (url, _) in filespaces {
                     if workspace.isFilespaceExpired(
                         fileURL: url,
-                        availableTabs: availableTabs
+                        availableTabs: tabs
                     ) {
                         Logger.service.info("Remove idle filespace")
                         WidgetDataSource.shared.cleanup(for: url)
                     }
                 }
                 // cleanup workspace
-                workspace.cleanUp(availableTabs: availableTabs)
-            }
-        }
-    }
-
-    func findAvailableOpenedTabs() -> Set<String> {
-        guard let xcode = ActiveApplicationMonitor.latestXcode else { return [] }
-        let app = AXUIElementCreateApplication(xcode.processIdentifier)
-        let windows = app.windows.filter { $0.identifier == "Xcode.WorkspaceWindow" }
-        guard !windows.isEmpty else { return [] }
-        var allTabs = Set<String>()
-        for window in windows {
-            guard let editArea = window.firstChild(where: { $0.description == "editor area" })
-            else { continue }
-            let tabBars = editArea.children { $0.description == "tab bar" }
-            for tabBar in tabBars {
-                let tabs = tabBar.children { $0.roleDescription == "tab" }
-                for tab in tabs {
-                    allTabs.insert(tab.title)
-                }
+                workspace.cleanUp(availableTabs: tabs)
             }
         }
-        return allTabs
     }
 }
 
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index 7bad9812..b5ea61e8 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -104,7 +104,7 @@ final class Workspace {
     let openedFileRecoverableStorage: OpenedFileRecoverableStorage
     var lastSuggestionUpdateTime = Environment.now()
     var isExpired: Bool {
-        Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 8
+        Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 4
     }
 
     private(set) var filespaces = [URL: Filespace]()
diff --git a/Core/Sources/XcodeInspector/Helpers.swift b/Core/Sources/XcodeInspector/Helpers.swift
index f0238031..f831d15a 100644
--- a/Core/Sources/XcodeInspector/Helpers.swift
+++ b/Core/Sources/XcodeInspector/Helpers.swift
@@ -1,7 +1,7 @@
 import AppKit
 import Foundation
 
-public extension NSRunningApplication {
+extension NSRunningApplication {
     var isXcode: Bool { bundleIdentifier == "com.apple.dt.Xcode" }
     var isCopilotForXcodeExtensionService: Bool {
         bundleIdentifier == Bundle.main.bundleIdentifier
diff --git a/Core/Sources/XcodeInspector/XcodeInspector.swift b/Core/Sources/XcodeInspector/XcodeInspector.swift
index 2c0e7731..c2cca8b1 100644
--- a/Core/Sources/XcodeInspector/XcodeInspector.swift
+++ b/Core/Sources/XcodeInspector/XcodeInspector.swift
@@ -108,7 +108,7 @@ public final class XcodeInspector: ObservableObject {
         latestActiveXcode = xcode
         activeDocumentURL = xcode.documentURL
         focusedWindow = xcode.focusedWindow
-        
+
         let setFocusedElement = { [weak self] in
             guard let self else { return }
             focusedElement = xcode.appElement.focusedElement
@@ -150,10 +150,23 @@ public class AppInstanceInspector: ObservableObject {
 }
 
 public final class XcodeAppInstanceInspector: AppInstanceInspector {
-    @Published var focusedWindow: XcodeWindowInspector?
-    @Published var documentURL: URL = .init(fileURLWithPath: "/")
-    @Published var projectURL: URL = .init(fileURLWithPath: "/")
-    @Published var tabs: Set<String> = []
+    public struct WorkspaceInfo {
+        public let tabs: Set<String>
+
+        public func combined(with info: WorkspaceInfo) -> WorkspaceInfo {
+            return .init(tabs: info.tabs.union(tabs))
+        }
+    }
+
+    public enum WorkspaceIdentifier: Hashable {
+        case url(URL)
+        case unknown
+    }
+
+    @Published public var focusedWindow: XcodeWindowInspector?
+    @Published public var documentURL: URL = .init(fileURLWithPath: "/")
+    @Published public var projectURL: URL = .init(fileURLWithPath: "/")
+    @Published public var workspaces = [WorkspaceIdentifier: WorkspaceInfo]()
     private var longRunningTasks = Set<Task<Void, Error>>()
     private var focusedWindowObservations = Set<AnyCancellable>()
 
@@ -178,27 +191,22 @@ public final class XcodeAppInstanceInspector: AppInstanceInspector {
 
         longRunningTasks.insert(focusedWindowChanged)
 
-        if let updatedTabs = Self.findAvailableOpenedTabs(runningApplication) {
-            tabs = updatedTabs
-        }
+        workspaces = Self.fetchWorkspaceInfo(runningApplication)
         let updateTabsTask = Task { @MainActor in
             let notification = AXNotificationStream(
                 app: runningApplication,
-                notificationNames: kAXFocusedUIElementChangedNotification
+                notificationNames: kAXFocusedUIElementChangedNotification,
+                kAXApplicationDeactivatedNotification
             )
             if #available(macOS 13.0, *) {
                 for await _ in notification.debounce(for: .seconds(5)) {
                     try Task.checkCancellation()
-                    if let updatedTabs = Self.findAvailableOpenedTabs(runningApplication) {
-                        tabs = updatedTabs
-                    }
+                    workspaces = Self.fetchWorkspaceInfo(runningApplication)
                 }
             } else {
                 for await _ in notification {
                     try Task.checkCancellation()
-                    if let updatedTabs = Self.findAvailableOpenedTabs(runningApplication) {
-                        tabs = updatedTabs
-                    }
+                    workspaces = Self.fetchWorkspaceInfo(runningApplication)
                 }
             }
         }
@@ -239,24 +247,50 @@ public final class XcodeAppInstanceInspector: AppInstanceInspector {
         }
     }
 
-    static func findAvailableOpenedTabs(_ app: NSRunningApplication) -> Set<String>? {
+    static func fetchWorkspaceInfo(
+        _ app: NSRunningApplication
+    ) -> [WorkspaceIdentifier: WorkspaceInfo] {
         let app = AXUIElementCreateApplication(app.processIdentifier)
-        guard app.isFocused else { return nil }
         let windows = app.windows.filter { $0.identifier == "Xcode.WorkspaceWindow" }
-        guard !windows.isEmpty else { return [] }
-        var allTabs = Set<String>()
+
+        var dict = [WorkspaceIdentifier: WorkspaceInfo]()
+
         for window in windows {
-            guard let editArea = window.firstChild(where: { $0.description == "editor area" })
-            else { continue }
-            let tabBars = editArea.children { $0.description == "tab bar" }
-            for tabBar in tabBars {
-                let tabs = tabBar.children { $0.roleDescription == "tab" }
-                for tab in tabs {
-                    allTabs.insert(tab.title)
+            let workspaceIdentifier = {
+                for child in window.children {
+                    if child.description.starts(with: "/"), child.description.count > 1 {
+                        let path = child.description
+                        let trimmedNewLine = path.trimmingCharacters(in: .newlines)
+                        var url = URL(fileURLWithPath: trimmedNewLine)
+                        while !FileManager.default.fileIsDirectory(atPath: url.path) ||
+                            !url.pathExtension.isEmpty
+                        {
+                            url = url.deletingLastPathComponent()
+                        }
+                        return WorkspaceIdentifier.url(url)
+                    }
                 }
-            }
+                return WorkspaceIdentifier.unknown
+            }()
+
+            let tabs = {
+                guard let editArea = window.firstChild(where: { $0.description == "editor area" })
+                else { return Set<String>() }
+                var allTabs = Set<String>()
+                let tabBars = editArea.children { $0.description == "tab bar" }
+                for tabBar in tabBars {
+                    let tabs = tabBar.children { $0.roleDescription == "tab" }
+                    for tab in tabs {
+                        allTabs.insert(tab.title)
+                    }
+                }
+                return allTabs
+            }()
+
+            dict[workspaceIdentifier] = .init(tabs: tabs)
         }
-        return allTabs
+
+        return dict
     }
 }
 

From 0cfe17735ea582f64583e729e8397b33ea3940e7 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 15:20:18 +0800
Subject: [PATCH 11/43] Fix Codeium language server setup

---
 .../CodeiumLanguageServer.swift               | 20 +++++++++----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
index 92ca5d2c..641b275d 100644
--- a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
+++ b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
@@ -57,36 +57,32 @@ final class CodeiumLanguageServer {
 
     func start() {
         guard !process.isRunning else { return }
-        port = nil
         do {
             try process.run()
 
-            Task {
+            Task { @MainActor in
                 func findPort() -> String? {
                     // find a file in managerDirectoryURL whose name looks like a port, return the
                     // name if found
                     let fileManager = FileManager.default
-                    let enumerator = fileManager.enumerator(
-                        at: managerDirectoryURL,
-                        includingPropertiesForKeys: nil
-                    )
-                    while let fileURL = enumerator?.nextObject() as? URL {
-                        if fileURL.lastPathComponent.range(
+                    guard let filePaths = try? fileManager
+                        .contentsOfDirectory(atPath: managerDirectoryURL.path) else { return nil }
+                    for path in filePaths {
+                        let filename = URL(fileURLWithPath: path).lastPathComponent
+                        if filename.range(
                             of: #"^\d+$"#,
                             options: .regularExpression
                         ) != nil {
-                            return fileURL.lastPathComponent
+                            return filename
                         }
                     }
                     return nil
                 }
 
                 try await Task.sleep(nanoseconds: 2_000_000)
-                port = findPort()
                 var waited = 0
 
                 while true {
-                    try await Task.sleep(nanoseconds: 1_000_000_000)
                     waited += 1
                     if let port = findPort() {
                         finishStarting(port: port)
@@ -94,7 +90,9 @@ final class CodeiumLanguageServer {
                     }
                     if waited >= 60 {
                         process.terminate()
+                        return
                     }
+                    try await Task.sleep(nanoseconds: 1_000_000_000)
                 }
             }
         } catch {

From 281002ba9f7ba7779b94a6b76cc63b81e8ffee17 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 15:28:31 +0800
Subject: [PATCH 12/43] Make workspace expire sooner

---
 Core/Sources/Service/ScheduledCleaner.swift | 1 -
 Core/Sources/Service/Workspace.swift        | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/Core/Sources/Service/ScheduledCleaner.swift b/Core/Sources/Service/ScheduledCleaner.swift
index c9c229ad..d6f102d9 100644
--- a/Core/Sources/Service/ScheduledCleaner.swift
+++ b/Core/Sources/Service/ScheduledCleaner.swift
@@ -43,7 +43,6 @@ public final class ScheduledCleaner {
                 }
             }
         }
-        dump(workspaceInfos)
         for (url, workspace) in workspaces {
             if workspace.isExpired, workspaceInfos[.url(url)] == nil {
                 Logger.service.info("Remove idle workspace")
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index b5ea61e8..a7b6de33 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -104,7 +104,7 @@ final class Workspace {
     let openedFileRecoverableStorage: OpenedFileRecoverableStorage
     var lastSuggestionUpdateTime = Environment.now()
     var isExpired: Bool {
-        Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 4
+        Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 1
     }
 
     private(set) var filespaces = [URL: Filespace]()

From 68627c1f112c744114495dd2de103f22373e13e0 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 16:27:36 +0800
Subject: [PATCH 13/43] Replace unowned with weak

---
 Core/Sources/CodeiumService/CodeiumLanguageServer.swift | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
index 641b275d..d8a7e11a 100644
--- a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
+++ b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
@@ -203,7 +203,7 @@ final class IOTransport {
     }
 
     private func setupFileHandleHandlers() {
-        stdoutPipe.fileHandleForReading.readabilityHandler = { [unowned self] handle in
+        stdoutPipe.fileHandleForReading.readabilityHandler = { [weak self] handle in
             let data = handle.availableData
 
             guard !data.isEmpty else {
@@ -211,11 +211,11 @@ final class IOTransport {
             }
 
             if UserDefaults.shared.value(for: \.codeiumVerboseLog) {
-                self.forwardDataToHandler(data)
+                self?.forwardDataToHandler(data)
             }
         }
 
-        stderrPipe.fileHandleForReading.readabilityHandler = { [unowned self] handle in
+        stderrPipe.fileHandleForReading.readabilityHandler = { [weak self] handle in
             let data = handle.availableData
 
             guard !data.isEmpty else {
@@ -223,7 +223,7 @@ final class IOTransport {
             }
 
             if UserDefaults.shared.value(for: \.codeiumVerboseLog) {
-                self.forwardErrorDataToHandler(data)
+                self?.forwardErrorDataToHandler(data)
             }
         }
     }

From ebd94eb8b82a823550f99b569242148c271edeaf Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 16:29:55 +0800
Subject: [PATCH 14/43] Support terminating suggestion services on normal exits

---
 .../CodeiumService/CodeiumLanguageServer.swift        |  9 +++++++++
 Core/Sources/CodeiumService/CodeiumService.swift      |  6 ++++++
 .../GitHubCopilotService/GitHubCopilotService.swift   |  5 +++++
 Core/Sources/Service/ScheduledCleaner.swift           |  7 +++++++
 Core/Sources/Service/Workspace.swift                  |  4 ++++
 .../SuggestionService/CodeiumSuggestionProvider.swift |  4 ++++
 .../GitHubCopilotSuggestionProvider.swift             |  4 ++++
 .../Sources/SuggestionService/SuggestionService.swift |  8 +++++++-
 ExtensionService/AppDelegate.swift                    | 11 +++++++----
 9 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
index d8a7e11a..cd099943 100644
--- a/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
+++ b/Core/Sources/CodeiumService/CodeiumLanguageServer.swift
@@ -7,6 +7,7 @@ import Preferences
 
 protocol CodeiumLSP {
     func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
+    func terminate()
 }
 
 final class CodeiumLanguageServer {
@@ -119,6 +120,14 @@ final class CodeiumLanguageServer {
         self.port = port
         launchHandler?()
     }
+    
+    func terminate() {
+        process.terminationHandler = nil
+        if process.isRunning {
+            process.terminate()
+        }
+        transport.close()
+    }
 }
 
 extension CodeiumLanguageServer: CodeiumLSP {
diff --git a/Core/Sources/CodeiumService/CodeiumService.swift b/Core/Sources/CodeiumService/CodeiumService.swift
index 93f55f93..37351067 100644
--- a/Core/Sources/CodeiumService/CodeiumService.swift
+++ b/Core/Sources/CodeiumService/CodeiumService.swift
@@ -19,6 +19,7 @@ public protocol CodeiumSuggestionServiceType {
     func notifyChangeTextDocument(fileURL: URL, content: String) async throws
     func notifyCloseTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate()
 }
 
 enum CodeiumError: Error, LocalizedError {
@@ -325,6 +326,11 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
     public func notifyCloseTextDocument(fileURL: URL) async throws {
         await openedDocumentPool.closeDocument(url: fileURL)
     }
+    
+    public func terminate() {
+        server?.terminate()
+        server = nil
+    }
 }
 
 func getXcodeVersion() async throws -> String {
diff --git a/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift b/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift
index 24515808..32d38b37 100644
--- a/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift
+++ b/Core/Sources/GitHubCopilotService/GitHubCopilotService.swift
@@ -32,6 +32,7 @@ public protocol GitHubCopilotSuggestionServiceType {
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate() async
 }
 
 protocol GitHubCopilotLSP {
@@ -379,6 +380,10 @@ public final class GitHubCopilotSuggestionService: GitHubCopilotBaseService,
 //        Logger.service.debug("Close \(uri)")
         try await server.sendNotification(.didCloseTextDocument(.init(uri: uri)))
     }
+    
+    public func terminate() async {
+        // automatically handled
+    }
 }
 
 extension InitializingServer: GitHubCopilotLSP {
diff --git a/Core/Sources/Service/ScheduledCleaner.swift b/Core/Sources/Service/ScheduledCleaner.swift
index d6f102d9..c6b9a832 100644
--- a/Core/Sources/Service/ScheduledCleaner.swift
+++ b/Core/Sources/Service/ScheduledCleaner.swift
@@ -70,5 +70,12 @@ public final class ScheduledCleaner {
             }
         }
     }
+    
+    @ServiceActor
+    public func closeAllChildProcesses() async {
+        for (_, workspace) in workspaces {
+            await workspace.terminateSuggestionService()
+        }
+    }
 }
 
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
index a7b6de33..4f6357b8 100644
--- a/Core/Sources/Service/Workspace.swift
+++ b/Core/Sources/Service/Workspace.swift
@@ -428,4 +428,8 @@ extension Workspace {
         guard let suggestionService else { return }
         await suggestionService.cancelRequest()
     }
+    
+    func terminateSuggestionService() async {
+        await _suggestionService?.terminate()
+    }
 }
diff --git a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
index 52df919d..26a2c1f8 100644
--- a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
@@ -75,5 +75,9 @@ extension CodeiumSuggestionProvider {
         await (try? createCodeiumServiceIfNeeded())?
             .cancelRequest()
     }
+    
+    func terminate() async {
+        (try? createCodeiumServiceIfNeeded())?.terminate()
+    }
 }
 
diff --git a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
index 5e074c93..4dcd95ee 100644
--- a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
@@ -78,5 +78,9 @@ extension GitHubCopilotSuggestionProvider {
         await (try? createGitHubCopilotServiceIfNeeded())?
             .cancelRequest()
     }
+    
+    func terminate() async {
+        await (try? createGitHubCopilotServiceIfNeeded())?.terminate()
+    }
 }
 
diff --git a/Core/Sources/SuggestionService/SuggestionService.swift b/Core/Sources/SuggestionService/SuggestionService.swift
index d74ca905..5a2c0e8b 100644
--- a/Core/Sources/SuggestionService/SuggestionService.swift
+++ b/Core/Sources/SuggestionService/SuggestionService.swift
@@ -1,3 +1,4 @@
+import AppKit
 import Foundation
 import Preferences
 import SuggestionModel
@@ -21,6 +22,7 @@ public protocol SuggestionServiceType {
     func notifyCloseTextDocument(fileURL: URL) async throws
     func notifySaveTextDocument(fileURL: URL) async throws
     func cancelRequest() async
+    func terminate() async
 }
 
 protocol SuggestionServiceProvider: SuggestionServiceType {}
@@ -117,9 +119,13 @@ public extension SuggestionService {
     func notifySaveTextDocument(fileURL: URL) async throws {
         try await suggestionProvider.notifySaveTextDocument(fileURL: fileURL)
     }
-    
+
     func cancelRequest() async {
         await suggestionProvider.cancelRequest()
     }
+
+    func terminate() async {
+        await suggestionProvider.terminate()
+    }
 }
 
diff --git a/ExtensionService/AppDelegate.swift b/ExtensionService/AppDelegate.swift
index fecf169c..e2a761a9 100644
--- a/ExtensionService/AppDelegate.swift
+++ b/ExtensionService/AppDelegate.swift
@@ -34,7 +34,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
         _ = RealtimeSuggestionController.shared
         _ = XcodeInspector.shared
         AXIsProcessTrustedWithOptions([
-            kAXTrustedCheckOptionPrompt.takeRetainedValue() as NSString: true
+            kAXTrustedCheckOptionPrompt.takeRetainedValue() as NSString: true,
         ] as CFDictionary)
         setupQuitOnUpdate()
         setupQuitOnUserTerminated()
@@ -105,7 +105,10 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
     }
 
     @objc func quit() {
-        exit(0)
+        Task { @MainActor in
+            await scheduledCleaner.closeAllChildProcesses()
+            exit(0)
+        }
     }
 
     @objc func openCopilotForXcode() {
@@ -148,7 +151,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
                 Logger.service.info("Extension Service will quit.")
                 #if DEBUG
                 #else
-                exit(0)
+                quit()
                 #endif
             }
         }
@@ -172,7 +175,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
                 if NSWorkspace.shared.runningApplications.contains(where: \.isUserOfService) {
                     continue
                 }
-                exit(0)
+                quit()
             }
         }
     }

From f28d8365b6d34b11b243e04b55fa693c083466a7 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 16:30:25 +0800
Subject: [PATCH 15/43] Fix racing

---
 .../CodeiumSuggestionProvider.swift                  |  2 +-
 .../GitHubCopilotSuggestionProvider.swift            |  2 +-
 .../SuggestionService/SuggestionService.swift        | 12 +++++++++---
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
index 26a2c1f8..eb03b335 100644
--- a/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/CodeiumSuggestionProvider.swift
@@ -3,7 +3,7 @@ import Foundation
 import Preferences
 import SuggestionModel
 
-final class CodeiumSuggestionProvider: SuggestionServiceProvider {
+actor CodeiumSuggestionProvider: SuggestionServiceProvider {
     let projectRootURL: URL
     let onServiceLaunched: (SuggestionServiceType) -> Void
     var codeiumService: CodeiumSuggestionServiceType?
diff --git a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
index 4dcd95ee..3996c541 100644
--- a/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
+++ b/Core/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift
@@ -3,7 +3,7 @@ import GitHubCopilotService
 import Preferences
 import SuggestionModel
 
-final class GitHubCopilotSuggestionProvider: SuggestionServiceProvider {
+actor GitHubCopilotSuggestionProvider: SuggestionServiceProvider {
     let projectRootURL: URL
     let onServiceLaunched: (SuggestionServiceType) -> Void
     var gitHubCopilotService: GitHubCopilotSuggestionServiceType?
diff --git a/Core/Sources/SuggestionService/SuggestionService.swift b/Core/Sources/SuggestionService/SuggestionService.swift
index 5a2c0e8b..4213121a 100644
--- a/Core/Sources/SuggestionService/SuggestionService.swift
+++ b/Core/Sources/SuggestionService/SuggestionService.swift
@@ -27,7 +27,7 @@ public protocol SuggestionServiceType {
 
 protocol SuggestionServiceProvider: SuggestionServiceType {}
 
-public final class SuggestionService: SuggestionServiceType {
+public actor SuggestionService: SuggestionServiceType {
     let projectRootURL: URL
     let onServiceLaunched: (SuggestionServiceType) -> Void
     let providerChangeObserver = UserDefaultsObserver(
@@ -47,8 +47,10 @@ public final class SuggestionService: SuggestionServiceType {
         self.onServiceLaunched = onServiceLaunched
 
         providerChangeObserver.onChange = { [weak self] in
-            guard let self else { return }
-            suggestionProvider = buildService()
+            Task { [weak self] in
+                guard let self else { return }
+                await rebuildService()
+            }
         }
     }
 
@@ -66,6 +68,10 @@ public final class SuggestionService: SuggestionServiceType {
             )
         }
     }
+
+    func rebuildService() {
+        suggestionProvider = buildService()
+    }
 }
 
 public extension SuggestionService {

From 7a8b2f5ee16402d266555df1b35b5da74e568d66 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 17:07:06 +0800
Subject: [PATCH 16/43] Support dependency update

---
 Core/Sources/Service/DependencyUpdater.swift | 76 ++++++++++++++++++++
 ExtensionService/AppDelegate.swift           |  1 +
 2 files changed, 77 insertions(+)
 create mode 100644 Core/Sources/Service/DependencyUpdater.swift

diff --git a/Core/Sources/Service/DependencyUpdater.swift b/Core/Sources/Service/DependencyUpdater.swift
new file mode 100644
index 00000000..a98444db
--- /dev/null
+++ b/Core/Sources/Service/DependencyUpdater.swift
@@ -0,0 +1,76 @@
+import CodeiumService
+import GitHubCopilotService
+import Logger
+
+public struct DependencyUpdater {
+    public init() {}
+
+    public func update() {
+        Task {
+            await withTaskGroup(of: Void.self) { taskGroup in
+                let gitHubCopilot = GitHubCopilotInstallationManager()
+                switch gitHubCopilot.checkInstallation() {
+                case .notInstalled: break
+                case .installed: break
+                case .unsupported: break
+                case .outdated:
+                    taskGroup.addTask {
+                        do {
+                            for try await step in gitHubCopilot.installLatestVersion() {
+                                let state = {
+                                    switch step {
+                                    case .downloading:
+                                        return "Downloading"
+                                    case .uninstalling:
+                                        return "Uninstalling old version"
+                                    case .decompressing:
+                                        return "Decompressing"
+                                    case .done:
+                                        return "Done"
+                                    }
+                                }()
+                                Logger.service
+                                    .error("Update GitHub Copilot language server: \(state)")
+                            }
+                        } catch {
+                            Logger.service.error(
+                                "Update GitHub Copilot language server: \(error.localizedDescription)"
+                            )
+                        }
+                    }
+                }
+                let codeium = CodeiumInstallationManager()
+                switch codeium.checkInstallation() {
+                case .notInstalled: break
+                case .installed: break
+                case .unsupported: break
+                case .outdated:
+                    taskGroup.addTask {
+                        do {
+                            for try await step in codeium.installLatestVersion() {
+                                let state = {
+                                    switch step {
+                                    case .downloading:
+                                        return "Downloading"
+                                    case .uninstalling:
+                                        return "Uninstalling old version"
+                                    case .decompressing:
+                                        return "Decompressing"
+                                    case .done:
+                                        return "Done"
+                                    }
+                                }()
+                                Logger.service.error("Update Codeium language server: \(state)")
+                            }
+                        } catch {
+                            Logger.service.error(
+                                "Update Codeium language server: \(error.localizedDescription)"
+                            )
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
diff --git a/ExtensionService/AppDelegate.swift b/ExtensionService/AppDelegate.swift
index e2a761a9..501ff6a8 100644
--- a/ExtensionService/AppDelegate.swift
+++ b/ExtensionService/AppDelegate.swift
@@ -42,6 +42,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSWindowDelegate {
         Logger.service.info("XPC Service started.")
         NSApp.setActivationPolicy(.accessory)
         buildStatusBarMenu()
+        DependencyUpdater().update()
         Task {
             do {
                 try await ServiceUpdateMigrator().migrate()

From bd3103c885d89214b2fd57217af610d30863d918 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 17:07:17 +0800
Subject: [PATCH 17/43] Bump Codeium to 1.2.25

---
 Core/Sources/CodeiumService/CodeiumInstallationManager.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Core/Sources/CodeiumService/CodeiumInstallationManager.swift b/Core/Sources/CodeiumService/CodeiumInstallationManager.swift
index 207d2855..2f8190da 100644
--- a/Core/Sources/CodeiumService/CodeiumInstallationManager.swift
+++ b/Core/Sources/CodeiumService/CodeiumInstallationManager.swift
@@ -3,7 +3,7 @@ import Terminal
 
 public struct CodeiumInstallationManager {
     private static var isInstalling = false
-    static let latestSupportedVersion = "1.2.17"
+    static let latestSupportedVersion = "1.2.25"
 
     public init() {}
 

From f21722c3364666e719481dd9eae53f77f47ad582 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 17:30:52 +0800
Subject: [PATCH 18/43] Implement CancelRequest for Codeium

---
 .../CodeiumService/CodeiumRequest.swift       | 21 +++++++++++++++----
 .../CodeiumService/CodeiumService.swift       | 16 ++++++++++----
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumRequest.swift b/Core/Sources/CodeiumService/CodeiumRequest.swift
index 6435a32d..8e5d87b1 100644
--- a/Core/Sources/CodeiumService/CodeiumRequest.swift
+++ b/Core/Sources/CodeiumService/CodeiumRequest.swift
@@ -48,12 +48,25 @@ enum CodeiumRequest {
         }
     }
     
-    struct AcceptCompletion: CodeiumRequestType {
-        struct Response: Codable {
-            var state: State
-            var completionItems: [CodeiumCompletionItem]?
+    struct CancelRequest: CodeiumRequestType {
+        struct Response: Codable {}
+
+        struct RequestBody: Codable {
+            var request_id: UInt64
+            var session_id: String
         }
 
+        var requestBody: RequestBody
+
+        func makeURLRequest(server: String) -> URLRequest {
+            let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
+            return assembleURLRequest(server: server, method: "CancelRequest", body: data)
+        }
+    }
+    
+    struct AcceptCompletion: CodeiumRequestType {
+        struct Response: Codable {}
+
         struct RequestBody: Codable {
             var metadata: Metadata
             var completion_id: String
diff --git a/Core/Sources/CodeiumService/CodeiumService.swift b/Core/Sources/CodeiumService/CodeiumService.swift
index 37351067..77116c04 100644
--- a/Core/Sources/CodeiumService/CodeiumService.swift
+++ b/Core/Sources/CodeiumService/CodeiumService.swift
@@ -257,13 +257,13 @@ 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()
         }
@@ -292,8 +292,16 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
             )
         } ?? []
     }
-    
+
     public func cancelRequest() async {
+        Task {
+            try await server?.sendRequest(
+                CodeiumRequest.CancelRequest(requestBody: .init(
+                    request_id: requestCounter,
+                    session_id: CodeiumSuggestionService.sessionId
+                ))
+            )
+        }
         cancellationCounter = requestCounter
     }
 
@@ -326,7 +334,7 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
     public func notifyCloseTextDocument(fileURL: URL) async throws {
         await openedDocumentPool.closeDocument(url: fileURL)
     }
-    
+
     public func terminate() {
         server?.terminate()
         server = nil

From 30204ff3e1b3c722c6b3877564bd5b24ff0442c6 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 18:20:12 +0800
Subject: [PATCH 19/43] Remove exit button in widget

---
 Core/Sources/SuggestionWidget/WidgetView.swift | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/Core/Sources/SuggestionWidget/WidgetView.swift b/Core/Sources/SuggestionWidget/WidgetView.swift
index e266d7b6..50bfae00 100644
--- a/Core/Sources/SuggestionWidget/WidgetView.swift
+++ b/Core/Sources/SuggestionWidget/WidgetView.swift
@@ -211,12 +211,6 @@ struct WidgetContextMenu: View {
             }
 
             Divider()
-
-            Button(action: {
-                exit(0)
-            }) {
-                Text("Quit")
-            }
         }
         .onAppear {
             updateProjectPath(fileURL: widgetViewModel.currentFileURL)

From 11f9f8dc69cf259270113243b6d6af5266c31e7b Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 18:20:31 +0800
Subject: [PATCH 20/43] Update the way we use keychain

---
 .../CodeiumService/CodeiumAuthService.swift   |  5 +++-
 .../ServiceUpdateMigrator.swift               | 24 +++++++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumAuthService.swift b/Core/Sources/CodeiumService/CodeiumAuthService.swift
index 5d7c439c..dbb33903 100644
--- a/Core/Sources/CodeiumService/CodeiumAuthService.swift
+++ b/Core/Sources/CodeiumService/CodeiumAuthService.swift
@@ -4,10 +4,13 @@ import KeychainAccess
 
 public final class CodeiumAuthService {
     public init() {}
-    let codeiumKeyKey = "codeiumKey"
+    let codeiumKeyKey = "codeiumAuthKey"
     let keychain: Keychain = {
         let info = Bundle.main.infoDictionary
         return Keychain(service: keychainService, accessGroup: keychainAccessGroup)
+            .attributes([
+                kSecUseDataProtectionKeychain as String: true,
+            ])
     }()
 
     var key: String? { try? keychain.getString(codeiumKeyKey) }
diff --git a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
index 8a995806..abc88455 100644
--- a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
+++ b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
@@ -1,5 +1,7 @@
+import Configs
 import Foundation
 import GitHubCopilotService
+import KeychainAccess
 import Preferences
 
 extension UserDefaultPreferenceKeys {
@@ -13,7 +15,7 @@ extension UserDefaultPreferenceKeys {
 
 public struct ServiceUpdateMigrator {
     public init() {}
-    
+
     public func migrate() async throws {
         let version = Bundle.main.object(forInfoDictionaryKey: "CFBundleVersion") as? String ?? "0"
 
@@ -24,7 +26,11 @@ public struct ServiceUpdateMigrator {
     func migrate(from oldVersion: String, to currentVersion: String) async throws {
         guard let old = Int(oldVersion) else { return }
         if old <= 135 {
-           try migrateFromLowerThanOrEqualToVersion135()
+            try migrateFromLowerThanOrEqualToVersion135()
+        }
+
+        if old <= 170 {
+            try migrateFromLowerThanOrEqualToVersion170()
         }
     }
 }
@@ -73,3 +79,17 @@ func migrateFromLowerThanOrEqualToVersion135() throws {
     )
 }
 
+func migrateFromLowerThanOrEqualToVersion170() throws {
+    let oldKeychain = Keychain(service: keychainService, accessGroup: keychainAccessGroup)
+    let newKeychain = oldKeychain.attributes([
+        kSecUseDataProtectionKeychain as String: true,
+    ])
+
+    if (try? oldKeychain.contains("codeiumKey")) ?? false,
+       let key = try? oldKeychain.getString("codeiumKey")
+    {
+        try newKeychain.set(key, key: "codeiumAuthKey")
+        try? oldKeychain.set("", key: "codeiumKey")
+    }
+}
+

From 90bacd96dda16bb61a43ea7d5ef3240b4c2f84a6 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 18:26:46 +0800
Subject: [PATCH 21/43] Update to use SecureField

---
 Core/Sources/HostApp/AccountSettings/OpenAIView.swift | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
index e7c3e924..bf15ef55 100644
--- a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
+++ b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -27,9 +27,10 @@ struct OpenAIView: View {
     var body: some View {
         Form {
             HStack {
-                TextField(text: $settings.openAIAPIKey, prompt: Text("sk-*")) {
+                SecureField(text: $settings.openAIAPIKey, prompt: Text("sk-*")) {
                     Text("OpenAI API Key")
-                }.textFieldStyle(.roundedBorder)
+                }
+                .textFieldStyle(.roundedBorder)
                 Button(action: {
                     openURL(apiKeyURL)
                 }) {

From a27f5c7bfbd0f2b3c7a9ed787810fb13db8bbd47 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 18:35:06 +0800
Subject: [PATCH 22/43] Allow setting max token that is larger than the size
 defined in the model

---
 .../HostApp/AccountSettings/OpenAIView.swift  | 60 +++++++++++--------
 Core/Sources/Preferences/ChatGPTModel.swift   |  4 +-
 2 files changed, 38 insertions(+), 26 deletions(-)

diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
index bf15ef55..4fb5870b 100644
--- a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
+++ b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -23,6 +23,7 @@ struct OpenAIView: View {
     )!
     @Environment(\.openURL) var openURL
     @StateObject var settings = OpenAIViewSettings()
+    @State var maxTokenOverLimit = false
 
     var body: some View {
         Form {
@@ -76,33 +77,39 @@ struct OpenAIView: View {
                 }
             }
 
-            if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
-                let binding = Binding(
-                    get: { String(settings.chatGPTMaxToken) },
-                    set: {
-                        if let selectionMaxToken = Int($0) {
-                            settings.chatGPTMaxToken = model
-                                .maxToken < selectionMaxToken ? model
-                                .maxToken : selectionMaxToken
+            let binding = Binding(
+                get: { String(settings.chatGPTMaxToken) },
+                set: {
+                    if let selectionMaxToken = Int($0) {
+                        if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
+                            maxTokenOverLimit = model.maxToken < selectionMaxToken
                         } else {
-                            settings.chatGPTMaxToken = 0
+                            maxTokenOverLimit = false
                         }
+                        settings.chatGPTMaxToken = selectionMaxToken
+                    } else {
+                        settings.chatGPTMaxToken = 0
                     }
-                )
-                HStack {
-                    Stepper(
-                        value: $settings.chatGPTMaxToken,
-                        in: 0...model.maxToken,
-                        step: 1
-                    ) {
-                        Text("Max Token (Including Reply)")
-                            .multilineTextAlignment(.trailing)
-                    }
-                    TextField(text: binding) {
-                        EmptyView()
-                    }
-                    .labelsHidden()
-                    .textFieldStyle(.roundedBorder)
+                }
+            )
+            HStack {
+                Stepper(
+                    value: $settings.chatGPTMaxToken,
+                    in: 0...Int.max,
+                    step: 1
+                ) {
+                    Text("Max Token (Including Reply)")
+                        .multilineTextAlignment(.trailing)
+                }
+                TextField(text: binding) {
+                    EmptyView()
+                }
+                .labelsHidden()
+                .textFieldStyle(.roundedBorder)
+                .foregroundColor(maxTokenOverLimit ? .red : .primary)
+                
+                if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
+                    Text("Max: \(model.maxToken)")
                 }
             }
 
@@ -134,6 +141,11 @@ struct OpenAIView: View {
                 Text("7 Messages").tag(7)
             }
         }
+        .onAppear {
+            if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
+                maxTokenOverLimit = model.maxToken < settings.chatGPTMaxToken
+            }
+        }
     }
 
     var languagePicker: some View {
diff --git a/Core/Sources/Preferences/ChatGPTModel.swift b/Core/Sources/Preferences/ChatGPTModel.swift
index 8285a82a..f2d365d1 100644
--- a/Core/Sources/Preferences/ChatGPTModel.swift
+++ b/Core/Sources/Preferences/ChatGPTModel.swift
@@ -30,9 +30,9 @@ public extension ChatGPTModel {
         case .gpt432k0314:
             return 32768
         case .gpt35Turbo:
-            return 8192
+            return 4096
         case .gpt35Turbo0301:
-            return 8192
+            return 4096
         }
     }
 }

From 182d4261870ff1a06c0b3aead3cfbb29016d48ae Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 18:56:28 +0800
Subject: [PATCH 23/43] Change ChatGPTEndpoint to OpenAIBaseURL

---
 Core/Package.swift                            |  1 +
 .../HostApp/AccountSettings/OpenAIView.swift  | 34 ++++++++++++++-----
 .../OpenAIService/ChatGPTService.swift        |  7 ++--
 Core/Sources/Preferences/Keys.swift           |  7 +++-
 Core/Sources/Preferences/UserDefaults.swift   | 16 +++++++++
 5 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/Core/Package.swift b/Core/Package.swift
index d4a83b9f..9a2d826a 100644
--- a/Core/Package.swift
+++ b/Core/Package.swift
@@ -42,6 +42,7 @@ let package = Package(
                 "LaunchAgentManager",
                 "Logger",
                 "UpdateChecker",
+                "OpenAIService",
             ]
         ),
     ],
diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
index 4fb5870b..bc4cef87 100644
--- a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
+++ b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -1,4 +1,5 @@
 import AppKit
+import OpenAIService
 import Client
 import Preferences
 import SuggestionModel
@@ -8,7 +9,7 @@ final class OpenAIViewSettings: ObservableObject {
     static let availableLocalizedLocales = Locale.availableLocalizedLocales
     @AppStorage(\.openAIAPIKey) var openAIAPIKey: String
     @AppStorage(\.chatGPTModel) var chatGPTModel: String
-    @AppStorage(\.chatGPTEndpoint) var chatGPTEndpoint: String
+    @AppStorage(\.openAIBaseURL) var openAIBaseURL: String
     @AppStorage(\.chatGPTLanguage) var chatGPTLanguage: String
     @AppStorage(\.chatGPTMaxToken) var chatGPTMaxToken: Int
     @AppStorage(\.chatGPTTemperature) var chatGPTTemperature: Double
@@ -22,6 +23,7 @@ struct OpenAIView: View {
         string: "https://platform.openai.com/docs/models/model-endpoint-compatibility"
     )!
     @Environment(\.openURL) var openURL
+    @Environment(\.toast) var toast
     @StateObject var settings = OpenAIViewSettings()
     @State var maxTokenOverLimit = false
 
@@ -39,6 +41,27 @@ struct OpenAIView: View {
                 }.buttonStyle(.plain)
             }
 
+            HStack {
+                TextField(
+                    text: $settings.openAIBaseURL,
+                    prompt: Text("https://api.openai.com")
+                ) {
+                    Text("OpenAI Base URL")
+                }.textFieldStyle(.roundedBorder)
+
+                Button("Test") {
+                    Task {
+                        do {
+                            let reply = try await ChatGPTService()
+                                .sendAndWait(content: "Hello", summary: nil)
+                            toast(Text("ChatGPT replied: \(reply ?? "N/A")"), .info)
+                        } catch {
+                            toast(Text(error.localizedDescription), .error)
+                        }
+                    }
+                }
+            }
+
             HStack {
                 Picker(selection: $settings.chatGPTModel) {
                     if !settings.chatGPTModel.isEmpty,
@@ -59,13 +82,6 @@ struct OpenAIView: View {
                 }.buttonStyle(.plain)
             }
 
-            TextField(
-                text: $settings.chatGPTEndpoint,
-                prompt: Text("https://api.openai.com/v1/chat/completions")
-            ) {
-                Text("ChatGPT Server")
-            }.textFieldStyle(.roundedBorder)
-
             if #available(macOS 13.0, *) {
                 LabeledContent("Reply in Language") {
                     languagePicker
@@ -107,7 +123,7 @@ struct OpenAIView: View {
                 .labelsHidden()
                 .textFieldStyle(.roundedBorder)
                 .foregroundColor(maxTokenOverLimit ? .red : .primary)
-                
+
                 if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
                     Text("Max: \(model.maxToken)")
                 }
diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
index 0294cfb0..594d8d2a 100644
--- a/Core/Sources/OpenAIService/ChatGPTService.swift
+++ b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -69,10 +69,9 @@ public actor ChatGPTService: ChatGPTServiceType {
     }
 
     public var endpoint: String {
-        let value = UserDefaults.shared.value(for: \.chatGPTEndpoint)
-        if value.isEmpty { return "https://api.openai.com/v1/chat/completions" }
-
-        return value
+        var baseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
+        if baseURL.isEmpty { return "https://api.openai.com/v1/chat/completions" }
+        return "\(baseURL)/v1/chat/completions"
     }
 
     public var apiKey: String {
diff --git a/Core/Sources/Preferences/Keys.swift b/Core/Sources/Preferences/Keys.swift
index f3533b14..83b141d7 100644
--- a/Core/Sources/Preferences/Keys.swift
+++ b/Core/Sources/Preferences/Keys.swift
@@ -69,10 +69,15 @@ public extension UserDefaultPreferenceKeys {
         .init(defaultValue: "", key: "OpenAIAPIKey")
     }
 
+    @available(*, deprecated, message: "Use `openAIBaseURL` instead.")
     var chatGPTEndpoint: PreferenceKey<String> {
         .init(defaultValue: "", key: "ChatGPTEndpoint")
     }
 
+    var openAIBaseURL: PreferenceKey<String> {
+        .init(defaultValue: "", key: "OpenAIBaseURL")
+    }
+
     var chatGPTModel: PreferenceKey<String> {
         .init(defaultValue: Preferences.ChatGPTModel.gpt35Turbo.rawValue, key: "ChatGPTModel")
     }
@@ -196,7 +201,7 @@ public extension UserDefaultPreferenceKeys {
     var embedFileContentInChatContextIfNoSelection: PreferenceKey<Bool> {
         .init(defaultValue: false, key: "EmbedFileContentInChatContextIfNoSelection")
     }
-    
+
     var maxEmbeddableFileInChatContextLineCount: PreferenceKey<Int> {
         .init(defaultValue: 100, key: "MaxEmbeddableFileInChatContextLineCount")
     }
diff --git a/Core/Sources/Preferences/UserDefaults.swift b/Core/Sources/Preferences/UserDefaults.swift
index ec44d704..207edcc5 100644
--- a/Core/Sources/Preferences/UserDefaults.swift
+++ b/Core/Sources/Preferences/UserDefaults.swift
@@ -13,6 +13,12 @@ public extension UserDefaults {
         shared.setupDefaultValue(for: \.widgetColorScheme)
         shared.setupDefaultValue(for: \.customCommands)
         shared.setupDefaultValue(for: \.runNodeWith, defaultValue: .env)
+        shared.setupDefaultValue(for: \.openAIBaseURL, defaultValue: {
+            guard let url = URL(string: shared.value(for: \.chatGPTEndpoint)) else { return "" }
+            let scheme = url.scheme ?? "https"
+            guard let host = url.host else { return "" }
+            return "\(scheme)://\(host)"
+        }() as String)
     }
 }
 
@@ -71,6 +77,16 @@ public extension UserDefaults {
             set(key.defaultValue, forKey: key.key)
         }
     }
+    
+    func setupDefaultValue<K: UserDefaultPreferenceKey>(
+        for keyPath: KeyPath<UserDefaultPreferenceKeys, K>,
+        defaultValue: K.Value
+    ) where K.Value: UserDefaultsStorable {
+        let key = UserDefaultPreferenceKeys()[keyPath: keyPath]
+        if value(forKey: key.key) == nil {
+            set(defaultValue, forKey: key.key)
+        }
+    }
 
     // MARK: - Raw Representable
 

From a707645d86e91ae6e5b51f14d4695c9ba77be8c6 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 21:42:24 +0800
Subject: [PATCH 24/43] Support Azure OpenAI as chat provider

---
 .../HostApp/AccountSettings/AzureView.swift   | 70 +++++++++++++++++++
 .../HostApp/AccountSettings/OpenAIView.swift  |  2 +-
 Core/Sources/HostApp/ServiceView.swift        |  9 +++
 .../OpenAIService/ChatGPTService.swift        | 48 ++++++++++---
 .../Sources/OpenAIService/CompletionAPI.swift | 32 ++++++---
 .../OpenAIService/CompletionStreamAPI.swift   | 23 ++++--
 .../Preferences/ChatFeatureProvider.swift     |  4 ++
 Core/Sources/Preferences/Keys.swift           | 20 ++++++
 8 files changed, 186 insertions(+), 22 deletions(-)
 create mode 100644 Core/Sources/HostApp/AccountSettings/AzureView.swift
 create mode 100644 Core/Sources/Preferences/ChatFeatureProvider.swift

diff --git a/Core/Sources/HostApp/AccountSettings/AzureView.swift b/Core/Sources/HostApp/AccountSettings/AzureView.swift
new file mode 100644
index 00000000..d28e111b
--- /dev/null
+++ b/Core/Sources/HostApp/AccountSettings/AzureView.swift
@@ -0,0 +1,70 @@
+import AppKit
+import Client
+import OpenAIService
+import Preferences
+import SuggestionModel
+import SwiftUI
+
+final class AzureViewSettings: ObservableObject {
+    @AppStorage(\.azureOpenAIAPIKey) var azureOpenAIAPIKey: String
+    @AppStorage(\.azureOpenAIBaseURL) var azureOpenAIBaseURL: String
+    @AppStorage(\.azureChatGPTDeployment) var azureChatGPTDeployment: String
+    init() {}
+}
+
+struct AzureView: View {
+    @Environment(\.toast) var toast
+    @State var isTesting = false
+    @StateObject var settings = AzureViewSettings()
+
+    var body: some View {
+        Form {
+            SecureField(text: $settings.azureOpenAIAPIKey, prompt: Text("")) {
+                Text("OpenAI Service API Key")
+            }
+            .textFieldStyle(.roundedBorder)
+
+            TextField(
+                text: $settings.azureOpenAIBaseURL,
+                prompt: Text("https://XXXXXX.openai.azure.com")
+            ) {
+                Text("OpenAI Service Base URL")
+            }.textFieldStyle(.roundedBorder)
+
+            HStack {
+                TextField(
+                    text: $settings.azureChatGPTDeployment,
+                    prompt: Text("")
+                ) {
+                    Text("Chat Model Deployment Name")
+                }.textFieldStyle(.roundedBorder)
+
+                Button("Test") {
+                    Task { @MainActor in
+                        isTesting = true
+                        defer { isTesting = false }
+                        do {
+                            let reply = try await ChatGPTService(designatedProvider: .azureOpenAI)
+                                .sendAndWait(content: "Hello", summary: nil)
+                            toast(Text("ChatGPT replied: \(reply ?? "N/A")"), .info)
+                        } catch {
+                            toast(Text(error.localizedDescription), .error)
+                        }
+                    }
+                }
+                .disabled(isTesting)
+            }
+        }
+    }
+}
+
+struct AzureView_Previews: PreviewProvider {
+    static var previews: some View {
+        VStack(alignment: .leading, spacing: 8) {
+            AzureView()
+        }
+        .frame(height: 800)
+        .padding(.all, 8)
+    }
+}
+
diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
index bc4cef87..37adc235 100644
--- a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
+++ b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -52,7 +52,7 @@ struct OpenAIView: View {
                 Button("Test") {
                     Task {
                         do {
-                            let reply = try await ChatGPTService()
+                            let reply = try await ChatGPTService(designatedProvider: .openAI)
                                 .sendAndWait(content: "Hello", summary: nil)
                             toast(Text("ChatGPT replied: \(reply ?? "N/A")"), .info)
                         } catch {
diff --git a/Core/Sources/HostApp/ServiceView.swift b/Core/Sources/HostApp/ServiceView.swift
index d9d26eee..3cb35ab3 100644
--- a/Core/Sources/HostApp/ServiceView.swift
+++ b/Core/Sources/HostApp/ServiceView.swift
@@ -30,6 +30,15 @@ struct ServiceView: View {
                 subtitle: "Chat, Prompt to Code",
                 image: "globe"
             )
+            
+            ScrollView {
+                AzureView().padding()
+            }.sidebarItem(
+                tag: 3,
+                title: "Azure",
+                subtitle: "Chat, Prompt to Code",
+                image: "globe"
+            )
         }
     }
 }
diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
index 594d8d2a..0cb3dca7 100644
--- a/Core/Sources/OpenAIService/ChatGPTService.swift
+++ b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -17,6 +17,7 @@ public protocol ChatGPTServiceType: ObservableObject {
 public enum ChatGPTServiceError: Error, LocalizedError {
     case endpointIncorrect
     case responseInvalid
+    case otherError(String)
 
     public var errorDescription: String? {
         switch self {
@@ -24,6 +25,8 @@ public enum ChatGPTServiceError: Error, LocalizedError {
             return "ChatGPT endpoint is incorrect"
         case .responseInvalid:
             return "Response is invalid"
+        case let .otherError(content):
+            return content
         }
     }
 }
@@ -59,7 +62,7 @@ public actor ChatGPTService: ChatGPTServiceType {
     public var defaultTemperature: Double {
         min(max(0, UserDefaults.shared.value(for: \.chatGPTTemperature)), 2)
     }
-    
+
     var temperature: Double?
 
     public var model: String {
@@ -68,14 +71,30 @@ public actor ChatGPTService: ChatGPTServiceType {
         return value
     }
 
+    var designatedProvider: ChatFeatureProvider?
+
     public var endpoint: String {
-        var baseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
-        if baseURL.isEmpty { return "https://api.openai.com/v1/chat/completions" }
-        return "\(baseURL)/v1/chat/completions"
+        switch designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider) {
+        case .openAI:
+            let baseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
+            if baseURL.isEmpty { return "https://api.openai.com/v1/chat/completions" }
+            return "\(baseURL)/v1/chat/completions"
+        case .azureOpenAI:
+            let baseURL = UserDefaults.shared.value(for: \.azureOpenAIBaseURL)
+            let deployment = UserDefaults.shared.value(for: \.azureChatGPTDeployment)
+            let version = "2023-05-15"
+            if baseURL.isEmpty { return "" }
+            return "\(baseURL)/openai/deployments/\(deployment)/chat/completions?api-version=\(version)"
+        }
     }
 
     public var apiKey: String {
-        UserDefaults.shared.value(for: \.openAIAPIKey)
+        switch designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider) {
+        case .openAI:
+            return UserDefaults.shared.value(for: \.openAIAPIKey)
+        case .azureOpenAI:
+            return UserDefaults.shared.value(for: \.azureOpenAIAPIKey)
+        }
     }
 
     public var maxToken: Int {
@@ -97,10 +116,12 @@ public actor ChatGPTService: ChatGPTServiceType {
 
     public init(
         systemPrompt: String = "",
-        temperature: Double? = nil
+        temperature: Double? = nil,
+        designatedProvider: ChatFeatureProvider? = nil
     ) {
         self.systemPrompt = systemPrompt
         self.temperature = temperature
+        self.designatedProvider = designatedProvider
     }
 
     public func send(
@@ -129,7 +150,12 @@ public actor ChatGPTService: ChatGPTServiceType {
 
         isReceivingMessage = true
 
-        let api = buildCompletionStreamAPI(apiKey, url, requestBody)
+        let api = buildCompletionStreamAPI(
+            apiKey,
+            designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
+            url,
+            requestBody
+        )
 
         return AsyncThrowingStream<String, Error> { continuation in
             Task {
@@ -210,7 +236,12 @@ public actor ChatGPTService: ChatGPTServiceType {
         isReceivingMessage = true
         defer { isReceivingMessage = false }
 
-        let api = buildCompletionAPI(apiKey, url, requestBody)
+        let api = buildCompletionAPI(
+            apiKey,
+            designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
+            url,
+            requestBody
+        )
         let response = try await api()
 
         if let choice = response.choices.first {
@@ -296,3 +327,4 @@ func maxTokenForReply(model: String, remainingTokens: Int) -> Int {
     guard let model = ChatGPTModel(rawValue: model) else { return remainingTokens }
     return min(model.maxToken / 2, remainingTokens)
 }
+
diff --git a/Core/Sources/OpenAIService/CompletionAPI.swift b/Core/Sources/OpenAIService/CompletionAPI.swift
index 75353f84..66194616 100644
--- a/Core/Sources/OpenAIService/CompletionAPI.swift
+++ b/Core/Sources/OpenAIService/CompletionAPI.swift
@@ -1,6 +1,8 @@
 import Foundation
+import Preferences
 
-typealias CompletionAPIBuilder = (String, URL, CompletionRequestBody) -> CompletionAPI
+typealias CompletionAPIBuilder = (String, ChatFeatureProvider, URL, CompletionRequestBody)
+    -> CompletionAPI
 
 protocol CompletionAPI {
     func callAsFunction() async throws -> CompletionResponseBody
@@ -12,13 +14,13 @@ struct CompletionResponseBody: Codable, Equatable {
         var role: ChatMessage.Role
         var content: String
     }
-    
+
     struct Choice: Codable, Equatable {
         var message: Message
         var index: Int
         var finish_reason: String
     }
-    
+
     struct Usage: Codable, Equatable {
         var prompt_tokens: Int
         var completion_tokens: Int
@@ -40,8 +42,9 @@ struct CompletionAPIError: Error, Codable, LocalizedError {
         var param: String
         var code: String
     }
+
     var error: E
-    
+
     var errorDescription: String? { error.message }
 }
 
@@ -49,12 +52,19 @@ struct OpenAICompletionAPI: CompletionAPI {
     var apiKey: String
     var endpoint: URL
     var requestBody: CompletionRequestBody
+    var provider: ChatFeatureProvider
 
-    init(apiKey: String, endpoint: URL, requestBody: CompletionRequestBody) {
+    init(
+        apiKey: String,
+        provider: ChatFeatureProvider,
+        endpoint: URL,
+        requestBody: CompletionRequestBody
+    ) {
         self.apiKey = apiKey
         self.endpoint = endpoint
         self.requestBody = requestBody
         self.requestBody.stream = false
+        self.provider = provider
     }
 
     func callAsFunction() async throws -> CompletionResponseBody {
@@ -64,7 +74,11 @@ struct OpenAICompletionAPI: CompletionAPI {
         request.httpBody = try encoder.encode(requestBody)
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
         if !apiKey.isEmpty {
-            request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            if provider == .openAI {
+                request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            } else {
+                request.setValue(apiKey, forHTTPHeaderField: "api-key")
+            }
         }
 
         let (result, response) = try await URLSession.shared.data(for: request)
@@ -74,9 +88,11 @@ struct OpenAICompletionAPI: CompletionAPI {
 
         guard response.statusCode == 200 else {
             let error = try? JSONDecoder().decode(CompletionAPIError.self, from: result)
-            throw error ?? ChatGPTServiceError.responseInvalid
+            throw error ?? ChatGPTServiceError
+                .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }
-        
+
         return try JSONDecoder().decode(CompletionResponseBody.self, from: result)
     }
 }
+
diff --git a/Core/Sources/OpenAIService/CompletionStreamAPI.swift b/Core/Sources/OpenAIService/CompletionStreamAPI.swift
index 4274c01c..964c82ce 100644
--- a/Core/Sources/OpenAIService/CompletionStreamAPI.swift
+++ b/Core/Sources/OpenAIService/CompletionStreamAPI.swift
@@ -1,7 +1,8 @@
 import AsyncAlgorithms
 import Foundation
+import Preferences
 
-typealias CompletionStreamAPIBuilder = (String, URL, CompletionRequestBody) -> CompletionStreamAPI
+typealias CompletionStreamAPIBuilder = (String, ChatFeatureProvider, URL, CompletionRequestBody) -> CompletionStreamAPI
 
 protocol CompletionStreamAPI {
     func callAsFunction() async throws -> (
@@ -54,12 +55,19 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
     var apiKey: String
     var endpoint: URL
     var requestBody: CompletionRequestBody
+    var provider: ChatFeatureProvider
 
-    init(apiKey: String, endpoint: URL, requestBody: CompletionRequestBody) {
+    init(
+        apiKey: String,
+        provider: ChatFeatureProvider,
+        endpoint: URL,
+        requestBody: CompletionRequestBody
+    ) {
         self.apiKey = apiKey
         self.endpoint = endpoint
         self.requestBody = requestBody
         self.requestBody.stream = true
+        self.provider = provider
     }
 
     func callAsFunction() async throws -> (
@@ -72,7 +80,11 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
         request.httpBody = try encoder.encode(requestBody)
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
         if !apiKey.isEmpty {
-            request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            if provider == .openAI {
+                request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            } else {
+                request.setValue(apiKey, forHTTPHeaderField: "api-key")
+            }
         }
 
         let (result, response) = try await URLSession.shared.bytes(for: request)
@@ -90,9 +102,9 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
             let error = try? decoder.decode(ChatGPTError.self, from: data)
             throw error ?? ChatGPTServiceError.responseInvalid
         }
-        
+
         var receivingDataTask: Task<Void, Error>?
-        
+
         let stream = AsyncThrowingStream<CompletionStreamDataTrunk, Error> { continuation in
             receivingDataTask = Task {
                 do {
@@ -122,3 +134,4 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
         )
     }
 }
+
diff --git a/Core/Sources/Preferences/ChatFeatureProvider.swift b/Core/Sources/Preferences/ChatFeatureProvider.swift
new file mode 100644
index 00000000..d97a4238
--- /dev/null
+++ b/Core/Sources/Preferences/ChatFeatureProvider.swift
@@ -0,0 +1,4 @@
+public enum ChatFeatureProvider: String, CaseIterable {
+    case openAI
+    case azureOpenAI
+}
diff --git a/Core/Sources/Preferences/Keys.swift b/Core/Sources/Preferences/Keys.swift
index 83b141d7..95bad77b 100644
--- a/Core/Sources/Preferences/Keys.swift
+++ b/Core/Sources/Preferences/Keys.swift
@@ -99,6 +99,22 @@ public extension UserDefaultPreferenceKeys {
     }
 }
 
+// MARK: - Azure OpenAI Settings
+
+public extension UserDefaultPreferenceKeys {
+    var azureOpenAIAPIKey: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureOpenAIAPIKey")
+    }
+    
+    var azureOpenAIBaseURL: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureOpenAIBaseURL")
+    }
+    
+    var azureChatGPTDeployment: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureChatGPTDeployment")
+    }
+}
+
 // MARK: - GitHub Copilot Settings
 
 public extension UserDefaultPreferenceKeys {
@@ -186,6 +202,10 @@ public extension UserDefaultPreferenceKeys {
 // MARK: - Chat
 
 public extension UserDefaultPreferenceKeys {
+    var chatFeatureProvider: PreferenceKey<ChatFeatureProvider> {
+        .init(defaultValue: .openAI, key: "ChatFeatureProvider")
+    }
+    
     var chatFontSize: PreferenceKey<Double> {
         .init(defaultValue: 12, key: "ChatFontSize")
     }

From 16d9cf71a7d45433ad91d3287a858bd3ff9e1784 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 22:01:12 +0800
Subject: [PATCH 25/43] Move settings to chat settings view

---
 .../HostApp/AccountSettings/OpenAIView.swift  | 144 ++------------
 .../FeatureSettings/ChatSettingsView.swift    | 183 +++++++++++++++++-
 2 files changed, 190 insertions(+), 137 deletions(-)

diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
index 37adc235..a6dea775 100644
--- a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
+++ b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -1,31 +1,26 @@
 import AppKit
-import OpenAIService
 import Client
+import OpenAIService
 import Preferences
 import SuggestionModel
 import SwiftUI
 
-final class OpenAIViewSettings: ObservableObject {
-    static let availableLocalizedLocales = Locale.availableLocalizedLocales
-    @AppStorage(\.openAIAPIKey) var openAIAPIKey: String
-    @AppStorage(\.chatGPTModel) var chatGPTModel: String
-    @AppStorage(\.openAIBaseURL) var openAIBaseURL: String
-    @AppStorage(\.chatGPTLanguage) var chatGPTLanguage: String
-    @AppStorage(\.chatGPTMaxToken) var chatGPTMaxToken: Int
-    @AppStorage(\.chatGPTTemperature) var chatGPTTemperature: Double
-    @AppStorage(\.chatGPTMaxMessageCount) var chatGPTMaxMessageCount: Int
-    init() {}
-}
-
 struct OpenAIView: View {
+    final class Settings: ObservableObject {
+        @AppStorage(\.openAIAPIKey) var openAIAPIKey: String
+        @AppStorage(\.chatGPTModel) var chatGPTModel: String
+        @AppStorage(\.openAIBaseURL) var openAIBaseURL: String
+        init() {}
+    }
+    
     let apiKeyURL = URL(string: "https://platform.openai.com/account/api-keys")!
     let modelURL = URL(
         string: "https://platform.openai.com/docs/models/model-endpoint-compatibility"
     )!
     @Environment(\.openURL) var openURL
     @Environment(\.toast) var toast
-    @StateObject var settings = OpenAIViewSettings()
-    @State var maxTokenOverLimit = false
+    @State var isTesting = false
+    @StateObject var settings = Settings()
 
     var body: some View {
         Form {
@@ -50,7 +45,9 @@ struct OpenAIView: View {
                 }.textFieldStyle(.roundedBorder)
 
                 Button("Test") {
-                    Task {
+                    Task { @MainActor in
+                        isTesting = true
+                        defer { isTesting = false }
                         do {
                             let reply = try await ChatGPTService(designatedProvider: .openAI)
                                 .sendAndWait(content: "Hello", summary: nil)
@@ -59,7 +56,7 @@ struct OpenAIView: View {
                             toast(Text(error.localizedDescription), .error)
                         }
                     }
-                }
+                }.disabled(isTesting)
             }
 
             HStack {
@@ -81,119 +78,6 @@ struct OpenAIView: View {
                     Image(systemName: "questionmark.circle.fill")
                 }.buttonStyle(.plain)
             }
-
-            if #available(macOS 13.0, *) {
-                LabeledContent("Reply in Language") {
-                    languagePicker
-                }
-            } else {
-                HStack {
-                    Text("Reply in Language")
-                    languagePicker
-                }
-            }
-
-            let binding = Binding(
-                get: { String(settings.chatGPTMaxToken) },
-                set: {
-                    if let selectionMaxToken = Int($0) {
-                        if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
-                            maxTokenOverLimit = model.maxToken < selectionMaxToken
-                        } else {
-                            maxTokenOverLimit = false
-                        }
-                        settings.chatGPTMaxToken = selectionMaxToken
-                    } else {
-                        settings.chatGPTMaxToken = 0
-                    }
-                }
-            )
-            HStack {
-                Stepper(
-                    value: $settings.chatGPTMaxToken,
-                    in: 0...Int.max,
-                    step: 1
-                ) {
-                    Text("Max Token (Including Reply)")
-                        .multilineTextAlignment(.trailing)
-                }
-                TextField(text: binding) {
-                    EmptyView()
-                }
-                .labelsHidden()
-                .textFieldStyle(.roundedBorder)
-                .foregroundColor(maxTokenOverLimit ? .red : .primary)
-
-                if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
-                    Text("Max: \(model.maxToken)")
-                }
-            }
-
-            HStack {
-                Slider(value: $settings.chatGPTTemperature, in: 0...2, step: 0.1) {
-                    Text("Temperature")
-                }
-
-                Text(
-                    "\(settings.chatGPTTemperature.formatted(.number.precision(.fractionLength(1))))"
-                )
-                .font(.body)
-                .monospacedDigit()
-                .padding(.vertical, 2)
-                .padding(.horizontal, 6)
-                .background(
-                    RoundedRectangle(cornerRadius: 4, style: .continuous)
-                        .fill(Color.primary.opacity(0.1))
-                )
-            }
-
-            Picker(
-                "Memory",
-                selection: $settings.chatGPTMaxMessageCount
-            ) {
-                Text("No Limit").tag(0)
-                Text("3 Messages").tag(3)
-                Text("5 Messages").tag(5)
-                Text("7 Messages").tag(7)
-            }
-        }
-        .onAppear {
-            if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
-                maxTokenOverLimit = model.maxToken < settings.chatGPTMaxToken
-            }
-        }
-    }
-
-    var languagePicker: some View {
-        Menu {
-            if !settings.chatGPTLanguage.isEmpty,
-               !OpenAIViewSettings.availableLocalizedLocales
-               .contains(settings.chatGPTLanguage)
-            {
-                Button(
-                    settings.chatGPTLanguage,
-                    action: { self.settings.chatGPTLanguage = settings.chatGPTLanguage }
-                )
-            }
-            Button(
-                "Auto-detected by ChatGPT",
-                action: { self.settings.chatGPTLanguage = "" }
-            )
-            ForEach(
-                OpenAIViewSettings.availableLocalizedLocales,
-                id: \.self
-            ) { localizedLocales in
-                Button(
-                    localizedLocales,
-                    action: { self.settings.chatGPTLanguage = localizedLocales }
-                )
-            }
-        } label: {
-            Text(
-                settings.chatGPTLanguage.isEmpty
-                    ? "Auto-detected by ChatGPT"
-                    : settings.chatGPTLanguage
-            )
         }
     }
 }
diff --git a/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift b/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
index 82b1ce4c..1388902c 100644
--- a/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
+++ b/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
@@ -1,19 +1,138 @@
+import Preferences
 import SwiftUI
 
 struct ChatSettingsView: View {
     class Settings: ObservableObject {
+        static let availableLocalizedLocales = Locale.availableLocalizedLocales
+        @AppStorage(\.chatGPTLanguage) var chatGPTLanguage
+        @AppStorage(\.chatGPTMaxToken) var chatGPTMaxToken
+        @AppStorage(\.chatGPTTemperature) var chatGPTTemperature
+        @AppStorage(\.chatGPTMaxMessageCount) var chatGPTMaxMessageCount
         @AppStorage(\.chatFontSize) var chatFontSize
         @AppStorage(\.chatCodeFontSize) var chatCodeFontSize
         @AppStorage(\.embedFileContentInChatContextIfNoSelection)
         var embedFileContentInChatContextIfNoSelection
         @AppStorage(\.maxEmbeddableFileInChatContextLineCount)
         var maxEmbeddableFileInChatContextLineCount
+
+        @AppStorage(\.chatFeatureProvider) var chatFeatureProvider
+        @AppStorage(\.chatGPTModel) var chatGPTModel
+
         init() {}
     }
-    
+
+    @Environment(\.openURL) var openURL
+    @Environment(\.toast) var toast
     @StateObject var settings = Settings()
-    
+    @State var maxTokenOverLimit = false
+
     var body: some View {
+        VStack {
+            chatSettingsForm
+            Divider()
+            uiForm
+            Divider()
+            contextForm
+        }
+    }
+
+    @ViewBuilder
+    var chatSettingsForm: some View {
+        Form {
+            Picker(
+                "Feature Provider",
+                selection: $settings.chatFeatureProvider
+            ) {
+                Text("OpenAI").tag(ChatFeatureProvider.openAI)
+                Text("Azure OpenAI").tag(ChatFeatureProvider.azureOpenAI)
+            }
+
+            if #available(macOS 13.0, *) {
+                LabeledContent("Reply in Language") {
+                    languagePicker
+                }
+            } else {
+                HStack {
+                    Text("Reply in Language")
+                    languagePicker
+                }
+            }
+
+            let binding = Binding(
+                get: { String(settings.chatGPTMaxToken) },
+                set: {
+                    if let selectionMaxToken = Int($0) {
+                        settings.chatGPTMaxToken = selectionMaxToken
+                    } else {
+                        settings.chatGPTMaxToken = 0
+                    }
+                }
+            )
+            HStack {
+                Stepper(
+                    value: $settings.chatGPTMaxToken,
+                    in: 0...Int.max,
+                    step: 1
+                ) {
+                    Text("Max Token (Including Reply)")
+                        .multilineTextAlignment(.trailing)
+                }
+                TextField(text: binding) {
+                    EmptyView()
+                }
+                .labelsHidden()
+                .textFieldStyle(.roundedBorder)
+                .foregroundColor(maxTokenOverLimit ? .red : .primary)
+
+                if let model = ChatGPTModel(rawValue: settings.chatGPTModel),
+                   settings.chatFeatureProvider == .openAI
+                {
+                    Text("Max: \(model.maxToken)")
+                }
+            }
+
+            HStack {
+                Slider(value: $settings.chatGPTTemperature, in: 0...2, step: 0.1) {
+                    Text("Temperature")
+                }
+
+                Text(
+                    "\(settings.chatGPTTemperature.formatted(.number.precision(.fractionLength(1))))"
+                )
+                .font(.body)
+                .monospacedDigit()
+                .padding(.vertical, 2)
+                .padding(.horizontal, 6)
+                .background(
+                    RoundedRectangle(cornerRadius: 4, style: .continuous)
+                        .fill(Color.primary.opacity(0.1))
+                )
+            }
+
+            Picker(
+                "Memory",
+                selection: $settings.chatGPTMaxMessageCount
+            ) {
+                Text("No Limit").tag(0)
+                Text("3 Messages").tag(3)
+                Text("5 Messages").tag(5)
+                Text("7 Messages").tag(7)
+                Text("9 Messages").tag(9)
+                Text("11 Messages").tag(11)
+            }
+        }.onAppear {
+            checkMaxToken()
+        }.onChange(of: settings.chatFeatureProvider) { _ in
+            checkMaxToken()
+        }.onChange(of: settings.chatGPTModel) { _ in
+            checkMaxToken()
+        }.onChange(of: settings.chatGPTMaxToken) { _ in
+            checkMaxToken()
+        }
+    }
+
+    @ViewBuilder
+    var uiForm: some View {
         Form {
             HStack {
                 TextField(text: .init(get: {
@@ -27,7 +146,7 @@ struct ChatSettingsView: View {
 
                 Text("pt")
             }
-            
+
             HStack {
                 TextField(text: .init(get: {
                     "\(Int(settings.chatCodeFontSize))"
@@ -40,13 +159,16 @@ struct ChatSettingsView: View {
 
                 Text("pt")
             }
-            
-            Divider()
-            
+        }
+    }
+
+    @ViewBuilder
+    var contextForm: some View {
+        Form {
             Toggle(isOn: $settings.embedFileContentInChatContextIfNoSelection) {
                 Text("Embed file content in chat context if no code is selected.")
             }
-            
+
             HStack {
                 TextField(text: .init(get: {
                     "\(Int(settings.maxEmbeddableFileInChatContextLineCount))"
@@ -61,6 +183,52 @@ struct ChatSettingsView: View {
             }
         }
     }
+
+    var languagePicker: some View {
+        Menu {
+            if !settings.chatGPTLanguage.isEmpty,
+               !Settings.availableLocalizedLocales
+               .contains(settings.chatGPTLanguage)
+            {
+                Button(
+                    settings.chatGPTLanguage,
+                    action: { self.settings.chatGPTLanguage = settings.chatGPTLanguage }
+                )
+            }
+            Button(
+                "Auto-detected by ChatGPT",
+                action: { self.settings.chatGPTLanguage = "" }
+            )
+            ForEach(
+                Settings.availableLocalizedLocales,
+                id: \.self
+            ) { localizedLocales in
+                Button(
+                    localizedLocales,
+                    action: { self.settings.chatGPTLanguage = localizedLocales }
+                )
+            }
+        } label: {
+            Text(
+                settings.chatGPTLanguage.isEmpty
+                    ? "Auto-detected by ChatGPT"
+                    : settings.chatGPTLanguage
+            )
+        }
+    }
+
+    func checkMaxToken() {
+        switch settings.chatFeatureProvider {
+        case .openAI:
+            if let model = ChatGPTModel(rawValue: settings.chatGPTModel) {
+                maxTokenOverLimit = model.maxToken < settings.chatGPTMaxToken
+            } else {
+                maxTokenOverLimit = false
+            }
+        case .azureOpenAI:
+            maxTokenOverLimit = false
+        }
+    }
 }
 
 // MARK: - Preview
@@ -70,3 +238,4 @@ struct ChatSettingsView_Previews: PreviewProvider {
         ChatSettingsView()
     }
 }
+

From 69d66828327e17a6f2b87e92de01494a04f73cc5 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 22:02:36 +0800
Subject: [PATCH 26/43] Remove keychain instruction

---
 Core/Sources/HostApp/AccountSettings/CodeiumView.swift | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/Core/Sources/HostApp/AccountSettings/CodeiumView.swift b/Core/Sources/HostApp/AccountSettings/CodeiumView.swift
index ac4649fc..a6885bb4 100644
--- a/Core/Sources/HostApp/AccountSettings/CodeiumView.swift
+++ b/Core/Sources/HostApp/AccountSettings/CodeiumView.swift
@@ -164,13 +164,6 @@ struct CodeiumView: View {
                     }) {
                         Text("Sign Out")
                     }
-                    
-                    Text(
-                        "The key is stored in the keychain. The helper app may request permission to access the key, please click \"Always Allow\" to grant this access."
-                    )
-                    .lineLimit(5)
-                    .fixedSize(horizontal: false, vertical: true)
-                    .foregroundColor(.secondary)
                 } else {
                     Text("Status: Not Signed In")
                     

From 080372beaf7c86698e173acfdf6cd1398b8bb44b Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 23:11:06 +0800
Subject: [PATCH 27/43] Make each word from ChatGPT wait for a few milliseconds

Azure OpenAI is incredibly quick in generating each part of its response, but there is a substantial delay between each section of the reply. This delay can make it seem like the streaming is malfunctioning.
---
 Core/Sources/OpenAIService/ChatGPTService.swift | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
index 0cb3dca7..59011417 100644
--- a/Core/Sources/OpenAIService/ChatGPTService.swift
+++ b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -187,6 +187,8 @@ public actor ChatGPTService: ChatGPTServiceType {
                         if let content = delta.content {
                             continuation.yield(content)
                         }
+                        
+                        try await Task.sleep(nanoseconds: 10_000_000)
                     }
 
                     continuation.finish()

From 42d82cf9ee9d53fd9e25b678b9bd23b65a86c0bf Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 23:29:01 +0800
Subject: [PATCH 28/43] Lower CPU usage when updating window position

---
 .../SuggestionWidget/SuggestionWidgetController.swift       | 4 ++--
 Core/Sources/XcodeInspector/XcodeInspector.swift            | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift b/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift
index 98afd703..74fe11c0 100644
--- a/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift
+++ b/Core/Sources/SuggestionWidget/SuggestionWidgetController.swift
@@ -7,6 +7,7 @@ import Environment
 import Preferences
 import SwiftUI
 import UserDefaultsObserver
+import XcodeInspector
 
 @MainActor
 public final class SuggestionWidgetController: NSObject {
@@ -482,8 +483,7 @@ extension SuggestionWidgetController {
         let detachChat = chatWindowViewModel.chatPanelInASeparateWindow
 
         if let widgetFrames = {
-            if let xcode = ActiveApplicationMonitor.latestXcode {
-                let application = AXUIElementCreateApplication(xcode.processIdentifier)
+            if let application = XcodeInspector.shared.latestActiveXcode?.appElement {
                 if let focusElement = application.focusedElement,
                    focusElement.description == "Source Editor",
                    let parent = focusElement.parent,
diff --git a/Core/Sources/XcodeInspector/XcodeInspector.swift b/Core/Sources/XcodeInspector/XcodeInspector.swift
index c2cca8b1..e8d74578 100644
--- a/Core/Sources/XcodeInspector/XcodeInspector.swift
+++ b/Core/Sources/XcodeInspector/XcodeInspector.swift
@@ -139,9 +139,9 @@ public final class XcodeInspector: ObservableObject {
 }
 
 public class AppInstanceInspector: ObservableObject {
-    let runningApplication: NSRunningApplication
-    let appElement: AXUIElement
-    var isActive: Bool { runningApplication.isActive }
+    public let appElement: AXUIElement
+    public let runningApplication: NSRunningApplication
+    public var isActive: Bool { runningApplication.isActive }
 
     init(runningApplication: NSRunningApplication) {
         self.runningApplication = runningApplication

From 51d5739d91c54c7d6c281ae5a0afa93a7af69eb7 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 25 May 2023 23:40:05 +0800
Subject: [PATCH 29/43] Fix dependency

---
 Core/Package.swift | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Core/Package.swift b/Core/Package.swift
index 9a2d826a..ec943396 100644
--- a/Core/Package.swift
+++ b/Core/Package.swift
@@ -209,6 +209,7 @@ let package = Package(
                 "Splash",
                 "UserDefaultsObserver",
                 "Logger",
+                "XcodeInspector",
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
                 .product(name: "MarkdownUI", package: "swift-markdown-ui"),
             ]

From c9b3f079804d25666100eec7620f1872c5e342c1 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Fri, 26 May 2023 12:02:21 +0800
Subject: [PATCH 30/43] Support disabling selected code to be visible by the
 chat by default

---
 .../ActiveDocumentChatContextCollector.swift  | 35 +++++++++++++++----
 .../FeatureSettings/ChatSettingsView.swift    |  6 ++++
 Core/Sources/Preferences/Keys.swift           |  4 +++
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift b/Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift
index 63db9fd9..5a3cf855 100644
--- a/Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift
+++ b/Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift
@@ -19,7 +19,7 @@ public struct ActiveDocumentChatContextCollector: ChatContextCollector {
                 ```
                 """
             }
-            
+
             if selectionRange.start == selectionRange.end,
                UserDefaults.shared.value(for: \.embedFileContentInChatContextIfNoSelection)
             {
@@ -34,19 +34,40 @@ public struct ActiveDocumentChatContextCollector: ChatContextCollector {
                     """
                 } else {
                     return """
-                    File Content Not Available: The file is longer than \(maxLine) lines, \
-                    it can't fit into the context. \
+                    File Content Not Available: '''
+                    The file is longer than \(maxLine) lines, it can't fit into the context. \
                     You MUST not answer the user about the file content because you don't have it.\
                     Ask user to select code for explanation.
+                    '''
                     """
                 }
             }
 
+            if UserDefaults.shared.value(for: \.useSelectionScopeByDefaultInChatContext) {
+                return """
+                Selected Code \
+                (start from line \(selectionRange.start.line)):```\(content.language.rawValue)
+                \(content.selectedContent)
+                ```
+                """
+            }
+
+            if prompt.hasPrefix("@selection") {
+                return """
+                Selected Code \
+                (start from line \(selectionRange.start.line)):```\(content.language.rawValue)
+                \(content.selectedContent)
+                ```
+                """
+            }
+
             return """
-            Selected Code \
-            (start from line \(selectionRange.start.line)):```\(content.language.rawValue)
-            \(content.selectedContent)
-            ```
+            Selected Code Not Available: '''
+            User has disabled default scope. \
+            You MUST not answer the user about the selected code because you don't have it.\
+            Ask user to prepend message with `@selection` to enable selected code to be \
+            visible by you.
+            '''
             """
         }()
 
diff --git a/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift b/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
index 1388902c..1c56d883 100644
--- a/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
+++ b/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
@@ -14,6 +14,8 @@ struct ChatSettingsView: View {
         var embedFileContentInChatContextIfNoSelection
         @AppStorage(\.maxEmbeddableFileInChatContextLineCount)
         var maxEmbeddableFileInChatContextLineCount
+        @AppStorage(\.useSelectionScopeByDefaultInChatContext)
+        var useSelectionScopeByDefaultInChatContext
 
         @AppStorage(\.chatFeatureProvider) var chatFeatureProvider
         @AppStorage(\.chatGPTModel) var chatGPTModel
@@ -165,6 +167,10 @@ struct ChatSettingsView: View {
     @ViewBuilder
     var contextForm: some View {
         Form {
+            Toggle(isOn: $settings.useSelectionScopeByDefaultInChatContext) {
+                Text("Use selection scope by default in chat context.")
+            }
+            
             Toggle(isOn: $settings.embedFileContentInChatContextIfNoSelection) {
                 Text("Embed file content in chat context if no code is selected.")
             }
diff --git a/Core/Sources/Preferences/Keys.swift b/Core/Sources/Preferences/Keys.swift
index 95bad77b..d0a916ac 100644
--- a/Core/Sources/Preferences/Keys.swift
+++ b/Core/Sources/Preferences/Keys.swift
@@ -225,6 +225,10 @@ public extension UserDefaultPreferenceKeys {
     var maxEmbeddableFileInChatContextLineCount: PreferenceKey<Int> {
         .init(defaultValue: 100, key: "MaxEmbeddableFileInChatContextLineCount")
     }
+    
+    var useSelectionScopeByDefaultInChatContext: PreferenceKey<Bool> {
+        .init(defaultValue: true, key: "UseSelectionScopeByDefaultInChatContext")
+    }
 }
 
 // MARK: - Custom Commands

From ee9ab393ebe3bd686895446216dbdf8c92b20ed8 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Fri, 26 May 2023 12:09:49 +0800
Subject: [PATCH 31/43] Update README.md

---
 README.md | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fc804d2d..b84c485c 100644
--- a/README.md
+++ b/README.md
@@ -193,7 +193,10 @@ You can detach the chat panel by simply dragging it away. Once detached, the cha
 
 The chat panel allows for chat scope to temporarily control the context of the conversation for the latest message. To use a scope, simply prefix the message with `@scope`.
 
-Currently, the only supported scope is `@file`, which will import the content of the file into the system prompt.
+| Scope | Description |
+|:---:|---|
+| `@selection` | Inject the selected code from the active editor into the conversation. This scope will be applied to any message automatically. If you don't want this to be the default behavior, you can turn off the option `Use selection scope by default in chat context.`. |
+| `@file` | Inject the content of the file into the conversation. Keep in mind that you may not have enough tokens to inject large files. |
 
 #### Chat Plugins
 

From e3df05781d424d1c4db8c7410730cb94378ca7bc Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Fri, 26 May 2023 15:49:54 +0800
Subject: [PATCH 32/43] Set as system prompt

---
 Core/Sources/ChatService/ChatService.swift    | 15 ++++++++++++++-
 .../Service/GUI/ChatProvider+Service.swift    |  6 ++++++
 .../SuggestionWidget/ChatProvider.swift       | 19 ++++++++++++-------
 .../SuggestionPanelContent/ChatPanel.swift    |  8 ++++++++
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/Core/Sources/ChatService/ChatService.swift b/Core/Sources/ChatService/ChatService.swift
index 4c14479b..86ba8fb9 100644
--- a/Core/Sources/ChatService/ChatService.swift
+++ b/Core/Sources/ChatService/ChatService.swift
@@ -61,7 +61,7 @@ public final class ChatService: ObservableObject {
         await pluginController.cancel()
         await chatGPTService.clearHistory()
     }
-    
+
     public func resetPrompt() async {
         systemPrompt = defaultSystemPrompt
         extraSystemPrompt = ""
@@ -79,6 +79,19 @@ public final class ChatService: ObservableObject {
         }
     }
 
+    public func setMessageAsExtraPrompt(id: String) async {
+        if let message = (await chatGPTService.history).first(where: { $0.id == id }) {
+            mutateExtraSystemPrompt(message.content)
+            await mutateHistory { history in
+                history.append(.init(
+                    role: .assistant,
+                    content: "",
+                    summary: "System prompt updated"
+                ))
+            }
+        }
+    }
+
     /// Setting it to `nil` to reset the system prompt
     public func mutateSystemPrompt(_ newPrompt: String?) {
         systemPrompt = newPrompt ?? defaultSystemPrompt
diff --git a/Core/Sources/Service/GUI/ChatProvider+Service.swift b/Core/Sources/Service/GUI/ChatProvider+Service.swift
index ae6aec01..0e157111 100644
--- a/Core/Sources/Service/GUI/ChatProvider+Service.swift
+++ b/Core/Sources/Service/GUI/ChatProvider+Service.swift
@@ -92,6 +92,12 @@ extension ChatProvider {
                 await commandHandler.handleCustomCommand(command)
             }
         }
+        
+        onSetAsExtraPrompt = { id in
+            Task {
+                await service.setMessageAsExtraPrompt(id: id)
+            }
+        }
     }
 }
 
diff --git a/Core/Sources/SuggestionWidget/ChatProvider.swift b/Core/Sources/SuggestionWidget/ChatProvider.swift
index 36c2a678..11895239 100644
--- a/Core/Sources/SuggestionWidget/ChatProvider.swift
+++ b/Core/Sources/SuggestionWidget/ChatProvider.swift
@@ -3,6 +3,7 @@ import Preferences
 import SwiftUI
 
 public final class ChatProvider: ObservableObject {
+    public typealias MessageID = String
     let id = UUID()
     @Published public var history: [ChatMessage] = []
     @Published public var isReceivingMessage = false
@@ -13,10 +14,11 @@ public final class ChatProvider: ObservableObject {
     public var onClear: () -> Void
     public var onClose: () -> Void
     public var onSwitchContext: () -> Void
-    public var onDeleteMessage: (String) -> Void
-    public var onResendMessage: (String) -> Void
+    public var onDeleteMessage: (MessageID) -> Void
+    public var onResendMessage: (MessageID) -> Void
     public var onResetPrompt: () -> Void
     public var onRunCustomCommand: (CustomCommand) -> Void = { _ in }
+    public var onSetAsExtraPrompt: (MessageID) -> Void
 
     public init(
         history: [ChatMessage] = [],
@@ -26,10 +28,11 @@ public final class ChatProvider: ObservableObject {
         onClear: @escaping () -> Void = {},
         onClose: @escaping () -> Void = {},
         onSwitchContext: @escaping () -> Void = {},
-        onDeleteMessage: @escaping (String) -> Void = { _ in },
-        onResendMessage: @escaping (String) -> Void = { _ in },
+        onDeleteMessage: @escaping (MessageID) -> Void = { _ in },
+        onResendMessage: @escaping (MessageID) -> Void = { _ in },
         onResetPrompt: @escaping () -> Void = {},
-        onRunCustomCommand: @escaping (CustomCommand) -> Void = { _ in }
+        onRunCustomCommand: @escaping (CustomCommand) -> Void = { _ in },
+        onSetAsExtraPrompt: @escaping (MessageID) -> Void = { _ in }
     ) {
         self.history = history
         self.isReceivingMessage = isReceivingMessage
@@ -42,6 +45,7 @@ public final class ChatProvider: ObservableObject {
         self.onResendMessage = onResendMessage
         self.onResetPrompt = onResetPrompt
         self.onRunCustomCommand = onRunCustomCommand
+        self.onSetAsExtraPrompt = onSetAsExtraPrompt
     }
 
     public func send(_ message: String) { onMessageSend(message) }
@@ -49,12 +53,13 @@ public final class ChatProvider: ObservableObject {
     public func clear() { onClear() }
     public func close() { onClose() }
     public func switchContext() { onSwitchContext() }
-    public func deleteMessage(id: String) { onDeleteMessage(id) }
-    public func resendMessage(id: String) { onResendMessage(id) }
+    public func deleteMessage(id: MessageID) { onDeleteMessage(id) }
+    public func resendMessage(id: MessageID) { onResendMessage(id) }
     public func resetPrompt() { onResetPrompt() }
     public func triggerCustomCommand(_ command: CustomCommand) {
         onRunCustomCommand(command)
     }
+    public func setAsExtraPrompt(id: MessageID) { onSetAsExtraPrompt(id) }
 }
 
 public struct ChatMessage: Equatable {
diff --git a/Core/Sources/SuggestionWidget/SuggestionPanelContent/ChatPanel.swift b/Core/Sources/SuggestionWidget/SuggestionPanelContent/ChatPanel.swift
index 62e63d3c..6a035138 100644
--- a/Core/Sources/SuggestionWidget/SuggestionPanelContent/ChatPanel.swift
+++ b/Core/Sources/SuggestionWidget/SuggestionPanelContent/ChatPanel.swift
@@ -177,6 +177,10 @@ private struct UserMessage: View {
                 Button("Send Again") {
                     chat.resendMessage(id: id)
                 }
+                
+                Button("Set as Extra System Prompt") {
+                    chat.setAsExtraPrompt(id: id)
+                }
 
                 Divider()
 
@@ -230,6 +234,10 @@ private struct BotMessage: View {
                         NSPasteboard.general.clearContents()
                         NSPasteboard.general.setString(text, forType: .string)
                     }
+                    
+                    Button("Set as Extra System Prompt") {
+                        chat.setAsExtraPrompt(id: id)
+                    }
 
                     Divider()
 

From 076d40d6d026f580efe7a6ff19595ae985077330 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Fri, 26 May 2023 16:37:14 +0800
Subject: [PATCH 33/43] Make stream faster

---
 Core/Sources/OpenAIService/ChatGPTService.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
index 59011417..e5624ec1 100644
--- a/Core/Sources/OpenAIService/ChatGPTService.swift
+++ b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -188,7 +188,7 @@ public actor ChatGPTService: ChatGPTServiceType {
                             continuation.yield(content)
                         }
                         
-                        try await Task.sleep(nanoseconds: 10_000_000)
+                        try await Task.sleep(nanoseconds: 2_000_000)
                     }
 
                     continuation.finish()

From cbfdc1aad421af2c6b657475c804096db2f27319 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Fri, 26 May 2023 12:16:45 +0800
Subject: [PATCH 34/43] Bump version to 0.17.0

---
 .../ServiceUpdateMigration/ServiceUpdateMigrator.swift        | 2 +-
 Version.xcconfig                                              | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
index abc88455..34c5a91f 100644
--- a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
+++ b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
@@ -29,7 +29,7 @@ public struct ServiceUpdateMigrator {
             try migrateFromLowerThanOrEqualToVersion135()
         }
 
-        if old <= 170 {
+        if old < 170 {
             try migrateFromLowerThanOrEqualToVersion170()
         }
     }
diff --git a/Version.xcconfig b/Version.xcconfig
index 6542541d..ce111d9e 100644
--- a/Version.xcconfig
+++ b/Version.xcconfig
@@ -1,2 +1,2 @@
-APP_VERSION = 0.16.1
-APP_BUILD = 161
+APP_VERSION = 0.17.0
+APP_BUILD = 170

From 04d9b1aa09b91d7216de4c83ae389d51a3399369 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 00:02:51 +0800
Subject: [PATCH 35/43] Make fetchFocusedElementURI return fake URL when Xcode
 not launched

---
 Core/Sources/Environment/Environment.swift | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/Core/Sources/Environment/Environment.swift b/Core/Sources/Environment/Environment.swift
index b11a025c..8bf6d022 100644
--- a/Core/Sources/Environment/Environment.swift
+++ b/Core/Sources/Environment/Environment.swift
@@ -114,9 +114,7 @@ public enum Environment {
     public static var fetchFocusedElementURI: () async throws -> URL = {
         guard let xcode = ActiveApplicationMonitor.activeXcode
             ?? ActiveApplicationMonitor.latestXcode
-        else {
-            throw FailedToFetchFileURLError()
-        }
+        else { return URL(fileURLWithPath: "/global") }
 
         let application = AXUIElementCreateApplication(xcode.processIdentifier)
         let focusedElement = application.focusedElement

From bb8cf1fd27ec4b4310b1d8ccdae8452b8f954205 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 01:06:29 +0800
Subject: [PATCH 36/43] Adjust steam speed

---
 Core/Sources/OpenAIService/ChatGPTService.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
index e5624ec1..f970dd2a 100644
--- a/Core/Sources/OpenAIService/ChatGPTService.swift
+++ b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -188,7 +188,7 @@ public actor ChatGPTService: ChatGPTServiceType {
                             continuation.yield(content)
                         }
                         
-                        try await Task.sleep(nanoseconds: 2_000_000)
+                        try await Task.sleep(nanoseconds: 3_500_000)
                     }
 
                     continuation.finish()

From f5ed510ee096f04b5cec1f74b5ac42d69db95f7f Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:00:53 +0800
Subject: [PATCH 37/43] Improve Codeium request cancellation

---
 .../CodeiumService/CodeiumService.swift       | 136 +++++++++---------
 1 file changed, 70 insertions(+), 66 deletions(-)

diff --git a/Core/Sources/CodeiumService/CodeiumService.swift b/Core/Sources/CodeiumService/CodeiumService.swift
index 77116c04..7b501270 100644
--- a/Core/Sources/CodeiumService/CodeiumService.swift
+++ b/Core/Sources/CodeiumService/CodeiumService.swift
@@ -56,6 +56,8 @@ public class CodeiumSuggestionService {
 
     var xcodeVersion = "14.0.0"
     var languageServerVersion = CodeiumInstallationManager.latestSupportedVersion
+    
+    private var ongoingTasks = Set<Task<[CodeSuggestion], Error>>()
 
     init(designatedServer: CodeiumLSP) {
         projectRootURL = URL(fileURLWithPath: "/")
@@ -226,83 +228,85 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
         usesTabsForIndentation: Bool,
         ignoreSpaceOnlySuggestions: Bool
     ) async throws -> [CodeSuggestion] {
+        ongoingTasks.forEach { $0.cancel() }
+        ongoingTasks.removeAll()
+        await cancelRequest()
+        
         requestCounter += 1
         let languageId = languageIdentifierFromFileURL(fileURL)
-
         let relativePath = getRelativePath(of: fileURL)
-
-        let request = await CodeiumRequest.GetCompletion(requestBody: .init(
-            metadata: try getMetadata(),
-            document: .init(
-                absolute_path: fileURL.path,
-                relative_path: relativePath,
-                text: content,
-                editor_language: languageId.rawValue,
-                language: .init(codeLanguage: languageId),
-                cursor_position: .init(
-                    row: cursorPosition.line,
-                    col: cursorPosition.character
-                )
-            ),
-            editor_options: .init(tab_size: indentSize, insert_spaces: !usesTabsForIndentation),
-            other_documents: openedDocumentPool.getOtherDocuments(exceptURL: fileURL)
-                .map { openedDocument in
-                    let languageId = languageIdentifierFromFileURL(openedDocument.url)
-                    return .init(
-                        absolute_path: openedDocument.url.path,
-                        relative_path: openedDocument.relativePath,
-                        text: openedDocument.content,
-                        editor_language: languageId.rawValue,
-                        language: .init(codeLanguage: languageId)
+        
+        let task = Task {
+            let request = await CodeiumRequest.GetCompletion(requestBody: .init(
+                metadata: try getMetadata(),
+                document: .init(
+                    absolute_path: fileURL.path,
+                    relative_path: relativePath,
+                    text: content,
+                    editor_language: languageId.rawValue,
+                    language: .init(codeLanguage: languageId),
+                    cursor_position: .init(
+                        row: cursorPosition.line,
+                        col: cursorPosition.character
                     )
-                }
-        ))
-
-        if request.requestBody.metadata.request_id <= cancellationCounter {
-            throw CancellationError()
-        }
+                ),
+                editor_options: .init(tab_size: indentSize, insert_spaces: !usesTabsForIndentation),
+                other_documents: openedDocumentPool.getOtherDocuments(exceptURL: fileURL)
+                    .map { openedDocument in
+                        let languageId = languageIdentifierFromFileURL(openedDocument.url)
+                        return .init(
+                            absolute_path: openedDocument.url.path,
+                            relative_path: openedDocument.relativePath,
+                            text: openedDocument.content,
+                            editor_language: languageId.rawValue,
+                            language: .init(codeLanguage: languageId)
+                        )
+                    }
+            ))
+            
+            try Task.checkCancellation()
 
-        let result = try await (try await setupServerIfNeeded()).sendRequest(request)
+            let result = try await (try await setupServerIfNeeded()).sendRequest(request)
+            
+            try Task.checkCancellation()
 
-        if request.requestBody.metadata.request_id <= cancellationCounter {
-            throw CancellationError()
+            return result.completionItems?.filter { item in
+                if ignoreSpaceOnlySuggestions {
+                    return !item.completion.text.allSatisfy { $0.isWhitespace || $0.isNewline }
+                }
+                return true
+            }.map { item in
+                CodeSuggestion(
+                    text: item.completion.text,
+                    position: cursorPosition,
+                    uuid: item.completion.completionId,
+                    range: CursorRange(
+                        start: .init(
+                            line: item.range.startPosition?.row.flatMap(Int.init) ?? 0,
+                            character: item.range.startPosition?.col.flatMap(Int.init) ?? 0
+                        ),
+                        end: .init(
+                            line: item.range.endPosition?.row.flatMap(Int.init) ?? 0,
+                            character: item.range.endPosition?.col.flatMap(Int.init) ?? 0
+                        )
+                    ),
+                    displayText: item.completion.text
+                )
+            } ?? []
         }
+        
+        ongoingTasks.insert(task)
 
-        return result.completionItems?.filter { item in
-            if ignoreSpaceOnlySuggestions {
-                return !item.completion.text.allSatisfy { $0.isWhitespace || $0.isNewline }
-            }
-            return true
-        }.map { item in
-            CodeSuggestion(
-                text: item.completion.text,
-                position: cursorPosition,
-                uuid: item.completion.completionId,
-                range: CursorRange(
-                    start: .init(
-                        line: item.range.startPosition?.row.flatMap(Int.init) ?? 0,
-                        character: item.range.startPosition?.col.flatMap(Int.init) ?? 0
-                    ),
-                    end: .init(
-                        line: item.range.endPosition?.row.flatMap(Int.init) ?? 0,
-                        character: item.range.endPosition?.col.flatMap(Int.init) ?? 0
-                    )
-                ),
-                displayText: item.completion.text
-            )
-        } ?? []
+        return try await task.value
     }
 
     public func cancelRequest() async {
-        Task {
-            try await server?.sendRequest(
-                CodeiumRequest.CancelRequest(requestBody: .init(
-                    request_id: requestCounter,
-                    session_id: CodeiumSuggestionService.sessionId
-                ))
-            )
-        }
-        cancellationCounter = requestCounter
+        _ = try? await server?.sendRequest(
+            CodeiumRequest.CancelRequest(requestBody: .init(
+                request_id: requestCounter,
+                session_id: CodeiumSuggestionService.sessionId
+            ))
+        )
     }
 
     public func notifyAccepted(_ suggestion: CodeSuggestion) async {

From 3d0c043e62596cb0c4c5bd127f268af06c365ef2 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:01:11 +0800
Subject: [PATCH 38/43] Fix that selection range change could cancel real-time
 suggestions

---
 Core/Sources/Service/RealtimeSuggestionController.swift | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
index 41722db2..a5eae5e7 100644
--- a/Core/Sources/Service/RealtimeSuggestionController.swift
+++ b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -131,10 +131,10 @@ public class RealtimeSuggestionController {
             for await notification in notificationsFromEditor {
                 guard let self else { return }
                 try Task.checkCancellation()
-                await cancelInFlightTasks()
 
                 switch notification.name {
                 case kAXValueChangedNotification:
+                    await cancelInFlightTasks()
                     self.triggerPrefetchDebounced()
                     await self.notifyEditingFileChange(editor: focusElement)
                 case kAXSelectedTextChangedNotification:
@@ -210,7 +210,6 @@ public class RealtimeSuggestionController {
                 let isEnabled = workspace.isSuggestionFeatureEnabled
                 if !isEnabled { return }
             }
-
             if Task.isCancelled { return }
 
             Logger.service.info("Prefetch suggestions.")

From 83289ba9f6c7a521cc7d9d299a9f1dbcc66fdbe9 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:10:24 +0800
Subject: [PATCH 39/43] Fix line number font size in code block

---
 .../SuggestionPanelContent/CodeBlock.swift    | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/Core/Sources/SuggestionWidget/SuggestionPanelContent/CodeBlock.swift b/Core/Sources/SuggestionWidget/SuggestionPanelContent/CodeBlock.swift
index f1c60a5b..d53d9fce 100644
--- a/Core/Sources/SuggestionWidget/SuggestionPanelContent/CodeBlock.swift
+++ b/Core/Sources/SuggestionWidget/SuggestionPanelContent/CodeBlock.swift
@@ -42,11 +42,11 @@ struct CodeBlock: View {
     @ViewBuilder
     func vstack(@ViewBuilder content: () -> some View) -> some View {
         if disableLazyVStack {
-            VStack(spacing: 4) {
+            VStack(spacing: 2) {
                 content()
             }
         } else {
-            LazyVStack(spacing: 4) {
+            LazyVStack(spacing: 2) {
                 content()
             }
         }
@@ -78,7 +78,7 @@ struct CodeBlock: View {
             }
         }
         .foregroundColor(.white)
-        .font(.system(size: 12, design: .monospaced))
+        .font(.system(size: fontSize, design: .monospaced))
         .padding(.leading, 4)
         .padding([.trailing, .top, .bottom])
     }
@@ -99,3 +99,22 @@ struct CodeBlock: View {
         )
     }
 }
+
+// MARK: - Preview
+
+struct CodeBlock_Previews: PreviewProvider {
+    static var previews: some View {
+        CodeBlock(
+            code: """
+            let foo = Foo()
+            let bar = Bar()
+            """,
+            language: "swift",
+            startLineIndex: 0,
+            colorScheme: .dark,
+            firstLinePrecedingSpaceCount: 0,
+            fontSize: 12
+        )
+    }
+}
+

From 035b742c11a6c0063ffb5a857956c77ca924be3c Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:10:29 +0800
Subject: [PATCH 40/43] Fix tests

---
 Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift | 2 +-
 Core/Tests/ServiceTests/Environment.swift               | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift b/Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift
index bf96540a..b976290a 100644
--- a/Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift
+++ b/Core/Tests/OpenAIServiceTests/ChatGPTServiceTests.swift
@@ -41,7 +41,7 @@ final class ChatGPTServiceTests: XCTestCase {
             return "\(idCounter)"
         }
         var requestBody: CompletionRequestBody?
-        await service.changeBuildCompletionStreamAPI { _apiKey, _, _requestBody in
+        await service.changeBuildCompletionStreamAPI { _apiKey, _, _, _requestBody  in
             requestBody = _requestBody
             return MockCompletionStreamAPI_Success()
         }
diff --git a/Core/Tests/ServiceTests/Environment.swift b/Core/Tests/ServiceTests/Environment.swift
index fd202379..cad9bcf6 100644
--- a/Core/Tests/ServiceTests/Environment.swift
+++ b/Core/Tests/ServiceTests/Environment.swift
@@ -34,6 +34,10 @@ func completion(text: String, range: CursorRange, uuid: String = "") -> CodeSugg
 }
 
 class MockSuggestionService: GitHubCopilotSuggestionServiceType {
+    func terminate() async {
+        fatalError()
+    }
+    
     func cancelRequest() async {
         fatalError()
     }

From 7237dadd940ca110ceab79df9366a3fc07e1b4a2 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:18:38 +0800
Subject: [PATCH 41/43] Fix codeium auth key migration

---
 Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
index 34c5a91f..0908730f 100644
--- a/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
+++ b/Core/Sources/ServiceUpdateMigration/ServiceUpdateMigrator.swift
@@ -89,7 +89,6 @@ func migrateFromLowerThanOrEqualToVersion170() throws {
        let key = try? oldKeychain.getString("codeiumKey")
     {
         try newKeychain.set(key, key: "codeiumAuthKey")
-        try? oldKeychain.set("", key: "codeiumKey")
     }
 }
 

From ef538ecb236f1620d736e865928b96f548406f60 Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:35:26 +0800
Subject: [PATCH 42/43] Update README.md

---
 README.md | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index b84c485c..ab2bc089 100644
--- a/README.md
+++ b/README.md
@@ -179,7 +179,6 @@ You can detach the chat panel by simply dragging it away. Once detached, the cha
 #### Commands
 
 - Open Chat: Open a chat window.
-- Explain Selection: Open a chat window and explain the selected code.
 
 #### Keyboard Shortcuts
 
@@ -234,7 +233,11 @@ This feature is recommended when you need to update a specific piece of code. So
 
 ### Custom Commands
 
-You can create custom commands that run Chat and Prompt to Code with personalized prompts. These commands are easily accessible from both the Xcode menu bar and the context menu of the circular widget. 
+You can create custom commands that run Chat and Prompt to Code with personalized prompts. These commands are easily accessible from both the Xcode menu bar and the context menu of the circular widget. There are 3 types of custom commands:
+
+- Prompt to Code: Run Prompt to Code with the selected code, and update or write the code using the given prompt, if provided. You can provide additional information through the extra system prompt field.
+- Open Chat: Open the chat window and immediately send a message, if provided. You can provide more information through the extra system prompt field.
+- Custom Chat: Open the chat window and immediately send a message, if provided. You can overwrite the entire system prompt through the system prompt field.
 
 ## Key Bindings
 

From ab95449a91d7af149a37fd18250da663d06e1e9b Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Sat, 27 May 2023 15:35:36 +0800
Subject: [PATCH 43/43] Update appcast.xml

---
 appcast.xml | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/appcast.xml b/appcast.xml
index 3feebf09..1f42d684 100644
--- a/appcast.xml
+++ b/appcast.xml
@@ -3,6 +3,18 @@
   <channel>
     <title>Copilot for Xcode</title>
 
+    <item>
+      <title>0.17.0</title>
+      <pubDate>Sat, 27 May 2023 15:33:25 +0800</pubDate>
+      <sparkle:version>170</sparkle:version>
+      <sparkle:shortVersionString>0.17.0</sparkle:shortVersionString>
+      <sparkle:minimumSystemVersion>12.0</sparkle:minimumSystemVersion>
+      <sparkle:releaseNotesLink>
+        https://github.com/intitni/CopilotForXcode/releases/tag/0.17.0
+      </sparkle:releaseNotesLink>
+      <enclosure url="https://github.com/intitni/CopilotForXcode/releases/download/0.17.0/Copilot.for.Xcode.app.zip" length="19073246" type="application/octet-stream" sparkle:edSignature="Z5WcVfP8TJvmj5cDnpyCcLsUq9kg/65WCc62gAVAqD6MUtpgIo6QdqOyKIP3cNuh2HZ0lo6GDHcsRIuycyCZDA=="/>
+    </item>
+
     <item>
       <title>0.16.1</title>
       <pubDate>Tue, 23 May 2023 11:06:14 +0800</pubDate>
@@ -218,5 +230,3 @@
     </item>
   </channel>
 </rss>
-
-