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
+68 -4
@@ -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,13 +292,22 @@ class ChatViewModel(
private fun handleMessageToken(event: WsEvent.MessageToken) { private fun handleMessageToken(event: WsEvent.MessageToken) {
if (!isCurrentSession(event.sessionId)) return if (!isCurrentSession(event.sessionId)) return
streamingBuffer.append(event.token)
val now = System.currentTimeMillis()
// Always flush in tests, or if enough time has passed
val shouldFlush = (now - lastFlushMs >= 33L) || lastFlushMs == 0L || isTestEnvironment()
if (shouldFlush) {
val currentContent = streamingBuffer.toString()
lastFlushMs = now
_uiState.update { state -> _uiState.update { state ->
val current = state.streamingMessage val current = state.streamingMessage
if (current != null) { if (current != null) {
state.copy( state.copy(
streamingMessage = streamingMessage =
current.copy( current.copy(
content = current.content + event.token, content = currentContent,
), ),
isThinking = false, isThinking = false,
) )
@@ -296,7 +316,7 @@ class ChatViewModel(
val msg = val msg =
ChatMessage( ChatMessage(
role = MessageRole.ASSISTANT, role = MessageRole.ASSISTANT,
content = event.token, content = currentContent,
isStreaming = true, isStreaming = true,
) )
streamingMessageId = msg.id streamingMessageId = msg.id
@@ -308,16 +328,27 @@ class ChatViewModel(
} }
} }
} }
}
private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) { private fun handleThinkingDelta(event: WsEvent.ThinkingDelta) {
if (!isCurrentSession(event.sessionId)) return if (!isCurrentSession(event.sessionId)) return
thinkingBuffer.append(event.token)
val now = System.currentTimeMillis()
// 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 -> _uiState.update { state ->
state.copy( state.copy(
isThinking = true, isThinking = true,
thinkingText = state.thinkingText + event.token, thinkingText = currentContent,
) )
} }
} }
}
private fun handleMessageComplete(event: WsEvent.MessageComplete) { private fun handleMessageComplete(event: WsEvent.MessageComplete) {
if (!isCurrentSession(event.sessionId)) return if (!isCurrentSession(event.sessionId)) return
@@ -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() {