From 4e2fb05077290af85d1004ef8a1e9da0e34d22b8 Mon Sep 17 00:00:00 2001 From: M57 <41645758+Hy4ri@users.noreply.github.com> Date: Sat, 25 Jul 2026 13:03:46 +0300 Subject: [PATCH] fix(chat): preserve reasoning across streaming transitions (#709) Flush throttled reasoning before state-transition events (MessageStart/ Complete/Done/ToolStart) so a reasoning delta arriving just before a transition is not dropped. Also compare reasoningText in sameMessages so a reasoning-only restore actually re-merges. - ChatStreamingController.flushPendingReasoning() forces a synchronous flush - handleWsEvent calls it for the transition events - sameMessages now includes reasoningText in its equality check - tests: ChatStreamingControllerTest (flush + no-resurrect), reasoning survives MessageDone after an unflushed delta Co-authored-by: M57 --- .../ui/chat/ChatStreamingController.kt | 11 +++ .../hermescontrol/ui/chat/ChatViewModel.kt | 17 +++- .../ui/chat/ChatStreamingControllerTest.kt | 77 +++++++++++++++++++ 3 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatStreamingController.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatStreamingController.kt index cf0d775a..1b02b821 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatStreamingController.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatStreamingController.kt @@ -56,6 +56,17 @@ class ChatStreamingController( resetStreaming() } + /** + * Flushes any throttled reasoning buffer before a state-transition event + * (MessageStart / MessageComplete / MessageDone / ToolStart). Without this, + * a reasoning delta that arrived just before a transition could be dropped + * because the next flush only fires on the ~33ms timer — so we force it out + * synchronously so the finalized message carries the latest reasoning. + */ + fun flushPendingReasoning() { + flushReasoning() + } + fun handleMessageToken(event: WsEvent.MessageToken) { if (!isCurrentSession(event.sessionId)) return 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 b0ec848d..cfac70a0 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 @@ -341,6 +341,18 @@ class ChatViewModel( } private fun handleWsEvent(event: WsEvent) { + // Flush any throttled reasoning before a state transition so the + // finalized/orphan message carries the latest reasoning text. + when (event) { + is WsEvent.MessageStart, + is WsEvent.MessageComplete, + is WsEvent.MessageDone, + is WsEvent.ToolStart, + -> streamingController.flushPendingReasoning() + + else -> Unit + } + // First, let the reducer compute the new state and any effects val result = ChatWsEventReducer.reduce( @@ -1627,7 +1639,10 @@ class ChatViewModel( ): Boolean = left.size == right.size && left.zip(right).all { (a, b) -> - a.id == b.id && a.role == b.role && a.content == b.content + a.id == b.id && + a.role == b.role && + a.content == b.content && + a.reasoningText == b.reasoningText } // ── UI actions ─────────────────────────────────────────────────────── diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt new file mode 100644 index 00000000..8793af62 --- /dev/null +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt @@ -0,0 +1,77 @@ +package com.m57.hermescontrol.ui.chat + +import com.m57.hermescontrol.data.ws.WsEvent +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.test.runTest +import org.junit.Assert.assertEquals +import org.junit.Test + +class ChatStreamingControllerTest { + @Test + fun flushPendingReasoning_flushesThrottledReasoningOnTransition() = + runTest { + // isTestEnvironment = false so the ~33ms throttle is live. The FIRST delta + // always flushes (lastFlushMs == 0), but a SECOND delta within 33ms is held + // in the buffer and only pushed out by flushPendingReasoning() (called at + // streaming transitions) — that is the exact drop we guard against. + val uiState = MutableStateFlow(ChatUiState()) + val streamingState = + MutableStateFlow( + StreamingState( + streamingMessage = + ChatMessage( + role = MessageRole.ASSISTANT, + content = "", + isStreaming = true, + ), + ), + ) + val controller = controller(this, uiState, streamingState, isTestEnvironment = { false }) + + controller.handleReasoningDelta(WsEvent.ReasoningDelta("A", "session")) + // First delta flushes immediately. + assertEquals("A", streamingState.value.streamingMessage?.reasoningText) + + // Second delta arrives <33ms later -> throttled, still buffered, not yet shown. + controller.handleReasoningDelta(WsEvent.ReasoningDelta("B", "session")) + assertEquals("A", streamingState.value.streamingMessage?.reasoningText) + + // A streaming transition forces the buffered reasoning onto the message. + controller.flushPendingReasoning() + + assertEquals("AB", streamingState.value.streamingMessage?.reasoningText) + } + + @Test + fun resetStreaming_clearsReasoningBufferSoStaleReasoningDoesNotResurrect() = + runTest { + val uiState = MutableStateFlow(ChatUiState()) + val streamingState = MutableStateFlow(StreamingState()) + val controller = controller(this, uiState, streamingState, isTestEnvironment = { false }) + + // Buffer some reasoning, then reset (e.g. on MessageStart of a new turn). + controller.handleReasoningDelta(WsEvent.ReasoningDelta("Stale reasoning", "session")) + controller.resetStreaming() + + // A fresh reasoning delta + flush must carry ONLY the new value, proving the + // stale buffer was cleared and cannot be resurrected. + controller.handleReasoningDelta(WsEvent.ReasoningDelta("Fresh reasoning", "session")) + controller.flushPendingReasoning() + + assertEquals("Fresh reasoning", streamingState.value.reasoningText) + } + + private fun controller( + scope: CoroutineScope, + uiState: MutableStateFlow, + streamingState: MutableStateFlow, + isTestEnvironment: () -> Boolean, + ) = ChatStreamingController( + scope = scope, + uiState = uiState, + streamingState = streamingState, + isCurrentSession = { true }, + isTestEnvironment = isTestEnvironment, + ) +}