krypto/kaizen-app
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Isolate streaming recomposition to active bubble only

Browse source 
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.
This commit is contained in:
Bruno Deanoz 2026-06-24 10:15:24 +02:00
parent 06757d6f69
commit f746c603b0
2 changed files with 66 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
View file

@ -166,11 +166,13 @@ fun ChatScreen(
// ── Auto-scroll ───────────────────────────────────────────────────── // ── Auto-scroll ─────────────────────────────────────────────────────
val autoScrollKey by remember { LaunchedEffect(Unit) {
derivedStateOf { vm.messages.size to (vm.messages.lastOrNull()?.content?.length ?: 0) } snapshotFlow { vm.messages.size }.collect {
if (vm.messages.isNotEmpty()) listState.scrollToItem(vm.messages.lastIndex)
}
} }
LaunchedEffect(Unit) { LaunchedEffect(Unit) {
snapshotFlow { autoScrollKey }.collect { vm.liveStream.collect {
if (vm.messages.isNotEmpty()) listState.scrollToItem(vm.messages.lastIndex) if (vm.messages.isNotEmpty()) listState.scrollToItem(vm.messages.lastIndex)
} }
} }
@ -248,8 +250,19 @@ fun ChatScreen(
placementSpec = spring(stiffness = Spring.StiffnessMediumLow), placementSpec = spring(stiffness = Spring.StiffnessMediumLow),
fadeOutSpec = tween(200), 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( MessageRow(
message = message, message = displayMessage,
onDelete = if (vm.conversationId != null && !vm.isStreaming) { wireId -> onDelete = if (vm.conversationId != null && !vm.isStreaming) { wireId ->
vm.deleteMessage(wireId) vm.deleteMessage(wireId)
} else null, } else null,

84
app/src/main/java/dev/kaizen/app/chat/ChatScreenViewModel.kt
View file

@ -38,6 +38,9 @@ import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job import kotlinx.coroutines.Job
import kotlinx.coroutines.async import kotlinx.coroutines.async
import kotlinx.coroutines.delay 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.launch
import kotlinx.coroutines.withContext import kotlinx.coroutines.withContext
import java.io.File import java.io.File
@ -60,6 +63,10 @@ class ChatScreenViewModel(
var input by mutableStateOf("") var input by mutableStateOf("")
private var nextId = 0L private var nextId = 0L
private var streamJob: Job? = null private var streamJob: Job? = null
var streamingMessageId by mutableStateOf<Long?>(null)
private set
private val _liveStream = MutableStateFlow(dev.kaizen.app.net.StreamState())
val liveStream: StateFlow<dev.kaizen.app.net.StreamState> = _liveStream.asStateFlow()
// ── Conversation ──────────────────────────────────────────────────────── // ── Conversation ────────────────────────────────────────────────────────
@ -230,7 +237,9 @@ class ChatScreenViewModel(
val assistantId = nextId++ val assistantId = nextId++
val assistantWireId = java.util.UUID.randomUUID().toString() val assistantWireId = java.util.UUID.randomUUID().toString()
val history = messages.map { WireMessage(if (it.role == Role.User) "user" else "assistant", it.content) } 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)) messages.add(Message(assistantId, Role.Assistant, "", streaming = true, thinking = true, wireId = assistantWireId))
streamingMessageId = assistantId
val assistantIdx = messages.lastIndex val assistantIdx = messages.lastIndex
val wasNewConversation = conversationId == null val wasNewConversation = conversationId == null
@ -253,28 +262,30 @@ class ChatScreenViewModel(
topP = samplingPrefs.topP.takeIf { it.enabled }?.value, topP = samplingPrefs.topP.takeIf { it.enabled }?.value,
topK = samplingPrefs.topK.takeIf { it.enabled }?.value?.toInt(), topK = samplingPrefs.topK.takeIf { it.enabled }?.value?.toInt(),
).collect { state -> ).collect { state ->
if (assistantIdx >= messages.size) return@collect
if (!sawContent && state.content.isNotEmpty()) { if (!sawContent && state.content.isNotEmpty()) {
sawContent = true sawContent = true
haptics.responseStart() haptics.responseStart()
} }
messages[assistantIdx] = messages[assistantIdx].copy( _liveStream.value = state
thinking = if (sawContent) false else messages[assistantIdx].thinking, if (sawContent && assistantIdx < messages.size && messages[assistantIdx].thinking) {
content = state.content, messages[assistantIdx] = messages[assistantIdx].copy(thinking = false)
reasoning = state.reasoning, }
tools = state.tools,
query = state.query,
sources = state.sources,
)
if (state.inputTokens > 0 || state.outputTokens > 0) { if (state.inputTokens > 0 || state.outputTokens > 0) {
usedTokens = state.inputTokens + state.outputTokens usedTokens = state.inputTokens + state.outputTokens
} }
} }
val finalMsg = if (assistantIdx < messages.size) messages[assistantIdx] else null val lastState = _liveStream.value
val finalContent = finalMsg?.content ?: "" val finalContent = lastState.content
val finalReasoning = finalMsg?.reasoning?.takeIf { it.isNotBlank() } val finalReasoning = lastState.reasoning.takeIf { it.isNotBlank() }
val finalSources = finalMsg?.sources?.takeIf { it.isNotEmpty() } val finalSources = lastState.sources.takeIf { it.isNotEmpty() }
if (assistantIdx < messages.size) messages[assistantIdx] = messages[assistantIdx].copy(streaming = false, thinking = false) 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 -> val convId = convIdDeferred?.await()?.also { cid ->
conversationId = cid conversationId = cid
@ -345,21 +356,21 @@ class ChatScreenViewModel(
} }
} catch (e: Exception) { } catch (e: Exception) {
val valid = assistantIdx < messages.size val valid = assistantIdx < messages.size
val liveContent = _liveStream.value.content
if (e is kotlinx.coroutines.CancellationException) { 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) { } 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() session.logout()
} else if (valid) { } else if (valid) {
val existing = messages[assistantIdx].content
val errorSuffix = "\n\n⚠ ${chatErrorText(e)}" val errorSuffix = "\n\n⚠ ${chatErrorText(e)}"
messages[assistantIdx] = messages[assistantIdx].copy( messages[assistantIdx] = messages[assistantIdx].copy(
streaming = false, streaming = false, thinking = false,
thinking = false, content = if (liveContent.isNotBlank()) liveContent + errorSuffix else chatErrorText(e),
content = if (existing.isNotBlank()) existing + errorSuffix else chatErrorText(e),
) )
} }
} finally { } finally {
streamingMessageId = null
isStreaming = false isStreaming = false
streamJob = null streamJob = null
haptics.responseEnd() haptics.responseEnd()
@ -777,7 +788,9 @@ class ChatScreenViewModel(
val assistantId = nextId++ val assistantId = nextId++
val assistantWireId = java.util.UUID.randomUUID().toString() val assistantWireId = java.util.UUID.randomUUID().toString()
val history = messages.map { WireMessage(if (it.role == Role.User) "user" else "assistant", it.content) } 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)) messages.add(Message(assistantId, Role.Assistant, "", streaming = true, thinking = true, wireId = assistantWireId))
streamingMessageId = assistantId
val voiceAssistantIdx = messages.lastIndex val voiceAssistantIdx = messages.lastIndex
isStreaming = true isStreaming = true
haptics.click() haptics.click()
@ -794,31 +807,32 @@ class ChatScreenViewModel(
cfg.baseUrl, cfg.token, cfg.model, history, cfg.baseUrl, cfg.token, cfg.model, history,
attachments = listOf(att), attachments = listOf(att),
).collect { state -> ).collect { state ->
if (voiceAssistantIdx >= messages.size) return@collect
if (!sawContent && state.content.isNotEmpty()) { if (!sawContent && state.content.isNotEmpty()) {
sawContent = true; haptics.responseStart() sawContent = true; haptics.responseStart()
} }
messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy( _liveStream.value = state
thinking = if (sawContent) false else messages[voiceAssistantIdx].thinking, if (sawContent && voiceAssistantIdx < messages.size && messages[voiceAssistantIdx].thinking) {
content = state.content, messages[voiceAssistantIdx] = messages[voiceAssistantIdx].copy(thinking = false)
reasoning = state.reasoning, }
tools = state.tools,
sources = state.sources,
query = state.query,
streaming = true,
)
} }
} catch (_: kotlinx.coroutines.CancellationException) { } catch (_: kotlinx.coroutines.CancellationException) {
} catch (_: Exception) {} } 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 isStreaming = false
haptics.responseEnd() haptics.responseEnd()
val cid = convIdDeferred?.await() ?: conversationId ?: return@launch val cid = convIdDeferred?.await() ?: conversationId ?: return@launch
if (conversationId == null) conversationId = cid if (conversationId == null) conversationId = cid
val assistantMsg = if (voiceAssistantIdx < messages.size) messages[voiceAssistantIdx] else return@launch val finalContent = lastState.content
val finalContent = assistantMsg.content val finalReasoning = lastState.reasoning.takeIf { it.isNotBlank() }
val finalReasoning = assistantMsg.reasoning.takeIf { it.isNotBlank() } val finalSources = lastState.sources.takeIf { it.isNotEmpty() }
val finalSources = assistantMsg.sources.takeIf { it.isNotEmpty() }
val userSave = SaveMessage(id = userMsg.wireId, parentId = userParentId, role = "user", content = "", attachments = listOf(att)) 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 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)) val saved = KaizenApi.saveMessages(cfg.baseUrl, cfg.token, cid, listOf(userSave, assistantSave))