Log json payload if we don't have a mapping

Author: daymxn

FirebaseAI/Sources/AILog.swift

@@ -68,15 +68,16 @@ enum AILog {
     case fallbackValueUsed = 3017
     case urlMetadataUnrecognizedURLRetrievalStatus = 3018
     case liveSessionUnsupportedMessage = 3019
-    case liveSessionFailedToEncodeClientMessage = 3020
-    case liveSessionFailedToEncodeClientMessagePayload = 3021
-    case liveSessionFailedToSendClientMessage = 3022
-    case liveSessionUnexpectedResponse = 3023
-    case liveSessionGoingAwaySoon = 3024
-    case decodedMissingProtoDurationSuffix = 3025
-    case decodedInvalidProtoDurationString = 3026
-    case decodedInvalidProtoDurationSeconds = 3027
-    case decodedInvalidProtoDurationNanoseconds = 3028
+    case liveSessionUnsupportedMessagePayload = 3020
+    case liveSessionFailedToEncodeClientMessage = 3021
+    case liveSessionFailedToEncodeClientMessagePayload = 3022
+    case liveSessionFailedToSendClientMessage = 3023
+    case liveSessionUnexpectedResponse = 3024
+    case liveSessionGoingAwaySoon = 3025
+    case decodedMissingProtoDurationSuffix = 3026
+    case decodedInvalidProtoDurationString = 3027
+    case decodedInvalidProtoDurationSeconds = 3028
+    case decodedInvalidProtoDurationNanoseconds = 3029
 
     // SDK State Errors
     case generateContentResponseNoCandidates = 4000

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

@@ -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)
+      throw InvalidMessageTypeError()
     }
 
     usageMetadata = try container.decodeIfPresent(
@@ -99,3 +95,11 @@ extension BidiGenerateContentServerMessage: Decodable {
     )
   }
 }
+
+struct InvalidMessageTypeError: Error, Sendable, CustomNSError {
+  public var errorUserInfo: [String: Any] {
+    [
+      NSLocalizedDescriptionKey: "Missing server message type.",
+    ]
+  }
+}

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

@@ -206,7 +206,26 @@ actor LiveSessionService {
               from: message
             )
           } catch {
-            throw LiveSessionUnsupportedMessageError(underlyingError: error)
+            // only log the json if it wasn't a decoding error, but an unsupported message type
+            if error is InvalidMessageTypeError {
+              AILog.error(
+                code: .liveSessionUnsupportedMessage,
+                "The server sent a message that we don't currently have a mapping for."
+              )
+
+              AILog.debug(
+                code: .liveSessionUnsupportedMessagePayload,
+                message.encodeToJsonString() ?? "\(message)"
+              )
+            }
+
+            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 {
@@ -231,13 +250,6 @@ actor LiveSessionService {
             }
 
             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(
-              code: .liveSessionUnsupportedMessage,
-              "The server sent a message that we don't currently have a mapping for: \(response)"
-            )
           }
         }
       } catch {
@@ -348,3 +360,19 @@ actor LiveSessionService {
     return AsyncWebSocket(urlSession: urlSession, urlRequest: urlRequest)
   }
 }
+
+private extension Data {
+  /// Encodes this into a raw json string, with no regard to specific keys.
+  ///
+  /// Will return `nil` if this data doesn't represent a valid json object.
+  func encodeToJsonString() -> String? {
+    do {
+      let object = try JSONSerialization.jsonObject(with: self)
+      let data = try JSONSerialization.data(withJSONObject: object)
+
+      return String(data: data, encoding: .utf8)
+    } catch {
+      return nil
+    }
+  }
+}