diff --git a/.changeset/retained-jetbrains-markdown.md b/.changeset/retained-jetbrains-markdown.md new file mode 100644 index 00000000000..e5e6cb0b0c7 --- /dev/null +++ b/.changeset/retained-jetbrains-markdown.md @@ -0,0 +1,5 @@ +--- +"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. 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 bf1a3513555..29adf9e05f6 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 @@ -17,7 +17,9 @@ import ai.kilocode.client.session.model.PermissionRequestState import ai.kilocode.client.session.model.Question import ai.kilocode.client.session.model.QuestionItem import ai.kilocode.client.session.model.QuestionOption +import ai.kilocode.client.session.model.Reasoning import ai.kilocode.client.session.model.ToolCallRef +import ai.kilocode.client.session.model.Text import ai.kilocode.client.plugin.KiloPluginSettings import ai.kilocode.client.session.SessionRef import ai.kilocode.client.telemetry.Telemetry @@ -897,7 +899,8 @@ class SessionController( is ChatEventDto.PartDelta -> { if (event.field == "text") { - model.appendDelta(event.messageID, event.partID, event.delta) + val delta = glue(event.messageID, event.partID, event.delta) + if (delta.isNotEmpty()) model.appendDelta(event.messageID, event.partID, delta) } } @@ -977,6 +980,18 @@ 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 span = (minOf(cur.length, delta.length) downTo 1) + .firstOrNull { n -> cur.regionMatches(cur.length - n, delta, 0, n) } ?: 0 + return delta.substring(span) + } + private fun handleHidden(event: ChatEventDto): Boolean = when (event) { is ChatEventDto.Error, is ChatEventDto.PermissionAsked, diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt index 5b356d0ff5e..c8f4c665f03 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt @@ -157,6 +157,7 @@ class SessionModel { fun appendDelta(messageId: String, contentId: String, delta: String) { val msg = entries[messageId] ?: return val existing = msg.parts[contentId] + val created = existing == null if (existing != null) { val buf = when (existing) { is Text -> existing.content @@ -170,7 +171,7 @@ class SessionModel { msg.parts[contentId] = content fire(SessionModelEvent.ContentAdded(messageId, content)) } - fire(SessionModelEvent.ContentDelta(messageId, contentId, delta)) + fire(SessionModelEvent.ContentDelta(messageId, contentId, delta, created)) updateHeader() } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModelEvent.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModelEvent.kt index 83b79b45b5d..1ca9fe5a342 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModelEvent.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModelEvent.kt @@ -36,7 +36,12 @@ sealed class SessionModelEvent { data class ContentRemoved(val messageId: String, val contentId: String) : SessionModelEvent() { override fun toString() = "ContentRemoved $messageId/$contentId" } - data class ContentDelta(val messageId: String, val contentId: String, val delta: String) : SessionModelEvent() { + data class ContentDelta( + val messageId: String, + val contentId: String, + val delta: String, + val created: Boolean = false, + ) : SessionModelEvent() { override fun toString() = "ContentDelta $messageId/$contentId" } data class StateChanged(val state: SessionState) : SessionModelEvent() { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt index d2255359cac..7d84b95b204 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt @@ -96,14 +96,11 @@ class SessionMessageListPanel( } is SessionModelEvent.ContentDelta -> { - // Use the full current content from the model rather than - // an incremental append. This avoids the double-write that - // occurs when ContentAdded and ContentDelta both fire for - // the same first delta (the model auto-creates the content - // on first appendDelta and fires both events in sequence). + if (event.created) return@addListener + val handled = msgToView[event.messageId]?.appendDelta(event.contentId, event.delta) == true + if (handled) return@addListener val content = model.content(event.messageId, event.contentId) if (content != null) msgToView[event.messageId]?.upsertPart(content) - refresh() } is SessionModelEvent.HistoryLoaded -> rebuild() diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt index 68d48f9d9b3..cbec04f5efd 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt @@ -182,10 +182,11 @@ class MessageView( } /** Append a streaming delta to the renderer for [contentId]. */ - fun appendDelta(contentId: String, delta: String) { - val part = parts[contentId] ?: return + fun appendDelta(contentId: String, delta: String): Boolean { + val part = parts[contentId] ?: return false part.appendDelta(delta) refresh() + return true } /** Look up a renderer by part id. */ diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/MdViewHybrid.kt index ef5c402103c..e9740eaa409 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/MdViewHybrid.kt @@ -126,8 +126,10 @@ internal class MdViewHybrid( private val source = StringBuilder() private var style = style private var rendered = "" - private var block: Disposable? = null private var disposed = false + private val blocks = mutableListOf() + private var openFence: Fence? = null + private var stale = false private val extensions = listOf( AutolinkExtension.create(), @@ -284,6 +286,7 @@ internal class MdViewHybrid( if (disposed) return if (this.selection === selection) return this.selection = selection + clearBlocks() syncBlocks() } @@ -315,6 +318,16 @@ internal class MdViewHybrid( override fun append(delta: String) { if (disposed) return if (delta.isEmpty()) return + val fence = openFence + val view = blocks.lastOrNull() as? CodeView + if (fence != null && view != null && clean(fence.char, delta)) { + source.append(delta) + view.grow(delta) + stale = true + root.revalidate() + root.repaint() + return + } source.append(delta) syncBlocks() } @@ -324,6 +337,8 @@ internal class MdViewHybrid( if (source.isEmpty() && rendered.isEmpty() && root.componentCount == 0) return source.clear() rendered = "" + openFence = null + stale = false clearBlocks() root.revalidate() root.repaint() @@ -340,7 +355,14 @@ internal class MdViewHybrid( override fun markdown(): String = source.toString() - override fun html(): String = rendered + override fun html(): String { + if (!stale) return rendered + val out = project(source.toString()) + rendered = out.html + openFence = out.open + stale = false + return rendered + } override fun overrideSheet(): String = MdCommon.rules(opts()) @@ -354,6 +376,8 @@ internal class MdViewHybrid( listeners.clear() source.clear() rendered = "" + openFence = null + stale = false clearBlocks() } @@ -362,37 +386,66 @@ internal class MdViewHybrid( val opts = opts() root.isOpaque = opts.opaque if (opts.opaque) root.background = opts.background - syncBlocks() + for (view in blocks) view.style(opts) + root.revalidate() + root.repaint() } private fun syncBlocks() { if (disposed) return val text = source.toString() - val doc = parser.parse(text) - val body = renderer.render(doc) - rendered = body - resetBlocks() + val out = project(text) + rendered = out.html + openFence = out.open + stale = false + val next = out.blocks if (text.isEmpty()) { + openFence = null + clearBlocks() root.revalidate() root.repaint() return } - val visitor = Visitor() - doc.accept(visitor) + sync(next) root.revalidate() root.repaint() } - private fun resetBlocks() { - block?.let { Disposer.dispose(it) } - block = Disposer.newDisposable("Markdown blocks") + private fun clearBlocks() { + blocks.forEach { Disposer.dispose(it.disposable) } + blocks.clear() root.removeAll() } - private fun clearBlocks() { - block?.let { Disposer.dispose(it) } - block = null - root.removeAll() + private fun sync(next: List) { + var at = 0 + while (at < blocks.size && at < next.size) { + val view = blocks[at] + val desc = next[at] + if (!view.compatible(desc)) break + view.update(desc) + at++ + } + removeBlocks(at) + for (desc in next.drop(at)) addBlock(view(desc)) + } + + private fun clean(char: Char, delta: String): Boolean { + if (delta.contains(char)) return false + val start = source.lastIndexOf("\n") + 1 + for (idx in start until source.length) { + if (source[idx] == char) return false + } + return true + } + + private fun removeBlocks(start: Int) { + if (start >= blocks.size) return + val idx = if (start == 0) 0 else start * 2 - 1 + while (root.componentCount > idx) root.remove(root.componentCount - 1) + val stale = blocks.drop(start) + repeat(blocks.size - start) { blocks.removeAt(blocks.lastIndex) } + stale.forEach { Disposer.dispose(it.disposable) } } private fun addGap() { @@ -400,15 +453,23 @@ internal class MdViewHybrid( root.add(Box.createVerticalStrut(JBUI.scale(SessionUiStyle.View.Code.BLOCK_GAP))) } - private fun addBlock(component: JComponent) { + private fun addBlock(view: View) { addGap() - component.alignmentX = JComponent.LEFT_ALIGNMENT - root.add(component) + view.component.alignmentX = JComponent.LEFT_ALIGNMENT + blocks.add(view) + root.add(view.component) } - private fun htmlBlock(node: Node): JComponent { + private fun view(desc: Desc): View { + val disposable = Disposer.newDisposable("Markdown block") + return when (desc) { + is Desc.Html -> HtmlView(desc, htmlBlock(desc.body, disposable), disposable) + is Desc.Code -> CodeView(desc, codeBlock(desc.text, desc.file, disposable), disposable) + } + } + + private fun htmlBlock(body: String, disposable: Disposable): JBHtmlPane { val opts = opts() - val body = renderer.render(node) return JBHtmlPane( JBHtmlPaneStyleConfiguration { enableInlineCodeBackground = true @@ -422,7 +483,7 @@ internal class MdViewHybrid( isOpaque = opts.opaque background = opts.background text = "$body" - block?.let { selection?.register(this, it) } + selection?.register(this, disposable) addHyperlinkListener { e -> if (e.eventType != HyperlinkEvent.EventType.ACTIVATED) return@addHyperlinkListener val href = e.description ?: return@addHyperlinkListener @@ -432,24 +493,20 @@ internal class MdViewHybrid( } } - private fun codeBlock(text: String, lang: String?): JComponent { + private fun codeBlock(text: String, file: FileType, disposable: Disposable): JBScrollPane { val opts = opts() val value = text.trimEnd('\n') val field = runCatching { - CodeField(file(lang), opts, text).also { ed -> - block?.let { ed.setDisposedWith(it) } - block?.let { selection?.register(ed, it) } + CodeField(file, opts, text).also { ed -> + ed.setDisposedWith(disposable) + selection?.register(ed, disposable) } }.getOrElse { err -> LOG.warn("kind=markdown codeEditor=true failed message=${err.message}", err) - textArea(text, opts) + textArea(text, opts, disposable) } - val height = codeHeight(field, value) - val width = codeWidth(field, value) - field.preferredSize = Dimension(width, height) - field.minimumSize = Dimension(0, height) - field.maximumSize = Dimension(Int.MAX_VALUE, height) - return object : JBScrollPane(field) { + sizeCodeField(field, value) + val pane = object : JBScrollPane(field) { override fun doLayout() { super.doLayout() val view = viewport.view ?: return @@ -457,7 +514,14 @@ internal class MdViewHybrid( if (size.height <= 0 || view.height == size.height) return view.setSize(view.width.coerceAtLeast(size.width), size.height) } - }.apply { + } + styleCodePane(pane, opts) + sizeCodePane(pane, field) + return pane + } + + private fun styleCodePane(pane: JBScrollPane, opts: MdStyle) { + pane.apply { border = JBUI.Borders.customLine(opts.tableBorder, SessionUiStyle.View.Code.BORDER_WIDTH) viewportBorder = JBUI.Borders.empty( SessionUiStyle.View.Code.topPadding(), @@ -475,14 +539,26 @@ internal class MdViewHybrid( horizontalScrollBar.preferredSize = Dimension(0, JBUI.scale(SessionUiStyle.View.Code.SCROLLBAR_HEIGHT)) horizontalScrollBar.isOpaque = true verticalScrollBar.preferredSize = JBUI.emptySize() - val pad = viewportBorder.getBorderInsets(this) - val size = height + insets.top + insets.bottom + pad.top + pad.bottom + horizontalScrollBar.preferredSize.height - preferredSize = Dimension(0, size) - minimumSize = Dimension(0, size) - maximumSize = Dimension(Int.MAX_VALUE, size) } } + private fun sizeCodeField(component: JComponent, text: String) { + val height = codeHeight(component, text) + val width = codeWidth(component, text) + component.preferredSize = Dimension(width, height) + component.minimumSize = Dimension(0, height) + component.maximumSize = Dimension(Int.MAX_VALUE, height) + } + + private fun sizeCodePane(pane: JBScrollPane, component: JComponent) { + val pad = pane.viewportBorder.getBorderInsets(pane) + val height = component.preferredSize.height + pane.insets.top + pane.insets.bottom + + pad.top + pad.bottom + pane.horizontalScrollBar.preferredSize.height + pane.preferredSize = Dimension(0, height) + pane.minimumSize = Dimension(0, height) + pane.maximumSize = Dimension(Int.MAX_VALUE, height) + } + private fun codeWidth(component: JComponent, text: String): Int { val metrics = component.getFontMetrics(component.font) val width = text.lineSequence().maxOfOrNull { metrics.stringWidth(it) } ?: 0 @@ -503,18 +579,22 @@ internal class MdViewHybrid( return line * rows } - private fun textArea(text: String, opts: MdStyle) = JBTextArea(text.trimEnd('\n')).apply { + private fun textArea(text: String, opts: MdStyle, disposable: Disposable) = JBTextArea(text.trimEnd('\n')).apply { isEditable = false lineWrap = false - isOpaque = opts.opaque - background = opts.preBg - foreground = opts.preFg - font = Font(opts.codeFont, Font.PLAIN, style.editorSize) + styleTextArea(this, opts) border = JBUI.Borders.empty( SessionUiStyle.View.Code.VIEWPORT_TOP_PADDING, SessionUiStyle.View.Code.VIEWPORT_HORIZONTAL_PADDING, ) - block?.let { selection?.register(this, it) } + selection?.register(this, disposable) + } + + private fun styleTextArea(area: JBTextArea, opts: MdStyle) { + area.isOpaque = opts.opaque + area.background = opts.preBg + area.foreground = opts.preFg + area.font = Font(opts.codeFont, Font.PLAIN, style.editorSize) } private inner class CodeField(file: FileType, opts: MdStyle, value: String) : @@ -586,17 +666,247 @@ internal class MdViewHybrid( ) } + private fun collect(doc: Node): List { + val visitor = Visitor() + doc.accept(visitor) + return visitor.blocks + } + + private fun project(text: String): Projection { + val blocks = mutableListOf() + val html = StringBuilder() + val md = StringBuilder() + val lines = lines(text) + var trailing: Fence? = null + var idx = 0 + + fun flush() { + if (md.isEmpty()) return + val doc = parser.parse(md.toString()) + html.append(renderer.render(doc)) + blocks.addAll(collect(doc)) + md.clear() + } + + while (idx < lines.size) { + val line = lines[idx] + val open = opener(line.text) + if (open == null) { + val pending = idx == lines.lastIndex && pendingOpener(line.text) + if (pending) { + flush() + blocks.add(Desc.Code("", PlainTextFileType.INSTANCE)) + html.append(codeHtml("")) + } else { + md.append(line.text).append(line.end) + } + idx++ + continue + } + + flush() + idx++ + val code = StringBuilder() + var closed = false + var trimmed = false + while (idx < lines.size) { + val item = lines[idx] + val close = closer(item.text, open) + if (close) { + closed = true + idx++ + break + } + val partial = idx == lines.lastIndex && partialCloser(item.text, open) + if (partial) trimmed = true + if (!partial) code.append(item.text).append(item.end) + idx++ + } + val desc = Desc.Code(code.toString(), file(open.info)) + blocks.add(desc) + html.append(codeHtml(desc.text)) + trailing = if (!closed && !trimmed) open else null + } + + flush() + return Projection(html.toString(), blocks, trailing) + } + + private fun lines(text: String): List { + if (text.isEmpty()) return emptyList() + val lines = mutableListOf() + var start = 0 + while (start < text.length) { + val end = text.indexOf('\n', start) + if (end == -1) { + lines.add(Line(text.substring(start), "")) + break + } + lines.add(Line(text.substring(start, end), "\n")) + start = end + 1 + } + return lines + } + + private fun opener(text: String): Fence? { + val trimmed = text.dropWhile { it == ' ' } + val indent = text.length - trimmed.length + if (indent > 3) return null + val char = trimmed.firstOrNull() ?: return null + if (char != '`' && char != '~') return null + val size = trimmed.takeWhile { it == char }.length + if (size < 3) return null + val info = trimmed.drop(size).trim() + if (char == '`' && info.contains('`')) return null + return Fence(char, size, info) + } + + private fun closer(text: String, fence: Fence): Boolean { + val trimmed = text.dropWhile { it == ' ' } + val indent = text.length - trimmed.length + if (indent > 3) return false + val size = trimmed.takeWhile { it == fence.char }.length + if (size < fence.size) return false + return trimmed.drop(size).isBlank() + } + + private fun pendingOpener(text: String): Boolean { + val trimmed = text.dropWhile { it == ' ' } + val indent = text.length - trimmed.length + if (indent > 3) return false + val char = trimmed.firstOrNull() ?: return false + if (char != '`' && char != '~') return false + val size = trimmed.takeWhile { it == char }.length + if (size !in 1..2) return false + return trimmed.drop(size).isBlank() + } + + private fun partialCloser(text: String, fence: Fence): Boolean { + val trimmed = text.dropWhile { it == ' ' } + val indent = text.length - trimmed.length + if (indent > 3) return false + val size = trimmed.takeWhile { it == fence.char }.length + if (size !in 1 until fence.size) return false + return trimmed.drop(size).isBlank() + } + + private fun codeHtml(text: String): String = "
${escape(text)}
\n" + + private fun escape(text: String): String = text + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace("\"", """) + + private sealed class Desc { + data class Html(val body: String) : Desc() + data class Code(val text: String, val file: FileType) : Desc() + } + + private data class Projection(val html: String, val blocks: List, val open: Fence?) + + private data class Line(val text: String, val end: String) + + private data class Fence(val char: Char, val size: Int, val info: String) + + private abstract inner class View( + var desc: Desc, + val component: JComponent, + val disposable: Disposable, + ) { + abstract fun compatible(desc: Desc): Boolean + abstract fun update(desc: Desc) + abstract fun style(opts: MdStyle) + } + + private inner class HtmlView(desc: Desc.Html, private val pane: JBHtmlPane, disposable: Disposable) : + View(desc, pane, disposable) { + override fun compatible(desc: Desc) = desc is Desc.Html + + override fun update(desc: Desc) { + if (this.desc == desc) return + this.desc = desc + pane.text = "${(desc as Desc.Html).body}" + } + + override fun style(opts: MdStyle) { + pane.isOpaque = opts.opaque + pane.background = opts.background + val item = desc as Desc.Html + pane.text = "${item.body}" + } + } + + private inner class CodeView(desc: Desc.Code, private val pane: JBScrollPane, disposable: Disposable) : + View(desc, pane, disposable) { + override fun compatible(desc: Desc) = desc is Desc.Code && (this.desc as Desc.Code).file == desc.file + + override fun update(desc: Desc) { + if (this.desc == desc) return + this.desc = desc + val value = (desc as Desc.Code).text.trimEnd('\n') + val view = pane.viewport.view + when (view) { + is CodeField -> view.text = value + is JBTextArea -> view.text = value + } + if (view is JComponent) { + sizeCodeField(view, value) + sizeCodePane(pane, view) + } + } + + fun grow(delta: String) { + val item = desc as Desc.Code + val next = item.copy(text = item.text + delta) + desc = next + val value = next.text.trimEnd('\n') + val view = pane.viewport.view + when (view) { + is CodeField -> view.text = value + is JBTextArea -> view.text = value + } + if (view is JComponent) { + sizeCodeField(view, value) + sizeCodePane(pane, view) + } + } + + override fun style(opts: MdStyle) { + styleCodePane(pane, opts) + val view = pane.viewport.view + when (view) { + is CodeField -> { + view.font = Font(opts.codeFont, Font.PLAIN, style.editorSize) + view.background = opts.preBg + } + is JBTextArea -> styleTextArea(view, opts) + } + if (view is JComponent) { + val text = when (view) { + is CodeField -> view.text + is JBTextArea -> view.text + else -> "" + } + sizeCodeField(view, text) + sizeCodePane(pane, view) + } + } + } + private inner class Visitor : AbstractVisitor() { + val blocks = mutableListOf() + override fun visit(document: Document) { visitChildren(document) } override fun visit(code: FencedCodeBlock) { - addBlock(codeBlock(code.literal, code.info)) + blocks.add(Desc.Code(code.literal, file(code.info))) } override fun visit(code: IndentedCodeBlock) { - addBlock(codeBlock(code.literal, null)) + blocks.add(Desc.Code(code.literal, file(null))) } public override fun visitChildren(parent: Node) { @@ -604,7 +914,9 @@ internal class MdViewHybrid( while (child != null) { val next = child.next if (child is FencedCodeBlock || child is IndentedCodeBlock) child.accept(this) - if (child is Block && child !is FencedCodeBlock && child !is IndentedCodeBlock) addBlock(htmlBlock(child)) + if (child is Block && child !is FencedCodeBlock && child !is IndentedCodeBlock) { + blocks.add(Desc.Html(renderer.render(child))) + } child = next } } 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 75c2e6887ab..11d4e92548b 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 @@ -12,6 +12,7 @@ import ai.kilocode.rpc.dto.QuestionRequestDto import ai.kilocode.rpc.dto.SessionStatusDto import ai.kilocode.rpc.dto.TodoDto import ai.kilocode.rpc.dto.ToolRefDto +import kotlinx.coroutines.ExperimentalCoroutinesApi class SessionUpdateQueueTest : SessionControllerTestBase() { @@ -234,6 +235,137 @@ class SessionUpdateQueueTest : SessionControllerTestBase() { assertEquals(listOf("hello world"), delta.map { it.delta }) } + fun `test text snapshot covered delta is not duplicated`() { + 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.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "import")), flush = false) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "import"), flush = false) + settle() + flush() + + assertModel( + """ + assistant#msg1 + text#prt1: + import + """, + 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() + val m = controller("ses_test", flushMs = Long.MAX_VALUE) + flush() + + emit(ChatEventDto.MessageUpdated("ses_test", msg("msg1", "ses_test", "assistant"))) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "import"))) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "import java")), flush = false) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", " java.util.List;"), flush = false) + settle() + flush() + + assertModel( + """ + assistant#msg1 + text#prt1: + import java.util.List; + """, + m, + ) + } + + fun `test repeated snapshots then lagging merged delta is not duplicated`() { + 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.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "import"))) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "import java"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "import java"), flush = false) + settle() + flush() + + assertModel( + """ + assistant#msg1 + text#prt1: + import java + """, + m, + ) + } + + fun `test multi round snapshot delta interleave stays single copy`() { + 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.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "```java\nimport"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "import")) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "```java\nimport java"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", " java")) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "```java\nimport java.util.List;\n"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", ".util.List;\n")) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = "```java\nimport java.util.List;\n\npublic class StreamBasics {\n}\n```"))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", "\npublic class StreamBasics {\n}\n```")) + + assertModel( + """ + assistant#msg1 + text#prt1: + ```java + import java.util.List; + + public class StreamBasics { + } + ``` + """, + m, + ) + } + + fun `test per token snapshot plus delta interleave does not double text or code`() { + 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"))) + + // Reproduces the streamed screenshot: for every token the backend sends a full + // PartUpdated snapshot AND a matching incremental PartDelta. The old dedup doubled + // each token ("ReadRead", "inputinput.txt.txt"); glue must keep a single copy. + val tokens = listOf( + "**Python**", "\n\n", "Read", " a", " file", " line", " by", " line", ",", + " which", " is", " stream-like", " because", " it", " avoids", " loading", + " the", " whole", " file", " into", " memory", ":", "\n\n", + "```python\n", "with", " open", "(\"input.txt\",", " \"r\")", " as", " file:", "\n", + " for", " line", " in", " file:", "\n", " print", "(line.strip())", "\n", + "```", + ) + val sb = StringBuilder() + for (token in tokens) { + sb.append(token) + emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = sb.toString()))) + emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", token)) + } + settle() + flush() + + val text = (m.model.message("msg1")!!.parts["prt1"] as ai.kilocode.client.session.model.Text).content.toString() + assertEquals(sb.toString(), text) + } + fun `test visible controller flushes after cadence`() { appRpc.state.value = ai.kilocode.rpc.dto.KiloAppStateDto(ai.kilocode.rpc.dto.KiloAppStatusDto.READY) projectRpc.state.value = workspaceReady() @@ -346,15 +478,15 @@ class SessionUpdateQueueTest : SessionControllerTestBase() { assertEquals(SessionState.Idle, m.model.state) } + @OptIn(ExperimentalCoroutinesApi::class) fun `test condensed and raw controller end with same final state on large corpus`() { appRpc.state.value = ai.kilocode.rpc.dto.KiloAppStateDto(ai.kilocode.rpc.dto.KiloAppStatusDto.READY) projectRpc.state.value = workspaceReady() val events = corpus() - val condensed = runCorpus(events, true) - val raw = runCorpus(events, false) - val a = snapshot(condensed) - val b = snapshot(raw) + val a = snapshot(runCorpus(events, true)) + rpc.events.resetReplayCache() + val b = snapshot(runCorpus(events, false)) if (a != b) fail("condensed=\n$a\nraw=\n$b") assertEquals(SessionState.Idle, a.state) @@ -362,7 +494,7 @@ class SessionUpdateQueueTest : SessionControllerTestBase() { assertTrue(a.body.contains("assistant#msg2")) assertTrue(a.body.contains("diff: src/A.kt src/B.kt")) assertTrue(a.body.contains("todo: [completed] ship feature")) - assertEquals(4, a.compacted) + assertEquals(2, a.compacted) } fun `test update hooks run on EDT around queued model batch`() { diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanelTest.kt index 6d5a067309d..588d4e95167 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanelTest.kt @@ -26,6 +26,7 @@ import com.intellij.openapi.Disposable import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase import java.awt.Container +import javax.swing.JPanel /** * Tests for [SessionMessageListPanel] — structural and index integrity. @@ -186,6 +187,31 @@ class SessionMessageListPanelTest : BasePlatformTestCase() { assertEquals("hello world", tv.markdown()) } + fun `test ContentDelta preserves TextView and markdown component`() { + model.upsertMessage(msg("a1", "assistant")) + model.updateContent("a1", part("p1", "a1", "text", text = "first\n\nsecond")) + val mv = panel.findMessage("a1")!! + val tv = mv.part("p1") as TextView + val comp = tv.md.component + val first = (comp as JPanel).components.first() + + model.appendDelta("a1", "p1", " more") + + assertSame(tv, mv.part("p1")) + assertSame(comp, tv.md.component) + assertSame(first, comp.components.first()) + assertEquals("first\n\nsecond more", tv.markdown()) + } + + fun `test created ContentDelta is not double applied`() { + model.upsertMessage(msg("a1", "assistant")) + + model.appendDelta("a1", "p1", "hello") + + val tv = panel.findMessage("a1")!!.part("p1") as TextView + assertEquals("hello", tv.markdown()) + } + fun `test ContentRemoved removes PartView from MessageView`() { model.upsertMessage(msg("a1", "assistant")) model.updateContent("a1", part("p1", "a1", "text", text = "x")) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewHybridTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewHybridTest.kt index 47f4e4368c6..0ff05c04f9a 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewHybridTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewHybridTest.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.fileTypes.UnknownFileType import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.ui.EditorTextField +import com.intellij.ui.components.JBHtmlPane import com.intellij.ui.components.JBScrollPane import com.intellij.util.ui.UIUtil import java.awt.BorderLayout @@ -128,6 +129,146 @@ class MdViewHybridTest : BasePlatformTestCase() { assertTrue("streamed editor should not be clipped vertically", editor.height >= editor.preferredSize.height) } + fun `test appending later html block preserves earlier block component`() { + view.set("first\n\nsecond") + val first = htmls().first() + + view.append(" more") + + assertSame(first, htmls().first()) + assertTrue(view.html().contains("second more")) + } + + fun `test streaming fenced code block preserves editor component`() { + view.append("```java\nclass Example {\n") + val pane = scrolls().single() + val editor = editors().single() + val initial = pane.preferredSize.height + + view.append("void method() {}\n}\n```") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("class Example {\nvoid method() {}\n}", editor.text) + assertTrue(scrolls().single().preferredSize.height >= initial) + } + + fun `test streaming code body append reuses editor and stays correct`() { + view.append("```java\nclass A {\n") + val pane = scrolls().single() + val editor = editors().single() + + view.append(" void run() {\n") + view.append(" }\n") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("class A {\n void run() {\n }", editor.text) + assertEquals("```java\nclass A {\n void run() {\n }\n", view.markdown()) + assertTrue(view.html().contains("class A")) + + view.append("}\n```") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("class A {\n void run() {\n }\n}", editor.text) + assertEquals("```java\nclass A {\n void run() {\n }\n}\n```", view.markdown()) + } + + fun `test streaming code body append keeps html in sync`() { + view.append("```java\n") + view.append("if (a < b) {\n") + + assertTrue(view.html().contains("a < b")) + assertEquals("if (a < b) {", editors().single().text) + } + + fun `test streaming partial fence opener renders code block without raw marker`() { + view.append("`") + + val pane = scrolls().single() + val editor = editors().single() + + assertEquals("", editor.text) + assertFalse(view.html().contains("`")) + + view.append("`") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("", editor.text) + assertFalse(view.html().contains("``")) + } + + fun `test streaming language prefix stays out of code text`() { + view.append("```") + view.append("p") + view.append("ython\nprint(1)\n") + + assertEquals("print(1)", editors().single().text) + assertFalse(view.html().contains("```")) + assertFalse(view.html().contains("python")) + } + + fun `test streaming partial closing fence stays out of code text`() { + view.append("```python\nprint(1)\n") + val pane = scrolls().single() + val editor = editors().single() + + view.append("`") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("print(1)", editor.text) + + view.append("`") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("print(1)", editor.text) + } + + fun `test completed closing fence preserves code block and renders following markdown`() { + view.append("```python\nprint(1)\n``") + val pane = scrolls().single() + val editor = editors().single() + + view.append("`\n\nafter") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals("print(1)", editor.text) + assertEquals(1, scrolls().size) + assertEquals(1, htmls().size) + assertTrue(view.html().contains("after")) + } + + fun `test appending new block preserves existing code block component`() { + view.set("```kotlin\nval value = 1\n```") + val pane = scrolls().single() + val editor = editors().single() + + view.append("\n\nhello") + + assertSame(pane, scrolls().single()) + assertSame(editor, editors().single()) + assertEquals(1, scrolls().size) + assertEquals(1, htmls().size) + } + + fun `test incompatible suffix replacement preserves prefix block`() { + view.set("before\n\n```kotlin\nval value = 1\n```") + val prefix = htmls().single() + val editor = editors().single().getEditor(true)!! + + view.set("before\n\nafter") + drainEdt() + + assertSame(prefix, htmls().first()) + assertTrue(editor.isDisposed) + assertTrue(scrolls().isEmpty()) + } + fun `test java code block with blank lines fits vertically`() { val code = """ import java.util.List; @@ -277,6 +418,8 @@ class MdViewHybridTest : BasePlatformTestCase() { private fun scrolls(): List = (view.component as JPanel).components.filterIsInstance() + private fun htmls(): List = (view.component as JPanel).components.filterIsInstance() + private fun editors(): List = scrolls().mapNotNull { it.viewport.view as? EditorTextField } private fun type(ext: String): FileType {