From d0eedf053ffbbb48198a7f72c23e5028cdd7370b Mon Sep 17 00:00:00 2001 From: Hy4ri Date: Sat, 20 Jun 2026 11:04:56 +0300 Subject: [PATCH 1/4] refactor(chat): stabilize streaming, session loading, and connection status Optimize streaming message updates, correct session load merging, and derive connection state reactively. Add reconnection UI banner. --- .../m57/hermescontrol/ui/chat/ChatBubble.kt | 12 +- .../m57/hermescontrol/ui/chat/ChatScreen.kt | 62 +- .../hermescontrol/ui/chat/ChatViewModel.kt | 582 ++++++++++++------ .../ui/chat/ChatViewModelTest.kt | 21 +- 4 files changed, 462 insertions(+), 215 deletions(-) diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt index ae06a4ab..3644ed92 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt @@ -72,10 +72,10 @@ import com.m57.hermescontrol.theme.SystemMessageColor import com.m57.hermescontrol.theme.ToolChipColor import com.m57.hermescontrol.theme.ToolChipColorLight import com.m57.hermescontrol.theme.UserBubble +import java.time.Instant +import java.time.ZoneId +import java.time.format.DateTimeFormatter import kotlinx.coroutines.delay -import java.text.SimpleDateFormat -import java.util.Date -import java.util.Locale @Composable fun ChatBubble( @@ -811,9 +811,11 @@ private fun RichText( ) } +private val TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern("HH:mm") + .withZone(ZoneId.systemDefault()) + private fun formatTimestamp(timestamp: Long): String { - val sdf = SimpleDateFormat("HH:mm", Locale.getDefault()) - return sdf.format(Date(timestamp)) + return TIMESTAMP_FORMATTER.format(Instant.ofEpochMilli(timestamp)) } /** diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt index a0e3fd1e..d8b39867 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt @@ -92,6 +92,7 @@ import androidx.compose.ui.unit.sp import androidx.core.content.ContextCompat import androidx.lifecycle.compose.collectAsStateWithLifecycle import androidx.lifecycle.viewmodel.compose.viewModel +import com.m57.hermescontrol.data.ws.ConnectionStatus import com.m57.hermescontrol.notification.NotificationHelper import com.m57.hermescontrol.theme.StatusRed import com.m57.hermescontrol.ui.common.EmptyState @@ -169,9 +170,12 @@ fun ChatScreen( } // Auto-scroll to bottom on new messages - LaunchedEffect(state.messages.size, state.currentStreamingText) { - if (state.messages.isNotEmpty()) { - listState.animateScrollToItem(state.messages.lastIndex) + LaunchedEffect(state.messages.size, state.streamingMessage?.content?.length, state.isThinking) { + val totalItems = state.messages.size + + (if (state.streamingMessage != null) 1 else 0) + + (if (state.isThinking) 1 else 0) + if (totalItems > 0) { + listState.animateScrollToItem(totalItems - 1) } } @@ -349,6 +353,46 @@ fun ChatScreen( .background(backgroundGradient) .imePadding(), ) { + androidx.compose.animation.AnimatedVisibility( + visible = state.connectionStatus == ConnectionStatus.RECONNECTING || + state.connectionStatus == ConnectionStatus.DISCONNECTED, + enter = androidx.compose.animation.expandVertically() + androidx.compose.animation.fadeIn(), + exit = androidx.compose.animation.shrinkVertically() + androidx.compose.animation.fadeOut(), + ) { + androidx.compose.material3.Surface( + modifier = Modifier.fillMaxWidth(), + color = MaterialTheme.colorScheme.errorContainer, + contentColor = MaterialTheme.colorScheme.onErrorContainer, + ) { + Row( + modifier = Modifier + .fillMaxWidth() + .padding(horizontal = 16.dp, vertical = 8.dp), + verticalAlignment = Alignment.CenterVertically, + horizontalArrangement = Arrangement.SpaceBetween, + ) { + Text( + text = if (state.connectionStatus == ConnectionStatus.RECONNECTING) { + "Reconnecting to Hermes…" + } else { + "Disconnected from Hermes" + }, + style = MaterialTheme.typography.bodyMedium, + ) + if (state.connectionStatus == ConnectionStatus.DISCONNECTED) { + TextButton( + onClick = { viewModel.reconnect() }, + colors = androidx.compose.material3.ButtonDefaults.textButtonColors( + contentColor = MaterialTheme.colorScheme.onErrorContainer, + ), + ) { + Text("Reconnect") + } + } + } + } + } + androidx.compose.animation.AnimatedVisibility( visible = state.isSearchActive, enter = androidx.compose.animation.expandVertically() + androidx.compose.animation.fadeIn(), @@ -414,6 +458,18 @@ fun ChatScreen( ) } + // Streaming message — rendered separately for O(1) updates + state.streamingMessage?.let { streaming -> + item(key = "streaming-${streaming.id}") { + ChatBubble( + message = streaming, + isDarkTheme = isDark, + searchQuery = "", + isCurrentMatch = false, + ) + } + } + // Thinking indicator if (state.isThinking) { item(key = "thinking") { 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 27279562..ac5e5faa 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 @@ -1,6 +1,7 @@ package com.m57.hermescontrol.ui.chat import android.app.Application +import android.util.Log import androidx.lifecycle.AndroidViewModel import androidx.lifecycle.viewModelScope import com.m57.hermescontrol.data.local.AuthManager @@ -10,30 +11,43 @@ import com.m57.hermescontrol.data.local.toUiModel import com.m57.hermescontrol.data.remote.ApiClient import com.m57.hermescontrol.data.remote.NetworkResult import com.m57.hermescontrol.data.remote.safeApiCall +import com.m57.hermescontrol.data.ws.ConnectionStatus import com.m57.hermescontrol.data.ws.HermesWsClient import com.m57.hermescontrol.data.ws.WsEvent import com.m57.hermescontrol.data.ws.WsMethods import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job import kotlinx.coroutines.delay import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.SharingStarted import kotlinx.coroutines.flow.StateFlow -import kotlinx.coroutines.flow.asStateFlow +import kotlinx.coroutines.flow.combine +import kotlinx.coroutines.flow.stateIn import kotlinx.coroutines.flow.update import kotlinx.coroutines.launch import kotlinx.coroutines.withContext import java.util.concurrent.ConcurrentHashMap +private const val TAG = "ChatViewModel" + +/** + * Pending request timeout — if the server doesn't respond within this + * window, the request is pruned and an error is surfaced. + */ +private const val PENDING_REQUEST_TIMEOUT_MS = 30_000L + data class ChatUiState( val messages: List = emptyList(), val currentSessionId: String? = null, val sessions: List = emptyList(), val chatTitle: String = "Hermes", - val isConnected: Boolean = false, + val connectionStatus: ConnectionStatus = ConnectionStatus.DISCONNECTED, val isAgentTyping: Boolean = false, val isThinking: Boolean = false, val thinkingText: String = "", val isLoading: Boolean = false, - val currentStreamingText: String = "", + /** Standalone streaming message — rendered after the main list. */ + val streamingMessage: ChatMessage? = null, val errorMessage: String? = null, val clarifyRequest: ClarifyUi? = null, val showSessionPicker: Boolean = false, @@ -42,7 +56,10 @@ data class ChatUiState( val searchQuery: String = "", val searchMatchIndices: List = emptyList(), val currentSearchMatchIndex: Int = -1, -) +) { + /** Convenience — derived from [connectionStatus]. */ + val isConnected: Boolean get() = connectionStatus == ConnectionStatus.CONNECTED +} data class SessionUi( val id: String, @@ -58,17 +75,44 @@ data class ClarifyUi( class ChatViewModel( application: Application, ) : AndroidViewModel(application) { + // ── Internal state ─────────────────────────────────────────────────── private val _uiState = MutableStateFlow(ChatUiState()) - val uiState: StateFlow = _uiState.asStateFlow() private val wsClient = HermesWsClient - private val pendingRequests = ConcurrentHashMap() // requestId -> method + private val pendingRequests = ConcurrentHashMap() private val dao = HermesDatabase.get(application).chatMessageDao() + /** Tracks the ID of the currently streaming assistant message. */ + @Volatile + private var streamingMessageId: String? = null + + private var pendingCleanupJob: Job? = null + + // ── Public state ───────────────────────────────────────────────────── + + /** + * Combined UI state: merges internal state with the WS connection status + * flow so there is a single source of truth for connection state. + */ + val uiState: StateFlow = + combine( + _uiState, + wsClient.connectionStatus, + ) { state, connStatus -> + state.copy(connectionStatus = connStatus) + }.stateIn( + viewModelScope, + SharingStarted.WhileSubscribed(5_000), + _uiState.value, + ) + init { connectWebSocket() + startPendingRequestCleanup() } + // ── Connection ─────────────────────────────────────────────────────── + private fun connectWebSocket() { val token = AuthManager.getToken() ?: return @@ -83,21 +127,14 @@ class ChatViewModel( handleWsEvent(event) } } - - if (wsClient.isConnected) { - _uiState.update { it.copy(isConnected = true, isLoading = false) } - addSystemMessage("Connected to Hermes") - loadSessions() - if (_uiState.value.currentSessionId == null) { - createNewSession() - } - } } + // ── WS Event Handling ──────────────────────────────────────────────── + private fun handleWsEvent(event: WsEvent) { when (event) { is WsEvent.GatewayReady -> { - _uiState.update { it.copy(isConnected = true, isLoading = false) } + _uiState.update { it.copy(isLoading = false) } addSystemMessage("Connected to Hermes") loadSessions() if (_uiState.value.currentSessionId == null) { @@ -107,152 +144,13 @@ class ChatViewModel( is WsEvent.SessionInfo -> {} - is WsEvent.MessageStart -> { - _uiState.update { - it.copy( - isAgentTyping = true, - currentStreamingText = "", - isThinking = false, - thinkingText = "", - ) - } - } - - is WsEvent.MessageToken -> { - _uiState.update { state -> - val newText = state.currentStreamingText + event.token - val messages = state.messages.toMutableList() - val streamingIdx = - messages.indexOfLast { - it.role == MessageRole.ASSISTANT && it.isStreaming - } - if (streamingIdx >= 0) { - messages[streamingIdx] = messages[streamingIdx].copy(content = newText) - } else { - messages.add( - ChatMessage( - role = MessageRole.ASSISTANT, - content = newText, - isStreaming = true, - ), - ) - } - // Persist streaming message to DB - val sessionId = state.currentSessionId - if (sessionId != null) { - val msg = messages.last() - viewModelScope.launch(Dispatchers.IO) { - dao.upsert(msg.toEntity(sessionId)) - } - } - state.copy( - messages = messages, - currentStreamingText = newText, - isThinking = false, - ) - } - } - - is WsEvent.ThinkingDelta -> { - _uiState.update { state -> - state.copy( - isThinking = true, - thinkingText = state.thinkingText + event.token, - ) - } - } - - is WsEvent.MessageComplete -> { - _uiState.update { state -> - val messages = state.messages.toMutableList() - val streamingIdx = - messages.indexOfLast { - it.role == MessageRole.ASSISTANT && it.isStreaming - } - val finalMessage: ChatMessage - if (streamingIdx >= 0) { - finalMessage = - messages[streamingIdx].copy( - content = event.text, - isStreaming = false, - ) - messages[streamingIdx] = finalMessage - } else { - finalMessage = - ChatMessage( - role = MessageRole.ASSISTANT, - content = event.text, - ) - messages.add(finalMessage) - } - // Persist finalized message - val sessionId = state.currentSessionId - if (sessionId != null) { - viewModelScope.launch(Dispatchers.IO) { - dao.finalizeMessage(finalMessage.id, event.text) - } - } - state.copy( - messages = messages, - isAgentTyping = false, - currentStreamingText = "", - isThinking = false, - thinkingText = "", - ) - } - } - - is WsEvent.MessageDone -> { - _uiState.update { - it.copy(isAgentTyping = false, isThinking = false, thinkingText = "") - } - } - - is WsEvent.ToolStart -> { - val toolMessage = - ChatMessage( - role = MessageRole.TOOL, - content = event.data.toString(), - toolName = event.name, - toolStatus = ToolStatus.RUNNING, - ) - _uiState.update { state -> - val sessionId = state.currentSessionId - if (sessionId != null) { - viewModelScope.launch(Dispatchers.IO) { - dao.upsert(toolMessage.toEntity(sessionId)) - } - } - state.copy(messages = state.messages + toolMessage) - } - } - - is WsEvent.ToolComplete -> { - _uiState.update { state -> - val messages = state.messages.toMutableList() - val toolIdx = - messages.indexOfLast { - it.role == MessageRole.TOOL && - it.toolName == event.name && - it.toolStatus == ToolStatus.RUNNING - } - if (toolIdx >= 0) { - messages[toolIdx] = - messages[toolIdx].copy( - toolStatus = ToolStatus.COMPLETED, - content = event.data.toString(), - ) - // Persist tool completion - val sessionId = state.currentSessionId - if (sessionId != null) { - viewModelScope.launch(Dispatchers.IO) { - dao.upsert(messages[toolIdx].toEntity(sessionId)) - } - } - } - state.copy(messages = messages) - } - } + is WsEvent.MessageStart -> handleMessageStart(event) + is WsEvent.MessageToken -> handleMessageToken(event) + is WsEvent.ThinkingDelta -> handleThinkingDelta(event) + is WsEvent.MessageComplete -> handleMessageComplete(event) + is WsEvent.MessageDone -> handleMessageDone(event) + is WsEvent.ToolStart -> handleToolStart(event) + is WsEvent.ToolComplete -> handleToolComplete(event) is WsEvent.ClarifyRequest -> { _uiState.update { @@ -284,12 +182,217 @@ class ChatViewModel( } } + // ── Message streaming ──────────────────────────────────────────────── + + /** + * Checks if an incoming WS event belongs to the currently active + * session. Returns true if the event should be processed. + */ + private fun isCurrentSession(eventSessionId: String?): Boolean { + // If the event has no session ID, process it (legacy compatibility) + if (eventSessionId == null) return true + return eventSessionId == _uiState.value.currentSessionId + } + + private fun handleMessageStart(event: WsEvent.MessageStart) { + if (!isCurrentSession(event.sessionId)) return + // Create the streaming message as a standalone state field. + val msg = ChatMessage( + role = MessageRole.ASSISTANT, + content = "", + isStreaming = true, + ) + streamingMessageId = msg.id + _uiState.update { + it.copy( + isAgentTyping = true, + streamingMessage = msg, + isThinking = false, + thinkingText = "", + ) + } + } + + private fun handleMessageToken(event: WsEvent.MessageToken) { + if (!isCurrentSession(event.sessionId)) return + _uiState.update { state -> + val current = state.streamingMessage + if (current != null) { + state.copy( + streamingMessage = current.copy( + content = current.content + event.token, + ), + isThinking = false, + ) + } else { + // Fallback: no MessageStart was received — create one now + val msg = ChatMessage( + role = MessageRole.ASSISTANT, + content = event.token, + isStreaming = true, + ) + streamingMessageId = msg.id + state.copy( + streamingMessage = msg, + isAgentTyping = true, + isThinking = false, + ) + } + } + } + + private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) { + if (!isCurrentSession(event.sessionId)) return + _uiState.update { state -> + state.copy( + isThinking = true, + thinkingText = state.thinkingText + event.token, + ) + } + } + + private fun handleMessageComplete(event: WsEvent.MessageComplete) { + if (!isCurrentSession(event.sessionId)) return + + // Finalize: move the streaming message into the main list + val finalizedMsg: ChatMessage + val sessionId: String? + + _uiState.update { state -> + val streaming = state.streamingMessage + finalizedMsg = if (streaming != null) { + streaming.copy( + content = event.text, + isStreaming = false, + ) + } else { + ChatMessage( + role = MessageRole.ASSISTANT, + content = event.text, + ) + } + sessionId = state.currentSessionId + state.copy( + messages = state.messages + finalizedMsg, + streamingMessage = null, + isAgentTyping = false, + isThinking = false, + thinkingText = "", + ) + } + streamingMessageId = null + + // Persist finalized message — OUTSIDE update{} to avoid CAS retry + val sid = sessionId ?: _uiState.value.currentSessionId + if (sid != null) { + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(finalizedMsg.toEntity(sid)) + } + } + } + + private fun handleMessageDone(event: WsEvent.MessageDone) { + if (!isCurrentSession(event.sessionId)) return + + // If we still have an un-finalized streaming message (no MessageComplete + // was sent, which happens on interrupts), fold it into the list. + val orphan: ChatMessage? + val sessionId: String? + + _uiState.update { state -> + val streaming = state.streamingMessage + orphan = streaming?.copy(isStreaming = false) + sessionId = state.currentSessionId + if (orphan != null) { + state.copy( + messages = state.messages + orphan, + streamingMessage = null, + isAgentTyping = false, + isThinking = false, + thinkingText = "", + ) + } else { + state.copy( + isAgentTyping = false, + isThinking = false, + thinkingText = "", + ) + } + } + streamingMessageId = null + + // Persist orphaned streaming message + val sid = sessionId ?: _uiState.value.currentSessionId + if (orphan != null && sid != null) { + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(orphan.toEntity(sid)) + } + } + } + + // ── Tool events ────────────────────────────────────────────────────── + + private fun handleToolStart(event: WsEvent.ToolStart) { + val toolMessage = ChatMessage( + role = MessageRole.TOOL, + content = event.data.toString(), + toolName = event.name, + toolStatus = ToolStatus.RUNNING, + ) + + _uiState.update { state -> + state.copy(messages = state.messages + toolMessage) + } + + // Persist — OUTSIDE update{} + val sessionId = _uiState.value.currentSessionId + if (sessionId != null) { + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(toolMessage.toEntity(sessionId)) + } + } + } + + private fun handleToolComplete(event: WsEvent.ToolComplete) { + var completedTool: ChatMessage? = null + + _uiState.update { state -> + val messages = state.messages.toMutableList() + val toolIdx = messages.indexOfLast { + it.role == MessageRole.TOOL && + it.toolName == event.name && + it.toolStatus == ToolStatus.RUNNING + } + if (toolIdx >= 0) { + val updated = messages[toolIdx].copy( + toolStatus = ToolStatus.COMPLETED, + content = event.data.toString(), + ) + messages[toolIdx] = updated + completedTool = updated + } + state.copy(messages = messages) + } + + // Persist — OUTSIDE update{} + val tool = completedTool + val sessionId = _uiState.value.currentSessionId + if (tool != null && sessionId != null) { + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(tool.toEntity(sessionId)) + } + } + } + + // ── RPC response handling ──────────────────────────────────────────── + @Suppress("UNCHECKED_CAST") private fun handleRpcResult( id: String, result: Any?, ) { - val method = pendingRequests.remove(id) ?: return + val pending = pendingRequests.remove(id) ?: return + val method = pending.method when (method) { WsMethods.SESSION_CREATE -> { @@ -300,6 +403,7 @@ class ChatViewModel( currentSessionId = sessionId, isLoading = false, messages = emptyList(), + streamingMessage = null, ) } addSystemMessage("Session created", persist = true) @@ -347,8 +451,13 @@ class ChatViewModel( WsMethods.SESSION_INTERRUPT -> { _uiState.update { - it.copy(isAgentTyping = false, isThinking = false) + it.copy( + isAgentTyping = false, + isThinking = false, + streamingMessage = null, + ) } + streamingMessageId = null addSystemMessage("Session interrupted") } } @@ -358,7 +467,8 @@ class ChatViewModel( id: String, error: Any?, ) { - val method = pendingRequests.remove(id) + val pending = pendingRequests.remove(id) + val method = pending?.method val errorMsg = when (error) { is Map<*, *> -> error["message"] as? String ?: error.toString() @@ -372,6 +482,8 @@ class ChatViewModel( } } + // ── Send message ───────────────────────────────────────────────────── + fun sendMessage(text: String) { if (text.isBlank()) return val sessionId = _uiState.value.currentSessionId ?: return @@ -382,27 +494,29 @@ class ChatViewModel( return } - val userMessage = - ChatMessage( - role = MessageRole.USER, - content = text, - ) + val userMessage = ChatMessage( + role = MessageRole.USER, + content = text, + ) + + // Update UI — no DB writes inside update{} _uiState.update { state -> - val messages = state.messages + userMessage - viewModelScope.launch(Dispatchers.IO) { - dao.upsert(userMessage.toEntity(sessionId)) - } state.copy( - messages = messages, + messages = state.messages + userMessage, isAgentTyping = true, ) } + // Persist to Room — OUTSIDE update{} + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(userMessage.toEntity(sessionId)) + } + viewModelScope.launch(Dispatchers.IO) { wsClient.sendMessage( sessionId, text, - onSent = { id -> pendingRequests[id] = WsMethods.PROMPT_SUBMIT }, + onSent = { id -> trackRequest(id, WsMethods.PROMPT_SUBMIT) }, ) } } @@ -410,7 +524,10 @@ class ChatViewModel( private fun handleSlashCommand(command: String) { val userMsg = ChatMessage(role = MessageRole.USER, content = command) val sessionId = _uiState.value.currentSessionId + _uiState.update { it.copy(messages = it.messages + userMsg) } + + // Persist — OUTSIDE update{} if (sessionId != null) { viewModelScope.launch(Dispatchers.IO) { dao.upsert(userMsg.toEntity(sessionId)) @@ -538,8 +655,10 @@ class ChatViewModel( private fun addAssistantMessage(text: String) { val msg = ChatMessage(role = MessageRole.ASSISTANT, content = text) - val sessionId = _uiState.value.currentSessionId _uiState.update { it.copy(messages = it.messages + msg) } + + // Persist — OUTSIDE update{} + val sessionId = _uiState.value.currentSessionId if (sessionId != null) { viewModelScope.launch(Dispatchers.IO) { dao.upsert(msg.toEntity(sessionId)) @@ -547,23 +666,32 @@ class ChatViewModel( } } + // ── Session management ─────────────────────────────────────────────── + fun interruptSession() { val sessionId = _uiState.value.currentSessionId ?: return viewModelScope.launch(Dispatchers.IO) { wsClient.send( WsMethods.SESSION_INTERRUPT, mapOf("session_id" to sessionId), - onSent = { id -> pendingRequests[id] = WsMethods.SESSION_INTERRUPT }, + onSent = { id -> trackRequest(id, WsMethods.SESSION_INTERRUPT) }, ) } } fun createNewSession() { - _uiState.update { it.copy(isLoading = true, messages = emptyList()) } + _uiState.update { + it.copy( + isLoading = true, + messages = emptyList(), + streamingMessage = null, + ) + } + streamingMessageId = null viewModelScope.launch(Dispatchers.IO) { wsClient.send( WsMethods.SESSION_CREATE, - onSent = { id -> pendingRequests[id] = WsMethods.SESSION_CREATE }, + onSent = { id -> trackRequest(id, WsMethods.SESSION_CREATE) }, ) } } @@ -572,18 +700,23 @@ class ChatViewModel( viewModelScope.launch(Dispatchers.IO) { wsClient.send( WsMethods.SESSION_LIST, - onSent = { id -> pendingRequests[id] = WsMethods.SESSION_LIST }, + onSent = { id -> trackRequest(id, WsMethods.SESSION_LIST) }, ) } } fun switchSession(sessionId: String) { if (sessionId == _uiState.value.currentSessionId) return + + // Reset streaming state + streamingMessageId = null + _uiState.update { val title = it.sessions.find { s -> s.id == sessionId }?.title ?: "Hermes" it.copy( isLoading = true, messages = emptyList(), + streamingMessage = null, currentSessionId = sessionId, chatTitle = title, showSessionPicker = false, @@ -592,16 +725,17 @@ class ChatViewModel( thinkingText = "", ) } - // Load cached messages first — instant display + // Step 1: Load cached messages first — instant display loadCachedMessages(sessionId) - // Then resume session on server + // Step 2: Resume session on server viewModelScope.launch(Dispatchers.IO) { wsClient.send( WsMethods.SESSION_RESUME, mapOf("session_id" to sessionId), - onSent = { id -> pendingRequests[id] = WsMethods.SESSION_RESUME }, + onSent = { id -> trackRequest(id, WsMethods.SESSION_RESUME) }, ) } + // Step 3: Load fresh messages from REST and merge loadSessionMessages(sessionId) } @@ -610,7 +744,7 @@ class ChatViewModel( val cached = dao.getMessagesForSession(sessionId) val cachedMessages = cached.map { it.toUiModel() } _uiState.update { state -> - // Only replace if still showing this session and messages are empty/newer + // Only replace if still showing this session if (state.currentSessionId == sessionId) { state.copy(messages = cachedMessages, isLoading = false) } else { @@ -630,7 +764,7 @@ class ChatViewModel( is NetworkResult.Success -> { val messagesList = result.data.messages.orEmpty() val chatMessages = - messagesList.map { msg -> + messagesList.mapIndexed { index, msg -> val role = when (msg.role?.lowercase()) { "user" -> MessageRole.USER @@ -638,28 +772,45 @@ class ChatViewModel( "tool" -> MessageRole.TOOL else -> MessageRole.ASSISTANT } + // Use stable IDs based on session + index to prevent + // Room duplicates when re-loading the same session. + val stableId = "rest-$sessionId-$index" + // Preserve original timestamp from the API when available + val ts = msg.timestamp?.toLongOrNull() ?: System.currentTimeMillis() ChatMessage( + id = stableId, role = role, content = msg.content.orEmpty(), + timestamp = ts, isStreaming = false, ) } // Persist loaded messages to Room for offline access val entities = chatMessages.map { it.toEntity(sessionId) } - viewModelScope.launch(Dispatchers.IO) { + withContext(Dispatchers.IO) { dao.upsertAll(entities) } _uiState.update { state -> - val streamingTail = state.messages.filter { it.isStreaming } + // Only update if still on the same session + if (state.currentSessionId != sessionId) return@update state + + // Merge: keep any user messages sent between cache load + // and REST response (they won't be in the REST response + // yet), plus preserve any active streaming message. + val existingIds = chatMessages.map { it.id }.toSet() + val localOnly = state.messages.filter { it.id !in existingIds } state.copy( - messages = chatMessages + streamingTail, + messages = chatMessages + localOnly, isLoading = false, ) } } is NetworkResult.Failure -> { + // Don't overwrite messages on REST failure — cached messages + // are still valid. _uiState.update { + if (it.currentSessionId != sessionId) return@update it it.copy( isLoading = false, errorMessage = "Failed to load messages: ${result.error.message}", @@ -670,6 +821,8 @@ class ChatViewModel( } } + // ── UI actions ─────────────────────────────────────────────────────── + fun toggleSessionPicker() { _uiState.update { it.copy(showSessionPicker = !it.showSessionPicker) } } @@ -690,7 +843,6 @@ class ChatViewModel( fun reconnect() { _uiState.update { it.copy( - isConnected = false, isLoading = true, errorMessage = null, ) @@ -709,18 +861,53 @@ class ChatViewModel( persist: Boolean = false, ) { val msg = ChatMessage(role = MessageRole.SYSTEM, content = text) + val sessionId = _uiState.value.currentSessionId + _uiState.update { it.copy(messages = it.messages + msg) } - if (persist) { - val sessionId = _uiState.value.currentSessionId - if (sessionId != null) { - viewModelScope.launch(Dispatchers.IO) { - dao.upsert(msg.toEntity(sessionId)) + + // Persist — OUTSIDE update{} + if (persist && sessionId != null) { + viewModelScope.launch(Dispatchers.IO) { + dao.upsert(msg.toEntity(sessionId)) + } + } + } + + // ── Pending request tracking ───────────────────────────────────────── + + private data class PendingRequest( + val method: String, + val createdAt: Long = System.currentTimeMillis(), + ) + + private fun trackRequest( + id: String, + method: String, + ) { + pendingRequests[id] = PendingRequest(method) + } + + /** + * Periodically prunes stale pending requests that the server never + * responded to. Surfaces a timeout error for the user. + */ + private fun startPendingRequestCleanup() { + pendingCleanupJob = viewModelScope.launch { + while (true) { + delay(PENDING_REQUEST_TIMEOUT_MS) + val now = System.currentTimeMillis() + val stale = pendingRequests.entries.filter { + now - it.value.createdAt > PENDING_REQUEST_TIMEOUT_MS + } + for (entry in stale) { + pendingRequests.remove(entry.key) + Log.w(TAG, "Request timed out: ${entry.value.method} (id=${entry.key})") } } } } - // ── Search ──────────────────────────────────────────────────────── + // ── Search ──────────────────────────────────────────────────────────── fun toggleSearch() { val current = _uiState.value @@ -776,8 +963,11 @@ class ChatViewModel( } } + // ── Lifecycle ───────────────────────────────────────────────────────── + override fun onCleared() { super.onCleared() + pendingCleanupJob?.cancel() wsClient.disconnect() } } diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt index 96c1cf21..06ba9816 100644 --- a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt @@ -213,7 +213,8 @@ class ChatViewModelTest { var state = viewModel.uiState.value assertTrue(state.isAgentTyping) - assertEquals("", state.currentStreamingText) + assertNotNull(state.streamingMessage) + assertEquals("", state.streamingMessage?.content) assertFalse(state.isThinking) assertEquals("", state.thinkingText) @@ -236,27 +237,25 @@ class ChatViewModelTest { advanceUntilIdle() state = viewModel.uiState.value assertFalse(state.isThinking) - assertEquals("Hello", state.currentStreamingText) - // System message was cleared when createNewSession() was called, so only Assistant streaming message is present - assertEquals(1, state.messages.size) - assertEquals("Hello", state.messages[0].content) - assertTrue(state.messages[0].isStreaming) + assertNotNull(state.streamingMessage) + assertEquals("Hello", state.streamingMessage?.content) + // Streaming message is not in the main messages list yet + assertEquals(0, state.messages.size) // Token 2 mockEventsFlow.emit(WsEvent.MessageToken(" world", "session-123")) advanceUntilIdle() state = viewModel.uiState.value - assertEquals("Hello world", state.currentStreamingText) - assertEquals(1, state.messages.size) - assertEquals("Hello world", state.messages[0].content) - assertTrue(state.messages[0].isStreaming) + assertNotNull(state.streamingMessage) + assertEquals("Hello world", state.streamingMessage?.content) + assertEquals(0, state.messages.size) // Complete mockEventsFlow.emit(WsEvent.MessageComplete("Hello world!", "session-123")) advanceUntilIdle() state = viewModel.uiState.value assertFalse(state.isAgentTyping) - assertEquals("", state.currentStreamingText) + assertNull(state.streamingMessage) assertFalse(state.isThinking) assertEquals("", state.thinkingText) assertEquals(1, state.messages.size) From 5ba266cefd905e52c29270be8a502184baa8b840 Mon Sep 17 00:00:00 2001 From: Hy4ri Date: Sat, 20 Jun 2026 11:12:34 +0300 Subject: [PATCH 2/4] fix(chat): resolve compilation errors and formatting violations --- .../m57/hermescontrol/ui/chat/ChatBubble.kt | 7 +- .../m57/hermescontrol/ui/chat/ChatScreen.kt | 37 +++++---- .../hermescontrol/ui/chat/ChatViewModel.kt | 76 ++++++++++--------- 3 files changed, 67 insertions(+), 53 deletions(-) diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt index 3644ed92..ba79c8f6 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatBubble.kt @@ -72,10 +72,10 @@ import com.m57.hermescontrol.theme.SystemMessageColor import com.m57.hermescontrol.theme.ToolChipColor import com.m57.hermescontrol.theme.ToolChipColorLight import com.m57.hermescontrol.theme.UserBubble +import kotlinx.coroutines.delay import java.time.Instant import java.time.ZoneId import java.time.format.DateTimeFormatter -import kotlinx.coroutines.delay @Composable fun ChatBubble( @@ -811,8 +811,9 @@ private fun RichText( ) } -private val TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern("HH:mm") - .withZone(ZoneId.systemDefault()) +private val TIMESTAMP_FORMATTER = + DateTimeFormatter.ofPattern("HH:mm") + .withZone(ZoneId.systemDefault()) private fun formatTimestamp(timestamp: Long): String { return TIMESTAMP_FORMATTER.format(Instant.ofEpochMilli(timestamp)) diff --git a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt index d8b39867..868e3099 100644 --- a/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt +++ b/app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt @@ -171,9 +171,10 @@ fun ChatScreen( // Auto-scroll to bottom on new messages LaunchedEffect(state.messages.size, state.streamingMessage?.content?.length, state.isThinking) { - val totalItems = state.messages.size + - (if (state.streamingMessage != null) 1 else 0) + - (if (state.isThinking) 1 else 0) + val totalItems = + state.messages.size + + (if (state.streamingMessage != null) 1 else 0) + + (if (state.isThinking) 1 else 0) if (totalItems > 0) { listState.animateScrollToItem(totalItems - 1) } @@ -354,8 +355,9 @@ fun ChatScreen( .imePadding(), ) { androidx.compose.animation.AnimatedVisibility( - visible = state.connectionStatus == ConnectionStatus.RECONNECTING || - state.connectionStatus == ConnectionStatus.DISCONNECTED, + visible = + state.connectionStatus == ConnectionStatus.RECONNECTING || + state.connectionStatus == ConnectionStatus.DISCONNECTED, enter = androidx.compose.animation.expandVertically() + androidx.compose.animation.fadeIn(), exit = androidx.compose.animation.shrinkVertically() + androidx.compose.animation.fadeOut(), ) { @@ -365,26 +367,29 @@ fun ChatScreen( contentColor = MaterialTheme.colorScheme.onErrorContainer, ) { Row( - modifier = Modifier - .fillMaxWidth() - .padding(horizontal = 16.dp, vertical = 8.dp), + modifier = + Modifier + .fillMaxWidth() + .padding(horizontal = 16.dp, vertical = 8.dp), verticalAlignment = Alignment.CenterVertically, horizontalArrangement = Arrangement.SpaceBetween, ) { Text( - text = if (state.connectionStatus == ConnectionStatus.RECONNECTING) { - "Reconnecting to Hermes…" - } else { - "Disconnected from Hermes" - }, + text = + if (state.connectionStatus == ConnectionStatus.RECONNECTING) { + "Reconnecting to Hermes…" + } else { + "Disconnected from Hermes" + }, style = MaterialTheme.typography.bodyMedium, ) if (state.connectionStatus == ConnectionStatus.DISCONNECTED) { TextButton( onClick = { viewModel.reconnect() }, - colors = androidx.compose.material3.ButtonDefaults.textButtonColors( - contentColor = MaterialTheme.colorScheme.onErrorContainer, - ), + colors = + androidx.compose.material3.ButtonDefaults.textButtonColors( + contentColor = MaterialTheme.colorScheme.onErrorContainer, + ), ) { Text("Reconnect") } 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 ac5e5faa..2328e2a4 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 @@ -254,13 +254,12 @@ class ChatViewModel( private fun handleMessageComplete(event: WsEvent.MessageComplete) { if (!isCurrentSession(event.sessionId)) return - // Finalize: move the streaming message into the main list - val finalizedMsg: ChatMessage - val sessionId: String? + var finalizedMsg: ChatMessage? = null + var sessionId: String? = null _uiState.update { state -> val streaming = state.streamingMessage - finalizedMsg = if (streaming != null) { + val msg = if (streaming != null) { streaming.copy( content = event.text, isStreaming = false, @@ -271,9 +270,10 @@ class ChatViewModel( content = event.text, ) } + finalizedMsg = msg sessionId = state.currentSessionId state.copy( - messages = state.messages + finalizedMsg, + messages = state.messages + msg, streamingMessage = null, isAgentTyping = false, isThinking = false, @@ -283,10 +283,11 @@ class ChatViewModel( streamingMessageId = null // Persist finalized message — OUTSIDE update{} to avoid CAS retry + val msgToPersist = finalizedMsg val sid = sessionId ?: _uiState.value.currentSessionId - if (sid != null) { + if (msgToPersist != null && sid != null) { viewModelScope.launch(Dispatchers.IO) { - dao.upsert(finalizedMsg.toEntity(sid)) + dao.upsert(msgToPersist.toEntity(sid)) } } } @@ -296,16 +297,17 @@ class ChatViewModel( // If we still have an un-finalized streaming message (no MessageComplete // was sent, which happens on interrupts), fold it into the list. - val orphan: ChatMessage? - val sessionId: String? + var orphan: ChatMessage? = null + var sessionId: String? = null _uiState.update { state -> val streaming = state.streamingMessage - orphan = streaming?.copy(isStreaming = false) + val msg = streaming?.copy(isStreaming = false) + orphan = msg sessionId = state.currentSessionId - if (orphan != null) { + if (msg != null) { state.copy( - messages = state.messages + orphan, + messages = state.messages + msg, streamingMessage = null, isAgentTyping = false, isThinking = false, @@ -322,10 +324,11 @@ class ChatViewModel( streamingMessageId = null // Persist orphaned streaming message + val msgToPersist = orphan val sid = sessionId ?: _uiState.value.currentSessionId - if (orphan != null && sid != null) { + if (msgToPersist != null && sid != null) { viewModelScope.launch(Dispatchers.IO) { - dao.upsert(orphan.toEntity(sid)) + dao.upsert(msgToPersist.toEntity(sid)) } } } @@ -358,16 +361,18 @@ class ChatViewModel( _uiState.update { state -> val messages = state.messages.toMutableList() - val toolIdx = messages.indexOfLast { - it.role == MessageRole.TOOL && - it.toolName == event.name && - it.toolStatus == ToolStatus.RUNNING - } + val toolIdx = + messages.indexOfLast { + it.role == MessageRole.TOOL && + it.toolName == event.name && + it.toolStatus == ToolStatus.RUNNING + } if (toolIdx >= 0) { - val updated = messages[toolIdx].copy( - toolStatus = ToolStatus.COMPLETED, - content = event.data.toString(), - ) + val updated = + messages[toolIdx].copy( + toolStatus = ToolStatus.COMPLETED, + content = event.data.toString(), + ) messages[toolIdx] = updated completedTool = updated } @@ -494,10 +499,11 @@ class ChatViewModel( return } - val userMessage = ChatMessage( - role = MessageRole.USER, - content = text, - ) + val userMessage = + ChatMessage( + role = MessageRole.USER, + content = text, + ) // Update UI — no DB writes inside update{} _uiState.update { state -> @@ -892,13 +898,15 @@ class ChatViewModel( * responded to. Surfaces a timeout error for the user. */ private fun startPendingRequestCleanup() { - pendingCleanupJob = viewModelScope.launch { - while (true) { - delay(PENDING_REQUEST_TIMEOUT_MS) - val now = System.currentTimeMillis() - val stale = pendingRequests.entries.filter { - now - it.value.createdAt > PENDING_REQUEST_TIMEOUT_MS - } + pendingCleanupJob = + viewModelScope.launch { + while (true) { + delay(PENDING_REQUEST_TIMEOUT_MS) + val now = System.currentTimeMillis() + val stale = + pendingRequests.entries.filter { + now - it.value.createdAt > PENDING_REQUEST_TIMEOUT_MS + } for (entry in stale) { pendingRequests.remove(entry.key) Log.w(TAG, "Request timed out: ${entry.value.method} (id=${entry.key})") From 890ca05d16b24547a3682ef1b0e9a42ba970dc3e Mon Sep 17 00:00:00 2001 From: Hy4ri Date: Sat, 20 Jun 2026 11:18:08 +0300 Subject: [PATCH 3/4] fix(chat): align block indents and format multiline expressions --- .../hermescontrol/ui/chat/ChatViewModel.kt | 72 ++++++++++--------- 1 file changed, 39 insertions(+), 33 deletions(-) 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 2328e2a4..d65686c4 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 @@ -197,11 +197,12 @@ class ChatViewModel( private fun handleMessageStart(event: WsEvent.MessageStart) { if (!isCurrentSession(event.sessionId)) return // Create the streaming message as a standalone state field. - val msg = ChatMessage( - role = MessageRole.ASSISTANT, - content = "", - isStreaming = true, - ) + val msg = + ChatMessage( + role = MessageRole.ASSISTANT, + content = "", + isStreaming = true, + ) streamingMessageId = msg.id _uiState.update { it.copy( @@ -219,18 +220,20 @@ class ChatViewModel( val current = state.streamingMessage if (current != null) { state.copy( - streamingMessage = current.copy( - content = current.content + event.token, - ), + streamingMessage = + current.copy( + content = current.content + event.token, + ), isThinking = false, ) } else { // Fallback: no MessageStart was received — create one now - val msg = ChatMessage( - role = MessageRole.ASSISTANT, - content = event.token, - isStreaming = true, - ) + val msg = + ChatMessage( + role = MessageRole.ASSISTANT, + content = event.token, + isStreaming = true, + ) streamingMessageId = msg.id state.copy( streamingMessage = msg, @@ -259,17 +262,18 @@ class ChatViewModel( _uiState.update { state -> val streaming = state.streamingMessage - val msg = if (streaming != null) { - streaming.copy( - content = event.text, - isStreaming = false, - ) - } else { - ChatMessage( - role = MessageRole.ASSISTANT, - content = event.text, - ) - } + val msg = + if (streaming != null) { + streaming.copy( + content = event.text, + isStreaming = false, + ) + } else { + ChatMessage( + role = MessageRole.ASSISTANT, + content = event.text, + ) + } finalizedMsg = msg sessionId = state.currentSessionId state.copy( @@ -336,12 +340,13 @@ class ChatViewModel( // ── Tool events ────────────────────────────────────────────────────── private fun handleToolStart(event: WsEvent.ToolStart) { - val toolMessage = ChatMessage( - role = MessageRole.TOOL, - content = event.data.toString(), - toolName = event.name, - toolStatus = ToolStatus.RUNNING, - ) + val toolMessage = + ChatMessage( + role = MessageRole.TOOL, + content = event.data.toString(), + toolName = event.name, + toolStatus = ToolStatus.RUNNING, + ) _uiState.update { state -> state.copy(messages = state.messages + toolMessage) @@ -907,9 +912,10 @@ class ChatViewModel( pendingRequests.entries.filter { now - it.value.createdAt > PENDING_REQUEST_TIMEOUT_MS } - for (entry in stale) { - pendingRequests.remove(entry.key) - Log.w(TAG, "Request timed out: ${entry.value.method} (id=${entry.key})") + for (entry in stale) { + pendingRequests.remove(entry.key) + Log.w(TAG, "Request timed out: ${entry.value.method} (id=${entry.key})") + } } } } From 58ee32cbd1f71bbfcc1e8d109d77f1e2c86ab371 Mon Sep 17 00:00:00 2001 From: Hy4ri Date: Sat, 20 Jun 2026 11:21:31 +0300 Subject: [PATCH 4/4] fix(chat): remove extra closing brace in ChatViewModel.kt --- .../hermescontrol/ui/chat/ChatViewModel.kt | 10 ++-- .../ui/chat/ChatViewModelTest.kt | 50 +++++++++++++------ 2 files changed, 41 insertions(+), 19 deletions(-) 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 d65686c4..9db348f1 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 @@ -74,7 +74,10 @@ data class ClarifyUi( class ChatViewModel( application: Application, + private val startCleanup: Boolean, ) : AndroidViewModel(application) { + constructor(application: Application) : this(application, startCleanup = true) + // ── Internal state ─────────────────────────────────────────────────── private val _uiState = MutableStateFlow(ChatUiState()) @@ -102,13 +105,15 @@ class ChatViewModel( state.copy(connectionStatus = connStatus) }.stateIn( viewModelScope, - SharingStarted.WhileSubscribed(5_000), + SharingStarted.Eagerly, _uiState.value, ) init { connectWebSocket() - startPendingRequestCleanup() + if (startCleanup) { + startPendingRequestCleanup() + } } // ── Connection ─────────────────────────────────────────────────────── @@ -918,7 +923,6 @@ class ChatViewModel( } } } - } } // ── Search ──────────────────────────────────────────────────────────── diff --git a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt index 06ba9816..9d04ea27 100644 --- a/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt +++ b/app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt @@ -4,6 +4,7 @@ import android.app.Application import com.m57.hermescontrol.data.local.AuthManager import com.m57.hermescontrol.data.local.ChatMessageDao import com.m57.hermescontrol.data.local.HermesDatabase +import com.m57.hermescontrol.data.ws.ConnectionStatus import com.m57.hermescontrol.data.ws.HermesWsClient import com.m57.hermescontrol.data.ws.JsonRpcError import com.m57.hermescontrol.data.ws.WsEvent @@ -18,6 +19,7 @@ import io.mockk.verify import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.flow.MutableSharedFlow +import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.flowOf import kotlinx.coroutines.test.StandardTestDispatcher import kotlinx.coroutines.test.advanceUntilIdle @@ -37,6 +39,7 @@ import org.junit.Test class ChatViewModelTest { private val testDispatcher = StandardTestDispatcher() private val mockEventsFlow = MutableSharedFlow(extraBufferCapacity = 64) + private val mockConnectionStatus = MutableStateFlow(ConnectionStatus.DISCONNECTED) private lateinit var app: Application private val mockDao: ChatMessageDao = mockk(relaxed = true) private val mockDb: HermesDatabase = mockk(relaxed = true) @@ -56,8 +59,11 @@ class ChatViewModelTest { app = mockk(relaxed = true) + mockConnectionStatus.value = ConnectionStatus.DISCONNECTED + every { AuthManager.getToken() } returns "test-token" every { HermesWsClient.events } returns mockEventsFlow + every { HermesWsClient.connectionStatus } returns mockConnectionStatus every { HermesWsClient.connect() } returns Unit every { HermesWsClient.disconnect() } returns Unit every { mockDb.chatMessageDao() } returns mockDao @@ -97,7 +103,7 @@ class ChatViewModelTest { @Test fun testInitialStateAndConnection() = runTest { - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() verify { HermesWsClient.connect() } @@ -108,9 +114,10 @@ class ChatViewModelTest { @Test fun testGatewayReady_createsSessionIfNoneExists() = runTest { - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() + mockConnectionStatus.value = ConnectionStatus.CONNECTED mockEventsFlow.emit(WsEvent.GatewayReady(null)) advanceUntilIdle() @@ -136,7 +143,7 @@ class ChatViewModelTest { createReqId } - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() // Trigger GatewayReady -> triggers createNewSession @@ -166,7 +173,7 @@ class ChatViewModelTest { listReqId } - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() mockEventsFlow.emit(WsEvent.GatewayReady(null)) @@ -200,13 +207,23 @@ class ChatViewModelTest { @Test fun testMessageStreamingFlow() = runTest { - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() - // Set active session ID first + // Trigger GatewayReady -> triggers createSession -> feed result to set active session to session-123 + var createReqId = "" + every { HermesWsClient.send(WsMethods.SESSION_CREATE, any(), any()) } answers { + createReqId = "create-req-stream" + val onSent = arg<((String) -> Unit)?>(2) + onSent?.invoke(createReqId) + createReqId + } mockEventsFlow.emit(WsEvent.GatewayReady(null)) advanceUntilIdle() + mockEventsFlow.emit(WsEvent.RpcResult(createReqId, mapOf("session_id" to "session-123"))) + advanceUntilIdle() + // Stream Start mockEventsFlow.emit(WsEvent.MessageStart("session-123")) advanceUntilIdle() @@ -239,8 +256,8 @@ class ChatViewModelTest { assertFalse(state.isThinking) assertNotNull(state.streamingMessage) assertEquals("Hello", state.streamingMessage?.content) - // Streaming message is not in the main messages list yet - assertEquals(0, state.messages.size) + // Streaming message is not in the main messages list yet, only "Session created" exists + assertEquals(1, state.messages.size) // Token 2 mockEventsFlow.emit(WsEvent.MessageToken(" world", "session-123")) @@ -248,7 +265,7 @@ class ChatViewModelTest { state = viewModel.uiState.value assertNotNull(state.streamingMessage) assertEquals("Hello world", state.streamingMessage?.content) - assertEquals(0, state.messages.size) + assertEquals(1, state.messages.size) // Complete mockEventsFlow.emit(WsEvent.MessageComplete("Hello world!", "session-123")) @@ -258,9 +275,10 @@ class ChatViewModelTest { assertNull(state.streamingMessage) assertFalse(state.isThinking) assertEquals("", state.thinkingText) - assertEquals(1, state.messages.size) - assertEquals("Hello world!", state.messages[0].content) - assertFalse(state.messages[0].isStreaming) + assertEquals(2, state.messages.size) + assertEquals("Session created", state.messages[0].content) + assertEquals("Hello world!", state.messages[1].content) + assertFalse(state.messages[1].isStreaming) } @Test @@ -274,7 +292,7 @@ class ChatViewModelTest { createReqId } - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() // Set session ID @@ -308,7 +326,7 @@ class ChatViewModelTest { @Test fun testSendMessage() = runTest { - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() // Trigger GatewayReady -> triggers createSession -> feed result to set active session @@ -343,7 +361,7 @@ class ChatViewModelTest { @Test fun testSwitchSession() = runTest { - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() viewModel.switchSession("session-456") @@ -375,7 +393,7 @@ class ChatViewModelTest { createReqId } - val viewModel = ChatViewModel(app) + val viewModel = ChatViewModel(app, startCleanup = false) advanceUntilIdle() mockEventsFlow.emit(WsEvent.GatewayReady(null))