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 1c716b70..44534ca5 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 @@ -1676,17 +1676,6 @@ class ChatViewModel( userMessage: ChatMessage? = null, ) { val dispatchGeneration = sessionGeneration - // Capture before attachment work or prompt submission can advance the - // cumulative counters. The reducer preserves this baseline across - // every assistant segment in a tool loop. - if (!_streamingState.value.turnUsageBaselineCaptured) { - _streamingState.update { - it.copy( - turnUsageBaseline = _uiState.value.sessionUsage, - turnUsageBaselineCaptured = true, - ) - } - } AuthManager.setLastOpenedSessionId(storageSessionId) val msgToPersist = userMessage ?: ChatMessage( @@ -1826,6 +1815,7 @@ class ChatViewModel( if (dispatchGeneration == sessionGeneration) { ActiveSessionHolder.set(agentSessionId, storageSessionId) } + captureTurnUsageBaselineIfNeeded() if (wasStreaming && attachments.isEmpty()) { wsClient.sendRedirect( agentSessionId, @@ -1859,6 +1849,20 @@ class ChatViewModel( } } + private fun captureTurnUsageBaselineIfNeeded() { + if (_streamingState.value.turnUsageBaselineCaptured) return + _streamingState.update { + if (it.turnUsageBaselineCaptured) { + it + } else { + it.copy( + turnUsageBaseline = _uiState.value.sessionUsage, + turnUsageBaselineCaptured = true, + ) + } + } + } + /** Snapshot one URI to private cache while enforcing the outbound frame limit. */ private suspend fun prepareAttachment(attachment: Attachment): PrepareAttachmentResult { var encodedFile: File? = null diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducer.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducer.kt index 60135e48..5b2f047b 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducer.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducer.kt @@ -15,6 +15,8 @@ private fun List.upsertById(message: ChatMessage): List 0.0 } + internal fun applyUsageSnapshot( state: ChatUiState, snapshot: com.m57.hermescontrol.data.model.UsageSnapshotResponse, @@ -432,7 +434,7 @@ object ChatWsEventReducer { streamingState = StreamingState(), ) } - val tps = (turnUsage?.avgTps ?: usageState.latestTps)?.takeIf { it.isFinite() && it > 0.0 } + val tps = turnUsage?.avgTps.validTpsOrNull() ?: usageState.latestTps.validTpsOrNull() val tokenCount = if (turnUsage?.outputTokens != null) { turnUsage.outputTokens.toIntOrNullSafely() 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 35ff224e..df49b884 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 @@ -1883,7 +1883,7 @@ class ChatViewModelTest { @Test fun sendMessage_oversizedDuringSnapshotLeavesNoPersistedGhostMessage() = runTest { - val (viewModel, _) = createViewModelWithSession() + val (viewModel, sessionId) = createViewModelWithSession() val persistedBeforeSend = fakeRepo.dao.count() val uriString = "content://unknown/growing-file" val mockUri = mockk() @@ -1920,6 +1920,7 @@ class ChatViewModelTest { viewModel.sendMessage("Inspect changing file") advanceUntilIdle() + assertFalse(viewModel.streamingState.value.turnUsageBaselineCaptured) assertEquals(persistedBeforeSend, fakeRepo.dao.count()) assertFalse( fakeRepo.dao @@ -1931,6 +1932,24 @@ class ChatViewModelTest { viewModel.uiState.value.messages .any { it.content == "Inspect changing file" }, ) + + mockEventsFlow.emit( + WsEvent.SessionUsage( + data = mapOf("usage" to mapOf("output" to 1300L)), + sessionId = sessionId, + ), + ) + advanceUntilIdle() + viewModel.removeAttachment(0) + viewModel.sendMessage("Retry without attachment") + advanceUntilIdle() + + assertEquals( + 1300L, + viewModel.streamingState.value.turnUsageBaseline + ?.outputTokens, + ) + assertTrue(viewModel.streamingState.value.turnUsageBaselineCaptured) } @Test @@ -4191,6 +4210,7 @@ class ChatViewModelTest { advanceUntilIdle() viewModel.sendMessage("prompt") + advanceUntilIdle() assertEquals( 1000L, diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducerTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducerTest.kt index 1cdeac1c..7faa6810 100644 --- a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducerTest.kt +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducerTest.kt @@ -4,6 +4,7 @@ import com.m57.hermescontrol.data.model.UsageSnapshotResponse import com.m57.hermescontrol.data.ws.WsEvent import org.junit.Assert.assertEquals import org.junit.Assert.assertFalse +import org.junit.Assert.assertNull import org.junit.Assert.assertTrue import org.junit.Test @@ -1185,6 +1186,46 @@ class ChatWsEventReducerTest { assertTrue(!result.streamingState.turnUsageBaselineCaptured) } + @Test + fun testMessageComplete_invalidFinalTpsFallsBackToLatestValidTps() { + fun complete( + finalTps: Double?, + latestTps: Double?, + ): Double? { + val state = + ChatUiState( + currentSessionId = "session-1", + latestTps = latestTps, + sessionUsage = UsageSnapshotResponse(outputTokens = 1000), + ) + val start = + ChatWsEventReducer.reduce( + state, + StreamingState(), + WsEvent.MessageStart("session-1"), + "session-1", + ) + val usage = mutableMapOf("output" to 1250L) + if (finalTps != null) usage["avg_tps"] = finalTps + return ChatWsEventReducer + .reduce( + state, + start.streamingState, + WsEvent.MessageComplete("reply", "session-1", rawPayload = mapOf("usage" to usage)), + "session-1", + ).state.messages + .single() + .tps + } + + assertEquals(42.5, complete(42.5, 30.0) ?: 0.0, 0.001) + assertEquals(37.5, complete(0.0, 37.5) ?: 0.0, 0.001) + assertEquals(37.5, complete(-2.0, 37.5) ?: 0.0, 0.001) + assertEquals(37.5, complete(null, 37.5) ?: 0.0, 0.001) + assertNull(complete(-2.0, null)) + assertNull(complete(0.0, -1.0)) + } + @Test fun testMessageComplete_withoutBaselineFallsBackToVisibleTextEstimate() { val result =