Merge a9c16ec into 525f293

Author: VinayGuthal

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

@@ -164,7 +164,7 @@ internal constructor(
     "wss://firebasevertexai.googleapis.com/ws/google.firebase.vertexai.v1beta.LlmBidiService/BidiGenerateContent/locations/$location?key=$key"
 
   suspend fun getWebSocketSession(location: String): ClientWebSocketSession =
-    client.webSocketSession(getBidiEndpoint(location))
+    client.webSocketSession(getBidiEndpoint(location)) { applyCommonHeaders() }
 
   fun generateContentStream(
     request: GenerateContentRequest
@@ -191,12 +191,7 @@ internal constructor(
       throw FirebaseCommonAIException.from(e)
     }
 
-  private fun HttpRequestBuilder.applyCommonConfiguration(request: Request) {
-    when (request) {
-      is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
-      is CountTokensRequest -> setBody<CountTokensRequest>(request)
-      is GenerateImageRequest -> setBody<GenerateImageRequest>(request)
-    }
+  private fun HttpRequestBuilder.applyCommonHeaders() {
     contentType(ContentType.Application.Json)
     header("x-goog-api-key", key)
     header("x-goog-api-client", apiClient)
@@ -205,6 +200,14 @@ internal constructor(
       header("X-Firebase-AppVersion", appVersion)
     }
   }
+  private fun HttpRequestBuilder.applyCommonConfiguration(request: Request) {
+    when (request) {
+      is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
+      is CountTokensRequest -> setBody<CountTokensRequest>(request)
+      is GenerateImageRequest -> setBody<GenerateImageRequest>(request)
+    }
+    applyCommonHeaders()
+  }
 
   private suspend fun HttpRequestBuilder.applyHeaderProvider() {
     if (headerProvider != null) {