kchro3 · kchro3 · Jan 2, 2024 · Jan 2, 2024 · Jan 2, 2024
diff --git 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 = "<group>"; };
 		2B9A89862B3CE1D300041856 /* CanFocusOnElement.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CanFocusOnElement.swift; sourceTree = "<group>"; };
 		2B9A898A2B4122DC00041856 /* UIElementVisitor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UIElementVisitor.swift; sourceTree = "<group>"; };
+		2B9A898E2B439BAC00041856 /* Task+Extension.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Task+Extension.swift"; sourceTree = "<group>"; };
 		2BA3C2362AADAD9A00537F95 /* SpecialCopyActor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SpecialCopyActor.swift; sourceTree = "<group>"; };
 		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 = "<group>"; };
@@ -592,6 +594,7 @@
 				2B3435022B1EA33500423EE8 /* String+HTMLParser.swift */,
 				2BE7BB8C2B258C4400164F88 /* AXUIElement+Extension.swift */,
 				2BDDB98A2B27DDFF00D52BF0 /* String+XMLMarkdown.swift */,
+				2B9A898E2B439BAC00041856 /* Task+Extension.swift */,
 			);
 			path = Extensions;
 			sourceTree = "<group>";
@@ -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/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
@@ -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(),
@@ -282,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")!,
@@ -320,6 +326,7 @@ class ClientManager: ObservableObject, CanGetUIElements {
                         completion: completion
                     )
 
+                    try Task.checkCancellation()
                     guard let stream = stream else {
                         self?.logger.debug("Failed to get stream")
                         return
@@ -334,7 +341,9 @@ class ClientManager: ObservableObject, CanGetUIElements {
                 }
             }
 
-            DispatchQueue.main.async {
+            await MainActor.run {
+                print("finishing client manager task")
+
                 self?.currentStreamingTask = nil
                 self?.isExecuting = false
             }
@@ -426,6 +435,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...")
@@ -462,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
@@ -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<C>(
+        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
@@ -24,7 +24,6 @@ extension FunctionManager {
             return
         }
 
-        try Task.checkCancellation()
         await modalManager.closeModal()
 
         guard let url = NSWorkspace.shared.urlForApplication(withBundleIdentifier: bundleIdentifier) else {
@@ -33,10 +32,10 @@ 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 await Task.sleepSafe(for: .seconds(2))
 
         let newAppContext = try await fetchAppContext()
         let (newUIElement, newElementMap) = getUIElements(appContext: newAppContext)
@@ -52,6 +51,7 @@ extension FunctionManager {
             return
         }
 
+        try Task.checkCancellation()
         await modalManager.showModal()
         await modalManager.appendTool(
             "Updated state: \(serializedUIElement)",
@@ -65,14 +65,6 @@ extension FunctionManager {
             apps: appInfo?.apps ?? [:]
         )
 
-        try Task.checkCancellation()
-
-        Task {
-            do {
-                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
@@ -27,30 +27,28 @@ 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 await simulateGoToFile()
-        try await Task.sleep(for: .seconds(1))
+        try await Task.sleepSafe(for: .seconds(1))
 
         // Paste filename to file search field
         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 await simulateEnter()
-        try await Task.sleep(for: .seconds(1))
+        try await Task.sleepSafe(for: .seconds(1))
+
         try await simulateEnter()
-        try await Task.sleep(for: .seconds(2))
+        try await Task.sleepSafe(for: .seconds(2))
 
         await modalManager.showModal()
 
@@ -77,12 +75,6 @@ extension FunctionManager {
             apps: appInfo?.apps ?? [:]
         )
 
-        Task {
-            do {
-                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
@@ -23,8 +23,9 @@ extension FunctionManager {
             appContext: appInfo?.appContext
         )
 
+        try Task.checkCancellation()
         try await openURL(url)
-        try await Task.sleep(for: .seconds(5))
+        try await Task.sleepSafe(for: .seconds(5))
 
         let (newUIElement, newElementMap) = getUIElements(appContext: appInfo?.appContext)
         if let serializedUIElement = newUIElement?.serialize(
@@ -49,13 +50,7 @@ extension FunctionManager {
             apps: appInfo?.apps ?? [:]
         )
 
-        Task {
-            do {
-                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 {
@@ -66,6 +61,7 @@ extension FunctionManager {
             return
         }
 
+        try Task.checkCancellation()
         if url == "<current>" {
             await modalManager.appendFunction(
                 "Scraping current page...",
@@ -79,8 +75,11 @@ extension FunctionManager {
                 app.activate(options: [.activateIgnoringOtherApps])
             }
 
+            try Task.checkCancellation()
             await modalManager.closeModal()
-            try await Task.sleep(for: .seconds(1))
+
+            try await Task.sleepSafe(for: .seconds(1))
+
             try await simulateSelectAll()
             try await simulateCopy()
         } else {
@@ -91,12 +90,18 @@ extension FunctionManager {
             )
 
             try await openURL(url)
+            try Task.checkCancellation()
+
             await modalManager.closeModal()
-            try await Task.sleep(for: .seconds(5))
+            try await Task.sleepSafe(for: .seconds(5))
+
             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 +144,6 @@ extension FunctionManager {
             }
         }
 
-        try await modalManager.continueReplying()
+        modalManager.continueReplying()
     }
 }