From a959149d44c34a0c685bcf7b1c9c5f7edcb77520 Mon Sep 17 00:00:00 2001 From: kchro3 Date: Mon, 1 Jan 2024 16:17:31 -0800 Subject: [PATCH 1/3] add some possible fixes --- TypeaheadAI/AppState.swift | 4 ++-- TypeaheadAI/ClientManager.swift | 4 +++- .../Functions/FunctionManager+PerformUIAction.swift | 2 +- TypeaheadAI/Functions/FunctionManager.swift | 2 +- TypeaheadAI/Views/Modal/ConversationView.swift | 2 +- TypeaheadAI/WindowManagers/ModalManager.swift | 11 ++++++----- 6 files changed, 14 insertions(+), 11 deletions(-) diff --git a/TypeaheadAI/AppState.swift b/TypeaheadAI/AppState.swift index 8fa59bd..6ab5c34 100644 --- a/TypeaheadAI/AppState.swift +++ b/TypeaheadAI/AppState.swift @@ -117,7 +117,7 @@ final class AppState: ObservableObject { Task { do { try await self.specialCopyActor?.specialCopy() - DispatchQueue.main.async { + await MainActor.run { NotificationCenter.default.post(name: .smartCopyPerformed, object: nil) } } catch { @@ -130,7 +130,7 @@ final class AppState: ObservableObject { Task { do { try await specialPasteActor?.specialPaste() - DispatchQueue.main.async { + await MainActor.run { NotificationCenter.default.post(name: .smartPastePerformed, object: nil) } } catch { diff --git a/TypeaheadAI/ClientManager.swift b/TypeaheadAI/ClientManager.swift index c971269..ecd269c 100644 --- a/TypeaheadAI/ClientManager.swift +++ b/TypeaheadAI/ClientManager.swift @@ -334,7 +334,7 @@ class ClientManager: ObservableObject, CanGetUIElements { } } - DispatchQueue.main.async { + await MainActor.run { self?.currentStreamingTask = nil self?.isExecuting = false } @@ -426,6 +426,8 @@ class ClientManager: ObservableObject, CanGetUIElements { // In the future, we can think about how to support completions with a full response, but worry about that later. var bufferedPayload: ChunkPayload = ChunkPayload(finishReason: nil) for try await line in data.lines { + try Task.checkCancellation() + guard let data = line.data(using: .utf8), let response = try? decoder.decode(ChunkPayload.self, from: data) else { let error = ClientManagerError.serverError("Failed to parse response...") diff --git a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift index ef835f7..255caf4 100644 --- a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift +++ b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift @@ -157,7 +157,7 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { apps: appInfo?.apps ?? [:] ) - DispatchQueue.main.async { + await MainActor.run { modalManager.cachedAppInfo = newAppInfo } diff --git a/TypeaheadAI/Functions/FunctionManager.swift b/TypeaheadAI/Functions/FunctionManager.swift index 94eef3f..9a264df 100644 --- a/TypeaheadAI/Functions/FunctionManager.swift +++ b/TypeaheadAI/Functions/FunctionManager.swift @@ -206,7 +206,7 @@ class FunctionManager: ObservableObject, modalManager.setError("Function \(functionCall.name) not supported", appContext: appContext) } - DispatchQueue.main.async { + await MainActor.run { self?.currentTask = nil self?.isExecuting = false } diff --git a/TypeaheadAI/Views/Modal/ConversationView.swift b/TypeaheadAI/Views/Modal/ConversationView.swift index 75211cf..6c77b29 100644 --- a/TypeaheadAI/Views/Modal/ConversationView.swift +++ b/TypeaheadAI/Views/Modal/ConversationView.swift @@ -59,7 +59,7 @@ struct ConversationView: View { Task { modalManager.cancelTasks() try await modalManager.rewindTo(index: index) - try await modalManager.replyToUserMessage() + try modalManager.replyToUserMessage() } }, onTruncate: { diff --git a/TypeaheadAI/WindowManagers/ModalManager.swift b/TypeaheadAI/WindowManagers/ModalManager.swift index 22cdede..7e07b11 100644 --- a/TypeaheadAI/WindowManagers/ModalManager.swift +++ b/TypeaheadAI/WindowManagers/ModalManager.swift @@ -674,12 +674,13 @@ class ModalManager: ObservableObject { /// Reply to the user /// If refresh, then pop the previous message before responding. - @MainActor - func replyToUserMessage() async throws { - isPending = true - userIntents = nil - + func replyToUserMessage() throws { Task { + await MainActor.run { + isPending = true + userIntents = nil + } + try await self.clientManager?.refine( messages: self.messages, streamHandler: defaultStreamHandler, From aadf3a69b9a118a8aa313fe0e4345c84b2f15506 Mon Sep 17 00:00:00 2001 From: kchro3 Date: Mon, 1 Jan 2024 17:17:03 -0800 Subject: [PATCH 2/3] added a ton of task check cancellations --- TypeaheadAI/ClientManager.swift | 3 +++ .../FunctionManager+OpenApplication.swift | 8 +++---- .../Functions/FunctionManager+OpenFile.swift | 8 +++++++ .../Functions/FunctionManager+OpenURL.swift | 16 +++++++++++++ .../FunctionManager+PerformUIAction.swift | 23 +++++++++++++------ .../Functions/FunctionManager+SaveFile.swift | 14 +++++++---- TypeaheadAI/Functions/FunctionManager.swift | 2 +- TypeaheadAI/WindowManagers/ModalManager.swift | 1 + 8 files changed, 59 insertions(+), 16 deletions(-) diff --git a/TypeaheadAI/ClientManager.swift b/TypeaheadAI/ClientManager.swift index ecd269c..8fe2721 100644 --- a/TypeaheadAI/ClientManager.swift +++ b/TypeaheadAI/ClientManager.swift @@ -221,6 +221,7 @@ class ClientManager: ObservableObject, CanGetUIElements { .flatMap { $0.quickActionId } .flatMap { self.promptManager?.getById($0) } + try Task.checkCancellation() if let quickAction = quickAction, let appContext = appInfo?.appContext, let copiedText = copiedText { @@ -233,6 +234,7 @@ class ClientManager: ObservableObject, CanGetUIElements { ) } + try Task.checkCancellation() await self.sendStreamRequest( id: UUID(), username: NSUserName(), @@ -320,6 +322,7 @@ class ClientManager: ObservableObject, CanGetUIElements { completion: completion ) + try Task.checkCancellation() guard let stream = stream else { self?.logger.debug("Failed to get stream") return diff --git a/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift b/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift index 5ff0999..fcc285d 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift @@ -24,7 +24,6 @@ extension FunctionManager { return } - try Task.checkCancellation() await modalManager.closeModal() guard let url = NSWorkspace.shared.urlForApplication(withBundleIdentifier: bundleIdentifier) else { @@ -33,11 +32,12 @@ extension FunctionManager { return } + try Task.checkCancellation() // Activate the app, bringing it to the foreground NSWorkspace.shared.open(url) try await Task.sleep(for: .seconds(2)) - try Task.checkCancellation() + try Task.checkCancellation() let newAppContext = try await fetchAppContext() let (newUIElement, newElementMap) = getUIElements(appContext: newAppContext) guard let serializedUIElement = newUIElement?.serialize( @@ -52,6 +52,7 @@ extension FunctionManager { return } + try Task.checkCancellation() await modalManager.showModal() await modalManager.appendTool( "Updated state: \(serializedUIElement)", @@ -65,10 +66,9 @@ extension FunctionManager { apps: appInfo?.apps ?? [:] ) - try Task.checkCancellation() - Task { do { + try Task.checkCancellation() try await modalManager.continueReplying(appInfo: newAppInfo) } catch { await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) diff --git a/TypeaheadAI/Functions/FunctionManager+OpenFile.swift b/TypeaheadAI/Functions/FunctionManager+OpenFile.swift index 85c1248..de25e68 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenFile.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenFile.swift @@ -39,21 +39,28 @@ extension FunctionManager { try Task.checkCancellation() // Select the file + try Task.checkCancellation() try await simulateGoToFile() try await Task.sleep(for: .seconds(1)) // Paste filename to file search field + try Task.checkCancellation() try await simulatePaste() try await Task.sleep(for: .seconds(1)) // Enter twice to pick and attach file + try Task.checkCancellation() try await simulateEnter() try await Task.sleep(for: .seconds(1)) + + try Task.checkCancellation() try await simulateEnter() try await Task.sleep(for: .seconds(2)) + try Task.checkCancellation() await modalManager.showModal() + try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] @@ -79,6 +86,7 @@ extension FunctionManager { Task { do { + try Task.checkCancellation() try await modalManager.continueReplying(appInfo: newAppInfo) } catch { await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) diff --git a/TypeaheadAI/Functions/FunctionManager+OpenURL.swift b/TypeaheadAI/Functions/FunctionManager+OpenURL.swift index ab9cd26..5f27fd5 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenURL.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenURL.swift @@ -23,9 +23,11 @@ extension FunctionManager { appContext: appInfo?.appContext ) + try Task.checkCancellation() try await openURL(url) try await Task.sleep(for: .seconds(5)) + try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] @@ -51,6 +53,7 @@ extension FunctionManager { Task { do { + try Task.checkCancellation() try await modalManager.continueReplying(appInfo: newAppInfo) } catch { await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) @@ -66,6 +69,7 @@ extension FunctionManager { return } + try Task.checkCancellation() if url == "" { await modalManager.appendFunction( "Scraping current page...", @@ -79,8 +83,12 @@ extension FunctionManager { app.activate(options: [.activateIgnoringOtherApps]) } + try Task.checkCancellation() await modalManager.closeModal() + try await Task.sleep(for: .seconds(1)) + + try Task.checkCancellation() try await simulateSelectAll() try await simulateCopy() } else { @@ -91,12 +99,19 @@ extension FunctionManager { ) try await openURL(url) + try Task.checkCancellation() + await modalManager.closeModal() try await Task.sleep(for: .seconds(5)) + + try Task.checkCancellation() try await simulateSelectAll() + try Task.checkCancellation() try await simulateCopy() + try Task.checkCancellation() try await simulateClose() + try Task.checkCancellation() if let bundleIdentifier = appContext?.bundleIdentifier, let app = NSRunningApplication.runningApplications(withBundleIdentifier: bundleIdentifier).first { // Activate the app, bringing it to the foreground @@ -139,6 +154,7 @@ extension FunctionManager { } } + try Task.checkCancellation() try await modalManager.continueReplying() } } diff --git a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift index 255caf4..6605147 100644 --- a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift +++ b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift @@ -51,11 +51,9 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { appContext: appInfo?.appContext ) - try await Task.sleep(for: .seconds(3)) - try Task.checkCancellation() - await modalManager.closeModal() + try Task.checkCancellation() if let bundleIdentifier = appContext?.bundleIdentifier, let app = NSRunningApplication.runningApplications(withBundleIdentifier: bundleIdentifier).first { // Activate the app, bringing it to the foreground @@ -63,7 +61,6 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { } try Task.checkCancellation() - guard let axElement = elementMap[action.id] else { // TERMINATE on invalid action await modalManager.showModal() @@ -71,11 +68,12 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { return } + try Task.checkCancellation() _ = AXUIElementPerformAction(axElement, "AXScrollToVisible" as CFString) try await Task.sleep(for: .milliseconds(100)) - try Task.checkCancellation() do { + try Task.checkCancellation() try await focus(on: axElement) } catch { // TERMINATE on failure @@ -84,6 +82,7 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { return } + try Task.checkCancellation() if let inputText = action.inputText, let role = axElement.stringValue(forAttribute: kAXRoleAttribute) { if role == "AXComboBox" { if let parent = axElement.parent(), @@ -106,10 +105,11 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { NSPasteboard.general.clearContents() NSPasteboard.general.setString(inputText, forType: .string) try await Task.sleep(for: .milliseconds(100)) - try Task.checkCancellation() + try Task.checkCancellation() try await simulatePaste() + try Task.checkCancellation() if action.pressEnter ?? false { try await simulateEnter() } @@ -118,30 +118,38 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { NSPasteboard.general.clearContents() NSPasteboard.general.setString(inputText, forType: .string) try await Task.sleep(for: .milliseconds(100)) + try Task.checkCancellation() - try await simulateSelectAll() + + try Task.checkCancellation() try await simulatePaste() if action.pressEnter ?? false { + try Task.checkCancellation() try await simulateEnter() } } } + try Task.checkCancellation() try await Task.sleep(for: .seconds(2)) + await modalManager.showModal() + try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] ) { + try Task.checkCancellation() await modalManager.appendTool( "Updated state: \(serializedUIElement)", functionCall: functionCall, appContext: appInfo?.appContext ) } else { + try Task.checkCancellation() await modalManager.appendToolError( "Could not capture app state", functionCall: functionCall, @@ -163,6 +171,7 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { Task { do { + try Task.checkCancellation() try await modalManager.continueReplying(appInfo: newAppInfo) } catch { await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) diff --git a/TypeaheadAI/Functions/FunctionManager+SaveFile.swift b/TypeaheadAI/Functions/FunctionManager+SaveFile.swift index f1c498c..3bde3ba 100644 --- a/TypeaheadAI/Functions/FunctionManager+SaveFile.swift +++ b/TypeaheadAI/Functions/FunctionManager+SaveFile.swift @@ -39,18 +39,22 @@ extension FunctionManager { NSPasteboard.general.clearContents() NSPasteboard.general.setString(file, forType: .string) try await Task.sleep(for: .milliseconds(100)) - try Task.checkCancellation() // Open file viewer + try Task.checkCancellation() try await simulateGoToFile() try await Task.sleep(for: .milliseconds(500)) // Paste filename to file viewer + try Task.checkCancellation() try await simulatePaste() // Enter to save as filename + try Task.checkCancellation() try await simulateEnter() try await Task.sleep(for: .milliseconds(500)) + + try Task.checkCancellation() try await simulateEnter() try await Task.sleep(for: .seconds(1)) @@ -58,25 +62,26 @@ extension FunctionManager { if let replaceButton = savePanel.findFirst(condition: { $0.stringValue(forAttribute: kAXIdentifierAttribute) == "action-button-1" }) { - print("found replace button") + try Task.checkCancellation() _ = AXUIElementPerformAction(replaceButton, "AXPress" as CFString) try await Task.sleep(for: .seconds(1)) - } else { - print("no replace button found") } await modalManager.showModal() + try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] ) { + try Task.checkCancellation() await modalManager.appendTool( "Updated state: \(serializedUIElement)", functionCall: functionCall, appContext: appInfo?.appContext ) } else { + try Task.checkCancellation() await modalManager.appendToolError( "Could not capture app state", functionCall: functionCall, @@ -92,6 +97,7 @@ extension FunctionManager { Task { do { + try Task.checkCancellation() try await modalManager.continueReplying(appInfo: newAppInfo) } catch { await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) diff --git a/TypeaheadAI/Functions/FunctionManager.swift b/TypeaheadAI/Functions/FunctionManager.swift index 9a264df..3291fd5 100644 --- a/TypeaheadAI/Functions/FunctionManager.swift +++ b/TypeaheadAI/Functions/FunctionManager.swift @@ -158,7 +158,7 @@ class FunctionManager: ObservableObject, } isExecuting = true - currentTask = Task.init { [weak self] in + currentTask = Task { [weak self] in switch functionCall.name { case "open_application": do { diff --git a/TypeaheadAI/WindowManagers/ModalManager.swift b/TypeaheadAI/WindowManagers/ModalManager.swift index 7e07b11..51a7462 100644 --- a/TypeaheadAI/WindowManagers/ModalManager.swift +++ b/TypeaheadAI/WindowManagers/ModalManager.swift @@ -732,6 +732,7 @@ class ModalManager: ObservableObject { userIntents = nil Task { + try Task.checkCancellation() try await self.clientManager?.refine( messages: self.messages, prevAppInfo: appInfo, From 7adf5d36c6f4ebb0d68fa8123ec5a31c30e3ab27 Mon Sep 17 00:00:00 2001 From: kchro3 Date: Tue, 2 Jan 2024 11:29:33 -0800 Subject: [PATCH 3/3] update --- TypeaheadAI.xcodeproj/project.pbxproj | 4 +++ TypeaheadAI/ClientManager.swift | 7 +++++ TypeaheadAI/Extensions/Task+Extension.swift | 22 ++++++++++++++ .../FunctionManager+OpenApplication.swift | 12 ++------ .../Functions/FunctionManager+OpenFile.swift | 30 +++++-------------- .../Functions/FunctionManager+OpenURL.swift | 21 ++++--------- .../FunctionManager+PerformUIAction.swift | 28 ++++------------- .../Functions/FunctionManager+SaveFile.swift | 27 +++++------------ TypeaheadAI/Functions/FunctionManager.swift | 18 +++++++++++ TypeaheadAI/Traits/CanFocusOnElement.swift | 5 ++-- TypeaheadAI/Traits/CanSimulateClose.swift | 4 +-- TypeaheadAI/Traits/CanSimulateCopy.swift | 4 +-- TypeaheadAI/Traits/CanSimulateEnter.swift | 4 +-- TypeaheadAI/Traits/CanSimulateGoToFile.swift | 4 +-- TypeaheadAI/Traits/CanSimulatePaste.swift | 4 +-- TypeaheadAI/Traits/CanSimulateSelectAll.swift | 4 +-- TypeaheadAI/Views/Modal/ModalFooterView.swift | 4 --- TypeaheadAI/Views/Modal/ModalView.swift | 2 +- TypeaheadAI/WindowManagers/ModalManager.swift | 13 ++++---- 19 files changed, 100 insertions(+), 117 deletions(-) create mode 100644 TypeaheadAI/Extensions/Task+Extension.swift diff --git a/TypeaheadAI.xcodeproj/project.pbxproj b/TypeaheadAI.xcodeproj/project.pbxproj index 5699433..17e8b44 100644 --- a/TypeaheadAI.xcodeproj/project.pbxproj +++ b/TypeaheadAI.xcodeproj/project.pbxproj @@ -79,6 +79,7 @@ 2B9A89852B3CDC9700041856 /* AXSavePanelVisitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2B9A89842B3CDC9700041856 /* AXSavePanelVisitor.swift */; }; 2B9A89872B3CE1D300041856 /* CanFocusOnElement.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2B9A89862B3CE1D300041856 /* CanFocusOnElement.swift */; }; 2B9A898B2B4122DC00041856 /* UIElementVisitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2B9A898A2B4122DC00041856 /* UIElementVisitor.swift */; }; + 2B9A898F2B439BAC00041856 /* Task+Extension.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2B9A898E2B439BAC00041856 /* Task+Extension.swift */; }; 2BA3C2352AADAC5700537F95 /* llama in Frameworks */ = {isa = PBXBuildFile; productRef = 2BA3C2342AADAC5700537F95 /* llama */; }; 2BA3C2372AADAD9A00537F95 /* SpecialCopyActor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2BA3C2362AADAD9A00537F95 /* SpecialCopyActor.swift */; }; 2BA7F0792A9ABBA8003D38BA /* TypeaheadAIApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2BA7F0782A9ABBA8003D38BA /* TypeaheadAIApp.swift */; }; @@ -216,6 +217,7 @@ 2B9A89842B3CDC9700041856 /* AXSavePanelVisitor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AXSavePanelVisitor.swift; sourceTree = ""; }; 2B9A89862B3CE1D300041856 /* CanFocusOnElement.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CanFocusOnElement.swift; sourceTree = ""; }; 2B9A898A2B4122DC00041856 /* UIElementVisitor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UIElementVisitor.swift; sourceTree = ""; }; + 2B9A898E2B439BAC00041856 /* Task+Extension.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Task+Extension.swift"; sourceTree = ""; }; 2BA3C2362AADAD9A00537F95 /* SpecialCopyActor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SpecialCopyActor.swift; sourceTree = ""; }; 2BA7F0752A9ABBA8003D38BA /* TypeaheadAI.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TypeaheadAI.app; sourceTree = BUILT_PRODUCTS_DIR; }; 2BA7F0782A9ABBA8003D38BA /* TypeaheadAIApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TypeaheadAIApp.swift; sourceTree = ""; }; @@ -592,6 +594,7 @@ 2B3435022B1EA33500423EE8 /* String+HTMLParser.swift */, 2BE7BB8C2B258C4400164F88 /* AXUIElement+Extension.swift */, 2BDDB98A2B27DDFF00D52BF0 /* String+XMLMarkdown.swift */, + 2B9A898E2B439BAC00041856 /* Task+Extension.swift */, ); path = Extensions; sourceTree = ""; @@ -802,6 +805,7 @@ 2B4BDB8B2ACC281E00E55D78 /* IntentManager.swift in Sources */, 2B9A89872B3CE1D300041856 /* CanFocusOnElement.swift in Sources */, 2B7D35862B01B16A00E85AEF /* OutroOnboardingView.swift in Sources */, + 2B9A898F2B439BAC00041856 /* Task+Extension.swift in Sources */, 2BAA97912B0366C400EC6A63 /* OnboardingWindowManager.swift in Sources */, 2BAFDB662AF58E0B009C8370 /* ConversationView.swift in Sources */, 2BAA97892B02DF6E00EC6A63 /* QuickActionExplanationOnboardingView.swift in Sources */, diff --git a/TypeaheadAI/ClientManager.swift b/TypeaheadAI/ClientManager.swift index 8fe2721..2e4dbc8 100644 --- a/TypeaheadAI/ClientManager.swift +++ b/TypeaheadAI/ClientManager.swift @@ -284,7 +284,11 @@ class ClientManager: ObservableObject, CanGetUIElements { ) async { cancelStreamingTask() isExecuting = true + print("starting stream request") + currentStreamingTask = Task.init { [weak self] in + print("starting client manager task") + let uuid = try? await self?.supabaseManager?.client.auth.session.user.id let payload = RequestPayload( uuid: uuid ?? UUID(uuidString: "00000000-0000-0000-0000-000000000000")!, @@ -338,6 +342,8 @@ class ClientManager: ObservableObject, CanGetUIElements { } await MainActor.run { + print("finishing client manager task") + self?.currentStreamingTask = nil self?.isExecuting = false } @@ -467,6 +473,7 @@ class ClientManager: ObservableObject, CanGetUIElements { @MainActor func cancelStreamingTask() { + print("cancelling client manager task") currentStreamingTask?.cancel() currentStreamingTask = nil isExecuting = false diff --git a/TypeaheadAI/Extensions/Task+Extension.swift b/TypeaheadAI/Extensions/Task+Extension.swift new file mode 100644 index 0000000..e4000d7 --- /dev/null +++ b/TypeaheadAI/Extensions/Task+Extension.swift @@ -0,0 +1,22 @@ +// +// Task+Extension.swift +// TypeaheadAI +// +// Created by Jeff Hara on 1/1/24. +// + +import Foundation + +extension Task where Success == Never, Failure == Never { + + /// Helper function to check for cancellations before and after sleeping + static func sleepSafe( + for duration: C.Instant.Duration, + tolerance: C.Instant.Duration? = nil, + clock: C = ContinuousClock() + ) async throws where C : Clock { + try Task.checkCancellation() + try await Task.sleep(for: duration, tolerance: tolerance, clock: clock) + try Task.checkCancellation() + } +} diff --git a/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift b/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift index fcc285d..8f91de8 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenApplication.swift @@ -35,9 +35,8 @@ extension FunctionManager { try Task.checkCancellation() // Activate the app, bringing it to the foreground NSWorkspace.shared.open(url) - try await Task.sleep(for: .seconds(2)) + try await Task.sleepSafe(for: .seconds(2)) - try Task.checkCancellation() let newAppContext = try await fetchAppContext() let (newUIElement, newElementMap) = getUIElements(appContext: newAppContext) guard let serializedUIElement = newUIElement?.serialize( @@ -66,13 +65,6 @@ extension FunctionManager { apps: appInfo?.apps ?? [:] ) - Task { - do { - try Task.checkCancellation() - try await modalManager.continueReplying(appInfo: newAppInfo) - } catch { - await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) - } - } + modalManager.continueReplying(appInfo: newAppInfo) } } diff --git a/TypeaheadAI/Functions/FunctionManager+OpenFile.swift b/TypeaheadAI/Functions/FunctionManager+OpenFile.swift index de25e68..49a390d 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenFile.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenFile.swift @@ -27,40 +27,31 @@ extension FunctionManager { let app = NSRunningApplication.runningApplications(withBundleIdentifier: bundleIdentifier).first { // Activate the app, bringing it to the foreground app.activate(options: [.activateIgnoringOtherApps]) - try await Task.sleep(for: .milliseconds(100)) + try await Task.sleepSafe(for: .milliseconds(100)) } - try Task.checkCancellation() - // Copy filename to clipboard NSPasteboard.general.clearContents() NSPasteboard.general.setString(file, forType: .string) - try await Task.sleep(for: .milliseconds(100)) - try Task.checkCancellation() + try await Task.sleepSafe(for: .milliseconds(100)) // Select the file - try Task.checkCancellation() try await simulateGoToFile() - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) // Paste filename to file search field - try Task.checkCancellation() try await simulatePaste() - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) // Enter twice to pick and attach file - try Task.checkCancellation() try await simulateEnter() - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) - try Task.checkCancellation() try await simulateEnter() - try await Task.sleep(for: .seconds(2)) + try await Task.sleepSafe(for: .seconds(2)) - try Task.checkCancellation() await modalManager.showModal() - try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] @@ -84,13 +75,6 @@ extension FunctionManager { apps: appInfo?.apps ?? [:] ) - Task { - do { - try Task.checkCancellation() - try await modalManager.continueReplying(appInfo: newAppInfo) - } catch { - await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) - } - } + modalManager.continueReplying(appInfo: newAppInfo) } } diff --git a/TypeaheadAI/Functions/FunctionManager+OpenURL.swift b/TypeaheadAI/Functions/FunctionManager+OpenURL.swift index 5f27fd5..85c79c3 100644 --- a/TypeaheadAI/Functions/FunctionManager+OpenURL.swift +++ b/TypeaheadAI/Functions/FunctionManager+OpenURL.swift @@ -25,9 +25,8 @@ extension FunctionManager { try Task.checkCancellation() try await openURL(url) - try await Task.sleep(for: .seconds(5)) + try await Task.sleepSafe(for: .seconds(5)) - try Task.checkCancellation() let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext) if let serializedUIElement = newUIElement?.serialize( excludedActions: ["AXShowMenu", "AXScrollToVisible", "AXCancel", "AXRaise"] @@ -51,14 +50,7 @@ extension FunctionManager { apps: appInfo?.apps ?? [:] ) - Task { - do { - try Task.checkCancellation() - try await modalManager.continueReplying(appInfo: newAppInfo) - } catch { - await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) - } - } + modalManager.continueReplying(appInfo: newAppInfo) } func openAndScrapeURL(_ functionCall: FunctionCall, appInfo: AppInfo?, modalManager: ModalManager) async throws { @@ -86,9 +78,8 @@ extension FunctionManager { try Task.checkCancellation() await modalManager.closeModal() - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) - try Task.checkCancellation() try await simulateSelectAll() try await simulateCopy() } else { @@ -102,9 +93,8 @@ extension FunctionManager { try Task.checkCancellation() await modalManager.closeModal() - try await Task.sleep(for: .seconds(5)) + try await Task.sleepSafe(for: .seconds(5)) - try Task.checkCancellation() try await simulateSelectAll() try Task.checkCancellation() try await simulateCopy() @@ -154,7 +144,6 @@ extension FunctionManager { } } - try Task.checkCancellation() - try await modalManager.continueReplying() + modalManager.continueReplying() } } diff --git a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift index 6605147..47d3662 100644 --- a/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift +++ b/TypeaheadAI/Functions/FunctionManager+PerformUIAction.swift @@ -70,10 +70,9 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { try Task.checkCancellation() _ = AXUIElementPerformAction(axElement, "AXScrollToVisible" as CFString) - try await Task.sleep(for: .milliseconds(100)) + try await Task.sleepSafe(for: .milliseconds(100)) do { - try Task.checkCancellation() try await focus(on: axElement) } catch { // TERMINATE on failure @@ -104,12 +103,10 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { } else { NSPasteboard.general.clearContents() NSPasteboard.general.setString(inputText, forType: .string) - try await Task.sleep(for: .milliseconds(100)) + try await Task.sleepSafe(for: .milliseconds(100)) - try Task.checkCancellation() try await simulatePaste() - try Task.checkCancellation() if action.pressEnter ?? false { try await simulateEnter() } @@ -117,9 +114,8 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { } else { NSPasteboard.general.clearContents() NSPasteboard.general.setString(inputText, forType: .string) - try await Task.sleep(for: .milliseconds(100)) - - try Task.checkCancellation() + try await Task.sleepSafe(for: .milliseconds(100)) + try await simulateSelectAll() try Task.checkCancellation() @@ -132,8 +128,7 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { } } - try Task.checkCancellation() - try await Task.sleep(for: .seconds(2)) + try await Task.sleepSafe(for: .seconds(2)) await modalManager.showModal() @@ -165,18 +160,7 @@ extension FunctionManager: CanSimulateEnter, CanGetUIElements { apps: appInfo?.apps ?? [:] ) - await MainActor.run { - modalManager.cachedAppInfo = newAppInfo - } - - Task { - do { - try Task.checkCancellation() - try await modalManager.continueReplying(appInfo: newAppInfo) - } catch { - await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) - } - } + modalManager.continueReplying(appInfo: newAppInfo) } /// Recursively traverse an AXList to select an option that matches the expected value. diff --git a/TypeaheadAI/Functions/FunctionManager+SaveFile.swift b/TypeaheadAI/Functions/FunctionManager+SaveFile.swift index 3bde3ba..5a24aee 100644 --- a/TypeaheadAI/Functions/FunctionManager+SaveFile.swift +++ b/TypeaheadAI/Functions/FunctionManager+SaveFile.swift @@ -30,33 +30,27 @@ extension FunctionManager { let app = NSRunningApplication.runningApplications(withBundleIdentifier: bundleIdentifier).first { // Activate the app, bringing it to the foreground app.activate(options: [.activateIgnoringOtherApps]) - try await Task.sleep(for: .milliseconds(100)) + try await Task.sleepSafe(for: .milliseconds(100)) } - try Task.checkCancellation() - // Copy filename to clipboard NSPasteboard.general.clearContents() NSPasteboard.general.setString(file, forType: .string) - try await Task.sleep(for: .milliseconds(100)) + try await Task.sleepSafe(for: .milliseconds(100)) // Open file viewer - try Task.checkCancellation() try await simulateGoToFile() - try await Task.sleep(for: .milliseconds(500)) + try await Task.sleepSafe(for: .milliseconds(500)) // Paste filename to file viewer - try Task.checkCancellation() try await simulatePaste() // Enter to save as filename - try Task.checkCancellation() try await simulateEnter() - try await Task.sleep(for: .milliseconds(500)) + try await Task.sleepSafe(for: .milliseconds(500)) - try Task.checkCancellation() try await simulateEnter() - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) // Check if there's a "Replace" dialog if let replaceButton = savePanel.findFirst(condition: { @@ -64,7 +58,7 @@ extension FunctionManager { }) { try Task.checkCancellation() _ = AXUIElementPerformAction(replaceButton, "AXPress" as CFString) - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) } await modalManager.showModal() @@ -95,13 +89,6 @@ extension FunctionManager { apps: appInfo?.apps ?? [:] ) - Task { - do { - try Task.checkCancellation() - try await modalManager.continueReplying(appInfo: newAppInfo) - } catch { - await modalManager.setError(error.localizedDescription, appContext: appInfo?.appContext) - } - } + modalManager.continueReplying(appInfo: newAppInfo) } } diff --git a/TypeaheadAI/Functions/FunctionManager.swift b/TypeaheadAI/Functions/FunctionManager.swift index 3291fd5..aeee3e3 100644 --- a/TypeaheadAI/Functions/FunctionManager.swift +++ b/TypeaheadAI/Functions/FunctionManager.swift @@ -149,6 +149,7 @@ class FunctionManager: ObservableObject, currentTask?.cancel() currentTask = nil isExecuting = false + print("starting parse-and-call") let appContext = appInfo?.appContext guard let jsonData = jsonString.data(using: .utf8), @@ -159,10 +160,15 @@ class FunctionManager: ObservableObject, isExecuting = true currentTask = Task { [weak self] in + + print("starting parse-and-call task") + switch functionCall.name { case "open_application": do { try await self?.openApplication(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when opening application...", appContext: appContext) } @@ -170,6 +176,8 @@ class FunctionManager: ObservableObject, case "open_url": do { try await self?.openURL(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when opening url...", appContext: appContext) } @@ -177,6 +185,8 @@ class FunctionManager: ObservableObject, case "perform_ui_action": do { try await self?.performUIAction(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when interacting with UI...", appContext: appContext) } @@ -184,6 +194,8 @@ class FunctionManager: ObservableObject, case "open_and_scrape_url": do { try await self?.openAndScrapeURL(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when scraping URL...", appContext: appContext) } @@ -191,6 +203,8 @@ class FunctionManager: ObservableObject, case "open_file": do { try await self?.openFile(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when opening file...", appContext: appContext) } @@ -198,6 +212,8 @@ class FunctionManager: ObservableObject, case "save_file": do { try await self?.saveFile(functionCall, appInfo: appInfo, modalManager: modalManager) + } catch _ as CancellationError { + modalManager.setError("Task was cancelled.", appContext: appContext) } catch { modalManager.setError("Failed when saving file...", appContext: appContext) } @@ -206,6 +222,7 @@ class FunctionManager: ObservableObject, modalManager.setError("Function \(functionCall.name) not supported", appContext: appContext) } + print("finishing parse-and-call task") await MainActor.run { self?.currentTask = nil self?.isExecuting = false @@ -215,6 +232,7 @@ class FunctionManager: ObservableObject, @MainActor func cancelTask() { + print("cancelling parse-and-call task") currentTask?.cancel() currentTask = nil isExecuting = false diff --git a/TypeaheadAI/Traits/CanFocusOnElement.swift b/TypeaheadAI/Traits/CanFocusOnElement.swift index a13f4cc..d596dd5 100644 --- a/TypeaheadAI/Traits/CanFocusOnElement.swift +++ b/TypeaheadAI/Traits/CanFocusOnElement.swift @@ -19,7 +19,7 @@ extension CanFocusOnElement { result = AXUIElementPerformAction(axElement, "AXPress" as CFString) if result == .cannotComplete { // Retry the action after one second - try await Task.sleep(for: .seconds(1)) + try await Task.sleepSafe(for: .seconds(1)) result = AXUIElementPerformAction(axElement, "AXPress" as CFString) } } else if let size = axElement.sizeValue(forAttribute: kAXSizeAttribute), @@ -33,8 +33,7 @@ extension CanFocusOnElement { result = .actionUnsupported } - try await Task.sleep(for: .milliseconds(100)) - try Task.checkCancellation() + try await Task.sleepSafe(for: .milliseconds(100)) guard result == .success else { print("Action failed (code: \(result?.rawValue ?? -1))") throw NSError() diff --git a/TypeaheadAI/Traits/CanSimulateClose.swift b/TypeaheadAI/Traits/CanSimulateClose.swift index c6809a5..007dc0e 100644 --- a/TypeaheadAI/Traits/CanSimulateClose.swift +++ b/TypeaheadAI/Traits/CanSimulateClose.swift @@ -22,8 +22,8 @@ extension CanSimulateClose { keyUp.flags = [.maskCommand] keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) } } diff --git a/TypeaheadAI/Traits/CanSimulateCopy.swift b/TypeaheadAI/Traits/CanSimulateCopy.swift index 419624d..948658a 100644 --- a/TypeaheadAI/Traits/CanSimulateCopy.swift +++ b/TypeaheadAI/Traits/CanSimulateCopy.swift @@ -28,9 +28,9 @@ extension CanSimulateCopy { let changeCount = NSPasteboard.general.changeCount keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) if changeCount == NSPasteboard.general.changeCount { throw CanSimulateCopyError.noChangesDetected } diff --git a/TypeaheadAI/Traits/CanSimulateEnter.swift b/TypeaheadAI/Traits/CanSimulateEnter.swift index 2dcb2c8..fac241f 100644 --- a/TypeaheadAI/Traits/CanSimulateEnter.swift +++ b/TypeaheadAI/Traits/CanSimulateEnter.swift @@ -23,8 +23,8 @@ extension CanSimulateEnter { keyUp.flags = [] keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) } } diff --git a/TypeaheadAI/Traits/CanSimulateGoToFile.swift b/TypeaheadAI/Traits/CanSimulateGoToFile.swift index 78894e0..f3b78c6 100644 --- a/TypeaheadAI/Traits/CanSimulateGoToFile.swift +++ b/TypeaheadAI/Traits/CanSimulateGoToFile.swift @@ -22,8 +22,8 @@ extension CanSimulateGoToFile { keyUp.flags = [.maskShift, .maskCommand] keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) } } diff --git a/TypeaheadAI/Traits/CanSimulatePaste.swift b/TypeaheadAI/Traits/CanSimulatePaste.swift index 6b154ff..3898ac2 100644 --- a/TypeaheadAI/Traits/CanSimulatePaste.swift +++ b/TypeaheadAI/Traits/CanSimulatePaste.swift @@ -28,8 +28,8 @@ extension CanSimulatePaste { } keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) } } diff --git a/TypeaheadAI/Traits/CanSimulateSelectAll.swift b/TypeaheadAI/Traits/CanSimulateSelectAll.swift index dbcdf6c..43ccd6e 100644 --- a/TypeaheadAI/Traits/CanSimulateSelectAll.swift +++ b/TypeaheadAI/Traits/CanSimulateSelectAll.swift @@ -22,8 +22,8 @@ extension CanSimulateSelectAll { keyUp.flags = [.maskCommand] keyDown.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(20)) + try await Task.sleepSafe(for: .milliseconds(20)) keyUp.post(tap: .cghidEventTap) - try await Task.sleep(for: .milliseconds(200)) + try await Task.sleepSafe(for: .milliseconds(200)) } } diff --git a/TypeaheadAI/Views/Modal/ModalFooterView.swift b/TypeaheadAI/Views/Modal/ModalFooterView.swift index 99ba7c4..5bfdb8c 100644 --- a/TypeaheadAI/Views/Modal/ModalFooterView.swift +++ b/TypeaheadAI/Views/Modal/ModalFooterView.swift @@ -47,10 +47,6 @@ struct ModalFooterView: View { await modalManager.addUserMessage(text, appContext: nil) } } - } else if let cachedAppInfo = modalManager.cachedAppInfo { - Task { - try await modalManager.continueReplying(appInfo: cachedAppInfo) - } } } .padding(.vertical, 5) diff --git a/TypeaheadAI/Views/Modal/ModalView.swift b/TypeaheadAI/Views/Modal/ModalView.swift index 9f2691c..6663b7e 100644 --- a/TypeaheadAI/Views/Modal/ModalView.swift +++ b/TypeaheadAI/Views/Modal/ModalView.swift @@ -102,7 +102,7 @@ struct ModalView: View { Task { modalManager.forceRefresh() modalManager.closeModal() - try await Task.sleep(for: .seconds(3)) + try await Task.sleepSafe(for: .seconds(3)) try await modalManager.specialRecordActor?.specialRecord() } } diff --git a/TypeaheadAI/WindowManagers/ModalManager.swift b/TypeaheadAI/WindowManagers/ModalManager.swift index 51a7462..a34d4eb 100644 --- a/TypeaheadAI/WindowManagers/ModalManager.swift +++ b/TypeaheadAI/WindowManagers/ModalManager.swift @@ -17,7 +17,6 @@ extension Notification.Name { class ModalManager: ObservableObject { private let context: NSManagedObjectContext - var cachedAppInfo: AppInfo? = nil @Published var messages: [Message] @Published var userIntents: [String]? @@ -681,6 +680,7 @@ class ModalManager: ObservableObject { userIntents = nil } + try Task.checkCancellation() try await self.clientManager?.refine( messages: self.messages, streamHandler: defaultStreamHandler, @@ -726,12 +726,13 @@ class ModalManager: ObservableObject { } } - @MainActor - func continueReplying(appInfo: AppInfo? = nil) async throws { - isPending = true - userIntents = nil - + func continueReplying(appInfo: AppInfo? = nil) { Task { + await MainActor.run { + isPending = true + userIntents = nil + } + try Task.checkCancellation() try await self.clientManager?.refine( messages: self.messages,