fix: task receive loop

fal-ai · Dec 1, 2023 · 57cad55 · 57cad55
1 parent 06edf92
commit 57cad55
Showing 1 changed file with 19 additions and 6 deletions.
diff --git a/Sources/FalClient/Realtime.swift b/Sources/FalClient/Realtime.swift
@@ -19,6 +19,7 @@ public enum FalRealtimeError: Error {
     case connectionError
     case unauthorized
     case invalidResult
+    case serviceError(type: String, reason: String)
 }
 
 public class RealtimeConnection<Input> {
@@ -165,25 +166,26 @@ class WebSocketConnection: NSObject, URLSessionWebSocketDelegate {
             switch incomingMessage {
             case let .success(message):
                 do {
+                    self?.receiveMessage()
                     let data = try message.data()
                     guard let parsedMessage = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
                         self?.onError(FalRealtimeError.invalidResult)
                         return
                     }
                     if isSuccessResult(parsedMessage) {
                         self?.onMessage(data)
+                        return
+                    }
+                    if let error = getError(parsedMessage) {
+                        self?.onError(error)
+                        return
                     }
-//                    if (parsedMessage["status"] as? String != "error") {
-//                        self?.task?.cancel()
-//                    }
-
                 } catch {
                     self?.onError(error)
                 }
             case let .failure(error):
                 self?.onError(error)
             }
-            self?.receiveMessage()
         }
     }
 
@@ -251,7 +253,18 @@ public protocol Realtime {
 }
 
 func isSuccessResult(_ message: [String: Any]) -> Bool {
-    return message["status"] as? String != "error" && message["type"] as? String != "x-fal-message"
+    return message["status"] as? String != "error"
+        && message["type"] as? String != "x-fal-message"
+        && message["type"] as? String != "x-fal-error"
+}
+
+func getError(_ message: [String: Any]) -> FalRealtimeError? {
+    if message["type"] as? String != "x-fal-error",
+       let error = message["error"] as? String,
+       let reason = message["reason"] as? String {
+        return FalRealtimeError.serviceError(type: error, reason: reason)
+    }
+    return nil
 }
 
 extension URLSessionWebSocketTask.Message {