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 new file mode 100644 index 00000000..9579eb08 --- /dev/null +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatStreamingController.kt @@ -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, + private val streamingState: MutableStateFlow, + 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, + ) + } + } + } +} 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 ec68b068..f0caa339 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 @@ -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(