refactor(chat): extract streaming state into ChatStreamingController (#489)

* refactor(chat): extract streaming state into ChatStreamingController

Moves streaming buffers, streamingMessageId, and the two token/delta
handlers out of ChatViewModel into a dedicated controller. Mirrors the
ChatSearchDelegate seam; _streamingState stays VM-owned and the WS
reducer is untouched. Behavior-preserving, ChatScreen unchanged.

* refactor(chat): drop write-only streamingMessageId in ChatStreamingController

Sourcery flagged streamingMessageId as write-only. Verified via git grep
that it was write-only in origin/main too (zero readers anywhere in the
repo), so removing it is behavior-preserving. beginStreamingMessage() now
delegates its buffer reset to resetStreaming() to remove the duplication
Sourcery noted. Kept the injected scope param — it mirrors the
ChatSearchDelegate seam (that delegate uses scope.launch) and keeps the
controller constructor shape consistent.
This commit is contained in:
M57 authored and GitHub committed 2026-07-07 23:57:13 +03:00
1 parent 3aa40414e6
commit 9508204574
2 files changed
+130 -94

No files matched your search

@@ -0,0 +1,112 @@
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.flow.update
/**
* 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.
*/
class ChatStreamingController(
private val scope: CoroutineScope,
private val uiState: MutableStateFlow<ChatUiState>,
private val streamingState: MutableStateFlow<StreamingState>,
private val isCurrentSession: (String?) -> Boolean,
private val isTestEnvironment: () -> Boolean,
) {
private val streamingBuffer = java.lang.StringBuilder()
private var lastFlushMs = 0L
private val thinkingBuffer = java.lang.StringBuilder()
private var lastThinkingFlushMs = 0L
/**
* 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() {
streamingBuffer.clear()
thinkingBuffer.clear()
lastFlushMs = 0L
lastThinkingFlushMs = 0L
}
/** Resets buffers and starts a fresh streaming message (called on MessageStart). */
fun beginStreamingMessage() {
resetStreaming()
}
fun handleMessageToken(event: WsEvent.MessageToken) {
if (!isCurrentSession(event.sessionId)) return
streamingBuffer.append(event.token)
val now = System.currentTimeMillis()
// Always flush in tests, or if enough time has passed
val shouldFlush = (now - lastFlushMs >= 33L) || lastFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = streamingBuffer.toString()
lastFlushMs = now
streamingState.update { state ->
val current = state.streamingMessage
if (current != null) {
state.copy(
streamingMessage =
current.copy(
content = currentContent,
),
isThinking = false,
)
} else {
// Fallback: no MessageStart was received — create one now
val msg =
ChatMessage(
role = MessageRole.ASSISTANT,
content = currentContent,
isStreaming = true,
)
state.copy(
streamingMessage = msg,
isThinking = false,
)
}
}
uiState.update { it.copy(isAgentTyping = true) }
}
}
fun handleThinkingDelta(event: WsEvent.ThinkingDelta) {
if (!isCurrentSession(event.sessionId)) return
thinkingBuffer.append(event.token)
val now = System.currentTimeMillis()
// Always flush in tests, or if enough time has passed
val shouldFlush =
(now - lastThinkingFlushMs >= 33L) || lastThinkingFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = thinkingBuffer.toString()
lastThinkingFlushMs = now
streamingState.update { state ->
state.copy(
isThinking = true,
thinkingText = currentContent,
)
}
}
}
}
@@ -120,18 +120,17 @@ class ChatViewModel(
)
private val attachmentsDelegate = ChatAttachmentsDelegate(uiState = _uiState)
/** Tracks the ID of the currently streaming assistant message. */
@Volatile
private var streamingMessageId: String? = null
private val streamingController =
ChatStreamingController(
scope = viewModelScope,
uiState = _uiState,
streamingState = _streamingState,
isCurrentSession = { sessionId -> isCurrentSession(sessionId) },
isTestEnvironment = { isTestEnvironment() },
)
private var pendingCleanupJob: Job? = null
private val streamingBuffer = java.lang.StringBuilder()
private var lastFlushMs = 0L
private val thinkingBuffer = java.lang.StringBuilder()
private var lastThinkingFlushMs = 0L
// ── Public state ─────────────────────────────────────────────────────
/**
@@ -287,41 +286,29 @@ class ChatViewModel(
}
is WsEvent.MessageToken -> {
handleMessageToken(event)
streamingController.handleMessageToken(event)
}
is WsEvent.ThinkingDelta -> {
handleThinkingDelta(event)
streamingController.handleThinkingDelta(event)
}
is WsEvent.MessageStart -> {
streamingMessageId =
java.util.UUID
.randomUUID()
.toString()
streamingBuffer.clear()
thinkingBuffer.clear()
lastFlushMs = 0L
lastThinkingFlushMs = 0L
streamingController.beginStreamingMessage()
}
is WsEvent.MessageComplete -> {
// Buffers cleared before reduce; ViewModel resets them after
streamingMessageId = null
streamingBuffer.clear()
thinkingBuffer.clear()
streamingController.resetStreaming()
}
is WsEvent.MessageDone -> {
streamingMessageId = null
streamingBuffer.clear()
thinkingBuffer.clear()
streamingController.resetStreaming()
}
is WsEvent.ToolStart -> {
// Reset streaming state when a tool starts
streamingBuffer.clear()
thinkingBuffer.clear()
streamingController.resetStreaming()
}
is WsEvent.RpcResult -> {
@@ -343,6 +330,7 @@ class ChatViewModel(
)
}
_streamingState.update { StreamingState() }
streamingController.resetStreaming()
}
is WsEvent.ApprovalRequest -> {
@@ -365,66 +353,6 @@ class ChatViewModel(
return eventSessionId == _uiState.value.currentSessionId
}
private fun handleMessageToken(event: WsEvent.MessageToken) {
if (!isCurrentSession(event.sessionId)) return
streamingBuffer.append(event.token)
val now = System.currentTimeMillis()
// Always flush in tests, or if enough time has passed
val shouldFlush = (now - lastFlushMs >= 33L) || lastFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = streamingBuffer.toString()
lastFlushMs = now
_streamingState.update { state ->
val current = state.streamingMessage
if (current != null) {
state.copy(
streamingMessage =
current.copy(
content = currentContent,
),
isThinking = false,
)
} else {
// Fallback: no MessageStart was received — create one now
val msg =
ChatMessage(
role = MessageRole.ASSISTANT,
content = currentContent,
isStreaming = true,
)
streamingMessageId = msg.id
state.copy(
streamingMessage = msg,
isThinking = false,
)
}
}
_uiState.update { it.copy(isAgentTyping = true) }
}
}
private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) {
if (!isCurrentSession(event.sessionId)) return
thinkingBuffer.append(event.token)
val now = System.currentTimeMillis()
// Always flush in tests, or if enough time has passed
val shouldFlush = (now - lastThinkingFlushMs >= 33L) || lastThinkingFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = thinkingBuffer.toString()
lastThinkingFlushMs = now
_streamingState.update { state ->
state.copy(
isThinking = true,
thinkingText = currentContent,
)
}
}
}
// ── RPC response handling ────────────────────────────────────────────
@Suppress("UNCHECKED_CAST")
@@ -496,15 +424,13 @@ class ChatViewModel(
}
WsMethods.SESSION_INTERRUPT -> {
streamingBuffer.clear()
thinkingBuffer.clear()
_uiState.update {
it.copy(
isAgentTyping = false,
)
}
_streamingState.update { StreamingState() }
streamingMessageId = null
streamingController.resetStreaming()
addSystemMessage("Session interrupted")
}
@@ -852,7 +778,7 @@ class ChatViewModel(
)
}
_streamingState.update { StreamingState() }
streamingMessageId = null
streamingController.resetStreaming()
viewModelScope.launch(Dispatchers.IO) {
wsClient.send(
WsMethods.SESSION_CREATE,
@@ -915,9 +841,7 @@ class ChatViewModel(
if (sessionId == _uiState.value.currentSessionId) return
// Reset streaming state
streamingMessageId = null
streamingBuffer.clear()
thinkingBuffer.clear()
streamingController.resetStreaming()
_uiState.update {
val title = it.sessions.find { s -> s.id == sessionId }?.title ?: "Hermes"
it.copy(