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 <hy4ri@users.noreply.github.com>
This commit is contained in:
M57andM57 authored and GitHub committed 2026-07-25 13:03:46 +03:00
1 parent 223f9b676c
commit 4e2fb05077
3 files changed
+104 -1

No files matched your search

@@ -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<ChatUiState>,
streamingState: MutableStateFlow<StreamingState>,
isTestEnvironment: () -> Boolean,
) = ChatStreamingController(
scope = scope,
uiState = uiState,
streamingState = streamingState,
isCurrentSession = { true },
isTestEnvironment = isTestEnvironment,
)
}