fix the chat loading behavior

Author: thatfiredev

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/feature/text/AudioSummarizationViewModel.kt

@@ -30,10 +30,8 @@ class AudioSummarizationViewModel : ChatViewModel() {
             }
         )
 
-        _uiState.value = ChatUiState.Success(
-            messages = chatHistory.map { UiChatMessage(it) },
-            attachments = emptyList()
-        )
+        _messages.value = chatHistory.map { UiChatMessage(it) }
+        _uiState.value = ChatUiState.Success
 
         val generativeModel = Firebase.ai.generativeModel(
             modelName = "gemini-2.5-flash"

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/feature/text/ChatViewModel.kt

@@ -17,14 +17,15 @@ import kotlinx.coroutines.launch
 @OptIn(PublicPreviewAPI::class)
 abstract class ChatViewModel : ViewModel() {
 
-    protected val _uiState: MutableStateFlow<ChatUiState> =
-        MutableStateFlow(
-            ChatUiState.Success(
-            messages = emptyList(),
-            attachments = emptyList()
-        ))
+    protected val _uiState = MutableStateFlow<ChatUiState>(ChatUiState.Success)
     val uiState: StateFlow<ChatUiState> = _uiState.asStateFlow()
 
+    protected val _messages = MutableStateFlow<List<UiChatMessage>>(emptyList())
+    val messages: StateFlow<List<UiChatMessage>> = _messages.asStateFlow()
+
+    protected val _attachments = MutableStateFlow<List<Attachment>>(emptyList())
+    val attachments: StateFlow<List<Attachment>> = _attachments.asStateFlow()
+
     abstract val initialPrompt: String
 
     // Builder for the next message
@@ -35,20 +36,16 @@ abstract class ChatViewModel : ViewModel() {
      * Handles adding the message to the UI and setting the loading state.
      */
     fun sendMessage(userMessage: String) {
-        val uiStateValue = _uiState.value
-        if (uiStateValue !is ChatUiState.Success) return
-
         val prompt = contentBuilder
             .text(userMessage)
             .build()
 
-        val updatedMessages = uiStateValue.messages + UiChatMessage(prompt)
-        _uiState.value = uiStateValue.copy(messages = updatedMessages)
+        _messages.value = _messages.value + UiChatMessage(prompt)
 
         viewModelScope.launch {
             _uiState.value = ChatUiState.Loading
             try {
-                performSendMessage(prompt, updatedMessages)
+                performSendMessage(prompt, _messages.value)
             } catch (e: Exception) {
                 _uiState.value = ChatUiState.Error(e.localizedMessage ?: "Unknown error")
             } finally {
@@ -83,10 +80,9 @@ abstract class ChatViewModel : ViewModel() {
                 "Could not display the response because it was missing required attribution components."
             )
         } else {
-            _uiState.value = ChatUiState.Success(
-                messages = currentMessages + UiChatMessage(candidate.content, candidate.groundingMetadata),
-                attachments = emptyList()
-            )
+            _messages.value = currentMessages + UiChatMessage(candidate.content, candidate.groundingMetadata)
+            _attachments.value = emptyList()
+            _uiState.value = ChatUiState.Success
         }
     }
 
@@ -95,19 +91,14 @@ abstract class ChatViewModel : ViewModel() {
         mimeType: String?,
         fileName: String? = "Unnamed file"
     ) {
-        val uiStateValue = _uiState.value
-        if (uiStateValue !is ChatUiState.Success) return
-
         if (mimeType?.contains("image") == true) {
             // images should be attached as ImageParts
             contentBuilder.image(decodeBitmapFromImage(fileInBytes))
         } else {
             contentBuilder.inlineData(fileInBytes, mimeType ?: "text/plain")
         }
-        
-        _uiState.value = uiStateValue.copy(
-            attachments = uiStateValue.attachments + Attachment(fileName ?: "Unnamed file")
-        )
+
+        _attachments.value = _attachments.value + Attachment(fileName ?: "Unnamed file")
     }
 
     protected fun decodeBitmapFromImage(input: ByteArray) =

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/feature/text/TravelTipsViewModel.kt

@@ -59,10 +59,8 @@ class TravelTipsViewModel : ChatViewModel() {
                 }
             ))
 
-        _uiState.value = ChatUiState.Success(
-            messages = chat.history.map { UiChatMessage(it) },
-            attachments = emptyList()
-        )
+        _messages.value = chat.history.map { UiChatMessage(it) }
+        _uiState.value = ChatUiState.Success
     }
 
     override suspend fun performSendMessage(prompt: Content, currentMessages: List<UiChatMessage>) {

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/feature/text/VideoSummarizationViewModel.kt

@@ -32,10 +32,8 @@ class VideoSummarizationViewModel : ChatViewModel() {
             }
         )
 
-        _uiState.value = ChatUiState.Success(
-            messages = chatHistory.map { UiChatMessage(it) },
-            attachments = emptyList()
-        )
+        _messages.value = chatHistory.map { UiChatMessage(it) }
+        _uiState.value = ChatUiState.Success
 
         val generativeModel = Firebase.ai.generativeModel(
             modelName = "gemini-2.5-flash"

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/ui/ChatScreen.kt

@@ -87,6 +87,8 @@ fun ChatScreen(
     chatViewModel: ChatViewModel
 ) {
     val uiState by chatViewModel.uiState.collectAsStateWithLifecycle()
+    val messages by chatViewModel.messages.collectAsStateWithLifecycle()
+    val attachments by chatViewModel.attachments.collectAsStateWithLifecycle()
 
     val initialPrompt: String = chatViewModel.initialPrompt
 
@@ -97,27 +99,13 @@ fun ChatScreen(
         modifier = Modifier
             .fillMaxSize()
     ) {
-        when (val state = uiState) {
-            is ChatUiState.Success -> {
-                ChatList(
-                    state.messages,
-                    listState,
-                    modifier = Modifier
-                        .fillMaxSize()
-                        .weight(0.5f)
-                )
-            }
-            is ChatUiState.Loading -> {
-//                Box(modifier = Modifier.weight(0.5f).fillMaxSize(), contentAlignment = Alignment.Center) {
-//                    LinearProgressIndicator()
-//                }
-            }
-            is ChatUiState.Error -> {
-                Box(modifier = Modifier.weight(0.5f).fillMaxSize(), contentAlignment = Alignment.Center) {
-                    Text(text = state.message, color = MaterialTheme.colorScheme.error)
-                }
-            }
-        }
+        ChatList(
+            messages,
+            listState,
+            modifier = Modifier
+                .fillMaxSize()
+                .weight(0.5f)
+        )
 
         Box(
             contentAlignment = Alignment.BottomCenter
@@ -148,9 +136,7 @@ fun ChatScreen(
                         )
                     }
                 }
-                (uiState as? ChatUiState.Success)?.let {
-                    AttachmentsList(it.attachments)
-                }
+                AttachmentsList(attachments)
                 val context = LocalContext.current
                 val contentResolver = context.contentResolver
                 MessageInput(

firebase-ai/app/src/main/java/com/google/firebase/quickstart/ai/ui/ChatUiState.kt

@@ -22,9 +22,6 @@ data class UiChatMessage(
 
 sealed interface ChatUiState {
     data object Loading : ChatUiState
-    data class Success(
-        val messages: List<UiChatMessage> = emptyList(),
-        val attachments: List<Attachment> = emptyList(),
-    ) : ChatUiState
+    data object Success : ChatUiState
     data class Error(val message: String) : ChatUiState
 }