Add an error for lost connection

Author: daymxn

FirebaseAI/Sources/Types/Internal/Live/AsyncWebSocket.swift

@@ -71,14 +71,30 @@ final class AsyncWebSocket: NSObject, @unchecked Sendable, URLSessionWebSocketDe
           let message = try await webSocketTask.receive()
           continuation.yield(message)
         } catch {
-          close(code: webSocketTask.closeCode, reason: webSocketTask.closeReason)
+          if let error = webSocketTask.error as? NSError {
+            close(
+              code: webSocketTask.closeCode,
+              reason: webSocketTask.closeReason,
+              underlyingError: error
+            )
+          } else {
+            close(code: webSocketTask.closeCode, reason: webSocketTask.closeReason)
+          }
         }
       }
     }
   }
 
-  private func close(code: URLSessionWebSocketTask.CloseCode, reason: Data?) {
-    let error = WebSocketClosedError(closeCode: code, closeReason: reason)
+  private func close(
+    code: URLSessionWebSocketTask.CloseCode,
+    reason: Data?,
+    underlyingError: Error? = nil
+  ) {
+    let error = WebSocketClosedError(
+      closeCode: code,
+      closeReason: reason,
+      underlyingError: underlyingError
+    )
     closeError.withLock {
       $0 = error
     }
@@ -110,21 +126,27 @@ private extension URLSessionWebSocketTask {
 ///
 /// See the `closeReason` for why, or the `errorCode` for the corresponding
 /// `URLSessionWebSocketTask.CloseCode`.
+///
+/// In some cases, the `NSUnderlyingErrorKey` key may be populated with an
+/// error for additional context.
 struct WebSocketClosedError: Error, Sendable, CustomNSError {
   let closeCode: URLSessionWebSocketTask.CloseCode
   let closeReason: String
+  let underlyingError: Error?
 
-  init(closeCode: URLSessionWebSocketTask.CloseCode, closeReason: Data?) {
+  init(closeCode: URLSessionWebSocketTask.CloseCode, closeReason: Data?, underlyingError: Error? = nil) {
     self.closeCode = closeCode
     self.closeReason = closeReason
       .flatMap { String(data: $0, encoding: .utf8) } ?? "Unknown reason."
+    self.underlyingError = underlyingError
   }
 
   var errorCode: Int { closeCode.rawValue }
 
   var errorUserInfo: [String: Any] {
     [
       NSLocalizedDescriptionKey: "WebSocket closed with code \(closeCode.rawValue). Reason: \(closeReason)",
+      NSUnderlyingErrorKey: underlyingError as Any,
     ]
   }
 }

FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift

@@ -244,9 +244,14 @@ actor LiveSessionService {
         if let error = error as? WebSocketClosedError {
           // only raise an error if the session didn't close normally (ie; the user calling close)
           if error.closeCode != .goingAway {
-            let error = LiveSessionUnexpectedClosureError(underlyingError: error)
+            let closureError: Error
+            if let error = error.underlyingError as? NSError, error.domain == NSURLErrorDomain, error.code == NSURLErrorNetworkConnectionLost {
+              closureError = LiveSessionLostConnectionError(underlyingError: error)
+            } else {
+              closureError = LiveSessionUnexpectedClosureError(underlyingError: error)
+            }
             close()
-            responseContinuation.finish(throwing: error)
+            responseContinuation.finish(throwing: closureError)
           }
         } else {
           // an error occurred outside the websocket, so it's likely not closed
@@ -365,10 +370,27 @@ public struct LiveSessionUnsupportedMessageError: Error, Sendable, CustomNSError
   }
 }
 
+/// The live session was closed, because the network connection was lost.
+///
+/// Check the `NSUnderlyingErrorKey` entry in ``errorUserInfo`` for the error that caused this.
+public struct LiveSessionLostConnectionError: Error, Sendable, CustomNSError {
+  let underlyingError: Error
+
+  init(underlyingError: Error) {
+    self.underlyingError = underlyingError
+  }
+
+  public var errorUserInfo: [String: Any] {
+    [
+      NSLocalizedDescriptionKey: "The live session lost connection to the server. Cause: \(underlyingError.localizedDescription)",
+      NSUnderlyingErrorKey: underlyingError,
+    ]
+  }
+}
+
 /// The live session was closed, but not for a reason the SDK expected.
 ///
 /// Check the `NSUnderlyingErrorKey` entry in ``errorUserInfo`` for the error that caused this.
-// TODO: two common causes I can think of are api limits and network issues. I wonder if we can catch these somehow, as they seem common enough to surface as actual errors.
 public struct LiveSessionUnexpectedClosureError: Error, Sendable, CustomNSError {
   let underlyingError: WebSocketClosedError