Migrate to protocol instead of enum

Author: daymxn

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

@@ -40,7 +40,7 @@ struct BidiGenerateContentServerMessage: Sendable {
   }
 
   /// The message type.
-  let messageType: MessageType
+  let messageType: MessageType?
 
   /// Usage metadata about the response(s).
   let usageMetadata: GenerateContentResponse.UsageMetadata?
@@ -86,11 +86,7 @@ extension BidiGenerateContentServerMessage: Decodable {
     } else if let goAway = try container.decodeIfPresent(GoAway.self, forKey: .goAway) {
       messageType = .goAway(goAway)
     } else {
-      let context = DecodingError.Context(
-        codingPath: decoder.codingPath,
-        debugDescription: "Could not decode server message."
-      )
-      throw DecodingError.dataCorrupted(context)
+      messageType = nil
     }
 
     usageMetadata = try container.decodeIfPresent(

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

@@ -206,10 +206,18 @@ actor LiveSessionService {
               from: message
             )
           } catch {
-            throw LiveSessionUnsupportedMessageError(underlyingError: error)
+            let error = LiveSessionUnsupportedMessageError(underlyingError: error)
+
+            // if we've already finished setting up, then only surface the error through responses
+            // otherwise, make the setup task error as well
+            if !resumed {
+              setupComplete.resume(throwing: error)
+            }
+            throw error
           }
 
