Isolate streaming recomposition to active bubble only
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
parent
06757d6f69
commit
f746c603b0
2 changed files with 66 additions and 39 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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<Long?>(null)
|
||||
private set
|
||||
private val _liveStream = MutableStateFlow(dev.kaizen.app.net.StreamState())
|
||||
val liveStream: StateFlow<dev.kaizen.app.net.StreamState> = _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))
|
||||
|
|
|
|||
Loading…
Reference in a new issue