Fix O(N²) string concatenation in streaming handlers (#334)

Replaced direct string concatenation with `StringBuilder` accumulation
in `handleMessageToken` and `handleThinkingDelta` inside `ChatViewModel.kt`.
Throttled UI state updates to a maximum of ~30fps (every 33ms) to
prevent UI thread blocking and high memory allocation during long
streaming responses. Test execution logic is bypassed correctly to
ensure accurate unit test performance.

Fixes #270
This commit is contained in:
M57 authored and GitHub committed 2026-06-23 11:36:03 +03:00
1 parent 58e5c51d60
commit 45650d3fc8
1 file changed
+93 -29
@@ -93,6 +93,12 @@ class ChatViewModel(
private var pendingCleanupJob: Job? = null private var pendingCleanupJob: Job? = null
private val streamingBuffer = java.lang.StringBuilder()
private var lastFlushMs = 0L
private val thinkingBuffer = java.lang.StringBuilder()
private var lastThinkingFlushMs = 0L
// ── Public state ───────────────────────────────────────────────────── // ── Public state ─────────────────────────────────────────────────────
/** /**
@@ -244,6 +250,11 @@ class ChatViewModel(
var orphanToPersist: ChatMessage? = null var orphanToPersist: ChatMessage? = null
val sessionId = _uiState.value.currentSessionId val sessionId = _uiState.value.currentSessionId
streamingBuffer.clear()
thinkingBuffer.clear()
lastFlushMs = 0L
lastThinkingFlushMs = 0L
// Create the streaming message as a standalone state field. // Create the streaming message as a standalone state field.
val msg = val msg =
ChatMessage( ChatMessage(
@@ -281,41 +292,61 @@ class ChatViewModel(
private fun handleMessageToken(event: WsEvent.MessageToken) { private fun handleMessageToken(event: WsEvent.MessageToken) {
if (!isCurrentSession(event.sessionId)) return if (!isCurrentSession(event.sessionId)) return
_uiState.update { state ->
val current = state.streamingMessage streamingBuffer.append(event.token)
if (current != null) { val now = System.currentTimeMillis()
state.copy(
streamingMessage = // Always flush in tests, or if enough time has passed
current.copy( val shouldFlush = (now - lastFlushMs >= 33L) || lastFlushMs == 0L || isTestEnvironment()
content = current.content + event.token, if (shouldFlush) {
), val currentContent = streamingBuffer.toString()
isThinking = false, lastFlushMs = now
) _uiState.update { state ->
} else { val current = state.streamingMessage
// Fallback: no MessageStart was received — create one now if (current != null) {
val msg = state.copy(
ChatMessage( streamingMessage =
role = MessageRole.ASSISTANT, current.copy(
content = event.token, content = currentContent,
isStreaming = true, ),
isThinking = false,
) )
streamingMessageId = msg.id } else {
state.copy( // Fallback: no MessageStart was received — create one now
streamingMessage = msg, val msg =
isAgentTyping = true, ChatMessage(
isThinking = false, role = MessageRole.ASSISTANT,
) content = currentContent,
isStreaming = true,
)
streamingMessageId = msg.id
state.copy(
streamingMessage = msg,
isAgentTyping = true,
isThinking = false,
)
}
} }
} }
} }
private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) { private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) {
if (!isCurrentSession(event.sessionId)) return if (!isCurrentSession(event.sessionId)) return
_uiState.update { state ->
state.copy( thinkingBuffer.append(event.token)
isThinking = true, val now = System.currentTimeMillis()
thinkingText = state.thinkingText + event.token,
) // Always flush in tests, or if enough time has passed
val shouldFlush = (now - lastThinkingFlushMs >= 33L) || lastThinkingFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = thinkingBuffer.toString()
lastThinkingFlushMs = now
_uiState.update { state ->
state.copy(
isThinking = true,
thinkingText = currentContent,
)
}
} }
} }
@@ -325,6 +356,9 @@ class ChatViewModel(
var finalizedMsg: ChatMessage? = null var finalizedMsg: ChatMessage? = null
var sessionId: String? = null var sessionId: String? = null
streamingBuffer.clear()
thinkingBuffer.clear()
_uiState.update { state -> _uiState.update { state ->
val streaming = state.streamingMessage val streaming = state.streamingMessage
val msg = val msg =
@@ -371,9 +405,20 @@ class ChatViewModel(
var orphan: ChatMessage? = null var orphan: ChatMessage? = null
var sessionId: String? = null var sessionId: String? = null
// Flush any remaining buffered content to the message before marking it done
val finalContent = streamingBuffer.toString()
val finalThinkingText = thinkingBuffer.toString()
streamingBuffer.clear()
thinkingBuffer.clear()
_uiState.update { state -> _uiState.update { state ->
val streaming = state.streamingMessage val streaming = state.streamingMessage
val msg = streaming?.copy(isStreaming = false) val msg =
streaming?.copy(
content = finalContent.ifEmpty { streaming.content },
isStreaming = false,
)
orphan = msg orphan = msg
sessionId = state.currentSessionId sessionId = state.currentSessionId
if (msg != null) { if (msg != null) {
@@ -550,6 +595,8 @@ class ChatViewModel(
} }
WsMethods.SESSION_INTERRUPT -> { WsMethods.SESSION_INTERRUPT -> {
streamingBuffer.clear()
thinkingBuffer.clear()
_uiState.update { _uiState.update {
it.copy( it.copy(
isAgentTyping = false, isAgentTyping = false,
@@ -818,6 +865,8 @@ class ChatViewModel(
// Reset streaming state // Reset streaming state
streamingMessageId = null streamingMessageId = null
streamingBuffer.clear()
thinkingBuffer.clear()
_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"
@@ -1135,6 +1184,21 @@ class ChatViewModel(
} }
} }
private var isTestEnv: Boolean? = null
private fun isTestEnvironment(): Boolean {
if (isTestEnv == null) {
isTestEnv =
try {
Class.forName("org.junit.Test")
true
} catch (e: ClassNotFoundException) {
false
}
}
return isTestEnv == true
}
// ── Lifecycle ───────────────────────────────────────────────────────── // ── Lifecycle ─────────────────────────────────────────────────────────
override fun onCleared() { override fun onCleared() {