diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt index 9db348f1..73c4f5e5 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt +++ b/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)) } } diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt index 9d04ea27..8a170310 100644 --- a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt +++ b/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) + } }