fix(jetbrains): preserve streamed chat markdown

.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.

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<SseEvent>(extraBufferCapacity = 64)
     val events: SharedFlow<SseEvent> = _events.asSharedFlow()
+    private val queue = Channel<SseEvent>(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")

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?) {}
+    }
 }

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<PartKey, String>()
+
+    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,

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()