From d0eedf053ffbbb48198a7f72c23e5028cdd7370b Mon Sep 17 00:00:00 2001 From: Hy4ri Date: Sat, 20 Jun 2026 11:04:56 +0300 Subject: [PATCH] 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)