add close reason

Author: VinayGuthal

firebase-ai/CHANGELOG.md

@@ -1,6 +1,7 @@
 # Unreleased
 
 - [changed] **Breaking Change**: Removed the `candidateCount` option from `LiveGenerationConfig`
+- [changed] Added better error messages to `ServiceConnectionHandshakeFailedException`
 
 # 17.3.0
 

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

@@ -116,22 +116,29 @@ internal constructor(
     val data: String = Json.encodeToString(clientMessage)
     try {
       val webSession = controller.getWebSocketSession(location)
-      webSession.send(Frame.Text(data))
-      val receivedJsonStr = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
-      val receivedJson = JSON.parseToJsonElement(receivedJsonStr)
+      try {
+        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")
+        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)
       }
     } catch (e: ClosedReceiveChannelException) {
-      throw ServiceConnectionHandshakeFailedException("Channel was closed by the server", e)
+      throw ServiceConnectionHandshakeFailedException("Channel was closed by the server.", e)
     }
   }
 

firebase-ai/src/main/kotlin/com/google/firebase/ai/common/APIController.kt

@@ -38,7 +38,7 @@ import io.ktor.client.engine.HttpClientEngine
 import io.ktor.client.engine.okhttp.OkHttp
 import io.ktor.client.plugins.HttpTimeout
 import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
-import io.ktor.client.plugins.websocket.ClientWebSocketSession
+import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
 import io.ktor.client.plugins.websocket.WebSockets
 import io.ktor.client.plugins.websocket.webSocketSession
 import io.ktor.client.request.HttpRequestBuilder
@@ -174,7 +174,7 @@ internal constructor(
         "wss://firebasevertexai.googleapis.com/ws/google.firebase.vertexai.v1beta.GenerativeService/BidiGenerateContent?key=$key"
     }
 
-  suspend fun getWebSocketSession(location: String): ClientWebSocketSession =
+  suspend fun getWebSocketSession(location: String): DefaultClientWebSocketSession =
     client.webSocketSession(getBidiEndpoint(location)) { applyCommonHeaders() }
 
   fun generateContentStream(

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/LiveSession.kt

@@ -29,7 +29,7 @@ import com.google.firebase.ai.common.util.CancelledCoroutineScope
 import com.google.firebase.ai.common.util.accumulateUntil
 import com.google.firebase.ai.common.util.childJob
 import com.google.firebase.annotations.concurrent.Blocking
-import io.ktor.client.plugins.websocket.ClientWebSocketSession
+import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
@@ -59,7 +59,7 @@ import kotlinx.serialization.json.Json
 @OptIn(ExperimentalSerializationApi::class)
 public class LiveSession
 internal constructor(
-  private val session: ClientWebSocketSession,
+  private val session: DefaultClientWebSocketSession,
   @Blocking private val blockingDispatcher: CoroutineContext,
   private var audioHelper: AudioHelper? = null,
   private val firebaseApp: FirebaseApp,