Refactor live bidi (#6870)

Per [b/410070347](https://b.corp.google.com/issues/410070347),

This refactors our bidi model to be more thread safe, and take better
advantage of immutability and kotlin coroutines. This solves various
edge case issues, as well as makes it easier to maintain.

This PR also fixes the following:
- [b/410063693](https://b.corp.google.com/issues/410063693) -> Add
serialization for bidi responses
- [b/410064609](https://b.corp.google.com/issues/410064609) -> Retain
audio data if present when turn complete or interrupted
- [b/410069806](https://b.corp.google.com/issues/410069806) -> Use
blocking instead of background dispatcher for bidi
- [b/410841715](https://b.corp.google.com/issues/410841715) -> Catch
Bidi exceptions
- [b/410067576](https://b.corp.google.com/issues/410067576) -> Send
zeroed out audio data to server when not speaking

Author: daymxn

firebase-vertexai/CHANGELOG.md

@@ -9,6 +9,11 @@
 * [feature] Added support for `HarmBlockThreshold.OFF`. See the
   [model documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters#how_to_configure_content_filters){: .external}
   for more information.
+* [fixed] Improved thread usage when using a `LiveGenerativeModel`. (#6870)
+* [fixed] Fixed an issue with `LiveContentResponse` audio data not being present when the model was
+  interrupted or the turn completed. (#6870)
+* [fixed] Fixed an issue with `LiveSession` not converting exceptions to `FirebaseVertexAIException`. (#6870)
+
 
 # 16.3.0
 * [feature] Emits a warning when attempting to use an incompatible model with

firebase-vertexai/api.txt

@@ -629,7 +629,7 @@ package com.google.firebase.vertexai.type {
     method public suspend Object? send(String text, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public suspend Object? sendFunctionResponse(java.util.List<com.google.firebase.vertexai.type.FunctionResponsePart> functionList, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public suspend Object? sendMediaStream(java.util.List<com.google.firebase.vertexai.type.MediaData> mediaChunks, kotlin.coroutines.Continuation<? super kotlin.Unit>);
-    method public suspend Object? startAudioConversation(kotlin.jvm.functions.Function1<? super com.google.firebase.vertexai.type.FunctionCallPart,com.google.firebase.vertexai.type.FunctionResponsePart>? functionCallHandler = null, kotlin.coroutines.Continuation<? super kotlin.Unit>);
+    method @RequiresPermission(android.Manifest.permission.RECORD_AUDIO) public suspend Object? startAudioConversation(kotlin.jvm.functions.Function1<? super com.google.firebase.vertexai.type.FunctionCallPart,com.google.firebase.vertexai.type.FunctionResponsePart>? functionCallHandler = null, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public void stopAudioConversation();
     method public void stopReceiving();
   }

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

@@ -19,7 +19,7 @@ package com.google.firebase.vertexai
 import android.util.Log
 import com.google.firebase.Firebase
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.app
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
@@ -41,7 +41,7 @@ import kotlin.coroutines.CoroutineContext
 public class FirebaseVertexAI
 internal constructor(
   private val firebaseApp: FirebaseApp,
-  @Background private val backgroundDispatcher: CoroutineContext,
+  @Blocking private val blockingDispatcher: CoroutineContext,
   private val location: String,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
@@ -133,7 +133,7 @@ internal constructor(
       "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}",
       firebaseApp.options.apiKey,
       firebaseApp,
-      backgroundDispatcher,
+      blockingDispatcher,
       generationConfig,
       tools,
       systemInstruction,

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

@@ -18,7 +18,7 @@ package com.google.firebase.vertexai
 
 import androidx.annotation.GuardedBy
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
@@ -31,7 +31,7 @@ import kotlin.coroutines.CoroutineContext
  */
 internal class FirebaseVertexAIMultiResourceComponent(
   private val app: FirebaseApp,
-  @Background val backgroundDispatcher: CoroutineContext,
+  @Blocking val blockingDispatcher: CoroutineContext,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
 ) {
@@ -43,7 +43,7 @@ internal class FirebaseVertexAIMultiResourceComponent(
       instances[location]
         ?: FirebaseVertexAI(
             app,
-            backgroundDispatcher,
+            blockingDispatcher,
             location,
             appCheckProvider,
             internalAuthProvider

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

@@ -18,7 +18,7 @@ package com.google.firebase.vertexai
 
 import androidx.annotation.Keep
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.components.Component
@@ -41,13 +41,13 @@ internal class FirebaseVertexAIRegistrar : ComponentRegistrar {
       Component.builder(FirebaseVertexAIMultiResourceComponent::class.java)
         .name(LIBRARY_NAME)
         .add(Dependency.required(firebaseApp))
-        .add(Dependency.required(backgroundDispatcher))
+        .add(Dependency.required(blockingDispatcher))
         .add(Dependency.optionalProvider(appCheckInterop))
         .add(Dependency.optionalProvider(internalAuthProvider))
         .factory { container ->
           FirebaseVertexAIMultiResourceComponent(
             container[firebaseApp],
-            container.get(backgroundDispatcher),
+            container.get(blockingDispatcher),
             container.getProvider(appCheckInterop),
             container.getProvider(internalAuthProvider)
           )
@@ -62,7 +62,7 @@ internal class FirebaseVertexAIRegistrar : ComponentRegistrar {
     private val firebaseApp = unqualified(FirebaseApp::class.java)
     private val appCheckInterop = unqualified(InteropAppCheckTokenProvider::class.java)
     private val internalAuthProvider = unqualified(InternalAuthProvider::class.java)
-    private val backgroundDispatcher =
-      Qualified.qualified(Background::class.java, CoroutineDispatcher::class.java)
+    private val blockingDispatcher =
+      Qualified.qualified(Blocking::class.java, CoroutineDispatcher::class.java)
   }
 }

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

@@ -17,13 +17,14 @@
 package com.google.firebase.vertexai
 
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 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.common.JSON
 import com.google.firebase.vertexai.type.Content
+import com.google.firebase.vertexai.type.LiveClientSetupMessage
 import com.google.firebase.vertexai.type.LiveGenerationConfig
 import com.google.firebase.vertexai.type.LiveSession
 import com.google.firebase.vertexai.type.PublicPreviewAPI
@@ -38,6 +39,7 @@ import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
 
 /**
  * Represents a multimodal model (like Gemini) capable of real-time content generation based on
@@ -47,7 +49,7 @@ import kotlinx.serialization.json.Json
 public class LiveGenerativeModel
 internal constructor(
   private val modelName: String,
-  @Background private val backgroundDispatcher: CoroutineContext,
+  @Blocking private val blockingDispatcher: CoroutineContext,
   private val config: LiveGenerationConfig? = null,
   private val tools: List<Tool>? = null,
   private val systemInstruction: Content? = null,
@@ -58,7 +60,7 @@ internal constructor(
     modelName: String,
     apiKey: String,
     firebaseApp: FirebaseApp,
-    backgroundDispatcher: CoroutineContext,
+    blockingDispatcher: CoroutineContext,
     config: LiveGenerationConfig? = null,
     tools: List<Tool>? = null,
     systemInstruction: Content? = null,
@@ -68,7 +70,7 @@ internal constructor(
     internalAuthProvider: InternalAuthProvider? = null,
   ) : this(
     modelName,
-    backgroundDispatcher,
+    blockingDispatcher,
     config,
     tools,
     systemInstruction,
@@ -93,7 +95,7 @@ internal constructor(
   @OptIn(ExperimentalSerializationApi::class)
   public suspend fun connect(): LiveSession {
     val clientMessage =
-      BidiGenerateContentClientMessage(
+      LiveClientSetupMessage(
           modelName,
           config?.toInternal(),
           tools?.map { it.toInternal() },
@@ -104,10 +106,11 @@ internal constructor(
     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)
+      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)
       } else {
         webSession.close()
         throw ServiceConnectionHandshakeFailedException("Unable to connect to the server")

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

@@ -165,6 +165,7 @@ internal constructor(
 
   suspend fun getWebSocketSession(location: String): ClientWebSocketSession =
     client.webSocketSession(getBidiEndpoint(location))
+
   fun generateContentStream(
     request: GenerateContentRequest
   ): Flow<GenerateContentResponse.Internal> =

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/util/android.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.vertexai.common.util
+
+import android.media.AudioRecord
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.yield
+
+/**
+ * The minimum buffer size for this instance.
+ *
+ * The same as calling [AudioRecord.getMinBufferSize], except the params are pre-populated.
+ */
+internal val AudioRecord.minBufferSize: Int
+  get() = AudioRecord.getMinBufferSize(sampleRate, channelConfiguration, audioFormat)
+
+/**
+ * Reads from this [AudioRecord] and returns the data in a flow.
+ *
+ * Will yield when this instance is not recording.
+ */
+internal fun AudioRecord.readAsFlow() = flow {
+  val buffer = ByteArray(minBufferSize)
+
+  while (true) {
+    if (recordingState != AudioRecord.RECORDSTATE_RECORDING) {
+      yield()
+      continue
+    }
+
+    val bytesRead = read(buffer, 0, buffer.size)
+    if (bytesRead > 0) {
+      emit(buffer.copyOf(bytesRead))
+    }
+  }
+}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/util/kotlin.kt

@@ -16,7 +16,16 @@
 
 package com.google.firebase.vertexai.common.util
 
+import java.io.ByteArrayOutputStream
 import java.lang.reflect.Field
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.fold
 
 /**
  * Removes the last character from the [StringBuilder].
@@ -39,3 +48,56 @@ internal fun StringBuilder.removeLast(): StringBuilder =
  * ```
  */
 internal inline fun <reified T : Annotation> Field.getAnnotation() = getAnnotation(T::class.java)
+
+/**
+ * Collects bytes from this flow and doesn't emit them back until [minSize] is reached.
+ *
+ * For example:
+ * ```
+ * val byteArr = flowOf(byteArrayOf(1), byteArrayOf(2, 3, 4), byteArrayOf(5, 6, 7, 8))
+ * val expectedResult = listOf(byteArrayOf(1, 2, 3, 4), byteArrayOf( 5, 6, 7, 8))
+ *
+ * byteArr.accumulateUntil(4).toList() shouldContainExactly expectedResult
+ * ```
+ *
+ * @param minSize The minimum about of bytes the array should have before being sent down-stream
+ * @param emitLeftOvers If the flow completes and there are bytes left over that don't meet the
+ * [minSize], send them anyways.
+ */
+internal fun Flow<ByteArray>.accumulateUntil(
+  minSize: Int,
+  emitLeftOvers: Boolean = false
+): Flow<ByteArray> = flow {
+  val remaining =
+    fold(ByteArrayOutputStream()) { buffer, it ->
+      buffer.apply {
+        write(it, 0, it.size)
+        if (size() >= minSize) {
+          emit(toByteArray())
+          reset()
+        }
+      }
+    }
+
+  if (emitLeftOvers && remaining.size() > 0) {
+    emit(remaining.toByteArray())
+  }
+}
+
+/**
+ * Create a [Job] that is a child of the [currentCoroutineContext], if any.
+ *
+ * This is useful when you want a coroutine scope to be canceled when its parent scope is canceled,
+ * and you don't have full control over the parent scope, but you don't want the cancellation of the
+ * child to impact the parent.
+ *
+ * If the parent coroutine context does not have a job, an empty one will be created.
+ */
+internal suspend inline fun childJob() = Job(currentCoroutineContext()[Job] ?: Job())
+
+/**
+ * A constant value pointing to a cancelled [CoroutineScope].
+ *
+ * Useful when you want to initialize a mutable [CoroutineScope] in a canceled state.
+ */
+internal val CancelledCoroutineScope = CoroutineScope(EmptyCoroutineContext).apply { cancel() }