diff --git a/ai-logic/firebase-ai/CHANGELOG.md b/ai-logic/firebase-ai/CHANGELOG.md index f1f6cfd0267..55c0bad49a5 100644 --- a/ai-logic/firebase-ai/CHANGELOG.md +++ b/ai-logic/firebase-ai/CHANGELOG.md @@ -1,5 +1,6 @@ # Unreleased +- [changed] Adjusts `LiveSession.isClosed` to better reflect underlying connection closure state and not consume frames (#8511) - [changed] Replaced the `"function"` conversational role with `"user"` for function response content. (#8508) # 17.15.0 diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/LiveSession.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/LiveSession.kt index 9a2bb34442c..3e19a04fbb8 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/LiveSession.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/LiveSession.kt @@ -294,7 +294,10 @@ internal constructor( } /** Indicates whether the underlying websocket connection is active. */ - public fun isClosed(): Boolean = !(session.isActive && !session.incoming.tryReceive().isClosed) + public fun isClosed(): Boolean { + val currentSession = session + return !currentSession.isActive || currentSession.closeReason.isCompleted + } /** Indicates whether an audio conversation is being used for this session object. */ public fun isAudioConversationActive(): Boolean = (audioHelper != null) diff --git a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/kotlin.kt b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/kotlin.kt deleted file mode 100644 index 5187607cc3b..00000000000 --- a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/kotlin.kt +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright 2024 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.ai.common.util - -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.runBlocking - -/** - * Runs the given [block] using [runBlocking] on the current thread for side effect. - * - * Using this function is like [runBlocking] with default context (which runs the given block on the - * calling thread) but forces the return type to be `Unit`, which is helpful when implementing - * suspending tests as expression functions: - * ``` - * @Test - * fun myTest() = doBlocking {...} - * ``` - */ -internal fun doBlocking(block: suspend CoroutineScope.() -> Unit) { - runBlocking(block = block) -} diff --git a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/tests.kt b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/tests.kt index b4ae47926b5..8a23464c693 100644 --- a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/tests.kt +++ b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/common/util/tests.kt @@ -35,6 +35,8 @@ import io.ktor.http.HttpHeaders import io.ktor.http.HttpStatusCode import io.ktor.http.headersOf import io.ktor.utils.io.ByteChannel +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.runBlocking import kotlinx.serialization.encodeToString import org.mockito.Mockito @@ -119,3 +121,18 @@ internal fun commonTest( ) CommonTestScope(channel, apiController).block() } + +/** + * Runs the given [block] using [runBlocking] on the current thread for side effect. + * + * Using this function is like [runBlocking] with default context (which runs the given block on the + * calling thread) but forces the return type to be `Unit`, which is helpful when implementing + * suspending tests as expression functions: + * ``` + * @Test + * fun myTest() = doBlocking {...} + * ``` + */ +internal fun doBlocking(block: suspend CoroutineScope.() -> Unit) { + runBlocking(block = block) +} diff --git a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/type/LiveSessionTest.kt b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/type/LiveSessionTest.kt new file mode 100644 index 00000000000..1ee843529ee --- /dev/null +++ b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/type/LiveSessionTest.kt @@ -0,0 +1,149 @@ +/* + * Copyright 2026 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.ai.type + +import com.google.firebase.FirebaseApp +import io.kotest.matchers.shouldBe +import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession +import io.ktor.websocket.CloseReason +import io.ktor.websocket.Frame +import io.mockk.coEvery +import io.mockk.every +import io.mockk.just +import io.mockk.mockk +import io.mockk.runs +import kotlin.coroutines.EmptyCoroutineContext +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner + +@RunWith(RobolectricTestRunner::class) +@OptIn(PublicPreviewAPI::class) +class LiveSessionTest { + + @Test(timeout = 10000) + fun testIsClosed_initiallyFalse() { + val mockSession = mockk() + val mockFirebaseApp = mockk() + val incomingChannel = Channel(Channel.UNLIMITED) + val closeReasonDeferred = CompletableDeferred() + + val job = Job() // Active job + every { mockSession.coroutineContext } returns EmptyCoroutineContext + job + every { mockSession.incoming } returns incomingChannel + every { mockSession.closeReason } returns closeReasonDeferred + + val liveSession = + LiveSession( + session = mockSession, + blockingDispatcher = Dispatchers.Unconfined, + firebaseApp = mockFirebaseApp + ) + + liveSession.isClosed() shouldBe false + } + + @Test(timeout = 10000) + fun testIsClosed_afterClose_returnsTrue() { + runBlocking { + val mockSession = mockk() + val mockFirebaseApp = mockk() + val incomingChannel = Channel(Channel.UNLIMITED) + val outgoingChannel = Channel(Channel.UNLIMITED) + val closeReasonDeferred = CompletableDeferred() + + val job = Job() // Active job + every { mockSession.coroutineContext } returns EmptyCoroutineContext + job + every { mockSession.incoming } returns incomingChannel + every { mockSession.outgoing } returns outgoingChannel + every { mockSession.closeReason } returns closeReasonDeferred + coEvery { mockSession.flush() } just runs + + // Mock send member function to delegate to outgoingChannel + coEvery { mockSession.send(any()) } coAnswers + { + val frame = firstArg() + outgoingChannel.send(frame) + } + + // Simulate Ktor behavior: sending close frame completes closeReason and cancels job + val monitorJob = launch { + for (frame in outgoingChannel) { + if (frame is Frame.Close) { + closeReasonDeferred.complete(CloseReason(CloseReason.Codes.NORMAL, "")) + job.cancel() + break + } + } + } + + val liveSession = + LiveSession( + session = mockSession, + blockingDispatcher = Dispatchers.Unconfined, + firebaseApp = mockFirebaseApp + ) + + liveSession.close() + monitorJob.join() + + liveSession.isClosed() shouldBe true + } + } + + @Test(timeout = 10000) + fun testIsClosed_serverClosedWithUnconsumedFrames_returnsTrue() { + runBlocking { + val mockSession = mockk() + val mockFirebaseApp = mockk() + val incomingChannel = Channel(Channel.UNLIMITED) + val closeReasonDeferred = CompletableDeferred() + + val job = Job() // Active job + every { mockSession.coroutineContext } returns EmptyCoroutineContext + job + every { mockSession.incoming } returns incomingChannel + every { mockSession.closeReason } returns closeReasonDeferred + + val liveSession = + LiveSession( + session = mockSession, + blockingDispatcher = Dispatchers.Unconfined, + firebaseApp = mockFirebaseApp + ) + + // Add some unconsumed frames to incoming channel + incomingChannel.send(Frame.Text("hello")) + + // Simulate server close: complete closeReason, cancel job, and close channel + closeReasonDeferred.complete(CloseReason(CloseReason.Codes.NORMAL, "")) + job.cancel() + incomingChannel.close() + + // The channel still has "hello" frame unconsumed, so it is not fully closed for receive yet + incomingChannel.isClosedForReceive shouldBe false + + // But the session should be considered closed because closeReason is completed + liveSession.isClosed() shouldBe true + } + } +}