diff --git a/.changeset/retained-jetbrains-markdown.md b/.changeset/retained-jetbrains-markdown.md index e5e6cb0b0c7..b01def7f0c5 100644 --- a/.changeset/retained-jetbrains-markdown.md +++ b/.changeset/retained-jetbrains-markdown.md @@ -2,4 +2,4 @@ "kilo-code": patch --- -Improve JetBrains chat streaming performance by retaining existing markdown and code block views while responses stream, and render code blocks without showing raw fence markers during streamed updates. +Improve JetBrains chat streaming performance by retaining existing markdown and code block views while responses stream, and keep streamed code fences intact without showing raw fence markers during updates. diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt index 4eccd0b4c74..c3349ce37bc 100644 --- a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt @@ -10,6 +10,7 @@ import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ensureActive import kotlinx.coroutines.Job +import kotlinx.coroutines.channels.Channel import kotlinx.coroutines.delay import kotlinx.coroutines.flow.MutableSharedFlow import kotlinx.coroutines.flow.MutableStateFlow @@ -88,6 +89,11 @@ class KiloConnectionService( private val _events = MutableSharedFlow(extraBufferCapacity = 64) val events: SharedFlow = _events.asSharedFlow() + private val queue = Channel(Channel.UNLIMITED) + private val lock = Any() + private val eventJob = cs.launch { + for (event in queue) _events.emit(event) + } /** Generated API client — null when disconnected. */ var api: DefaultApi? = null @@ -232,25 +238,31 @@ class KiloConnectionService( } private val listener = object : EventSourceListener() { - override fun onOpen(src: EventSource, response: Response) { + override fun onOpen(eventSource: EventSource, response: Response) { log.info("SSE: connected") setState(ConnectionState.Connected(port, password)) lastEvent.set(System.currentTimeMillis()) } - override fun onEvent(src: EventSource, id: String?, type: String?, data: String) { - lastEvent.set(System.currentTimeMillis()) - val kind = type ?: KiloCliDataParser.extractEventType(data) - log.debug { "evt=$kind bytes=${data.length} hasId=${id != null} ${ChatLogSummary.body(data)}" } - cs.launch { _events.emit(SseEvent(type = kind, data = data)) } + override fun onEvent(eventSource: EventSource, id: String?, type: String?, data: String) { + synchronized(lock) { + if (disposed) return@synchronized + lastEvent.set(System.currentTimeMillis()) + val kind = type ?: KiloCliDataParser.extractEventType(data) + log.debug { "evt=$kind bytes=${data.length} hasId=${id != null} ${ChatLogSummary.body(data)}" } + val result = queue.trySend(SseEvent(type = kind, data = data)) + if (result.isFailure && !disposed) { + log.warn("SSE: event queue rejected type=$kind", result.exceptionOrNull()) + } + } } - override fun onClosed(src: EventSource) { + override fun onClosed(eventSource: EventSource) { log.info("SSE: stream closed — scheduling reconnect") scheduleReconnect() } - override fun onFailure(src: EventSource, t: Throwable?, response: Response?) { + override fun onFailure(eventSource: EventSource, t: Throwable?, response: Response?) { val detail = when { t != null -> t.stackTraceToString() response != null -> response.body?.string() @@ -369,6 +381,8 @@ class KiloConnectionService( healthJob?.cancel() processJob?.cancel() reconnectJob?.cancel() + eventJob.cancel() + queue.close() close() _state.value = ConnectionState.Disconnected log.info("KiloConnectionService disposed") diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt index de84bfced22..3b62046dd36 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt @@ -6,15 +6,24 @@ import ai.kilocode.backend.app.KiloConnectionService import ai.kilocode.backend.testing.FakeCliServer import ai.kilocode.backend.testing.MockCliServer import ai.kilocode.backend.testing.TestLog +import ai.kilocode.log.KiloLog import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.async import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.flow.first +import kotlinx.coroutines.flow.take +import kotlinx.coroutines.flow.toList import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeout +import okhttp3.Request +import okhttp3.sse.EventSource +import okhttp3.sse.EventSourceListener +import java.util.concurrent.CountDownLatch +import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger import kotlin.test.AfterTest import kotlin.test.Test @@ -93,6 +102,61 @@ class KiloConnectionServiceTest { } } + @Test + fun `SSE events preserve callback order`() = runBlocking { + val svc = KiloConnectionService(scope, fake, {}, log) + svc.connect() + mock.awaitSseConnection() + + withTimeout(5_000) { + svc.state.first { it is ConnectionState.Connected } + } + + val count = 100 + val received = async { + withTimeout(5_000) { + svc.events.take(count).toList() + } + } + + delay(200) + repeat(count) { idx -> + mock.pushEvent("test.event", idx.toString()) + } + + assertEquals((0 until count).map { it.toString() }, received.await().map { it.data }) + } + + @Test + fun `SSE concurrent callbacks preserve callback order`() = runBlocking { + val blocked = BlockingLog() + val svc = KiloConnectionService(scope, fake, {}, blocked) + val field = KiloConnectionService::class.java.getDeclaredField("listener") + field.isAccessible = true + val listener = field.get(svc) as EventSourceListener + val source = object : EventSource { + override fun request(): Request = Request.Builder().url("http://127.0.0.1/global/event").build() + override fun cancel() {} + } + val received = scope.async { + withTimeout(5_000) { + svc.events.take(2).toList() + } + } + + delay(200) + val first = Thread { listener.onEvent(source, null, "first.event", "first") } + val second = Thread { listener.onEvent(source, null, "second.event", "second") } + first.start() + assertTrue(blocked.started.await(1, TimeUnit.SECONDS)) + second.start() + first.join() + second.join() + + assertEquals(listOf("first", "second"), received.await().map { it.data }) + svc.dispose() + } + @Test fun `SSE close triggers error state`() = runBlocking { val svc = KiloConnectionService(scope, fake, {}, log) @@ -208,4 +272,20 @@ class KiloConnectionServiceTest { svc.state.first { it !is ConnectionState.Connected } } } + + private class BlockingLog : KiloLog { + val started = CountDownLatch(1) + override var isDebugEnabled: Boolean = true + + override fun debug(block: () -> String) { + val msg = block() + if (!msg.contains("evt=first.event")) return + started.countDown() + Thread.sleep(250) + } + + override fun info(msg: String) {} + override fun warn(msg: String, t: Throwable?) {} + override fun error(msg: String, t: Throwable?) {} + } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt index 29adf9e05f6..a1c2076218e 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt @@ -145,6 +145,9 @@ class SessionController( private var prefModel: String? = null private var prefAgent: String? = null private var modelTime: Double? = null + private val snapshots = mutableMapOf() + + private data class PartKey(val messageId: String, val partId: String) val ready: Boolean get() = model.isReady() val autoApprove: Boolean get() = KiloPluginSettings.getAutoApprove() @@ -653,6 +656,7 @@ class SessionController( if (disposed) return@runEdt if (sid != id) return@runEdt updateModel { + snapshots.clear() this@SessionController.model.loadHistory(items) syncHistoryAgent(items) if (session != null) this@SessionController.model.setSession(session) @@ -701,6 +705,7 @@ class SessionController( ref = SessionRef.Local(session) setRecentSessionsState(RecentsState.Idle) updateModel { + snapshots.clear() this@SessionController.model.loadHistory(items) syncHistoryAgent(items) this@SessionController.model.setSession(session) @@ -890,7 +895,15 @@ class SessionController( is ChatEventDto.PartUpdated -> { partType = event.part.type tool = event.part.tool + val key = PartKey(event.part.messageID, event.part.id) + val prev = content(event.part.messageID, event.part.id) model.updateContent(event.part.messageID, event.part) + val next = content(event.part.messageID, event.part.id) + if (next != null && next != prev) { + snapshots[key] = next + } else { + snapshots.remove(key) + } if (model.state is SessionState.Busy) { model.setState(SessionState.Busy(status())) } @@ -905,6 +918,7 @@ class SessionController( } is ChatEventDto.PartRemoved -> { + snapshots.remove(PartKey(event.messageID, event.partID)) model.removeContent(event.messageID, event.partID) } @@ -938,6 +952,7 @@ class SessionController( } is ChatEventDto.MessageRemoved -> { + snapshots.keys.removeAll { it.messageId == event.messageID } model.removeMessage(event.messageID) } @@ -982,16 +997,24 @@ class SessionController( private fun glue(messageId: String, partId: String, delta: String): String { if (delta.isEmpty()) return delta - val cur = when (val content = model.content(messageId, partId)) { - is Text -> content.content - is Reasoning -> content.content - else -> return delta - } + val key = PartKey(messageId, partId) + val cur = snapshots[key] ?: return delta val span = (minOf(cur.length, delta.length) downTo 1) .firstOrNull { n -> cur.regionMatches(cur.length - n, delta, 0, n) } ?: 0 + if (span == delta.length) { + snapshots.remove(key) + return "" + } + snapshots.remove(key) return delta.substring(span) } + private fun content(messageId: String, partId: String): String? = when (val content = model.content(messageId, partId)) { + is Text -> content.content.toString() + is Reasoning -> content.content.toString() + else -> null + } + private fun handleHidden(event: ChatEventDto): Boolean = when (event) { is ChatEventDto.Error, is ChatEventDto.PermissionAsked, diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt index 11d4e92548b..1f08b695574 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt @@ -257,6 +257,54 @@ class SessionUpdateQueueTest : SessionControllerTestBase() { ) } + fun `test pure text deltas preserve incidental overlap`() { + appRpc.state.value = ai.kilocode.rpc.dto.KiloAppStateDto(ai.kilocode.rpc.dto.KiloAppStatusDto.READY) + projectRpc.state.value = workspaceReady() + val m = controller("ses_test", flushMs = Long.MAX_VALUE) + flush() + + emit(ChatEventDto.MessageUpdated("ses_test", msg("msg1", "ses_test", "assistant"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "hel")) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "lo")) + settle() + flush() + + assertModel( + """ + assistant#msg1 + text#prt1: + hello + """, + m, + ) + } + + fun `test pure text deltas preserve split closing fence`() { + appRpc.state.value = ai.kilocode.rpc.dto.KiloAppStateDto(ai.kilocode.rpc.dto.KiloAppStatusDto.READY) + projectRpc.state.value = workspaceReady() + val m = controller("ses_test", flushMs = Long.MAX_VALUE) + flush() + + emit(ChatEventDto.MessageUpdated("ses_test", msg("msg1", "ses_test", "assistant"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "```python\nprint(1)\n``")) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "`\n\nafter")) + settle() + flush() + + assertModel( + """ + assistant#msg1 + text#prt1: + ```python + print(1) + ``` + + after + """, + m, + ) + } + fun `test text snapshot covered prefix is trimmed from merged delta`() { appRpc.state.value = ai.kilocode.rpc.dto.KiloAppStateDto(ai.kilocode.rpc.dto.KiloAppStatusDto.READY) projectRpc.state.value = workspaceReady()