From 58794bee00f044bc7356503d00445414d9943432 Mon Sep 17 00:00:00 2001 From: m57 <41645758+Hy4ri@users.noreply.github.com> Date: Mon, 28 Sep 2026 15:05:58 +0300 Subject: [PATCH] refactor(chat): unify throttled streaming buffers (#1338) (#1358) --- .../ui/chat/ChatStreamingController.kt | 197 ++++-------------- .../hermescontrol/ui/chat/ChatViewModel.kt | 8 +- .../hermescontrol/ui/chat/StreamingBuffer.kt | 58 ++++++ .../ui/chat/ChatStreamingControllerTest.kt | 30 ++- .../ui/chat/StreamingBufferTest.kt | 56 +++++ 5 files changed, 177 insertions(+), 172 deletions(-) create mode 100644 app/src/main/java/com/m57/hermescontrol/ui/chat/StreamingBuffer.kt create mode 100644 app/src/test/java/com/m57/hermescontrol/ui/chat/StreamingBufferTest.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 9652cf5f..2aa0191f 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 @@ -2,26 +2,18 @@ package com.m57.hermescontrol.ui.chat import com.m57.hermescontrol.data.ws.WsEvent import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Job -import kotlinx.coroutines.delay import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.update -import kotlinx.coroutines.launch /** * Holds chat message-streaming state and logic, extracted from [ChatViewModel] * to keep the god-object focused on messaging/session concerns. * - * Behavior is identical to the original inline implementation: it owns the - * streaming buffers and the two token/delta handlers, and mutates the shared - * [uiState] + [streamingState] flows. The owning ViewModel keeps [streamingState] - * as the single source of truth for the reduced [StreamingState] (applied via the - * WS event reducer); this controller only writes into it for buffer-driven flushes. - * - * [isCurrentSession] and [isTestEnvironment] are injected so the controller - * stays free of ViewModel-specific context while preserving exact behavior. - * [scope] mirrors the [com.m57.hermescontrol.ui.chat.ChatSearchDelegate] seam - * and is reserved for future streaming coroutine work. + * Owns the three throttled streams and their synchronous transition flushes, and mutates + * the shared [uiState] + [streamingState] flows. The owning ViewModel keeps + * [streamingState] as the single source of truth for the reduced [StreamingState]. + * [isCurrentSession] and [isTestEnvironment] preserve the existing session + * filtering and immediate-flush test behavior. */ class ChatStreamingController( private val scope: CoroutineScope, @@ -34,36 +26,22 @@ class ChatStreamingController( const val FLUSH_INTERVAL_MS = 33L } - private val streamingBuffer = java.lang.StringBuilder() - private var lastFlushMs = 0L - private var trailingMessageFlushJob: Job? = null - - private val thinkingBuffer = java.lang.StringBuilder() - private var lastThinkingFlushMs = 0L - private var trailingThinkingFlushJob: Job? = null - - private val reasoningBuffer = java.lang.StringBuilder() - private var lastReasoningFlushMs = 0L - private var trailingReasoningFlushJob: Job? = null - private fun nowMs(): Long = System.nanoTime() / 1_000_000L + private val tokens = StreamingBuffer(scope, FLUSH_INTERVAL_MS, ::nowMs, ::flushTokens) + private val thinking = StreamingBuffer(scope, FLUSH_INTERVAL_MS, ::nowMs, ::flushThinking) + private val reasoning = StreamingBuffer(scope, FLUSH_INTERVAL_MS, ::nowMs, ::flushReasoning) + /** * Resets all streaming buffers. Centralizes the clear logic that was * previously scattered across MessageStart, MessageComplete, MessageDone, * ToolStart, ClarifyRequest, session switches, and interrupt handling. */ fun resetStreaming() { - trailingMessageFlushJob?.cancel() - trailingThinkingFlushJob?.cancel() - trailingReasoningFlushJob?.cancel() - flushReasoning() - streamingBuffer.clear() - thinkingBuffer.clear() - reasoningBuffer.clear() - lastFlushMs = 0L - lastThinkingFlushMs = 0L - lastReasoningFlushMs = 0L + reasoning.flushPending() + tokens.clear() + thinking.clear() + reasoning.clear() // Issue #755: the shared streaming state must not carry the previous // message's reasoning into the next one. The finalized message keeps // its own copy (persisted to Room), so clearing here is safe. @@ -84,19 +62,13 @@ class ChatStreamingController( * Used at tool.start: the reducer keeps the streaming message + its * reasoning alive across the tool call (issue #771), and the VM must * not undo that by resetting the shared streaming state. The buffers - * themselves are drained (flushPendingReasoning) before reduce, so a + * themselves are drained before reduce, so a * buffer-only clear is safe and prevents stale deltas from re-flushing. */ fun clearStreamingBuffers() { - trailingMessageFlushJob?.cancel() - trailingThinkingFlushJob?.cancel() - trailingReasoningFlushJob?.cancel() - streamingBuffer.clear() - thinkingBuffer.clear() - reasoningBuffer.clear() - lastFlushMs = 0L - lastThinkingFlushMs = 0L - lastReasoningFlushMs = 0L + tokens.clear() + thinking.clear() + reasoning.clear() } /** Resets buffers and starts a fresh streaming message (called on MessageStart). */ @@ -104,78 +76,31 @@ 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() { - trailingReasoningFlushJob?.cancel() - flushReasoning() - } - - /** - * Force-flushes any throttled TOKEN buffer onto the streaming message - * before a state-transition event (ToolStart / MessageComplete / - * MessageDone). `handleMessageToken` only flushes when >=33ms have passed - * since the last flush, so a delta that lands just before the transition - * can still be sitting in the buffer. The reducer seals the streaming - * message into a finalized orphan at tool.start (issue #842) — sealing a - * stale copy truncated the narration's tail (observed on-device: the - * orphan ended "...chickpea" while the REST row ended "...chickpea - * goodness:"), which then failed the logical-match against the clean - * REST row and re-added the commentary as a ghost bubble. Flushing the - * tokens synchronously first makes the sealed orphan complete. - */ - fun flushPendingTokens() { - trailingMessageFlushJob?.cancel() - val currentContent = streamingBuffer.toString() - if (currentContent.isEmpty()) return - lastFlushMs = nowMs() - streamingState.update { state -> - val current = state.streamingMessage ?: return@update state - state.copy( - streamingMessage = - current.copy( - content = currentContent, - reasoningText = - if (current.reasoningText.isNotBlank()) current.reasoningText else state.reasoningText, - ), - ) + /** Drain throttled reasoning and tokens before the reducer seals or resets the message. */ + fun flushPendingTransition() { + reasoning.flushPending() + tokens.flushPending { content -> + streamingState.update { state -> + val current = state.streamingMessage ?: return@update state + state.copy( + streamingMessage = + current.copy( + content = content, + reasoningText = + if (current.reasoningText.isNotBlank()) current.reasoningText else state.reasoningText, + ), + ) + } } } fun handleMessageToken(event: WsEvent.MessageToken) { if (!isCurrentSession(event.sessionId)) return - streamingBuffer.append(event.token) - val now = nowMs() - - // Always flush in tests, or if enough time has passed - val shouldFlush = (now - lastFlushMs >= FLUSH_INTERVAL_MS) || lastFlushMs == 0L || isTestEnvironment() - if (shouldFlush) { - trailingMessageFlushJob?.cancel() - flushTokens(now) - } else { - scheduleTrailingMessageFlush() - } + tokens.append(event.token, isTestEnvironment()) } - private fun scheduleTrailingMessageFlush() { - if (trailingMessageFlushJob?.isActive == true) return - trailingMessageFlushJob = - scope.launch { - delay(FLUSH_INTERVAL_MS) - flushTokens(nowMs()) - } - } - - private fun flushTokens(now: Long) { - val currentContent = streamingBuffer.toString() - if (currentContent.isEmpty()) return - lastFlushMs = now + private fun flushTokens(currentContent: String) { streamingState.update { state -> val current = state.streamingMessage if (current != null) { @@ -210,33 +135,10 @@ class ChatStreamingController( fun handleThinkingDelta(event: WsEvent.ThinkingDelta) { if (!isCurrentSession(event.sessionId)) return - thinkingBuffer.append(event.token) - val now = nowMs() - - // Always flush in tests, or if enough time has passed - val shouldFlush = - (now - lastThinkingFlushMs >= FLUSH_INTERVAL_MS) || lastThinkingFlushMs == 0L || isTestEnvironment() - if (shouldFlush) { - trailingThinkingFlushJob?.cancel() - flushThinking(now) - } else { - scheduleTrailingThinkingFlush() - } + thinking.append(event.token, isTestEnvironment()) } - private fun scheduleTrailingThinkingFlush() { - if (trailingThinkingFlushJob?.isActive == true) return - trailingThinkingFlushJob = - scope.launch { - delay(FLUSH_INTERVAL_MS) - flushThinking(nowMs()) - } - } - - private fun flushThinking(now: Long) { - val currentContent = thinkingBuffer.toString() - if (currentContent.isEmpty()) return - lastThinkingFlushMs = now + private fun flushThinking(currentContent: String) { streamingState.update { state -> state.copy( isThinking = true, @@ -248,33 +150,10 @@ class ChatStreamingController( fun handleReasoningDelta(event: WsEvent.ReasoningDelta) { if (!isCurrentSession(event.sessionId)) return - reasoningBuffer.append(event.token) - val now = nowMs() - - // Always flush in tests, or if enough time has passed - val shouldFlush = - (now - lastReasoningFlushMs >= FLUSH_INTERVAL_MS) || lastReasoningFlushMs == 0L || isTestEnvironment() - if (shouldFlush) { - trailingReasoningFlushJob?.cancel() - flushReasoning(now) - } else { - scheduleTrailingReasoningFlush() - } + reasoning.append(event.token, isTestEnvironment()) } - private fun scheduleTrailingReasoningFlush() { - if (trailingReasoningFlushJob?.isActive == true) return - trailingReasoningFlushJob = - scope.launch { - delay(FLUSH_INTERVAL_MS) - flushReasoning(nowMs()) - } - } - - private fun flushReasoning(now: Long = nowMs()) { - val currentContent = reasoningBuffer.toString() - if (currentContent.isEmpty()) return - lastReasoningFlushMs = now + private fun flushReasoning(currentContent: String) { streamingState.update { state -> state.copy( isReasoning = true, 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 69b4a1c4..14763576 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 @@ -1042,13 +1042,7 @@ class ChatViewModel( is WsEvent.MessageDone, is WsEvent.ToolStart, -> { - streamingController.flushPendingReasoning() - // Issue #842: the token buffer can hold deltas that landed - // <33ms before the transition. The reducer seals the - // streaming message into the orphan at tool.start — flush - // first so the seal carries the COMPLETE narration (a - // truncated seal fails the later REST dedupe and ghosts). - streamingController.flushPendingTokens() + streamingController.flushPendingTransition() } else -> {} diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/StreamingBuffer.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/StreamingBuffer.kt new file mode 100644 index 00000000..445151e2 --- /dev/null +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/StreamingBuffer.kt @@ -0,0 +1,58 @@ +package com.m57.hermescontrol.ui.chat + +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Job +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch + +/** One throttled streaming channel: cumulative text, last flush and trailing timer. */ +internal class StreamingBuffer( + private val scope: CoroutineScope, + private val intervalMs: Long, + private val nowMs: () -> Long, + private val flush: (String) -> Unit, +) { + private val content = StringBuilder() + private var lastFlushMs = 0L + private var trailingJob: Job? = null + + fun append( + delta: String, + flushImmediately: Boolean, + ) { + content.append(delta) + val now = nowMs() + if (flushImmediately || lastFlushMs == 0L || now - lastFlushMs >= intervalMs) { + trailingJob?.cancel() + flushNow(now) + } else if (trailingJob?.isActive != true) { + trailingJob = + scope.launch { + delay(intervalMs) + flushNow(nowMs()) + } + } + } + + /** Keep the cumulative text: subsequent deltas append to the same message. */ + fun flushPending(override: ((String) -> Unit)? = null) { + trailingJob?.cancel() + flushNow(nowMs(), override) + } + + fun clear() { + trailingJob?.cancel() + trailingJob = null + content.clear() + lastFlushMs = 0L + } + + private fun flushNow( + now: Long, + override: ((String) -> Unit)? = null, + ) { + if (content.isEmpty()) return + lastFlushMs = now + (override ?: flush)(content.toString()) + } +} 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 index 6efcd0a2..e6aaa5a9 100644 --- a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatStreamingControllerTest.kt @@ -9,11 +9,11 @@ import org.junit.Test class ChatStreamingControllerTest { @Test - fun flushPendingReasoning_flushesThrottledReasoningOnTransition() = + fun flushPendingTransition_flushesThrottledReasoning() = 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 + // in the buffer and only pushed out by flushPendingTransition() (called at // streaming transitions) — that is the exact drop we guard against. val uiState = MutableStateFlow(ChatUiState()) val streamingState = @@ -38,13 +38,13 @@ class ChatStreamingControllerTest { assertEquals("A", streamingState.value.streamingMessage?.reasoningText) // A streaming transition forces the buffered reasoning onto the message. - controller.flushPendingReasoning() + controller.flushPendingTransition() assertEquals("AB", streamingState.value.streamingMessage?.reasoningText) } @Test - fun flushPendingTokens_flushesThrottledTokensBeforeSeal() = + fun flushPendingTransition_flushesThrottledTokensBeforeSeal() = runTest { // Issue #842: a delta landing <33ms before tool.start stays in the // throttled buffer; flushing the tokens before the reducer seals the @@ -73,7 +73,7 @@ class ChatStreamingControllerTest { // The tool.start transition forces the buffered tokens out so the // sealed orphan carries the complete narration. - controller.flushPendingTokens() + controller.flushPendingTransition() assertEquals( "tool's loaded! 🔍 now searchin' for the best hummus", @@ -95,7 +95,7 @@ class ChatStreamingControllerTest { // 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() + controller.flushPendingTransition() assertEquals("Fresh reasoning", streamingState.value.reasoningText) } @@ -168,6 +168,24 @@ class ChatStreamingControllerTest { assertEquals("First Second", streamingState.value.streamingMessage?.content) } + @Test + fun thinkingDelta_trailingFlushAndResetDoNotLeakIntoNextTurn() = + runTest { + val uiState = MutableStateFlow(ChatUiState()) + val streamingState = MutableStateFlow(StreamingState()) + val controller = controller(this, uiState, streamingState, isTestEnvironment = { false }) + + controller.handleThinkingDelta(WsEvent.ThinkingDelta("Think", "session")) + controller.handleThinkingDelta(WsEvent.ThinkingDelta(" more", "session")) + assertEquals("Think", streamingState.value.thinkingText) + testScheduler.advanceTimeBy(35L) + assertEquals("Think more", streamingState.value.thinkingText) + + controller.resetStreaming() + controller.handleThinkingDelta(WsEvent.ThinkingDelta("Fresh", "session")) + assertEquals("Fresh", streamingState.value.thinkingText) + } + private fun controller( scope: CoroutineScope, uiState: MutableStateFlow, diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/StreamingBufferTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/StreamingBufferTest.kt new file mode 100644 index 00000000..8d501cac --- /dev/null +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/StreamingBufferTest.kt @@ -0,0 +1,56 @@ +package com.m57.hermescontrol.ui.chat + +import kotlinx.coroutines.test.runTest +import org.junit.Assert.assertEquals +import org.junit.Test + +class StreamingBufferTest { + @Test + fun channelsThrottleIndependentlyAndFlushCumulativeText() = + runTest { + var now = 100L + val tokens = mutableListOf() + val reasoning = mutableListOf() + val tokenBuffer = StreamingBuffer(this, 33L, { now }, tokens::add) + val reasoningBuffer = StreamingBuffer(this, 33L, { now }, reasoning::add) + + tokenBuffer.append("A", flushImmediately = false) + tokenBuffer.append("B", flushImmediately = false) + reasoningBuffer.append("R", flushImmediately = false) + assertEquals(listOf("A"), tokens) + assertEquals(listOf("R"), reasoning) + + now += 34 + testScheduler.advanceTimeBy(34L) + assertEquals(listOf("A", "AB"), tokens) + assertEquals(listOf("R"), reasoning) + } + + @Test + fun transitionFlushCancelsTrailingJobWithoutDroppingTheNextDelta() = + runTest { + val output = mutableListOf() + val buffer = StreamingBuffer(this, 33L, { 100L }, output::add) + buffer.append("A", flushImmediately = false) + buffer.append("B", flushImmediately = false) + buffer.flushPending() + testScheduler.advanceTimeBy(34L) + assertEquals(listOf("A", "AB"), output) + + buffer.append("C", flushImmediately = true) + assertEquals(listOf("A", "AB", "ABC"), output) + } + + @Test + fun clearCancelsTrailingJobAndStartsFresh() = + runTest { + val output = mutableListOf() + val buffer = StreamingBuffer(this, 33L, { 100L }, output::add) + buffer.append("Old", flushImmediately = false) + buffer.append(" tail", flushImmediately = false) + buffer.clear() + testScheduler.advanceTimeBy(34L) + buffer.append("Fresh", flushImmediately = false) + assertEquals(listOf("Old", "Fresh"), output) + } +}