update exceptions

Author: VinayGuthal

firebase-vertexai/api.txt

@@ -325,10 +325,6 @@ package com.google.firebase.vertexai.type {
     property public final kotlinx.serialization.json.JsonObject response;
   }
 
-  public final class GeminiConnectionHandshakeFailed extends com.google.firebase.vertexai.type.FirebaseVertexAIException {
-    ctor public GeminiConnectionHandshakeFailed();
-  }
-
   public final class GenerateContentResponse {
     ctor public GenerateContentResponse(java.util.List<com.google.firebase.vertexai.type.Candidate> candidates, com.google.firebase.vertexai.type.PromptFeedback? promptFeedback, com.google.firebase.vertexai.type.UsageMetadata? usageMetadata);
     method public java.util.List<com.google.firebase.vertexai.type.Candidate> getCandidates();
@@ -816,6 +812,10 @@ package com.google.firebase.vertexai.type {
   public final class ServerException extends com.google.firebase.vertexai.type.FirebaseVertexAIException {
   }
 
+  public final class ServiceConnectionHandshakeFailedException extends com.google.firebase.vertexai.type.FirebaseVertexAIException {
+    ctor public ServiceConnectionHandshakeFailedException(String message, Throwable? cause = null);
+  }
+
   public final class ServiceDisabledException extends com.google.firebase.vertexai.type.FirebaseVertexAIException {
   }
 

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

@@ -33,6 +33,7 @@ import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
 import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 
@@ -82,8 +83,8 @@ internal constructor(
   /**
    * Returns a LiveSession object using which you could send/receive messages from the server
    * @return LiveSession object created
-   * @throws [BidiServerHandshakeFailed] if the handshake with the server failed.
-   * @throws [ClosedReceiveChannelException] if the channel was closed by the server.
+   * @throws [ServiceConnectionHandshakeFailedException] if the client was not able to establish a
+   * connection with the server.
    */
   public suspend fun connect(): LiveSession {
     val clientMessage =
@@ -95,15 +96,19 @@ internal constructor(
         )
         .toInternal()
     val data: String = Json.encodeToString(clientMessage)
-    val webSession = controller.getWebSocketSession(location)
-    webSession.send(Frame.Text(data))
-    val receivedJson = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
-    // TODO: Try to decode the json instead of string matching.
-    return if (receivedJson.contains("setupComplete")) {
-      LiveSession(session = webSession, backgroundDispatcher = backgroundDispatcher)
-    } else {
-      webSession.close()
-      throw ServiceConnectionHandshakeFailedException()
+    try {
+      val webSession = controller.getWebSocketSession(location)
+      webSession.send(Frame.Text(data))
+      val receivedJson = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
+      // TODO: Try to decode the json instead of string matching.
+      return if (receivedJson.contains("setupComplete")) {
+        LiveSession(session = webSession, backgroundDispatcher = backgroundDispatcher)
+      } else {
+        webSession.close()
+        throw ServiceConnectionHandshakeFailedException("Unable to connect to the server")
+      }
+    } catch (e: ClosedReceiveChannelException) {
+      throw ServiceConnectionHandshakeFailedException("Channel was closed by the server", e)
     }
   }
 

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/java/LiveModelFutures.kt

@@ -21,7 +21,7 @@ import com.google.common.util.concurrent.ListenableFuture
 import com.google.firebase.vertexai.GenerativeModel
 import com.google.firebase.vertexai.LiveGenerativeModel
 import com.google.firebase.vertexai.type.LiveSession
-import java.nio.channels.ClosedChannelException
+import com.google.firebase.vertexai.type.ServiceConnectionHandshakeFailedException
 
 /**
  * Wrapper class providing Java compatible methods for [GenerativeModel].
@@ -33,7 +33,8 @@ public abstract class LiveModelFutures 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 [ServiceConnectionHandshakeFailedException] if the client was not able to establish a
+   * connection with the server.
    */
   public abstract fun connect(): ListenableFuture<LiveSession>
 

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

@@ -186,8 +186,8 @@ public class AudioRecordInitializationFailedException(message: String) :
   FirebaseVertexAIException(message)
 
 /** Handshake failed with the server */
-public class ServiceConnectionHandshakeFailedException :
-  FirebaseVertexAIException("Handshake failed with the server.")
+public class ServiceConnectionHandshakeFailedException(message: String, cause: Throwable? = null) :
+  FirebaseVertexAIException(message, cause)
 
 /** Catch all case for exceptions not explicitly expected. */
 public class UnknownException internal constructor(message: String, cause: Throwable? = null) :