Refactor to use async/await and remove URLSessionWebSocketDelegate

Author: andrewheard

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

@@ -15,7 +15,7 @@
 import Foundation
 
 /// Response message for BidiGenerateContent RPC call.
-public struct BidiGenerateContentServerMessage {
+public struct BidiGenerateContentServerMessage: Sendable {
   // TODO: Make this type `internal`
 
   /// The type of the message.

FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift

@@ -40,15 +40,15 @@ public final class LiveGenerativeModel {
     self.urlSession = urlSession
   }
 
-  public func connect() async throws -> LiveSession {
+  public func connect() -> LiveSession {
     let liveSession = LiveSession(
       modelResourceName: modelResourceName,
       generationConfig: generationConfig,
       url: webSocketURL(),
       urlSession: urlSession
     )
     print("Opening Live Session...")
-    try await liveSession.open()
+    liveSession.openConnection()
     return liveSession
   }
 

FirebaseAI/Sources/Types/Public/Live/LiveSession.swift

@@ -15,35 +15,18 @@
 import Foundation
 
 // TODO: Extract most of this file into a service class similar to `GenerativeAIService`.
-public final class LiveSession: NSObject, URLSessionWebSocketDelegate, URLSessionTaskDelegate {
-  private enum State {
-    case notConnected
-    case connecting
-    case setupSent
-    case ready
-    case closed
-  }
-
-  private enum WebSocketError: Error {
-    case connectionClosed
-  }
-
+public final class LiveSession: Sendable {
   let modelResourceName: String
   let generationConfig: LiveGenerationConfig?
   let webSocket: URLSessionWebSocketTask
 
-  // TODO: Refactor this property, potentially returning responses after `connect`.
   public let responses: AsyncThrowingStream<BidiGenerateContentServerMessage, Error>
+  private let responseContinuation: AsyncThrowingStream<BidiGenerateContentServerMessage, Error>
+    .Continuation
 
-  private var state: State = .notConnected
-  private var pendingMessages: [(String, CheckedContinuation<Void, Error>)] = []
   private let jsonEncoder = JSONEncoder()
   private let jsonDecoder = JSONDecoder()
 
-  // TODO: Properly wrap callback code using `withCheckedContinuation` or similar.
-  private let responseContinuation: AsyncThrowingStream<BidiGenerateContentServerMessage, Error>
-    .Continuation
-
   init(modelResourceName: String,
        generationConfig: LiveGenerationConfig?,
        url: URL,
@@ -54,166 +37,61 @@ public final class LiveSession: NSObject, URLSessionWebSocketDelegate, URLSessio
     (responses, responseContinuation) = AsyncThrowingStream.makeStream()
   }
 
-  func open() async throws {
-    guard state == .notConnected else {
-      print("Web socket is not in a valid state to be opened: \(state)")
-      return
-    }
-
-    state = .connecting
-    webSocket.delegate = self
-    webSocket.resume()
-
-    print("Opening websocket")
-  }
-
-  private func failPendingMessages(with error: Error) {
-    for (_, continuation) in pendingMessages {
-      continuation.resume(throwing: error)
-    }
-    pendingMessages.removeAll()
-    responseContinuation.finish(throwing: error)
+  deinit {
+    webSocket.cancel(with: .goingAway, reason: nil)
   }
 
-  private func processPendingMessages() {
-    for (message, continuation) in pendingMessages {
-      Task {
-        do {
-          try await send(message)
-          continuation.resume()
-        } catch {
-          continuation.resume(throwing: error)
-        }
-      }
-    }
-    pendingMessages.removeAll()
-  }
-
-  private func send(_ message: String) async throws {
+  public func sendMessage(_ message: String) async throws {
     let content = ModelContent(role: "user", parts: [message])
     let clientContent = BidiGenerateContentClientContent(turns: [content], turnComplete: true)
     let clientMessage = BidiGenerateContentClientMessage.clientContent(clientContent)
     let clientMessageData = try jsonEncoder.encode(clientMessage)
-    let clientMessageJSON = String(data: clientMessageData, encoding: .utf8)
-    print("Client Message JSON: \(clientMessageJSON)")
     try await webSocket.send(.data(clientMessageData))
-    setReceiveHandler()
   }
 
-  public func sendMessage(_ message: String) async throws {
-    if state == .ready {
-      try await send(message)
-    } else {
-      try await withCheckedThrowingContinuation { continuation in
-        pendingMessages.append((message, continuation))
-      }
+  func openConnection() {
+    webSocket.resume()
+    // TODO: Verify that this task gets cancelled on deinit
+    Task {
+      await startEventLoop()
     }
   }
 
-  public func urlSession(_ session: URLSession,
-                         webSocketTask: URLSessionWebSocketTask,
-                         didOpenWithProtocol protocol: String?) {
-    print("Web Socket opened.")
-
-    guard state == .connecting else {
-      print("Web socket is not in a valid state to be opened: \(state)")
-      return
+  private func startEventLoop() async {
+    defer {
+      webSocket.cancel(with: .goingAway, reason: nil)
     }
 
     do {
-      let setup = BidiGenerateContentSetup(
-        model: modelResourceName, generationConfig: generationConfig
-      )
-      let message = BidiGenerateContentClientMessage.setup(setup)
-      let messageData = try jsonEncoder.encode(message)
-      let messageJSON = String(data: messageData, encoding: .utf8)
-      print("JSON: \(messageJSON)")
-      webSocketTask.send(.data(messageData)) { error in
-        if let error {
-          print("Send Error: \(error)")
-          self.state = .closed
-          self.failPendingMessages(with: error)
-          return
-        }
-
-        self.state = .setupSent
-        self.setReceiveHandler()
-      }
-    } catch {
-      print(error)
-      state = .closed
-      failPendingMessages(with: error)
-    }
-  }
-
-  public func urlSession(_ session: URLSession,
-                         webSocketTask: URLSessionWebSocketTask,
-                         didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
-                         reason: Data?) {
-    print("Web Socket closed.")
-    state = .closed
-    failPendingMessages(with: WebSocketError.connectionClosed)
-    responseContinuation.finish()
-  }
-
-  func setReceiveHandler() {
-    guard state == .setupSent || state == .ready else {
-      print("Web socket is not in a valid state to receive messages: \(state)")
-      return
-    }
+      try await sendSetupMessage()
 
-    webSocket.receive { result in
-      do {
-        let message = try result.get()
+      while !Task.isCancelled {
+        let message = try await webSocket.receive()
         switch message {
         case let .string(string):
           print("Unexpected string response: \(string)")
-          self.setReceiveHandler()
         case let .data(data):
-          let response = try self.jsonDecoder.decode(
+          let response = try jsonDecoder.decode(
             BidiGenerateContentServerMessage.self,
             from: data
           )
-          let responseJSON = String(data: data, encoding: .utf8)
-
-          switch response.messageType {
-          case .setupComplete:
-            print("Setup Complete: \(responseJSON)")
-            self.state = .ready
-            self.processPendingMessages()
-          case .serverContent:
-            print("Server Content: \(responseJSON)")
-          case .toolCall:
-            // TODO: Tool calls not yet implemented
-            print("Tool Call: \(responseJSON)")
-          case .toolCallCancellation:
-            // TODO: Tool call cancellation not yet implemented
-            print("Tool Call Cancellation: \(responseJSON)")
-          case let .goAway(goAway):
-            if let timeLeft = goAway.timeLeft {
-              print("Server will disconnect in \(timeLeft) seconds.")
-            } else {
-              print("Server will disconnect soon.")
-            }
-          }
-
-          self.responseContinuation.yield(response)
-
-          if self.state == .closed {
-            print("Web socket is closed, not listening for more messages.")
-          } else {
-            self.setReceiveHandler()
-          }
+          responseContinuation.yield(response)
         @unknown default:
           print("Unknown message received")
-          self.setReceiveHandler()
         }
-      } catch {
-        // handle the error
-        print(error)
-        self.state = .closed
-        self.responseContinuation.finish(throwing: error)
       }
+    } catch {
+      responseContinuation.finish(throwing: error)
     }
+    responseContinuation.finish()
+  }
+
+  private func sendSetupMessage() async throws {
+    let setup = BidiGenerateContentSetup(
+      model: modelResourceName, generationConfig: generationConfig
+    )
+    let message = BidiGenerateContentClientMessage.setup(setup)
+    let messageData = try jsonEncoder.encode(message)
+    try await webSocket.send(.data(messageData))
   }
 }

FirebaseAI/Tests/TestApp/Sources/ContentView.swift

@@ -37,7 +37,7 @@ struct ContentView: View {
     .padding()
     .task {
       do {
-        let liveSession = try await liveModel.connect()
+        let liveSession = liveModel.connect()
         try await liveSession.sendMessage("Why is the sky blue?")
         for try await response in liveSession.responses {
           responses.append(String(describing: response))