-          if case .setupComplete = response.messageType {
+          switch response.messageType {
+          case .setupComplete:
             if resumed {
               AILog.debug(
                 code: .duplicateLiveSessionSetupComplete,
@@ -221,17 +229,33 @@ actor LiveSessionService {
               resumed = true
               setupComplete.resume()
             }
-          } else if let liveMessage = LiveServerMessage(from: response) {
-            if case let .goingAwayNotice(message) = liveMessage.messageType {
-              // TODO: (b/444045023) When auto session resumption is enabled, call `connect` again
-              AILog.debug(
-                code: .liveSessionGoingAwaySoon,
-                "Session expires in: \(message.goAway.timeLeft?.timeInterval ?? 0)"
-              )
-            }
-
-            responseContinuation.yield(liveMessage)
-          } else {
+          case let .goAway(goAway):
+            // TODO: (b/444045023) When auto session resumption is enabled, call `connect` again
+            AILog.debug(
+              code: .liveSessionGoingAwaySoon,
+              "Session expires in: \(goAway.timeLeft?.timeInterval ?? 0)"
+            )
+            
+            responseContinuation.yield(LiveServerGoingAwayNotice(
+              goAway,
+              usageMetadata: response.usageMetadata
+            ))
+          case let .serverContent(serverContent):
+            responseContinuation.yield(LiveServerContent(
+              serverContent,
+              usageMetadata: response.usageMetadata
+            ))
+          case let .toolCall(toolCall):
+            responseContinuation.yield(LiveServerToolCall(
+              toolCall,
+              usageMetadata: response.usageMetadata
+            ))
+          case let .toolCallCancellation(toolCallCancellation):
+            responseContinuation.yield(LiveServerToolCallCancellation(
+              toolCallCancellation,
+              usageMetadata: response.usageMetadata
+            ))
+          case .none:
             // we don't raise an error, since this allows us to add support internally but not
             // publicly. We still log it in debug though, in case it's not expected.
             AILog.debug(

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

@@ -19,7 +19,7 @@
 /// may choose to buffer and play it out in realtime.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerContent: Sendable {
+public struct LiveServerContent: LiveServerMessage {
   let serverContent: BidiGenerateContentServerContent
 
   /// The content that the model has generated as part of the current
@@ -77,7 +77,11 @@ public struct LiveServerContent: Sendable {
     serverContent.outputTranscription.map { LiveTranscription($0) }
   }
 
-  init(_ serverContent: BidiGenerateContentServerContent) {
+  public var usageMetadata: GenerateContentResponse.UsageMetadata?
+
+  init(_ serverContent: BidiGenerateContentServerContent,
+       usageMetadata: GenerateContentResponse.UsageMetadata?) {
     self.serverContent = serverContent
+    self.usageMetadata = usageMetadata
   }
 }

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

@@ -19,15 +19,18 @@ import Foundation
 /// To  learn more about session limits,  see the docs on [Maximum session duration](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/multimodal-live#maximum-session-duration)\.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerGoingAwayNotice: Sendable {
+public struct LiveServerGoingAwayNotice: LiveServerMessage {
   let goAway: GoAway
   /// The remaining time before the connection will be terminated as ABORTED.
   ///
   /// The minimal time returned here is specified differently together with
   /// the rate limits for a given model.
   public var timeLeft: TimeInterval? { goAway.timeLeft?.timeInterval }
 
-  init(_ goAway: GoAway) {
+  public var usageMetadata: GenerateContentResponse.UsageMetadata?
+
+  init(_ goAway: GoAway, usageMetadata: GenerateContentResponse.UsageMetadata?) {
     self.goAway = goAway
+    self.usageMetadata = usageMetadata
   }
 }

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

@@ -15,63 +15,7 @@
 /// Update from the  server, generated from the model in response to client messages.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerMessage: Sendable {
-  let serverMessage: BidiGenerateContentServerMessage
-
-  /// The type of message sent from the server.
-  public enum MessageType: Sendable {
-    /// Content generated by the model in response to client messages.
-    case content(LiveServerContent)
-
-    /// Request for the client to execute the provided functions.
-    case toolCall(LiveServerToolCall)
-
-    /// Notification for the client that a previously issued ``LiveServerToolCall`` should be
-    /// cancelled.
-    case toolCallCancellation(LiveServerToolCallCancellation)
-
-    /// Server will disconnect soon.
-    case goingAwayNotice(LiveServerGoingAwayNotice)
-  }
-
-  /// The actual message sent from the server.
-  public var messageType: MessageType
-
+public protocol LiveServerMessage: Sendable {
   /// Metadata on the usage of the cached content.
-  public var usageMetadata: GenerateContentResponse.UsageMetadata? { serverMessage.usageMetadata }
-}
-
-// MARK: - Internal parsing
-
-@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
-@available(watchOS, unavailable)
-extension LiveServerMessage {
-  init?(from serverMessage: BidiGenerateContentServerMessage) {
-    guard let messageType = LiveServerMessage.MessageType(from: serverMessage.messageType) else {
-      return nil
-    }
-
-    self.serverMessage = serverMessage
-    self.messageType = messageType
-  }
-}
-
-@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
-@available(watchOS, unavailable)
-extension LiveServerMessage.MessageType {
-  init?(from serverMessage: BidiGenerateContentServerMessage.MessageType) {
-    switch serverMessage {
-    case .setupComplete:
-      // this is handled internally, and should not be surfaced to users
-      return nil
-    case let .serverContent(msg):
-      self = .content(LiveServerContent(msg))
-    case let .toolCall(msg):
-      self = .toolCall(LiveServerToolCall(msg))
-    case let .toolCallCancellation(msg):
-      self = .toolCallCancellation(LiveServerToolCallCancellation(msg))
-    case let .goAway(msg):
-      self = .goingAwayNotice(LiveServerGoingAwayNotice(msg))
-    }
-  }
+  var usageMetadata: GenerateContentResponse.UsageMetadata? { get }
 }

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

@@ -18,15 +18,19 @@
 /// correspond to individual ``FunctionCallPart``s.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerToolCall: Sendable {
+public struct LiveServerToolCall: LiveServerMessage {
   let serverToolCall: BidiGenerateContentToolCall
 
   /// A list of ``FunctionCallPart`` to run and return responses for.
   public var functionCalls: [FunctionCallPart]? {
     serverToolCall.functionCalls?.map { FunctionCallPart($0) }
   }
 
-  init(_ serverToolCall: BidiGenerateContentToolCall) {
+  public var usageMetadata: GenerateContentResponse.UsageMetadata?
+
+  init(_ serverToolCall: BidiGenerateContentToolCall,
+       usageMetadata: GenerateContentResponse.UsageMetadata?) {
     self.serverToolCall = serverToolCall
+    self.usageMetadata = usageMetadata
   }
 }

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

@@ -18,13 +18,17 @@
 /// ``FunctionCallPart``s.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerToolCallCancellation: Sendable {
+public struct LiveServerToolCallCancellation: LiveServerMessage {
   let serverToolCallCancellation: BidiGenerateContentToolCallCancellation
   /// A list of `functionId`s matching the `functionId` provided in a previous
   /// ``LiveServerToolCall``, where only the provided `functionId`s should be cancelled.
   public var ids: [String]? { serverToolCallCancellation.ids }
 
-  init(_ serverToolCallCancellation: BidiGenerateContentToolCallCancellation) {
+  public var usageMetadata: GenerateContentResponse.UsageMetadata?
+
+  init(_ serverToolCallCancellation: BidiGenerateContentToolCallCancellation,
+       usageMetadata: GenerateContentResponse.UsageMetadata?) {
     self.serverToolCallCancellation = serverToolCallCancellation
+    self.usageMetadata = usageMetadata
   }
 }