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:
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(
|
||||
|
||||
Reference in new issue
Block a user