Revert to 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,7 +86,11 @@ extension BidiGenerateContentServerMessage: Decodable {
     } else if let goAway = try container.decodeIfPresent(GoAway.self, forKey: .goAway) {
       messageType = .goAway(goAway)
     } else {
-      messageType = nil
+      let context = DecodingError.Context(
+        codingPath: decoder.codingPath,
+        debugDescription: "Could not decode server message."
+      )
+      throw DecodingError.dataCorrupted(context)
     }
 
     usageMetadata = try container.decodeIfPresent(

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

@@ -206,18 +206,10 @@ actor LiveSessionService {
               from: message
             )
           } catch {
-            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
+            throw LiveSessionUnsupportedMessageError(underlyingError: error)
           }
 
-          switch response.messageType {
-          case .setupComplete:
+          if case .setupComplete = response.messageType {
             if resumed {
               AILog.debug(
                 code: .duplicateLiveSessionSetupComplete,
@@ -229,33 +221,17 @@ actor LiveSessionService {
               resumed = true
               setupComplete.resume()
             }
-          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)"
-            )
+          } 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(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:
+            responseContinuation.yield(liveMessage)
+          } else {
             // 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: LiveServerMessage {
+public struct LiveServerContent: Sendable {
   let serverContent: BidiGenerateContentServerContent
 
   /// The content that the model has generated as part of the current
@@ -76,11 +76,7 @@ public struct LiveServerContent: LiveServerMessage {
     serverContent.outputTranscription.map { LiveTranscription($0) }
   }
 
-  public var usageMetadata: GenerateContentResponse.UsageMetadata?
-
-  init(_ serverContent: BidiGenerateContentServerContent,
-       usageMetadata: GenerateContentResponse.UsageMetadata?) {
+  init(_ serverContent: BidiGenerateContentServerContent) {
     self.serverContent = serverContent
-    self.usageMetadata = usageMetadata
   }
 }

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

@@ -19,18 +19,15 @@ 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: LiveServerMessage {
+public struct LiveServerGoingAwayNotice: Sendable {
   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 }
 
-  public var usageMetadata: GenerateContentResponse.UsageMetadata?
-
-  init(_ goAway: GoAway, usageMetadata: GenerateContentResponse.UsageMetadata?) {
+  init(_ goAway: GoAway) {
     self.goAway = goAway
-    self.usageMetadata = usageMetadata
   }
 }

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

@@ -15,7 +15,63 @@
 /// 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 protocol LiveServerMessage: Sendable {
+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
+
   /// Metadata on the usage of the cached content.
-  var usageMetadata: GenerateContentResponse.UsageMetadata? { get }
+  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))
+    }
+  }
 }

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

@@ -18,19 +18,15 @@
 /// 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: LiveServerMessage {
+public struct LiveServerToolCall: Sendable {
   let serverToolCall: BidiGenerateContentToolCall
 
   /// A list of ``FunctionCallPart`` to run and return responses for.
   public var functionCalls: [FunctionCallPart]? {
     serverToolCall.functionCalls?.map { FunctionCallPart($0) }
   }
 
-  public var usageMetadata: GenerateContentResponse.UsageMetadata?
-
-  init(_ serverToolCall: BidiGenerateContentToolCall,
-       usageMetadata: GenerateContentResponse.UsageMetadata?) {
+  init(_ serverToolCall: BidiGenerateContentToolCall) {
     self.serverToolCall = serverToolCall
-    self.usageMetadata = usageMetadata
   }
 }

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

@@ -18,17 +18,13 @@
 /// ``FunctionCallPart``s.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
 @available(watchOS, unavailable)
-public struct LiveServerToolCallCancellation: LiveServerMessage {
+public struct LiveServerToolCallCancellation: Sendable {
   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 }
 
-  public var usageMetadata: GenerateContentResponse.UsageMetadata?
-
-  init(_ serverToolCallCancellation: BidiGenerateContentToolCallCancellation,
-       usageMetadata: GenerateContentResponse.UsageMetadata?) {
+  init(_ serverToolCallCancellation: BidiGenerateContentToolCallCancellation) {
     self.serverToolCallCancellation = serverToolCallCancellation
-    self.usageMetadata = usageMetadata
   }
 }