From f746c603b0b7c78c55a8db9c30faca2677eee053 Mon Sep 17 00:00:00 2001 From: Bruno Deanoz Date: Wed, 24 Jun 2026 10:15:24 +0200 Subject: [PATCH] Isolate streaming recomposition to active bubble only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit During streaming, content updates now flow through a StateFlow (liveStream) instead of mutating the SnapshotStateList on every chunk. Only the single MessageRow with id == streamingMessageId subscribes to the flow via collectAsState(). All other items in the LazyColumn stay completely static during streaming. Before: every chunk mutated messages[idx] → entire LazyColumn evaluated After: only the live bubble recomposes, list mutation happens once at stream end Auto-scroll split into two LaunchedEffects: one for list size changes, one for liveStream emissions. --- .../java/dev/kaizen/app/chat/ChatScreen.kt | 21 ++++- .../kaizen/app/chat/ChatScreenViewModel.kt | 84 +++++++++++-------- 2 files changed, 66 insertions(+), 39 deletions(-) diff --git a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt index da66542..afaf5e7 100644 --- a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt +++ b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt @@ -166,11 +166,13 @@ fun ChatScreen( // ── Auto-scroll ───────────────────────────────────────────────────── - val autoScrollKey by remember { - derivedStateOf { vm.messages.size to (vm.messages.lastOrNull()?.content?.length ?: 0) } + LaunchedEffect(Unit) { + snapshotFlow { vm.messages.size }.collect { + if (vm.messages.isNotEmpty()) listState.scrollToItem(vm.messages.lastIndex) + } } LaunchedEffect(Unit) { - snapshotFlow { autoScrollKey }.collect { + vm.liveStream.collect { if (vm.messages.isNotEmpty()) listState.scrollToItem(vm.messages.lastIndex) } } @@ -248,8 +250,19 @@ fun ChatScreen( placementSpec = spring(stiffness = Spring.StiffnessMediumLow), fadeOutSpec = tween(200), )) { + val isLive = message.id == vm.streamingMessageId + val displayMessage = if (isLive) { + val stream by vm.liveStream.collectAsState() + message.copy( + content = stream.content, + reasoning = stream.reasoning, + tools = stream.tools, + query = stream.query, + sources = stream.sources, + ) + } else message MessageRow( - message = message, + message = displayMessage, onDelete = if (vm.conversationId != null && !vm.isStreaming) { wireId -> vm.deleteMessage(wireId) } else null, diff --git a/app/src/main/java/dev/kaizen/app/chat/ChatScreenViewModel.kt b/app/src/main/java/dev/kaizen/app/chat/ChatScreenViewModel.kt index c7ce05f..5f32c62 100644 --- a/app/src/main/java/dev/kaizen/app/chat/ChatScreenViewModel.kt +++ b/app/src/main/java/dev/kaizen/app/chat/ChatScreenViewModel.kt @@ -38,6 +38,9 @@ import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.async import kotlinx.coroutines.delay +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.StateFlow +import kotlinx.coroutines.flow.asStateFlow import kotlinx.coroutines.launch import kotlinx.coroutines.withContext import java.io.File @@ -60,6 +63,10 @@ class ChatScreenViewModel( var input by mutableStateOf("") private var nextId = 0L private var streamJob: Job? = null + var streamingMessageId by mutableStateOf(null) + private set + private val _liveStream = MutableStateFlow(dev.kaizen.app.net.StreamState()) + val liveStream: StateFlow = _liveStream.asStateFlow() // ── Conversation ──────────────────────────────────────────────────────── @@ -230,7 +237,9 @@ class ChatScreenViewModel( val assistantId = nextId++ val assistantWireId = java.util.UUID.randomUUID().toString() val history = messages.map { WireMessage(if (it.role == Role.User) "user" else "assistant", it.content) } + _liveStream.value = dev.kaizen.app.net.StreamState() messages.add(Message(assistantId, Role.Assistant, "", streaming = true, thinking = true, wireId = assistantWireId)) + streamingMessageId = assistantId val assistantIdx = messages.lastIndex val wasNewConversation = conversationId == null @@ -253,28 +262,30 @@ class ChatScreenViewModel( topP = samplingPrefs.topP.takeIf { it.enabled }?.value, topK = samplingPrefs.topK.takeIf { it.enabled }?.value?.toInt(), ).collect { state -> - if (assistantIdx >= messages.size) return@collect if (!sawContent && state.content.isNotEmpty()) { sawContent = true haptics.responseStart() } - messages[assistantIdx] = messages[assistantIdx].copy( - thinking = if (sawContent) false else messages[assistantIdx].thinking, - content = state.content, - reasoning = state.reasoning, - tools = state.tools, - query = state.query, - sources = state.sources, - ) + _liveStream.value = state + if (sawContent && assistantIdx < messages.size && messages[assistantIdx].thinking) { + messages[assistantIdx] = messages[assistantIdx].copy(thinking = false) + } if (state.inputTokens > 0 || state.outputTokens > 0) { usedTokens = state.inputTokens + state.outputTokens } } - val finalMsg = if (assistantIdx < messages.size) messages[assistantIdx] else null - val finalContent = finalMsg?.content ?: "" - val finalReasoning = finalMsg?.reasoning?.takeIf { it.isNotBlank() } - val finalSources = finalMsg?.sources?.takeIf { it.isNotEmpty() } - if (assistantIdx < messages.size) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false) + val lastState = _liveStream.value + val finalContent = lastState.content + val finalReasoning = lastState.reasoning.takeIf { it.isNotBlank() } + val finalSources = lastState.sources.takeIf { it.isNotEmpty() } + if (assistantIdx < messages.size) { + messages[assistantIdx] = messages[assistantIdx].copy( + streaming = false, thinking = false, + content = finalContent, reasoning = lastState.reasoning, + tools = lastState.tools, query = lastState.query, sources = lastState.sources, + ) + } + streamingMessageId = null val convId = convIdDeferred?.await()?.also { cid -> conversationId = cid @@ -345,21 +356,21 @@ class ChatScreenViewModel( } } catch (e: Exception) { val valid = assistantIdx < messages.size + val liveContent = _liveStream.value.content if (e is kotlinx.coroutines.CancellationException) { - if (valid) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false) + if (valid) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false, content = liveContent, reasoning = _liveStream.value.reasoning, tools = _liveStream.value.tools, sources = _liveStream.value.sources) } else if (e is ChatHttpException && e.code == 401) { - if (valid) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false) + if (valid) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false, content = liveContent) session.logout() } else if (valid) { - val existing = messages[assistantIdx].content val errorSuffix = "\n\n⚠ ${chatErrorText(e)}" messages[assistantIdx] = messages[assistantIdx].copy( - streaming = false, - thinking = false, - content = if (existing.isNotBlank()) existing + errorSuffix else chatErrorText(e), + streaming = false, thinking = false, + content = if (liveContent.isNotBlank()) liveContent + errorSuffix else chatErrorText(e), ) } } finally { + streamingMessageId = null isStreaming = false streamJob = null haptics.responseEnd() @@ -777,7 +788,9 @@ class ChatScreenViewModel( val assistantId = nextId++ val assistantWireId = java.util.UUID.randomUUID().toString() val history = messages.map { WireMessage(if (it.role == Role.User) "user" else "assistant", it.content) } + _liveStream.value = dev.kaizen.app.net.StreamState() messages.add(Message(assistantId, Role.Assistant, "", streaming = true, thinking = true, wireId = assistantWireId)) + streamingMessageId = assistantId val voiceAssistantIdx = messages.lastIndex isStreaming = true haptics.click() @@ -794,31 +807,32 @@ class ChatScreenViewModel( cfg.baseUrl, cfg.token, cfg.model, history, attachments = listOf(att), ).collect { state -> - if (voiceAssistantIdx >= messages.size) return@collect if (!sawContent && state.content.isNotEmpty()) { sawContent = true; haptics.responseStart() } - messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy( - thinking = if (sawContent) false else messages[voiceAssistantIdx].thinking, - content = state.content, - reasoning = state.reasoning, - tools = state.tools, - sources = state.sources, - query = state.query, - streaming = true, - ) + _liveStream.value = state + if (sawContent && voiceAssistantIdx < messages.size && messages[voiceAssistantIdx].thinking) { + messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy(thinking = false) + } } } catch (_: kotlinx.coroutines.CancellationException) { } catch (_: Exception) {} - if (voiceAssistantIdx < messages.size) messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy(streaming = false, thinking = false) + val lastState = _liveStream.value + if (voiceAssistantIdx < messages.size) { + messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy( + streaming = false, thinking = false, + content = lastState.content, reasoning = lastState.reasoning, + tools = lastState.tools, sources = lastState.sources, query = lastState.query, + ) + } + streamingMessageId = null isStreaming = false haptics.responseEnd() val cid = convIdDeferred?.await() ?: conversationId ?: return@launch if (conversationId == null) conversationId = cid - val assistantMsg = if (voiceAssistantIdx < messages.size) messages[voiceAssistantIdx] else return@launch - val finalContent = assistantMsg.content - val finalReasoning = assistantMsg.reasoning.takeIf { it.isNotBlank() } - val finalSources = assistantMsg.sources.takeIf { it.isNotEmpty() } + val finalContent = lastState.content + val finalReasoning = lastState.reasoning.takeIf { it.isNotBlank() } + val finalSources = lastState.sources.takeIf { it.isNotEmpty() } val userSave = SaveMessage(id = userMsg.wireId, parentId = userParentId, role = "user", content = "", attachments = listOf(att)) val assistantSave = SaveMessage(id = assistantWireId, parentId = userMsg.wireId, role = "assistant", content = finalContent, model = cfg.model, reasoning = finalReasoning, sources = finalSources?.map { SearchSource(it.title, it.url, it.snippet) }) val saved = KaizenApi.saveMessages(cfg.baseUrl, cfg.token, cid, listOf(userSave, assistantSave))