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.
This commit is contained in:
Hy4ri committed 2026-06-20 11:04:56 +03:00
1 parent a2d299e136
commit d0eedf053f
4 files changed
+462 -215

No files matched your search

@@ -72,10 +72,10 @@ import com.m57.hermescontrol.theme.SystemMessageColor
import com.m57.hermescontrol.theme.ToolChipColor import com.m57.hermescontrol.theme.ToolChipColor
import com.m57.hermescontrol.theme.ToolChipColorLight import com.m57.hermescontrol.theme.ToolChipColorLight
import com.m57.hermescontrol.theme.UserBubble import com.m57.hermescontrol.theme.UserBubble
import java.time.Instant
import java.time.ZoneId
import java.time.format.DateTimeFormatter
import kotlinx.coroutines.delay import kotlinx.coroutines.delay
import java.text.SimpleDateFormat
import java.util.Date
import java.util.Locale
@Composable @Composable
fun ChatBubble( 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 { private fun formatTimestamp(timestamp: Long): String {
val sdf = SimpleDateFormat("HH:mm", Locale.getDefault()) return TIMESTAMP_FORMATTER.format(Instant.ofEpochMilli(timestamp))
return sdf.format(Date(timestamp))
} }
/** /**
@@ -92,6 +92,7 @@ import androidx.compose.ui.unit.sp
import androidx.core.content.ContextCompat import androidx.core.content.ContextCompat
import androidx.lifecycle.compose.collectAsStateWithLifecycle import androidx.lifecycle.compose.collectAsStateWithLifecycle
import androidx.lifecycle.viewmodel.compose.viewModel import androidx.lifecycle.viewmodel.compose.viewModel
import com.m57.hermescontrol.data.ws.ConnectionStatus
import com.m57.hermescontrol.notification.NotificationHelper import com.m57.hermescontrol.notification.NotificationHelper
import com.m57.hermescontrol.theme.StatusRed import com.m57.hermescontrol.theme.StatusRed
import com.m57.hermescontrol.ui.common.EmptyState import com.m57.hermescontrol.ui.common.EmptyState
@@ -169,9 +170,12 @@ fun ChatScreen(
} }
// Auto-scroll to bottom on new messages // Auto-scroll to bottom on new messages
LaunchedEffect(state.messages.size, state.currentStreamingText) { LaunchedEffect(state.messages.size, state.streamingMessage?.content?.length, state.isThinking) {
if (state.messages.isNotEmpty()) { val totalItems = state.messages.size +
listState.animateScrollToItem(state.messages.lastIndex) (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) .background(backgroundGradient)
.imePadding(), .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( androidx.compose.animation.AnimatedVisibility(
visible = state.isSearchActive, visible = state.isSearchActive,
enter = androidx.compose.animation.expandVertically() + androidx.compose.animation.fadeIn(), 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 // Thinking indicator
if (state.isThinking) { if (state.isThinking) {
item(key = "thinking") { item(key = "thinking") {
@@ -1,6 +1,7 @@
package com.m57.hermescontrol.ui.chat package com.m57.hermescontrol.ui.chat
import android.app.Application import android.app.Application
import android.util.Log
import androidx.lifecycle.AndroidViewModel import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope import androidx.lifecycle.viewModelScope
import com.m57.hermescontrol.data.local.AuthManager 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.ApiClient
import com.m57.hermescontrol.data.remote.NetworkResult import com.m57.hermescontrol.data.remote.NetworkResult
import com.m57.hermescontrol.data.remote.safeApiCall 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.HermesWsClient
import com.m57.hermescontrol.data.ws.WsEvent import com.m57.hermescontrol.data.ws.WsEvent
import com.m57.hermescontrol.data.ws.WsMethods import com.m57.hermescontrol.data.ws.WsMethods
import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.StateFlow 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.flow.update
import kotlinx.coroutines.launch import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext import kotlinx.coroutines.withContext
import java.util.concurrent.ConcurrentHashMap 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( data class ChatUiState(
val messages: List<ChatMessage> = emptyList(), val messages: List<ChatMessage> = emptyList(),
val currentSessionId: String? = null, val currentSessionId: String? = null,
val sessions: List<SessionUi> = emptyList(), val sessions: List<SessionUi> = emptyList(),
val chatTitle: String = "Hermes", val chatTitle: String = "Hermes",
val isConnected: Boolean = false, val connectionStatus: ConnectionStatus = ConnectionStatus.DISCONNECTED,
val isAgentTyping: Boolean = false, val isAgentTyping: Boolean = false,
val isThinking: Boolean = false, val isThinking: Boolean = false,
val thinkingText: String = "", val thinkingText: String = "",
val isLoading: Boolean = false, val isLoading: Boolean = false,
val currentStreamingText: String = "", /** Standalone streaming message — rendered after the main list. */
val streamingMessage: ChatMessage? = null,
val errorMessage: String? = null, val errorMessage: String? = null,
val clarifyRequest: ClarifyUi? = null, val clarifyRequest: ClarifyUi? = null,
val showSessionPicker: Boolean = false, val showSessionPicker: Boolean = false,
@@ -42,7 +56,10 @@ data class ChatUiState(
val searchQuery: String = "", val searchQuery: String = "",
val searchMatchIndices: List<Int> = emptyList(), val searchMatchIndices: List<Int> = emptyList(),
val currentSearchMatchIndex: Int = -1, val currentSearchMatchIndex: Int = -1,
) ) {
/** Convenience — derived from [connectionStatus]. */
val isConnected: Boolean get() = connectionStatus == ConnectionStatus.CONNECTED
}
data class SessionUi( data class SessionUi(
val id: String, val id: String,
@@ -58,17 +75,44 @@ data class ClarifyUi(
class ChatViewModel( class ChatViewModel(
application: Application, application: Application,
) : AndroidViewModel(application) { ) : AndroidViewModel(application) {
// ── Internal state ───────────────────────────────────────────────────
private val _uiState = MutableStateFlow(ChatUiState()) private val _uiState = MutableStateFlow(ChatUiState())
val uiState: StateFlow<ChatUiState> = _uiState.asStateFlow()
private val wsClient = HermesWsClient 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() 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.WhileSubscribed(5_000),
_uiState.value,
)
init { init {
connectWebSocket() connectWebSocket()
startPendingRequestCleanup()
} }
// ── Connection ───────────────────────────────────────────────────────
private fun connectWebSocket() { private fun connectWebSocket() {
val token = AuthManager.getToken() ?: return val token = AuthManager.getToken() ?: return
@@ -83,21 +127,14 @@ class ChatViewModel(
handleWsEvent(event) 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) { private fun handleWsEvent(event: WsEvent) {
when (event) { when (event) {
is WsEvent.GatewayReady -> { is WsEvent.GatewayReady -> {
_uiState.update { it.copy(isConnected = true, isLoading = false) } _uiState.update { it.copy(isLoading = false) }
addSystemMessage("Connected to Hermes") addSystemMessage("Connected to Hermes")
loadSessions() loadSessions()
if (_uiState.value.currentSessionId == null) { if (_uiState.value.currentSessionId == null) {
@@ -107,152 +144,13 @@ class ChatViewModel(
is WsEvent.SessionInfo -> {} is WsEvent.SessionInfo -> {}
is WsEvent.MessageStart -> { is WsEvent.MessageStart -> handleMessageStart(event)
_uiState.update { is WsEvent.MessageToken -> handleMessageToken(event)
it.copy( is WsEvent.ThinkingDelta -> handleThinkingDelta(event)
isAgentTyping = true, is WsEvent.MessageComplete -> handleMessageComplete(event)
currentStreamingText = "", is WsEvent.MessageDone -> handleMessageDone(event)
isThinking = false, is WsEvent.ToolStart -> handleToolStart(event)
thinkingText = "", is WsEvent.ToolComplete -> handleToolComplete(event)
)
}
}
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.ClarifyRequest -> { is WsEvent.ClarifyRequest -> {
_uiState.update { _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") @Suppress("UNCHECKED_CAST")
private fun handleRpcResult( private fun handleRpcResult(
id: String, id: String,
result: Any?, result: Any?,
) { ) {
val method = pendingRequests.remove(id) ?: return val pending = pendingRequests.remove(id) ?: return
val method = pending.method
when (method) { when (method) {
WsMethods.SESSION_CREATE -> { WsMethods.SESSION_CREATE -> {
@@ -300,6 +403,7 @@ class ChatViewModel(
currentSessionId = sessionId, currentSessionId = sessionId,
isLoading = false, isLoading = false,
messages = emptyList(), messages = emptyList(),
streamingMessage = null,
) )
} }
addSystemMessage("Session created", persist = true) addSystemMessage("Session created", persist = true)
@@ -347,8 +451,13 @@ class ChatViewModel(
WsMethods.SESSION_INTERRUPT -> { WsMethods.SESSION_INTERRUPT -> {
_uiState.update { _uiState.update {
it.copy(isAgentTyping = false, isThinking = false) it.copy(
isAgentTyping = false,
isThinking = false,
streamingMessage = null,
)
} }
streamingMessageId = null
addSystemMessage("Session interrupted") addSystemMessage("Session interrupted")
} }
} }
@@ -358,7 +467,8 @@ class ChatViewModel(
id: String, id: String,
error: Any?, error: Any?,
) { ) {
val method = pendingRequests.remove(id) val pending = pendingRequests.remove(id)
val method = pending?.method
val errorMsg = val errorMsg =
when (error) { when (error) {
is Map<*, *> -> error["message"] as? String ?: error.toString() is Map<*, *> -> error["message"] as? String ?: error.toString()
@@ -372,6 +482,8 @@ class ChatViewModel(
} }
} }
// ── Send message ─────────────────────────────────────────────────────
fun sendMessage(text: String) { fun sendMessage(text: String) {
if (text.isBlank()) return if (text.isBlank()) return
val sessionId = _uiState.value.currentSessionId ?: return val sessionId = _uiState.value.currentSessionId ?: return
@@ -382,27 +494,29 @@ class ChatViewModel(
return return
} }
val userMessage = val userMessage = ChatMessage(
ChatMessage( role = MessageRole.USER,
role = MessageRole.USER, content = text,
content = text, )
)
// Update UI — no DB writes inside update{}
_uiState.update { state -> _uiState.update { state ->
val messages = state.messages + userMessage
viewModelScope.launch(Dispatchers.IO) {
dao.upsert(userMessage.toEntity(sessionId))
}
state.copy( state.copy(
messages = messages, messages = state.messages + userMessage,
isAgentTyping = true, isAgentTyping = true,
) )
} }
// Persist to Room — OUTSIDE update{}
viewModelScope.launch(Dispatchers.IO) {
dao.upsert(userMessage.toEntity(sessionId))
}
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
wsClient.sendMessage( wsClient.sendMessage(
sessionId, sessionId,
text, 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) { private fun handleSlashCommand(command: String) {
val userMsg = ChatMessage(role = MessageRole.USER, content = command) val userMsg = ChatMessage(role = MessageRole.USER, content = command)
val sessionId = _uiState.value.currentSessionId val sessionId = _uiState.value.currentSessionId
_uiState.update { it.copy(messages = it.messages + userMsg) } _uiState.update { it.copy(messages = it.messages + userMsg) }
// Persist — OUTSIDE update{}
if (sessionId != null) { if (sessionId != null) {
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
dao.upsert(userMsg.toEntity(sessionId)) dao.upsert(userMsg.toEntity(sessionId))
@@ -538,8 +655,10 @@ class ChatViewModel(
private fun addAssistantMessage(text: String) { private fun addAssistantMessage(text: String) {
val msg = ChatMessage(role = MessageRole.ASSISTANT, content = text) val msg = ChatMessage(role = MessageRole.ASSISTANT, content = text)
val sessionId = _uiState.value.currentSessionId
_uiState.update { it.copy(messages = it.messages + msg) } _uiState.update { it.copy(messages = it.messages + msg) }
// Persist — OUTSIDE update{}
val sessionId = _uiState.value.currentSessionId
if (sessionId != null) { if (sessionId != null) {
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
dao.upsert(msg.toEntity(sessionId)) dao.upsert(msg.toEntity(sessionId))
@@ -547,23 +666,32 @@ class ChatViewModel(
} }
} }
// ── Session management ───────────────────────────────────────────────
fun interruptSession() { fun interruptSession() {
val sessionId = _uiState.value.currentSessionId ?: return val sessionId = _uiState.value.currentSessionId ?: return
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
wsClient.send( wsClient.send(
WsMethods.SESSION_INTERRUPT, WsMethods.SESSION_INTERRUPT,
mapOf("session_id" to sessionId), mapOf("session_id" to sessionId),
onSent = { id -> pendingRequests[id] = WsMethods.SESSION_INTERRUPT }, onSent = { id -> trackRequest(id, WsMethods.SESSION_INTERRUPT) },
) )
} }
} }
fun createNewSession() { 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) { viewModelScope.launch(Dispatchers.IO) {
wsClient.send( wsClient.send(
WsMethods.SESSION_CREATE, 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) { viewModelScope.launch(Dispatchers.IO) {
wsClient.send( wsClient.send(
WsMethods.SESSION_LIST, WsMethods.SESSION_LIST,
onSent = { id -> pendingRequests[id] = WsMethods.SESSION_LIST }, onSent = { id -> trackRequest(id, WsMethods.SESSION_LIST) },
) )
} }
} }
fun switchSession(sessionId: String) { fun switchSession(sessionId: String) {
if (sessionId == _uiState.value.currentSessionId) return if (sessionId == _uiState.value.currentSessionId) return
// Reset streaming state
streamingMessageId = null
_uiState.update { _uiState.update {
val title = it.sessions.find { s -> s.id == sessionId }?.title ?: "Hermes" val title = it.sessions.find { s -> s.id == sessionId }?.title ?: "Hermes"
it.copy( it.copy(
isLoading = true, isLoading = true,
messages = emptyList(), messages = emptyList(),
streamingMessage = null,
currentSessionId = sessionId, currentSessionId = sessionId,
chatTitle = title, chatTitle = title,
showSessionPicker = false, showSessionPicker = false,
@@ -592,16 +725,17 @@ class ChatViewModel(
thinkingText = "", thinkingText = "",
) )
} }
// Load cached messages first — instant display // Step 1: Load cached messages first — instant display
loadCachedMessages(sessionId) loadCachedMessages(sessionId)
// Then resume session on server // Step 2: Resume session on server
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
wsClient.send( wsClient.send(
WsMethods.SESSION_RESUME, WsMethods.SESSION_RESUME,
mapOf("session_id" to sessionId), 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) loadSessionMessages(sessionId)
} }
@@ -610,7 +744,7 @@ class ChatViewModel(
val cached = dao.getMessagesForSession(sessionId) val cached = dao.getMessagesForSession(sessionId)
val cachedMessages = cached.map { it.toUiModel() } val cachedMessages = cached.map { it.toUiModel() }
_uiState.update { state -> _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) { if (state.currentSessionId == sessionId) {
state.copy(messages = cachedMessages, isLoading = false) state.copy(messages = cachedMessages, isLoading = false)
} else { } else {
@@ -630,7 +764,7 @@ class ChatViewModel(
is NetworkResult.Success -> { is NetworkResult.Success -> {
val messagesList = result.data.messages.orEmpty() val messagesList = result.data.messages.orEmpty()
val chatMessages = val chatMessages =
messagesList.map { msg -> messagesList.mapIndexed { index, msg ->
val role = val role =
when (msg.role?.lowercase()) { when (msg.role?.lowercase()) {
"user" -> MessageRole.USER "user" -> MessageRole.USER
@@ -638,28 +772,45 @@ class ChatViewModel(
"tool" -> MessageRole.TOOL "tool" -> MessageRole.TOOL
else -> MessageRole.ASSISTANT 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( ChatMessage(
id = stableId,
role = role, role = role,
content = msg.content.orEmpty(), content = msg.content.orEmpty(),
timestamp = ts,
isStreaming = false, isStreaming = false,
) )
} }
// Persist loaded messages to Room for offline access // Persist loaded messages to Room for offline access
val entities = chatMessages.map { it.toEntity(sessionId) } val entities = chatMessages.map { it.toEntity(sessionId) }
viewModelScope.launch(Dispatchers.IO) { withContext(Dispatchers.IO) {
dao.upsertAll(entities) dao.upsertAll(entities)
} }
_uiState.update { state -> _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( state.copy(
messages = chatMessages + streamingTail, messages = chatMessages + localOnly,
isLoading = false, isLoading = false,
) )
} }
} }
is NetworkResult.Failure -> { is NetworkResult.Failure -> {
// Don't overwrite messages on REST failure — cached messages
// are still valid.
_uiState.update { _uiState.update {
if (it.currentSessionId != sessionId) return@update it
it.copy( it.copy(
isLoading = false, isLoading = false,
errorMessage = "Failed to load messages: ${result.error.message}", errorMessage = "Failed to load messages: ${result.error.message}",
@@ -670,6 +821,8 @@ class ChatViewModel(
} }
} }
// ── UI actions ───────────────────────────────────────────────────────
fun toggleSessionPicker() { fun toggleSessionPicker() {
_uiState.update { it.copy(showSessionPicker = !it.showSessionPicker) } _uiState.update { it.copy(showSessionPicker = !it.showSessionPicker) }
} }
@@ -690,7 +843,6 @@ class ChatViewModel(
fun reconnect() { fun reconnect() {
_uiState.update { _uiState.update {
it.copy( it.copy(
isConnected = false,
isLoading = true, isLoading = true,
errorMessage = null, errorMessage = null,
) )
@@ -709,18 +861,53 @@ class ChatViewModel(
persist: Boolean = false, persist: Boolean = false,
) { ) {
val msg = ChatMessage(role = MessageRole.SYSTEM, content = text) val msg = ChatMessage(role = MessageRole.SYSTEM, content = text)
val sessionId = _uiState.value.currentSessionId
_uiState.update { it.copy(messages = it.messages + msg) } _uiState.update { it.copy(messages = it.messages + msg) }
if (persist) {
val sessionId = _uiState.value.currentSessionId // Persist — OUTSIDE update{}
if (sessionId != null) { if (persist && sessionId != null) {
viewModelScope.launch(Dispatchers.IO) { viewModelScope.launch(Dispatchers.IO) {
dao.upsert(msg.toEntity(sessionId)) 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() { fun toggleSearch() {
val current = _uiState.value val current = _uiState.value
@@ -776,8 +963,11 @@ class ChatViewModel(
} }
} }
// ── Lifecycle ─────────────────────────────────────────────────────────
override fun onCleared() { override fun onCleared() {
super.onCleared() super.onCleared()
pendingCleanupJob?.cancel()
wsClient.disconnect() wsClient.disconnect()
} }
} }
@@ -213,7 +213,8 @@ class ChatViewModelTest {
var state = viewModel.uiState.value var state = viewModel.uiState.value
assertTrue(state.isAgentTyping) assertTrue(state.isAgentTyping)
assertEquals("", state.currentStreamingText) assertNotNull(state.streamingMessage)
assertEquals("", state.streamingMessage?.content)
assertFalse(state.isThinking) assertFalse(state.isThinking)
assertEquals("", state.thinkingText) assertEquals("", state.thinkingText)
@@ -236,27 +237,25 @@ class ChatViewModelTest {
advanceUntilIdle() advanceUntilIdle()
state = viewModel.uiState.value state = viewModel.uiState.value
assertFalse(state.isThinking) assertFalse(state.isThinking)
assertEquals("Hello", state.currentStreamingText) assertNotNull(state.streamingMessage)
// System message was cleared when createNewSession() was called, so only Assistant streaming message is present assertEquals("Hello", state.streamingMessage?.content)
assertEquals(1, state.messages.size) // Streaming message is not in the main messages list yet
assertEquals("Hello", state.messages[0].content) assertEquals(0, state.messages.size)
assertTrue(state.messages[0].isStreaming)
// Token 2 // Token 2
mockEventsFlow.emit(WsEvent.MessageToken(" world", "session-123")) mockEventsFlow.emit(WsEvent.MessageToken(" world", "session-123"))
advanceUntilIdle() advanceUntilIdle()
state = viewModel.uiState.value state = viewModel.uiState.value
assertEquals("Hello world", state.currentStreamingText) assertNotNull(state.streamingMessage)
assertEquals(1, state.messages.size) assertEquals("Hello world", state.streamingMessage?.content)
assertEquals("Hello world", state.messages[0].content) assertEquals(0, state.messages.size)
assertTrue(state.messages[0].isStreaming)
// Complete // Complete
mockEventsFlow.emit(WsEvent.MessageComplete("Hello world!", "session-123")) mockEventsFlow.emit(WsEvent.MessageComplete("Hello world!", "session-123"))
advanceUntilIdle() advanceUntilIdle()
state = viewModel.uiState.value state = viewModel.uiState.value
assertFalse(state.isAgentTyping) assertFalse(state.isAgentTyping)
assertEquals("", state.currentStreamingText) assertNull(state.streamingMessage)
assertFalse(state.isThinking) assertFalse(state.isThinking)
assertEquals("", state.thinkingText) assertEquals("", state.thinkingText)
assertEquals(1, state.messages.size) assertEquals(1, state.messages.size)