perf: optimize startup, database, and ui rendering (#405)

* perf: optimize startup, database, and ui rendering

* fix(#405): address code review feedback

* test: wait for async prefs initialization to prevent coroutine leaks

* test: stub mockViewModel.streamingState in ChatScreenTest to prevent ClassCastException on Android

app/src/main/java/com/m57/hermescontrol/data/local/AuthManager.kt

@@ -7,9 +7,14 @@ import androidx.security.crypto.MasterKeys
 import com.m57.hermescontrol.theme.BottomNavDisplayMode
 import com.m57.hermescontrol.theme.ThemePreference
 import com.m57.hermescontrol.theme.ThemePreset
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.async
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.runBlocking
 
 /**
  * Singleton that manages encrypted storage of the Hermes dashboard token
@@ -43,7 +48,7 @@ object AuthManager {
     private const val DEFAULT_TYPING_EFFECT_DELAY_MS = 30
 
     @Volatile
-    private var prefs: SharedPreferences? = null
+    private var prefsDeferred: Deferred<SharedPreferences>? = null
 
     private val _bottomNavItemsFlow = MutableStateFlow<List<String>>(emptyList())
     val bottomNavItemsFlow: StateFlow<List<String>> = _bottomNavItemsFlow.asStateFlow()
@@ -70,32 +75,54 @@ object AuthManager {
      * Call this once from Application.onCreate() or MainActivity.onCreate().
      */
     fun init(context: Context) {
-        if (prefs != null) return
+        if (prefsDeferred != null) return
         synchronized(this) {
-            if (prefs != null) return
-            val masterKey =
-                MasterKeys.getOrCreate(MasterKeys.AES256_GCM_SPEC)
+            if (prefsDeferred != null) return
+            prefsDeferred =
+                CoroutineScope(Dispatchers.IO).async {
+                    val masterKey =
+                        MasterKeys.getOrCreate(MasterKeys.AES256_GCM_SPEC)
 
-            prefs =
-                EncryptedSharedPreferences.create(
-                    PREFS_FILE,
-                    masterKey,
-                    context.applicationContext,
-                    EncryptedSharedPreferences.PrefKeyEncryptionScheme.AES256_SIV,
-                    EncryptedSharedPreferences.PrefValueEncryptionScheme.AES256_GCM,
-                )
+                    val p =
+                        EncryptedSharedPreferences.create(
+                            PREFS_FILE,
+                            masterKey,
+                            context.applicationContext,
+                            EncryptedSharedPreferences.PrefKeyEncryptionScheme.AES256_SIV,
+                            EncryptedSharedPreferences.PrefValueEncryptionScheme.AES256_GCM,
+                        )
 
-            _bottomNavItemsFlow.value = getBottomNavItems()
-            _themePreferenceFlow.value = getThemePreference()
-            _useDynamicColorsFlow.value = isUseDynamicColors()
-            _themePresetFlow.value = getThemePreset()
-            _bottomNavDisplayModeFlow.value = getBottomNavDisplayMode()
-            _tokenFlow.value = getToken()
+                    // Initialize flows with values from prefs once loaded
+                    _bottomNavItemsFlow.value = getBottomNavItemsInternal(p)
+                    _themePreferenceFlow.value = getThemePreferenceInternal(p)
+                    _useDynamicColorsFlow.value = isUseDynamicColorsInternal(p)
+                    _themePresetFlow.value = getThemePresetInternal(p)
+                    _bottomNavDisplayModeFlow.value = getBottomNavDisplayModeInternal(p)
+                    _tokenFlow.value = getTokenInternal(p)
+                    p
+                }
         }
     }
 
+    /**
+     * Retrieves the initialized [SharedPreferences] instance.
+     *
+     * WARNING: This method is synchronous and will block the caller thread (using [runBlocking])
+     * if the asynchronous initialization is still in progress. Callers should avoid invoking this
+     * on the main thread during early startup to prevent frame drops or potential ANRs.
+     *
+     * Times out and throws [IllegalStateException] if initialization takes longer than 2 seconds.
+     */
     private fun requirePrefs(): SharedPreferences =
-        prefs ?: throw IllegalStateException("AuthManager not initialized. Call init(context) first.")
+        runBlocking {
+            val deferred =
+                prefsDeferred ?: throw IllegalStateException(
+                    "AuthManager not initialized. Call init(context) first.",
+                )
+            kotlinx.coroutines.withTimeoutOrNull(2000) {
+                deferred.await()
+            } ?: throw IllegalStateException("AuthManager initialization timed out after 2 seconds.")
+        }
 
     fun setWsAuthParam(param: String) {
         requirePrefs().edit().putString(KEY_WS_AUTH_PARAM, param).apply()
@@ -209,6 +236,15 @@ object AuthManager {
         }
     }
 
+    private fun getTokenInternal(p: SharedPreferences): String? {
+        val selectedId = p.getString(KEY_SELECTED_PROFILE_ID, null)?.takeIf { it.isNotBlank() }
+        return if (selectedId != null) {
+            p.getString("token_$selectedId", null) ?: p.getString(KEY_TOKEN, null)
+        } else {
+            p.getString(KEY_TOKEN, null)
+        }
+    }
+
     fun setToken(token: String?) {
         val selectedId = getSelectedProfileId()
         if (selectedId != null) {
@@ -289,8 +325,10 @@ object AuthManager {
     // start reset the theme to SYSTEM. Fix: round-trip via EncryptedSharedPreferences
     // (the same store already used by token / host / port / auto_reconnect).
 
-    fun getThemePreference(): ThemePreference =
-        requirePrefs().getString(KEY_THEME_PREFERENCE, ThemePreference.SYSTEM.name)?.let { name ->
+    fun getThemePreference(): ThemePreference = getThemePreferenceInternal(requirePrefs())
+
+    private fun getThemePreferenceInternal(p: SharedPreferences): ThemePreference =
+        p.getString(KEY_THEME_PREFERENCE, ThemePreference.SYSTEM.name)?.let { name ->
             runCatching { ThemePreference.valueOf(name) }.getOrNull()
         } ?: ThemePreference.SYSTEM
 
@@ -299,15 +337,19 @@ object AuthManager {
         _themePreferenceFlow.value = theme
     }
 
-    fun isUseDynamicColors(): Boolean = requirePrefs().getBoolean(KEY_USE_DYNAMIC_COLORS, true)
+    fun isUseDynamicColors(): Boolean = isUseDynamicColorsInternal(requirePrefs())
+
+    private fun isUseDynamicColorsInternal(p: SharedPreferences): Boolean = p.getBoolean(KEY_USE_DYNAMIC_COLORS, true)
 
     fun setUseDynamicColors(value: Boolean) {
         requirePrefs().edit().putBoolean(KEY_USE_DYNAMIC_COLORS, value).apply()
         _useDynamicColorsFlow.value = value
     }
 
-    fun getThemePreset(): ThemePreset =
-        requirePrefs().getString(KEY_THEME_PRESET, ThemePreset.DEFAULT.name)?.let { name ->
+    fun getThemePreset(): ThemePreset = getThemePresetInternal(requirePrefs())
+
+    private fun getThemePresetInternal(p: SharedPreferences): ThemePreset =
+        p.getString(KEY_THEME_PRESET, ThemePreset.DEFAULT.name)?.let { name ->
             runCatching { ThemePreset.valueOf(name) }.getOrNull()
         } ?: ThemePreset.DEFAULT
 
@@ -337,8 +379,10 @@ object AuthManager {
         listOf("ChatScreen", "SkillsScreen", "CronJobsScreen", "SystemScreen", "SettingsScreen")
 
     /** Returns the list of selected bottom-nav item keys (data-object names). */
-    fun getBottomNavItems(): List<String> {
-        val raw = requirePrefs().getString(KEY_BOTTOM_NAV_ITEMS, null) ?: return DEFAULT_BOTTOM_NAV_ITEMS
+    fun getBottomNavItems(): List<String> = getBottomNavItemsInternal(requirePrefs())
+
+    private fun getBottomNavItemsInternal(p: SharedPreferences): List<String> {
+        val raw = p.getString(KEY_BOTTOM_NAV_ITEMS, null) ?: return DEFAULT_BOTTOM_NAV_ITEMS
         return raw.split(",").filter { it.isNotBlank() }
     }
 
@@ -366,8 +410,10 @@ object AuthManager {
 
     // ── Bottom Nav Display Mode ──────────────────────────────────────────
 
-    fun getBottomNavDisplayMode(): BottomNavDisplayMode =
-        requirePrefs().getString(KEY_BOTTOM_NAV_DISPLAY_MODE, BottomNavDisplayMode.ICON_AND_TEXT.name)?.let { name ->
+    fun getBottomNavDisplayMode(): BottomNavDisplayMode = getBottomNavDisplayModeInternal(requirePrefs())
+
+    private fun getBottomNavDisplayModeInternal(p: SharedPreferences): BottomNavDisplayMode =
+        p.getString(KEY_BOTTOM_NAV_DISPLAY_MODE, BottomNavDisplayMode.ICON_AND_TEXT.name)?.let { name ->
             runCatching { BottomNavDisplayMode.valueOf(name) }.getOrNull()
         } ?: BottomNavDisplayMode.ICON_AND_TEXT
 

app/src/main/java/com/m57/hermescontrol/data/local/ChatMessageEntity.kt

@@ -2,6 +2,7 @@ package com.m57.hermescontrol.data.local
 
 import androidx.room.ColumnInfo
 import androidx.room.Entity
+import androidx.room.Index
 import androidx.room.PrimaryKey
 
 /**
@@ -12,7 +13,12 @@ import androidx.room.PrimaryKey
  * thread. The [timestamp] field preserves original ordering even if Room
  * reorders internally.
  */
-@Entity(tableName = "chat_messages")
+@Entity(
+    tableName = "chat_messages",
+    indices = [
+        Index(value = ["session_id", "timestamp"]),
+    ],
+)
 data class ChatMessageEntity(
     @PrimaryKey
     val id: String,

app/src/main/java/com/m57/hermescontrol/data/local/HermesDatabase.kt

@@ -4,12 +4,14 @@ import android.content.Context
 import androidx.room.Database
 import androidx.room.Room
 import androidx.room.RoomDatabase
+import androidx.room.migration.Migration
+import androidx.sqlite.db.SupportSQLiteDatabase
 import net.zetetic.database.sqlcipher.SupportOpenHelperFactory
 import java.io.File
 
 @Database(
     entities = [ChatMessageEntity::class],
-    version = 2,
+    version = 3,
     exportSchema = true,
 )
 abstract class HermesDatabase : RoomDatabase() {
@@ -33,20 +35,31 @@ abstract class HermesDatabase : RoomDatabase() {
                 System.loadLibrary("sqlcipher")
                 val factory = SupportOpenHelperFactory(AuthManager.getDatabasePassword())
 
+                val migration2to3 =
+                    object : Migration(2, 3) {
+                        override fun migrate(db: SupportSQLiteDatabase) {
+                            db.execSQL(
+                                "CREATE INDEX IF NOT EXISTS `index_chat_messages_session_id_timestamp` " +
+                                    "ON `chat_messages` (`session_id`, `timestamp`)",
+                            )
+                        }
+                    }
+
                 instance ?: Room
                     .databaseBuilder(
                         context.applicationContext,
                         HermesDatabase::class.java,
                         "hermes_control.db",
                     ).openHelperFactory(factory)
+                    .addMigrations(migration2to3)
                     .fallbackToDestructiveMigration(false)
                     .build()
                     .also { instance = it }
             }
 
         /** Returns true if the database file starts with the SQLCipher magic header. */
-        private fun isSqlCipherDatabase(file: File): Boolean {
-            return try {
+        private fun isSqlCipherDatabase(file: File): Boolean =
+            try {
                 val header = ByteArray(16)
                 file.inputStream().use { it.read(header) }
                 // SQLCipher 4.x databases start with bytes that differ from
@@ -56,7 +69,6 @@ abstract class HermesDatabase : RoomDatabase() {
             } catch (_: Exception) {
                 false // if we can't read it, treat as plaintext and delete
             }
-        }
 
         /** For testing — inject a custom instance. */
         fun setForTest(db: HermesDatabase?) {

app/src/main/java/com/m57/hermescontrol/ui/chat/ChatScreen.kt

@@ -123,6 +123,7 @@ fun ChatScreen(
     viewModel: ChatViewModel = viewModel(),
 ) {
     val state by viewModel.uiState.collectAsStateWithLifecycle()
+    val streamingState by viewModel.streamingState.collectAsStateWithLifecycle()
     val listState = rememberLazyListState()
     val showScrollToBottom by remember {
         derivedStateOf {
@@ -142,8 +143,8 @@ fun ChatScreen(
         connectionStatus = state.connectionStatus,
         currentSessionId = state.currentSessionId,
         messages = state.messages,
-        streamingMessage = state.streamingMessage,
-        isThinking = state.isThinking,
+        streamingMessage = streamingState.streamingMessage,
+        isThinking = streamingState.isThinking,
         errorMessage = state.errorMessage,
         isSearchActive = state.isSearchActive,
         currentSearchMatchIndex = state.currentSearchMatchIndex,
@@ -246,9 +247,9 @@ fun ChatScreen(
             ) {
                 ChatMessageList(
                     messages = state.messages,
-                    streamingMessage = state.streamingMessage,
-                    isThinking = state.isThinking,
-                    thinkingText = state.thinkingText,
+                    streamingMessage = streamingState.streamingMessage,
+                    isThinking = streamingState.isThinking,
+                    thinkingText = streamingState.thinkingText,
                     isSearchActive = state.isSearchActive,
                     searchQuery = state.searchQuery,
                     currentSearchMatchIndex = state.currentSearchMatchIndex,
@@ -272,8 +273,8 @@ fun ChatScreen(
                     scrollScope = scrollScope,
                     listState = listState,
                     messages = state.messages,
-                    streamingMessage = state.streamingMessage,
-                    isThinking = state.isThinking,
+                    streamingMessage = streamingState.streamingMessage,
+                    isThinking = streamingState.isThinking,
                 )
             }
 
@@ -286,8 +287,8 @@ fun ChatScreen(
                     scrollScope.launch {
                         val totalItems =
                             state.messages.size +
-                                (if (state.streamingMessage != null) 1 else 0) +
-                                (if (state.isThinking) 1 else 0)
+                                (if (streamingState.streamingMessage != null) 1 else 0) +
+                                (if (streamingState.isThinking) 1 else 0)
                         if (totalItems > 0) {
                             listState.animateScrollToItem(totalItems - 1)
                         }
@@ -557,8 +558,7 @@ private fun ClarifyDialog(
                                             role = Role.Button,
                                             onClick = { typedText = option },
                                             onLongClick = { onOptionSelected(option) },
-                                        )
-                                        .padding(horizontal = 24.dp, vertical = 10.dp),
+                                        ).padding(horizontal = 24.dp, vertical = 10.dp),
                                 contentAlignment = Alignment.Center,
                             ) {
                                 Text(
@@ -865,10 +865,12 @@ private fun ChatLifecycleEffects(
                         viewModel.refreshSettings()
                         viewModel.refreshCurrentSession()
                     }
+
                     androidx.lifecycle.Lifecycle.Event.ON_STOP -> {
                         NotificationHelper.setAppForeground(context, false)
                         NotificationHelper.start(context)
                     }
+
                     else -> {}
                 }
             }

app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt

@@ -21,6 +21,7 @@ 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
@@ -87,6 +88,9 @@ class ChatViewModel(
     // ── Internal state ───────────────────────────────────────────────────
     private val _uiState = MutableStateFlow(ChatUiState())
 
+    private val _streamingState = MutableStateFlow(StreamingState())
+    val streamingState: StateFlow<StreamingState> = _streamingState.asStateFlow()
+
     private val wsClient = HermesWsClient
     private val pendingRequests = ConcurrentHashMap<String, PendingRequest>()
     private val repo =
@@ -165,10 +169,17 @@ class ChatViewModel(
 
     private fun handleWsEvent(event: WsEvent) {
         // First, let the reducer compute the new state and any effects
-        val result = ChatWsEventReducer.reduce(_uiState.value, event, _uiState.value.currentSessionId)
+        val result =
+            ChatWsEventReducer.reduce(
+                _uiState.value,
+                _streamingState.value,
+                event,
+                _uiState.value.currentSessionId,
+            )
 
         // Apply the new state
         _uiState.update { result.state }
+        _streamingState.update { result.streamingState }
 
         // Process side-effects from the reducer
         for (effect in result.effects) {
@@ -267,6 +278,7 @@ class ChatViewModel(
                         isAgentTyping = false,
                     )
                 }
+                _streamingState.update { StreamingState() }
             }
 
             is WsEvent.ApprovalRequest -> {
@@ -300,7 +312,7 @@ class ChatViewModel(
         if (shouldFlush) {
             val currentContent = streamingBuffer.toString()
             lastFlushMs = now
-            _uiState.update { state ->
+            _streamingState.update { state ->
                 val current = state.streamingMessage
                 if (current != null) {
                     state.copy(
@@ -321,11 +333,11 @@ class ChatViewModel(
                     streamingMessageId = msg.id
                     state.copy(
                         streamingMessage = msg,
-                        isAgentTyping = true,
                         isThinking = false,
                     )
                 }
             }
+            _uiState.update { it.copy(isAgentTyping = true) }
         }
     }
 
@@ -340,7 +352,7 @@ class ChatViewModel(
         if (shouldFlush) {
             val currentContent = thinkingBuffer.toString()
             lastThinkingFlushMs = now
-            _uiState.update { state ->
+            _streamingState.update { state ->
                 state.copy(
                     isThinking = true,
                     thinkingText = currentContent,
@@ -368,10 +380,10 @@ class ChatViewModel(
                         currentSessionId = sessionId,
                         isLoading = false,
                         messages = emptyList(),
-                        streamingMessage = null,
                         chatTitle = "Hermes",
                     )
                 }
+                _streamingState.update { StreamingState() }
                 addSystemMessage("Session created", persist = true)
                 loadSessions()
             }
@@ -422,10 +434,9 @@ class ChatViewModel(
                 _uiState.update {
                     it.copy(
                         isAgentTyping = false,
-                        isThinking = false,
-                        streamingMessage = null,
                     )
                 }
+                _streamingState.update { StreamingState() }
                 streamingMessageId = null
                 addSystemMessage("Session interrupted")
             }
@@ -728,10 +739,10 @@ class ChatViewModel(
             it.copy(
                 isLoading = true,
                 messages = emptyList(),
-                streamingMessage = null,
                 chatTitle = "Hermes",
             )
         }
+        _streamingState.update { StreamingState() }
         streamingMessageId = null
         viewModelScope.launch(Dispatchers.IO) {
             wsClient.send(
@@ -777,8 +788,8 @@ class ChatViewModel(
     @Suppress("UNCHECKED_CAST")
     private fun parseCommandCatalog(map: Map<*, *>): CommandCatalog? =
         try {
-            val json = gson.toJson(map)
-            gson.fromJson(json, CommandCatalog::class.java)
+            val json = OkHttpProvider.gson.toJson(map)
+            OkHttpProvider.gson.fromJson(json, CommandCatalog::class.java)
         } catch (e: Exception) {
             Log.e(TAG, "Failed to parse command catalog", e)
             null
@@ -791,21 +802,18 @@ class ChatViewModel(
         streamingMessageId = null
         streamingBuffer.clear()
         thinkingBuffer.clear()
-
         _uiState.update {
             val title = it.sessions.find { s -> s.id == sessionId }?.title ?: "Hermes"
             it.copy(
                 isLoading = true,
-                messages = emptyList(),
-                streamingMessage = null,
                 currentSessionId = sessionId,
+                messages = emptyList(),
                 chatTitle = title,
                 showSessionPicker = false,
                 isAgentTyping = false,
-                isThinking = false,
-                thinkingText = "",
             )
         }
+        _streamingState.update { StreamingState() }
         // Step 1: Load cached messages first — instant display
         loadCachedMessages(sessionId)
         // Step 2: Resume session on server
@@ -1184,6 +1192,5 @@ class ChatViewModel(
     }
 
     companion object {
-        private val gson = OkHttpProvider.gson
     }
 }

app/src/main/java/com/m57/hermescontrol/ui/chat/ChatWsEventReducer.kt

@@ -1,5 +1,6 @@
 package com.m57.hermescontrol.ui.chat
 
+import com.m57.hermescontrol.data.remote.OkHttpProvider
 import com.m57.hermescontrol.data.ws.WsEvent
 
 /**
@@ -15,6 +16,7 @@ import com.m57.hermescontrol.data.ws.WsEvent
 object ChatWsEventReducer {
     fun reduce(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent,
         currentSessionId: String? = null,
     ): ReducerResult {
@@ -31,60 +33,65 @@ object ChatWsEventReducer {
                 else -> null
             }
         if (eventSessionId != null && currentSessionId != null && eventSessionId != currentSessionId) {
-            return ReducerResult(state)
+            return ReducerResult(state = state, streamingState = streamingState)
         }
-        return reduceInternal(state, event)
+        return reduceInternal(state, streamingState, event)
     }
 
     private fun reduceInternal(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent,
     ): ReducerResult =
         when (event) {
-            is WsEvent.GatewayReady -> onGatewayReady(state)
+            is WsEvent.GatewayReady -> onGatewayReady(state, streamingState)
 
-            is WsEvent.MessageStart -> onMessageStart(state, event)
+            is WsEvent.MessageStart -> onMessageStart(state, streamingState, event)
 
-            is WsEvent.MessageToken -> onMessageToken(state, event)
+            is WsEvent.MessageToken -> onMessageToken(state, streamingState, event)
 
-            is WsEvent.ThinkingDelta -> onThinkingDelta(state, event)
+            is WsEvent.ThinkingDelta -> onThinkingDelta(state, streamingState, event)
 
-            is WsEvent.MessageComplete -> onMessageComplete(state, event)
+            is WsEvent.MessageComplete -> onMessageComplete(state, streamingState, event)
 
-            is WsEvent.MessageDone -> onMessageDone(state)
+            is WsEvent.MessageDone -> onMessageDone(state, streamingState)
 
-            is WsEvent.ToolStart -> onToolStart(state, event)
+            is WsEvent.ToolStart -> onToolStart(state, streamingState, event)
 
-            is WsEvent.ToolComplete -> onToolComplete(state, event)
+            is WsEvent.ToolComplete -> onToolComplete(state, streamingState, event)
 
-            is WsEvent.ClarifyRequest -> onClarifyRequest(state, event)
+            is WsEvent.ClarifyRequest -> onClarifyRequest(state, streamingState, event)
 
-            is WsEvent.RpcError -> onRpcError(state, event)
+            is WsEvent.RpcError -> onRpcError(state, streamingState, event)
 
-            is WsEvent.SessionUpdated -> onSessionUpdated(state)
+            is WsEvent.SessionUpdated -> onSessionUpdated(state, streamingState)
 
-            is WsEvent.StatusUpdate -> onStatusUpdate(state)
+            is WsEvent.StatusUpdate -> onStatusUpdate(state, streamingState)
 
-            is WsEvent.Unknown -> onUnknown(state)
+            is WsEvent.Unknown -> onUnknown(state, streamingState)
 
             // SessionInfo is a no-op in the original code
-            is WsEvent.SessionInfo -> ReducerResult(state)
+            is WsEvent.SessionInfo -> ReducerResult(state = state, streamingState = streamingState)
 
             // RpcResult is handled by the ViewModel (needs pending request context)
-            is WsEvent.RpcResult -> ReducerResult(state)
+            is WsEvent.RpcResult -> ReducerResult(state = state, streamingState = streamingState)
 
             // ApprovalRequest is handled by the ViewModel (needs active session + WS client)
-            is WsEvent.ApprovalRequest -> ReducerResult(state)
+            is WsEvent.ApprovalRequest -> ReducerResult(state = state, streamingState = streamingState)
         }
 
     // ── GatewayReady ──────────────────────────────────────────────────
 
-    private fun onGatewayReady(state: ChatUiState): ReducerResult = ReducerResult(state)
+    private fun onGatewayReady(
+        state: ChatUiState,
+        streamingState: StreamingState,
+    ): ReducerResult = ReducerResult(state = state, streamingState = streamingState)
 
     // ── MessageStart ──────────────────────────────────────────────────
 
     private fun onMessageStart(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.MessageStart,
     ): ReducerResult {
         val msg =
@@ -98,53 +105,52 @@ object ChatWsEventReducer {
         // Build new state: finalize any orphan streaming message, then set the new one
         var orphan: ChatMessage? = null
         val preState =
-            if (state.streamingMessage?.content?.isNotEmpty() == true) {
-                val finalized = state.streamingMessage!!.copy(isStreaming = false)
+            if (streamingState.streamingMessage?.content?.isNotEmpty() == true) {
+                val finalized = streamingState.streamingMessage.copy(isStreaming = false)
                 orphan = finalized
                 state.copy(
                     messages = state.messages + finalized,
-                    streamingMessage = null,
                     isAgentTyping = true,
-                    isThinking = false,
-                    thinkingText = "",
                 )
             } else {
                 state.copy(
                     isAgentTyping = true,
-                    isThinking = false,
-                    thinkingText = "",
                 )
             }
-        val newState = preState.copy(streamingMessage = msg)
+        val newStreamingState = StreamingState(streamingMessage = msg)
+        val newState = preState
         val sid = newState.currentSessionId
         if (orphan != null && sid != null) {
             effects.add(ReducerEffect.PersistMessage(orphan, sid))
         }
 
-        return ReducerResult(newState, effects)
+        return ReducerResult(state = newState, streamingState = newStreamingState, effects = effects)
     }
 
     // ── MessageToken ──────────────────────────────────────────────────
 
     private fun onMessageToken(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.MessageToken,
     ): ReducerResult {
         // Token accumulation is done by the ViewModel (streaming buffer + flush timer).
         // The reducer only signals that new content arrived — ViewModel owns the timer.
-        return ReducerResult(state)
+        return ReducerResult(state = state, streamingState = streamingState)
     }
 
     // ── ThinkingDelta ─────────────────────────────────────────────────
 
     private fun onThinkingDelta(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.ThinkingDelta,
     ): ReducerResult {
-        val currentContent = state.thinkingText + event.token
+        val currentContent = streamingState.thinkingText + event.token
         return ReducerResult(
-            state =
-                state.copy(
+            state = state,
+            streamingState =
+                streamingState.copy(
                     isThinking = true,
                     thinkingText = currentContent,
                 ),
@@ -155,9 +161,10 @@ object ChatWsEventReducer {
 
     private fun onMessageComplete(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.MessageComplete,
     ): ReducerResult {
-        val streaming = state.streamingMessage
+        val streaming = streamingState.streamingMessage
         val msg =
             streaming?.copy(
                 content = event.text ?: streaming.content,
@@ -176,10 +183,7 @@ object ChatWsEventReducer {
             state =
                 state.copy(
                     messages = state.messages + msg,
-                    streamingMessage = null,
                     isAgentTyping = false,
-                    isThinking = false,
-                    thinkingText = "",
                 ),
             effects = effects,
         )
@@ -187,8 +191,15 @@ object ChatWsEventReducer {
 
     // ── MessageDone ───────────────────────────────────────────────────
 
-    private fun onMessageDone(state: ChatUiState): ReducerResult {
-        val streaming = state.streamingMessage ?: return ReducerResult(state)
+    private fun onMessageDone(
+        state: ChatUiState,
+        streamingState: StreamingState,
+    ): ReducerResult {
+        val streaming =
+            streamingState.streamingMessage ?: return ReducerResult(
+                state = state,
+                streamingState = streamingState,
+            )
         val msg = streaming.copy(isStreaming = false)
         val effects = mutableListOf<ReducerEffect>()
         val sid = state.currentSessionId
@@ -199,10 +210,7 @@ object ChatWsEventReducer {
             state =
                 state.copy(
                     messages = state.messages + msg,
-                    streamingMessage = null,
                     isAgentTyping = false,
-                    isThinking = false,
-                    thinkingText = "",
                 ),
             effects = effects,
         )
@@ -212,13 +220,12 @@ object ChatWsEventReducer {
 
     private fun onToolStart(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.ToolStart,
     ): ReducerResult {
         val contentJson =
             event.data?.let {
-                com.m57.hermescontrol.data.remote.OkHttpProvider
-                    .gson
-                    .toJson(it)
+                OkHttpProvider.gson.toJson(it)
             } ?: ""
         val toolMessage =
             ChatMessage(
@@ -231,15 +238,14 @@ object ChatWsEventReducer {
         // Finalize any orphan streaming message
         var orphanToPersist: ChatMessage? = null
         val newState =
-            if (state.streamingMessage?.content?.isNotEmpty() == true) {
-                val finalized = state.streamingMessage!!.copy(isStreaming = false)
+            if (streamingState.streamingMessage?.content?.isNotEmpty() == true) {
+                val finalized = streamingState.streamingMessage.copy(isStreaming = false)
                 orphanToPersist = finalized
                 state.copy(
                     messages = state.messages + finalized,
-                    streamingMessage = null,
                 )
             } else {
-                state.copy(streamingMessage = null)
+                state
             }
         val sid = newState.currentSessionId
         val effects = mutableListOf<ReducerEffect>()
@@ -260,13 +266,12 @@ object ChatWsEventReducer {
 
     private fun onToolComplete(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.ToolComplete,
     ): ReducerResult {
         val contentJson =
             event.data?.let {
-                com.m57.hermescontrol.data.remote.OkHttpProvider
-                    .gson
-                    .toJson(it)
+                OkHttpProvider.gson.toJson(it)
             } ?: ""
         val messages = state.messages.toMutableList()
         val toolIdx =
@@ -275,7 +280,7 @@ object ChatWsEventReducer {
                     it.toolName == event.name &&
                     it.toolStatus == ToolStatus.RUNNING
             }
-        if (toolIdx < 0) return ReducerResult(state)
+        if (toolIdx < 0) return ReducerResult(state = state, streamingState = streamingState)
 
         val updated =
             messages[toolIdx].copy(
@@ -291,6 +296,7 @@ object ChatWsEventReducer {
         }
         return ReducerResult(
             state = state.copy(messages = messages),
+            streamingState = streamingState,
             effects = effects,
         )
     }
@@ -299,6 +305,7 @@ object ChatWsEventReducer {
 
     private fun onClarifyRequest(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.ClarifyRequest,
     ): ReducerResult =
         ReducerResult(
@@ -318,6 +325,7 @@ object ChatWsEventReducer {
 
     private fun onRpcError(
         state: ChatUiState,
+        streamingState: StreamingState,
         event: WsEvent.RpcError,
     ): ReducerResult =
         ReducerResult(
@@ -326,52 +334,37 @@ object ChatWsEventReducer {
                     isLoading = false,
                     errorMessage = "Error: ${event.error}",
                 ),
+            streamingState = streamingState,
         )
 
     // ── SessionUpdated / StatusUpdate / Unknown ───────────────────────
 
-    private fun onSessionUpdated(state: ChatUiState): ReducerResult = ReducerResult(state)
-
-    private fun onStatusUpdate(state: ChatUiState): ReducerResult = ReducerResult(state)
-
-    private fun onUnknown(state: ChatUiState): ReducerResult = ReducerResult(state)
-
-    // ── Helpers ───────────────────────────────────────────────────────
-
-    /**
-     * If there's a non-empty streaming message, finalize it into the message
-     * list and return it as an orphan that should be persisted.
-     */
-    private fun finalizeStreamingMessage(
+    private fun onSessionUpdated(
         state: ChatUiState,
-        newMessage: ChatMessage,
-    ): Pair<ChatUiState, ChatMessage?> {
-        val existing = state.streamingMessage
-        if (existing != null && existing.content.isNotEmpty()) {
-            val finalized = existing.copy(isStreaming = false)
-            return Pair(
-                state.copy(
-                    messages = state.messages + finalized,
-                    streamingMessage = newMessage,
-                ),
-                finalized,
-            )
-        }
-        return Pair(
-            state.copy(streamingMessage = newMessage),
-            null,
-        )
-    }
+        streamingState: StreamingState,
+    ): ReducerResult = ReducerResult(state = state, streamingState = streamingState)
+
+    private fun onStatusUpdate(
+        state: ChatUiState,
+        streamingState: StreamingState,
+    ): ReducerResult = ReducerResult(state = state, streamingState = streamingState)
+
+    private fun onUnknown(
+        state: ChatUiState,
+        streamingState: StreamingState,
+    ): ReducerResult = ReducerResult(state = state, streamingState = streamingState)
 }
 
 /**
  * Result of reducing one event.
  *
  * @param state The new UI state after applying the event.
+ * @param streamingState The new streaming state after applying the event.
  * @param effects Side-effects the ViewModel should execute.
  */
 data class ReducerResult(
     val state: ChatUiState,
+    val streamingState: StreamingState = StreamingState(),
     val effects: List<ReducerEffect> = emptyList(),
 )
 

app/src/main/java/com/m57/hermescontrol/ui/chat/StreamingState.kt

@@ -0,0 +1,7 @@
+package com.m57.hermescontrol.ui.chat
+
+data class StreamingState(
+    val streamingMessage: ChatMessage? = null,
+    val isThinking: Boolean = false,
+    val thinkingText: String = "",
+)