Merge pull request #151 from Hy4ri/refactor/chat-system-stability

refactor: stabilize chat system streaming, session loading, and connection state
This commit is contained in:
M57 authored and GitHub committed 2026-06-20 12:12:26 +03:00
commit 196883b5c1
4 files changed
+511 -222

No files matched your search

@@ -73,9 +73,9 @@ import com.m57.hermescontrol.theme.ToolChipColor
import com.m57.hermescontrol.theme.ToolChipColorLight
import com.m57.hermescontrol.theme.UserBubble
import kotlinx.coroutines.delay
import java.text.SimpleDateFormat
import java.util.Date
import java.util.Locale
import java.time.Instant
import java.time.ZoneId
import java.time.format.DateTimeFormatter
@Composable
fun ChatBubble(
@@ -811,9 +811,12 @@ 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))
}
/**
@@ -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,13 @@ 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 +354,50 @@ 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 +463,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") {
@@ -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<ChatMessage> = emptyList(),
val currentSessionId: String? = null,
val sessions: List<SessionUi> = 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<Int> = emptyList(),
val currentSearchMatchIndex: Int = -1,
)
) {
/** Convenience — derived from [connectionStatus]. */
val isConnected: Boolean get() = connectionStatus == ConnectionStatus.CONNECTED
}
data class SessionUi(
val id: String,
@@ -57,18 +74,50 @@ 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())
val uiState: StateFlow<ChatUiState> = _uiState.asStateFlow()
private val wsClient = HermesWsClient
private val pendingRequests = ConcurrentHashMap<String, String>() // requestId -> method
private val pendingRequests = ConcurrentHashMap<String, PendingRequest>()
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<ChatUiState> =
combine(
_uiState,
wsClient.connectionStatus,
) { state, connStatus ->
state.copy(connectionStatus = connStatus)
}.stateIn(
viewModelScope,
SharingStarted.Eagerly,
_uiState.value,
)
init {
connectWebSocket()
if (startCleanup) {
startPendingRequestCleanup()
}
}
// ── Connection ───────────────────────────────────────────────────────
private fun connectWebSocket() {
val token = AuthManager.getToken() ?: return
@@ -83,21 +132,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 +149,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 +187,227 @@ 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
var finalizedMsg: ChatMessage? = null
var sessionId: String? = null
_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,
)
}
finalizedMsg = msg
sessionId = state.currentSessionId
state.copy(
messages = state.messages + msg,
streamingMessage = null,
isAgentTyping = false,
isThinking = false,
thinkingText = "",
)
}
streamingMessageId = null
// Persist finalized message — OUTSIDE update{} to avoid CAS retry
val msgToPersist = finalizedMsg
val sid = sessionId ?: _uiState.value.currentSessionId
if (msgToPersist != null && sid != null) {
viewModelScope.launch(Dispatchers.IO) {
dao.upsert(msgToPersist.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.
var orphan: ChatMessage? = null
var sessionId: String? = null
_uiState.update { state ->
val streaming = state.streamingMessage
val msg = streaming?.copy(isStreaming = false)
orphan = msg
sessionId = state.currentSessionId
if (msg != null) {
state.copy(
messages = state.messages + msg,
streamingMessage = null,
isAgentTyping = false,
isThinking = false,
thinkingText = "",
)
} else {
state.copy(
isAgentTyping = false,
isThinking = false,
thinkingText = "",
)
}
}
streamingMessageId = null
// Persist orphaned streaming message
val msgToPersist = orphan
val sid = sessionId ?: _uiState.value.currentSessionId
if (msgToPersist != null && sid != null) {
viewModelScope.launch(Dispatchers.IO) {
dao.upsert(msgToPersist.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 +418,7 @@ class ChatViewModel(
currentSessionId = sessionId,
isLoading = false,
messages = emptyList(),
streamingMessage = null,
)
}
addSystemMessage("Session created", persist = true)
@@ -347,8 +466,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 +482,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 +497,8 @@ class ChatViewModel(
}
}
// ── Send message ─────────────────────────────────────────────────────
fun sendMessage(text: String) {
if (text.isBlank()) return
val sessionId = _uiState.value.currentSessionId ?: return
@@ -387,22 +514,25 @@ class ChatViewModel(
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 +540,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 +671,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 +682,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 +716,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 +741,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 +760,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 +780,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 +788,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 +837,8 @@ class ChatViewModel(
}
}
// ── UI actions ───────────────────────────────────────────────────────
fun toggleSessionPicker() {
_uiState.update { it.copy(showSessionPicker = !it.showSessionPicker) }
}
@@ -690,7 +859,6 @@ class ChatViewModel(
fun reconnect() {
_uiState.update {
it.copy(
isConnected = false,
isLoading = true,
errorMessage = null,
)
@@ -709,18 +877,55 @@ 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))
}
}
}
// ── Search ────────────────────────────────────────────────────────
// ── 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 ────────────────────────────────────────────────────────────
fun toggleSearch() {
val current = _uiState.value
@@ -776,8 +981,11 @@ class ChatViewModel(
}
}
// ── Lifecycle ─────────────────────────────────────────────────────────
override fun onCleared() {
super.onCleared()
pendingCleanupJob?.cancel()
wsClient.disconnect()
}
}
@@ -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<WsEvent>(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,20 +207,31 @@ 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()
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,32 +254,31 @@ 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
assertNotNull(state.streamingMessage)
assertEquals("Hello", state.streamingMessage?.content)
// Streaming message is not in the main messages list yet, only "Session created" exists
assertEquals(1, state.messages.size)
assertEquals("Hello", state.messages[0].content)
assertTrue(state.messages[0].isStreaming)
// Token 2
mockEventsFlow.emit(WsEvent.MessageToken(" world", "session-123"))
advanceUntilIdle()
state = viewModel.uiState.value
assertEquals("Hello world", state.currentStreamingText)
assertNotNull(state.streamingMessage)
assertEquals("Hello world", state.streamingMessage?.content)
assertEquals(1, state.messages.size)
assertEquals("Hello world", state.messages[0].content)
assertTrue(state.messages[0].isStreaming)
// 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)
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
@@ -275,7 +292,7 @@ class ChatViewModelTest {
createReqId
}
val viewModel = ChatViewModel(app)
val viewModel = ChatViewModel(app, startCleanup = false)
advanceUntilIdle()
// Set session ID
@@ -309,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
@@ -344,7 +361,7 @@ class ChatViewModelTest {
@Test
fun testSwitchSession() =
runTest {
val viewModel = ChatViewModel(app)
val viewModel = ChatViewModel(app, startCleanup = false)
advanceUntilIdle()
viewModel.switchSession("session-456")
@@ -376,7 +393,7 @@ class ChatViewModelTest {
createReqId
}
val viewModel = ChatViewModel(app)
val viewModel = ChatViewModel(app, startCleanup = false)
advanceUntilIdle()
mockEventsFlow.emit(WsEvent.GatewayReady(null))