Revert to previous content on error

Author: intitni

Core/Sources/PromptToCodeService/PromptToCodeService.swift

@@ -8,12 +8,14 @@ public final class PromptToCodeService: ObservableObject {
     var promptToCodeAPI: PromptToCodeAPI {
         designatedPromptToCodeAPI ?? OpenAIPromptToCodeAPI()
     }
+
     var runningAPI: PromptToCodeAPI?
 
     @Published public var oldCode: String?
     @Published public var code: String
     @Published public var isResponding: Bool = false
     @Published public var description: String = ""
+    public var oldDescription: String?
     public var canRevert: Bool { oldCode != nil }
     public var selectionRange: CursorRange
     public var language: CopilotLanguage
@@ -39,31 +41,48 @@ public final class PromptToCodeService: ObservableObject {
         runningAPI = api
         isResponding = true
         let toBemodified = code
+        oldDescription = description
         oldCode = code
         code = ""
         description = ""
         defer { isResponding = false }
-        let stream = try await api.modifyCode(
-            code: toBemodified,
-            language: language,
-            indentSize: indentSize,
-            usesTabsForIndentation: usesTabsForIndentation,
-            requirement: prompt
-        )
-        for try await fragment in stream {
-            Task { @MainActor in
-                code = fragment.code
-                description = fragment.description
+        do {
+            let stream = try await api.modifyCode(
+                code: toBemodified,
+                language: language,
+                indentSize: indentSize,
+                usesTabsForIndentation: usesTabsForIndentation,
+                requirement: prompt
+            )
+            for try await fragment in stream {
+                Task { @MainActor in
+                    code = fragment.code
+                    description = fragment.description
+                }
             }
+        } catch {
+            if let oldCode {
+                code = oldCode
+            }
+            if let oldDescription {
+                description = oldDescription
+            }
+            oldCode = nil
+            oldDescription = nil
+            throw error
         }
     }
 
     public func revert() {
         guard let oldCode = oldCode else { return }
         code = oldCode
+        if let oldDescription {
+            description = oldDescription
+        }
         self.oldCode = nil
+        self.oldDescription = nil
     }
-    
+
     public func generateCompletion() -> CopilotCompletion {
         .init(
             text: code,
@@ -73,7 +92,7 @@ public final class PromptToCodeService: ObservableObject {
             displayText: code
         )
     }
-    
+
     public func stopResponding() {
         runningAPI?.stopResponding()
         isResponding = false
@@ -88,19 +107,19 @@ protocol PromptToCodeAPI {
         usesTabsForIndentation: Bool,
         requirement: String
     ) async throws -> AsyncThrowingStream<(code: String, description: String), Error>
-    
+
     func stopResponding()
 }
 
 final class OpenAIPromptToCodeAPI: PromptToCodeAPI {
     var service: (any ChatGPTServiceType)?
-    
+
     func stopResponding() {
         Task {
             await service?.stopReceivingMessage()
         }
     }
-    
+
     func modifyCode(
         code: String,
         language: CopilotLanguage,
@@ -198,7 +217,7 @@ final class OpenAIPromptToCodeAPI: PromptToCodeAPI {
         }
 
         let description = extractDescriptionFromMarkdown(content, startIndex: endIndex)
-        
+
         return (code, description)
     }
 }
@@ -207,7 +226,7 @@ final class CopilotPromptToCodeAPI: PromptToCodeAPI {
     func stopResponding() {
         fatalError()
     }
-    
+
     func modifyCode(
         code: String,
         language: CopilotLanguage,