Add timeout to accept suggestion

Author: intitni

EditorExtension/AcceptSuggestionCommand.swift

@@ -13,13 +13,15 @@ class AcceptSuggestionCommand: NSObject, XCSourceEditorCommand, CommandType {
     ) {
         Task {
             do {
-                let service = try getService()
-                if let content = try await service.getSuggestionAcceptedCode(
-                    editorContent: .init(invocation)
-                ) {
-                    invocation.accept(content)
-                }
-                completionHandler(nil)
+                try await (Task(timeout: 7) {
+                    let service = try getService()
+                    if let content = try await service.getSuggestionAcceptedCode(
+                        editorContent: .init(invocation)
+                    ) {
+                        invocation.accept(content)
+                    }
+                    completionHandler(nil)
+                }.value)
             } catch is CancellationError {
                 completionHandler(nil)
             } catch {
@@ -28,3 +30,33 @@ class AcceptSuggestionCommand: NSObject, XCSourceEditorCommand, CommandType {
         }
     }
 }
+
+/// https://gist.github.com/swhitty/9be89dfe97dbb55c6ef0f916273bbb97
+extension Task where Failure == Error {
+    // Start a new Task with a timeout. If the timeout expires before the operation is
+    // completed then the task is cancelled and an error is thrown.
+    init(
+        priority: TaskPriority? = nil,
+        timeout: TimeInterval,
+        operation: @escaping @Sendable () async throws -> Success
+    ) {
+        self = Task(priority: priority) {
+            try await withThrowingTaskGroup(of: Success.self) { group -> Success in
+                group.addTask(operation: operation)
+                group.addTask {
+                    try await _Concurrency.Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
+                    throw TimeoutError()
+                }
+                guard let success = try await group.next() else {
+                    throw _Concurrency.CancellationError()
+                }
+                group.cancelAll()
+                return success
+            }
+        }
+    }
+}
+
+private struct TimeoutError: LocalizedError {
+    var errorDescription: String? = "Task timed out before completion"
+}