refactor(chat): unify throttled streaming buffers (#1338) (#1358)

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
+ * Owns the three throttled streams and their synchronous transition flushes, and mutates
- * streaming buffers and the two token/delta handlers, and mutates the shared
+ * the shared [uiState] + [streamingState] flows. The owning ViewModel keeps
- * [uiState] + [streamingState] flows. The owning ViewModel keeps [streamingState]
+ * [streamingState] as the single source of truth for the reduced [StreamingState].
- * as the single source of truth for the reduced [StreamingState] (applied via the
+ * [isCurrentSession] and [isTestEnvironment] preserve the existing session
- * WS event reducer); this controller only writes into it for buffer-driven flushes.
+ * filtering and immediate-flush test behavior.
- *
- * [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.
  */
 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()
+        reasoning.flushPending()
-        trailingThinkingFlushJob?.cancel()
+        tokens.clear()
-        trailingReasoningFlushJob?.cancel()
+        thinking.clear()
-        flushReasoning()
+        reasoning.clear()
-        streamingBuffer.clear()
-        thinkingBuffer.clear()
-        reasoningBuffer.clear()
-        lastFlushMs = 0L
-        lastThinkingFlushMs = 0L
-        lastReasoningFlushMs = 0L
         // 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()
+        tokens.clear()
-        trailingThinkingFlushJob?.cancel()
+        thinking.clear()
-        trailingReasoningFlushJob?.cancel()
+        reasoning.clear()
-        streamingBuffer.clear()
-        thinkingBuffer.clear()
-        reasoningBuffer.clear()
-        lastFlushMs = 0L
-        lastThinkingFlushMs = 0L
-        lastReasoningFlushMs = 0L
     }
 
     /** Resets buffers and starts a fresh streaming message (called on MessageStart). */
@@ -104,78 +76,31 @@ class ChatStreamingController(
         resetStreaming()
     }
 
-    /**
+    /** Drain throttled reasoning and tokens before the reducer seals or resets the message. */
-     * Flushes any throttled reasoning buffer before a state-transition event
+    fun flushPendingTransition() {
-     * (MessageStart / MessageComplete / MessageDone / ToolStart). Without this,
+        reasoning.flushPending()
-     * a reasoning delta that arrived just before a transition could be dropped
+        tokens.flushPending { content ->
-     * because the next flush only fires on the ~33ms timer — so we force it out
+            streamingState.update { state ->
-     * synchronously so the finalized message carries the latest reasoning.
+                val current = state.streamingMessage ?: return@update state
-     */
+                state.copy(
-    fun flushPendingReasoning() {
+                    streamingMessage =
-        trailingReasoningFlushJob?.cancel()
+                        current.copy(
-        flushReasoning()
+                            content = content,
-    }
+                            reasoningText =
-
+                                if (current.reasoningText.isNotBlank()) current.reasoningText else state.reasoningText,
-    /**
+                        ),
-     * 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,
-                    ),
-            )
         }
     }
 
     fun handleMessageToken(event: WsEvent.MessageToken) {
         if (!isCurrentSession(event.sessionId)) return
 
-        streamingBuffer.append(event.token)
+        tokens.append(event.token, isTestEnvironment())
-        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()
-        }
     }
 
-    private fun scheduleTrailingMessageFlush() {
+    private fun flushTokens(currentContent: String) {
-        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
         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)
+        thinking.append(event.token, isTestEnvironment())
-        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()
-        }
     }
 
-    private fun scheduleTrailingThinkingFlush() {
+    private fun flushThinking(currentContent: String) {
-        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
         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)
+        reasoning.append(event.token, isTestEnvironment())
-        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()
-        }
     }
 
-    private fun scheduleTrailingReasoningFlush() {
+    private fun flushReasoning(currentContent: String) {
-        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
         streamingState.update { state ->
             state.copy(
                 isReasoning = true,

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()
+                streamingController.flushPendingTransition()
-                // 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()
             }
 
             else -> {}

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())
+    }
+}

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<ChatUiState>,

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<String>()
+            val reasoning = mutableListOf<String>()
+            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<String>()
+            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<String>()
+            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)
+        }
+}