update live generative model

Author: VinayGuthal

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

@@ -32,6 +32,7 @@ import com.google.firebase.ai.type.Tool
 import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
+import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
@@ -114,31 +115,28 @@ internal constructor(
         )
         .toInternal()
     val data: String = Json.encodeToString(clientMessage)
+    var webSession: DefaultClientWebSocketSession? = null
     try {
-      val webSession = controller.getWebSocketSession(location)
-      try {
-        webSession.send(Frame.Text(data))
-        val receivedJsonStr = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
-        val receivedJson = JSON.parseToJsonElement(receivedJsonStr)
+      webSession = controller.getWebSocketSession(location)
+      webSession.send(Frame.Text(data))
+      val receivedJsonStr = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
+      val receivedJson = JSON.parseToJsonElement(receivedJsonStr)
 
-        return if (receivedJson is JsonObject && "setupComplete" in receivedJson) {
-          LiveSession(
-            session = webSession,
-            blockingDispatcher = blockingDispatcher,
-            firebaseApp = firebaseApp
-          )
-        } else {
-          webSession.close()
-          throw ServiceConnectionHandshakeFailedException("Unable to connect to the server")
-        }
-      } catch (e: ClosedReceiveChannelException) {
-        val reason = webSession.closeReason.await()
-        val message =
-          "Channel was closed by the server.${if(reason!=null) " Details: ${reason.message}" else "" }"
-        throw ServiceConnectionHandshakeFailedException(message, e)
+      return if (receivedJson is JsonObject && "setupComplete" in receivedJson) {
+        LiveSession(
+          session = webSession,
+          blockingDispatcher = blockingDispatcher,
+          firebaseApp = firebaseApp
+        )
+      } else {
+        webSession.close()
+        throw ServiceConnectionHandshakeFailedException("Unable to connect to the server")
       }
     } catch (e: ClosedReceiveChannelException) {
-      throw ServiceConnectionHandshakeFailedException("Channel was closed by the server.", e)
+      val reason = webSession?.closeReason?.await()
+      val message =
+        "Channel was closed by the server.${if(reason!=null) " Details: ${reason.message}" else "" }"
+      throw ServiceConnectionHandshakeFailedException(message, e)
     }
   }