update

Author: VinayGuthal

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/LiveGenerativeModel.kt

@@ -22,6 +22,7 @@ import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.vertexai.common.APIController
 import com.google.firebase.vertexai.common.AppCheckHeaderProvider
 import com.google.firebase.vertexai.type.BidiGenerateContentClientMessage
+import com.google.firebase.vertexai.type.BidiServerHandshakeFailed
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.LiveGenerationConfig
 import com.google.firebase.vertexai.type.LiveSession
@@ -33,8 +34,8 @@ import io.ktor.client.engine.okhttp.OkHttp
 import io.ktor.client.plugins.websocket.WebSockets
 import io.ktor.client.plugins.websocket.webSocketSession
 import io.ktor.websocket.Frame
+import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
-import java.nio.channels.ClosedChannelException
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 
@@ -84,9 +85,9 @@ internal constructor(
   /**
    * Returns a LiveSession object using which you could send/receive messages from the server
    * @return LiveSession object created. Returns null if the object cannot be created.
-   * @throws [ClosedChannelException] if channel was closed before creating a websocket connection.
+   * @throws [BidiServerHandshakeFailed] if the handshake with the server failed.
    */
-  public suspend fun connect(): LiveSession? {
+  public suspend fun connect(): LiveSession {
     val client = HttpClient(OkHttp) { install(WebSockets) }
 
     val bidiEndPoint = this.controller.getBidiEndpoint(location)
@@ -115,7 +116,8 @@ internal constructor(
     return if (shouldReturn) {
       LiveSession(session = webSession, isRecording = false)
     } else {
-      null
+      webSession.close()
+      throw BidiServerHandshakeFailed()
     }
   }
 

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/BidiGenerateContentSetup.kt renamed to firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/BidiGenerateContentClientMessage.kt

@@ -1,7 +1,9 @@
 package com.google.firebase.vertexai.type
 
+import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 
+@OptIn(ExperimentalSerializationApi::class)
 internal class BidiGenerateContentClientMessage(
   val model: String,
   val generationConfig: LiveGenerationConfig.Internal?,

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Exceptions.kt

@@ -189,6 +189,10 @@ public class AudioRecordInitializationFailedException :
 public class AudioRecordInvalidBufferSizeException :
   FirebaseVertexAIException("Audio Record buffer size is invalid")
 
+/** Handshake failed with the server */
+public class BidiServerHandshakeFailed :
+  FirebaseVertexAIException("Handshake failed with the server.")
+
 /** Catch all case for exceptions not explicitly expected. */
 public class UnknownException internal constructor(message: String, cause: Throwable? = null) :
   FirebaseVertexAIException(message, cause)