Merge pull request #152 from Hy4ri/fix/issue-146-lost-replies

fix(#146): finalize active streaming message before tool execution or new message start

app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt

@@ -201,6 +201,10 @@ class ChatViewModel(
 
     private fun handleMessageStart(event: WsEvent.MessageStart) {
         if (!isCurrentSession(event.sessionId)) return
+
+        var orphanToPersist: ChatMessage? = null
+        val sessionId = _uiState.value.currentSessionId
+
         // Create the streaming message as a standalone state field.
         val msg =
             ChatMessage(
@@ -209,14 +213,31 @@ class ChatViewModel(
                 isStreaming = true,
             )
         streamingMessageId = msg.id
-        _uiState.update {
+
-            it.copy(
+        _uiState.update { state ->
+            val messages = state.messages.toMutableList()
+            val existing = state.streamingMessage
+            if (existing != null && existing.content.isNotEmpty()) {
+                val finalized = existing.copy(isStreaming = false)
+                messages.add(finalized)
+                orphanToPersist = finalized
+            }
+            state.copy(
+                messages = messages,
                 isAgentTyping = true,
                 streamingMessage = msg,
                 isThinking = false,
                 thinkingText = "",
             )
         }
+
+        // Persist — OUTSIDE update{}
+        val orphan = orphanToPersist
+        if (orphan != null && sessionId != null) {
+            viewModelScope.launch(Dispatchers.IO) {
+                dao.upsert(orphan.toEntity(sessionId))
+            }
+        }
     }
 
     private fun handleMessageToken(event: WsEvent.MessageToken) {
@@ -345,6 +366,9 @@ class ChatViewModel(
     // ── Tool events ──────────────────────────────────────────────────────
 
     private fun handleToolStart(event: WsEvent.ToolStart) {
+        var orphanToPersist: ChatMessage? = null
+        val sessionId = _uiState.value.currentSessionId
+
         val toolMessage =
             ChatMessage(
                 role = MessageRole.TOOL,
@@ -354,13 +378,27 @@ class ChatViewModel(
             )
 
         _uiState.update { state ->
-            state.copy(messages = state.messages + toolMessage)
+            val messages = state.messages.toMutableList()
+            val existing = state.streamingMessage
+            if (existing != null && existing.content.isNotEmpty()) {
+                val finalized = existing.copy(isStreaming = false)
+                messages.add(finalized)
+                orphanToPersist = finalized
+            }
+            messages.add(toolMessage)
+            state.copy(
+                messages = messages,
+                streamingMessage = null,
+            )
         }
 
         // Persist — OUTSIDE update{}
-        val sessionId = _uiState.value.currentSessionId
+        val orphan = orphanToPersist
         if (sessionId != null) {
             viewModelScope.launch(Dispatchers.IO) {
+                if (orphan != null) {
+                    dao.upsert(orphan.toEntity(sessionId))
+                }
                 dao.upsert(toolMessage.toEntity(sessionId))
             }
         }

app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt

@@ -413,4 +413,85 @@ class ChatViewModelTest {
             assertTrue(state.errorMessage!!.contains("Internal error during creation"))
             assertTrue(state.errorMessage.contains(WsMethods.SESSION_CREATE))
         }
+
+    @Test
+    fun testToolExecution_finalizesPreviousStreamingMessage() =
+        runTest {
+            val viewModel = ChatViewModel(app, startCleanup = false)
+            advanceUntilIdle()
+
+            // Setup active session
+            var createReqId = ""
+            every { HermesWsClient.send(WsMethods.SESSION_CREATE, any(), any()) } answers {
+                createReqId = "create-req-tool-test"
+                val onSent = arg<((String) -> Unit)?>(2)
+                onSent?.invoke(createReqId)
+                createReqId
+            }
+            mockEventsFlow.emit(WsEvent.GatewayReady(null))
+            advanceUntilIdle()
+            mockEventsFlow.emit(WsEvent.RpcResult(createReqId, mapOf("session_id" to "session-123")))
+            advanceUntilIdle()
+
+            // Start typing some text
+            mockEventsFlow.emit(WsEvent.MessageStart("session-123"))
+            mockEventsFlow.emit(WsEvent.MessageToken("Calculating sum", "session-123"))
+            advanceUntilIdle()
+
+            // Start tool call
+            mockEventsFlow.emit(WsEvent.ToolStart("calculator", mapOf("input" to "2+2")))
+            advanceUntilIdle()
+
+            val state = viewModel.uiState.value
+            // Messages should be: 1. "Session created" 2. "Calculating sum" (finalized) 3. Tool bubble
+            assertEquals(3, state.messages.size)
+            assertEquals("Session created", state.messages[0].content)
+            assertEquals("Calculating sum", state.messages[1].content)
+            assertEquals(MessageRole.ASSISTANT, state.messages[1].role)
+            assertFalse(state.messages[1].isStreaming)
+            assertEquals(MessageRole.TOOL, state.messages[2].role)
+            assertNull(state.streamingMessage)
+        }
+
+    @Test
+    fun testMessageStart_finalizesPreviousStreamingMessage() =
+        runTest {
+            val viewModel = ChatViewModel(app, startCleanup = false)
+            advanceUntilIdle()
+
+            // Setup active session
+            var createReqId = ""
+            every { HermesWsClient.send(WsMethods.SESSION_CREATE, any(), any()) } answers {
+                createReqId = "create-req-msg-test"
+                val onSent = arg<((String) -> Unit)?>(2)
+                onSent?.invoke(createReqId)
+                createReqId
+            }
+            mockEventsFlow.emit(WsEvent.GatewayReady(null))
+            advanceUntilIdle()
+            mockEventsFlow.emit(WsEvent.RpcResult(createReqId, mapOf("session_id" to "session-123")))
+            advanceUntilIdle()
+
+            // Start typing message 1
+            mockEventsFlow.emit(WsEvent.MessageStart("session-123"))
+            mockEventsFlow.emit(WsEvent.MessageToken("First response segment", "session-123"))
+            advanceUntilIdle()
+
+            // Start typing message 2 without message 1 complete
+            mockEventsFlow.emit(WsEvent.MessageStart("session-123"))
+            mockEventsFlow.emit(WsEvent.MessageToken("Second response segment", "session-123"))
+            advanceUntilIdle()
+
+            val state = viewModel.uiState.value
+            // Messages should contain: 1. "Session created" 2. "First response segment" (finalized)
+            assertEquals(2, state.messages.size)
+            assertEquals("Session created", state.messages[0].content)
+            assertEquals("First response segment", state.messages[1].content)
+            assertFalse(state.messages[1].isStreaming)
+
+            // Active streaming message should be the second segment
+            assertNotNull(state.streamingMessage)
+            assertEquals("Second response segment", state.streamingMessage?.content)
+            assertTrue(state.streamingMessage?.isStreaming == true)
+        }
 }