From ccfd17ef13895c0170afcad2602410aa372708bd Mon Sep 17 00:00:00 2001 From: kirillk Date: Mon, 24 Aug 2026 14:16:17 -0400 Subject: [PATCH 01/11] feat(jetbrains): add mermaid diagram engine --- .../ai/kilocode/client/ui/diagram/Art.kt | 97 +++++ .../ai/kilocode/client/ui/diagram/Engine.kt | 40 ++ .../ai/kilocode/client/ui/diagram/Measure.kt | 13 + .../ai/kilocode/client/ui/diagram/Metrics.kt | 23 + .../ai/kilocode/client/ui/diagram/Type.kt | 40 ++ .../client/ui/diagram/mermaid/Flow.kt | 324 ++++++++++++++ .../client/ui/diagram/mermaid/FlowLayout.kt | 408 ++++++++++++++++++ .../client/ui/diagram/mermaid/FlowMarks.kt | 229 ++++++++++ .../client/ui/diagram/mermaid/Mermaid.kt | 58 +++ .../kilocode/client/ui/diagram/mermaid/Seq.kt | 182 ++++++++ .../client/ui/diagram/mermaid/SeqLayout.kt | 264 ++++++++++++ .../client/ui/diagram/mermaid/Source.kt | 93 ++++ .../kilocode/client/ui/diagram/CancelTest.kt | 50 +++ .../client/ui/diagram/ConformanceTest.kt | 78 ++++ .../client/ui/diagram/DiagramAsserts.kt | 115 +++++ .../kilocode/client/ui/diagram/ErrorTest.kt | 54 +++ .../kilocode/client/ui/diagram/FakeMeasure.kt | 31 ++ .../client/ui/diagram/FlowLayoutTest.kt | 88 ++++ .../client/ui/diagram/InvariantTest.kt | 42 ++ .../kilocode/client/ui/diagram/LimitsTest.kt | 54 +++ .../client/ui/diagram/SeqLayoutTest.kt | 102 +++++ .../client/ui/diagram/SerializeTest.kt | 57 +++ .../ai/kilocode/client/ui/diagram/TypeTest.kt | 43 ++ .../ui/diagram/mermaid/FlowParseTest.kt | 181 ++++++++ .../client/ui/diagram/mermaid/SeqParseTest.kt | 140 ++++++ .../client/ui/diagram/mermaid/SourceTest.kt | 74 ++++ .../src/test/resources/diagram/flow-basic.mmd | 6 + .../src/test/resources/diagram/flow-cycle.mmd | 9 + .../src/test/resources/diagram/flow-long.mmd | 7 + .../test/resources/diagram/flow-shapes.mmd | 17 + .../test/resources/diagram/flow-subgraph.mmd | 10 + .../src/test/resources/diagram/seq-basic.mmd | 7 + .../src/test/resources/diagram/seq-blocks.mmd | 16 + .../src/test/resources/diagram/seq-notes.mmd | 8 + 34 files changed, 2960 insertions(+) create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Art.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Measure.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Type.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ConformanceTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ErrorTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FakeMeasure.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/TypeTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-basic.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-cycle.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-long.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-shapes.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-subgraph.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-basic.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-blocks.mmd create mode 100644 packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-notes.mmd diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Art.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Art.kt new file mode 100644 index 00000000000..37d422b85aa --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Art.kt @@ -0,0 +1,97 @@ +package ai.kilocode.client.ui.diagram + +import kotlinx.serialization.Serializable +import kotlinx.serialization.SerialName +import kotlin.math.roundToInt + +@Serializable +internal data class Pt(val x: Double, val y: Double) { + override fun toString() = "${fmt(x)},${fmt(y)}" +} + +@Serializable +internal data class Rect(val x: Double, val y: Double, val w: Double, val h: Double) { + override fun toString() = "${fmt(x)},${fmt(y)} ${fmt(w)}x${fmt(h)}" +} + +@Serializable +internal data class Size(val w: Double, val h: Double) { + override fun toString() = "${fmt(w)}x${fmt(h)}" +} + +internal enum class Role { Surface, Border, Text, Muted, Accent, Note, Cluster, Line } + +internal enum class Head { None, Arrow, Open, Cross, Dot } + +internal enum class Anchor { TopLeft, Top, TopRight, Left, Center, Right, BottomLeft, Bottom, BottomRight } + +@Serializable +internal sealed interface Mark { + @Serializable + data class Box(val rect: Rect, val arc: Double, val fill: Role?, val line: Role?, val dash: Boolean = false) : Mark { + override fun toString() = "box $rect arc=${fmt(arc)} fill=${fill.name()} line=${line.name()} dash=$dash" + } + + @Serializable + data class Oval(val rect: Rect, val fill: Role?, val line: Role?) : Mark { + override fun toString() = "oval $rect fill=${fill.name()} line=${line.name()}" + } + + @Serializable + data class Poly(val points: List, val fill: Role?, val line: Role?) : Mark { + override fun toString() = "poly ${points.joinToString(" ")} fill=${fill.name()} line=${line.name()}" + } + + @Serializable + data class Edge( + val points: List, + val role: Role, + val dash: Boolean = false, + val thick: Boolean = false, + val head: Head = Head.None, + val tail: Head = Head.None, + ) : Mark { + override fun toString() = "edge ${points.joinToString(" ")} role=$role dash=$dash thick=$thick head=$head tail=$tail" + } + + @Serializable + data class Text(val text: String, val at: Pt, val anchor: Anchor, val role: Role, val bold: Boolean = false) : Mark { + override fun toString() = "text ${quote(text)} at=$at anchor=$anchor role=$role bold=$bold" + } + + @Serializable + data class Group(val id: String?, val marks: List) : Mark { + override fun toString() = buildString { + append("group ${id ?: "-"}") + for (mark in marks) append('\n').append(mark.toString().prependIndent(" ")) + } + } +} + +@Serializable +internal sealed interface Art + +@Serializable +internal data class Scene(@SerialName("diagram") val type: Type, val marks: List, val size: Size) : Art { + override fun toString() = buildString { + append("scene $type $size") + for (mark in marks) append('\n').append(mark) + } +} + +private fun Role?.name() = this?.name ?: "-" + +internal fun fmt(value: Double): String = value.roundToInt().toString() + +private fun quote(value: String) = buildString { + append('"') + for (char in value) { + when (char) { + '\\' -> append("\\\\") + '"' -> append("\\\"") + '\n' -> append("\\n") + else -> append(char) + } + } + append('"') +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt new file mode 100644 index 00000000000..06f205d595b --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt @@ -0,0 +1,40 @@ +package ai.kilocode.client.ui.diagram + +import kotlinx.serialization.Serializable + +internal interface Engine { + fun accepts(type: Type): Boolean + + /** Off-EDT, cancellable, no Swing/AWT/IntelliJ types in or out. */ + suspend fun draw(source: String, spec: Spec): Out +} + +internal sealed interface Out { + data class Ok(val art: Art) : Out + data class Err(val fault: Fault, val message: String, val line: Int? = null) : Out +} + +internal enum class Fault { Syntax, Unsupported, Limit, Internal } + +@Serializable +internal data class Spec( + val font: FontSpec, + val metrics: Metrics = Metrics(), + val limits: Limits = Limits(), +) + +@Serializable +internal data class FontSpec(val family: String, val size: Int, val bold: Boolean = false) + +@Serializable +internal data class Metrics( + val pad: Double = 8.0, + val gap: Double = 24.0, + val rank: Double = 48.0, + val line: Double = 1.0, + val arc: Double = 4.0, + val wrap: Double = 0.0, +) + +@Serializable +internal data class Limits(val nodes: Int = 400, val edges: Int = 800, val lines: Int = 2_000) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Measure.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Measure.kt new file mode 100644 index 00000000000..e1e3b5ac87e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Measure.kt @@ -0,0 +1,13 @@ +package ai.kilocode.client.ui.diagram + +/** + * Text measurement capability used by in-process engines. + * + * Implementations must be deterministic for a given [FontSpec]. Tests use a fake implementation + * so geometry snapshots do not depend on system fonts or CI image contents. + */ +internal interface Measure { + fun width(text: String, font: FontSpec): Double + fun height(font: FontSpec): Double + fun ascent(font: FontSpec): Double +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt new file mode 100644 index 00000000000..0489211e69e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt @@ -0,0 +1,23 @@ +package ai.kilocode.client.ui.diagram + +import java.awt.Font +import java.awt.FontMetrics +import java.awt.image.BufferedImage + +internal class AwtMeasure : Measure { + private val img = BufferedImage(1, 1, BufferedImage.TYPE_INT_ARGB) + private val g = img.createGraphics() + private val cache = linkedMapOf() + + override fun width(text: String, font: FontSpec) = metrics(font).stringWidth(text).toDouble() + override fun height(font: FontSpec) = metrics(font).height.toDouble() + override fun ascent(font: FontSpec) = metrics(font).ascent.toDouble() + + private fun metrics(font: FontSpec): FontMetrics { + cache[font]?.let { return it } + val style = if (font.bold) Font.BOLD else Font.PLAIN + val value = g.getFontMetrics(Font(font.family, style, font.size)) + cache[font] = value + return value + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Type.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Type.kt new file mode 100644 index 00000000000..22652deb265 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Type.kt @@ -0,0 +1,40 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Clean +import ai.kilocode.client.ui.diagram.mermaid.Source +import kotlinx.serialization.Serializable + +@Serializable +internal enum class Type { + Flowchart, + Sequence, + Class, + State, + Er, + Gantt, + Pie, + Unknown; + + companion object { + /** + * Detects the diagram type from preprocessed text, so frontmatter, `%%` comments and + * `%%{init}%%` directives cannot shift the answer. + */ + fun of(source: String): Type = of(Source.clean(source)) + + fun of(clean: Clean): Type { + val head = clean.lines.firstOrNull { it.text.isNotBlank() }?.text?.trim() ?: return Unknown + val token = head.takeWhile { !it.isWhitespace() }.lowercase() + return when (token) { + "graph", "flowchart" -> Flowchart + "sequencediagram" -> Sequence + "classdiagram", "classdiagram-v2" -> Class + "statediagram", "statediagram-v2" -> State + "erdiagram" -> Er + "gantt" -> Gantt + "pie" -> Pie + else -> Unknown + } + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt new file mode 100644 index 00000000000..c5c19536568 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt @@ -0,0 +1,324 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Head + +internal enum class Dir { Down, Up, Left, Right } + +internal enum class Shape { + Rect, + Round, + Stadium, + Subroutine, + Cylinder, + Circle, + Doubled, + Rhombus, + Hexagon, + Skew, + SkewAlt, + Trapezoid, + TrapezoidAlt, + Flag, +} + +internal enum class Link { Solid, Dotted, Thick } + +internal data class FlowNode( + val id: String, + val label: List, + val shape: Shape, + val index: Int, + val cluster: String?, +) + +internal data class FlowEdge( + val from: String, + val to: String, + val link: Link, + val head: Head, + val tail: Head, + val label: List, + val index: Int, +) + +internal data class Cluster(val id: String, val label: List, val parent: String?, val index: Int) + +internal data class Graph( + val dir: Dir, + val nodes: Map, + val edges: List, + val clusters: Map, +) + +internal sealed interface FlowOut { + data class Ok(val graph: Graph) : FlowOut + data class Err(val message: String, val line: Int) : FlowOut +} + +/** Line-oriented flowchart parser. Unknown statements are skipped rather than failing the diagram. */ +internal class Flow { + private val nodes = linkedMapOf() + private val edges = mutableListOf() + private val clusters = linkedMapOf() + private val stack = ArrayDeque() + private var dir = Dir.Down + + fun parse(clean: Clean): FlowOut { + var first = true + for (line in clean.lines) { + val text = line.text.trim() + if (text.isEmpty()) continue + if (first) { + first = false + if (header(text)) continue + } + val err = stmt(text, line.at) ?: continue + return FlowOut.Err(err, line.at) + } + if (stack.isNotEmpty()) { + return FlowOut.Err("subgraph is missing a matching end", clean.lines.lastOrNull()?.at ?: 1) + } + return FlowOut.Ok(Graph(dir, nodes, edges, clusters)) + } + + private fun header(text: String): Boolean { + val token = text.substringBefore(' ').lowercase() + if (token != "graph" && token != "flowchart") return false + dir = dirOf(text.substringAfter(' ', "").trim()) + return true + } + + private fun stmt(text: String, at: Int): String? { + val token = text.substringBefore(' ').substringBefore('[').lowercase() + if (token in SKIP) return null + if (token == "end") { + if (stack.isEmpty()) return "end without a matching subgraph" + stack.removeLast() + return null + } + if (token == "subgraph") return group(text) + return chain(text, at) + } + + private fun group(text: String): String? { + val rest = text.substringAfter("subgraph").trim() + val open = rest.indexOf('[') + val id = when { + rest.isEmpty() -> "sub${clusters.size + 1}" + open > 0 && rest.endsWith("]") -> rest.substring(0, open).trim() + else -> rest + } + val label = when { + rest.isEmpty() -> listOf(id) + open > 0 && rest.endsWith("]") -> Source.label(rest.substring(open + 1, rest.length - 1)) + else -> Source.label(rest) + } + if (clusters.containsKey(id)) return "duplicate subgraph id $id" + clusters[id] = Cluster(id, label, stack.lastOrNull(), clusters.size) + stack.addLast(id) + return null + } + + private fun chain(text: String, at: Int): String? { + val hits = hits(text) + if (hits.isEmpty()) { + refs(text) ?: return null + return null + } + val segs = split(text, hits) + val labels = pipes(segs, hits.size) + val groups = segs.map { refs(it) ?: return "expected a node on both sides of the link" } + for (idx in hits.indices) { + val hit = hits[idx] + val label = labels[idx] ?: hit.label + for (from in groups[idx]) { + for (to in groups[idx + 1]) { + edges.add(FlowEdge(from, to, hit.link, hit.head, hit.tail, label, edges.size)) + } + } + } + return null + } + + private fun split(text: String, hits: List): MutableList { + val segs = mutableListOf() + var cursor = 0 + for (hit in hits) { + segs.add(text.substring(cursor, hit.start)) + cursor = hit.end + } + segs.add(text.substring(cursor)) + return segs + } + + /** Pulls `|label|` off the front of each right-hand segment, attaching it to the preceding link. */ + private fun pipes(segs: MutableList, count: Int): Array?> { + val out = arrayOfNulls>(count) + for (idx in 1 until segs.size) { + val trimmed = segs[idx].trimStart() + if (!trimmed.startsWith("|")) continue + val close = trimmed.indexOf('|', 1) + if (close < 0) continue + out[idx - 1] = Source.label(trimmed.substring(1, close)) + segs[idx] = trimmed.substring(close + 1) + } + return out + } + + private fun refs(segment: String): List? { + val out = mutableListOf() + for (token in parts(segment)) { + val id = ref(token) ?: continue + out.add(id) + } + return out.ifEmpty { null } + } + + /** Splits `A & B` groups at bracket depth zero. */ + private fun parts(segment: String): List { + val out = mutableListOf() + var start = 0 + for (idx in segment.indices) { + if (segment[idx] != '&') continue + if (!Source.open(segment, idx)) continue + out.add(segment.substring(start, idx)) + start = idx + 1 + } + out.add(segment.substring(start)) + return out + } + + private fun ref(token: String): String? { + val text = classes(token.trim()) + if (text.isEmpty()) return null + val wrap = WRAPS.firstOrNull { fits(text, it) } + if (wrap == null) { + add(text, listOf(text), Shape.Rect) + return text + } + val open = text.indexOf(wrap.open) + val id = text.substring(0, open) + val body = text.substring(open + wrap.open.length, text.length - wrap.close.length) + add(id, Source.label(body), wrap.shape) + return id + } + + private fun fits(text: String, wrap: Wrap): Boolean { + val open = text.indexOf(wrap.open) + if (open <= 0) return false + if (!text.endsWith(wrap.close)) return false + return text.length >= open + wrap.open.length + wrap.close.length + } + + /** Drops a trailing `:::class` assignment; class styling is not modelled. */ + private fun classes(text: String): String { + val cut = text.lastIndexOf(":::") + if (cut <= 0) return text + val tail = text.substring(cut + 3) + if (tail.isEmpty() || !tail.all { it.isLetterOrDigit() || it == '_' || it == '-' }) return text + return text.substring(0, cut) + } + + private fun add(id: String, label: List, shape: Shape) { + val prior = nodes[id] + if (prior == null) { + nodes[id] = FlowNode(id, label, shape, nodes.size, stack.lastOrNull()) + return + } + val implicit = prior.label == listOf(prior.id) && prior.shape == Shape.Rect + if (!implicit || label == listOf(id)) return + nodes[id] = prior.copy(label = label, shape = shape) + } + + private fun hits(text: String): List { + val out = mutableListOf() + var idx = 0 + while (idx < text.length) { + if (!Source.open(text, idx)) { + idx++ + continue + } + val hit = edge(text, idx) + if (hit == null) { + idx++ + continue + } + val back = hit.start > 0 && text[hit.start - 1] == '<' + out.add(if (back) hit.copy(start = hit.start - 1, tail = Head.Arrow) else hit) + idx = hit.end + } + return out + } + + private fun edge(text: String, at: Int): Hit? { + if (text[at] != '-' && text[at] != '=') return null + var idx = at + while (idx < text.length && Source.rail(text[idx])) idx++ + val rail = text.substring(at, idx) + if (rail.length < 2) return null + val head = headOf(text.getOrNull(idx)) + if (head != Head.None) idx++ + if (head != Head.None || rail.length >= 3) { + return Hit(at, idx, linkOf(rail), head, Head.None, emptyList()) + } + val rest = text.substring(idx) + val match = RAIL.find(rest) ?: return Hit(at, idx, linkOf(rail), Head.None, Head.None, emptyList()) + val end = idx + match.range.last + 1 + val label = Source.label(rest.substring(0, match.range.first)) + return Hit(at, end, linkOf(rail + match.value), headOf(match.value.last()), Head.None, label) + } + + private data class Hit( + val start: Int, + val end: Int, + val link: Link, + val head: Head, + val tail: Head, + val label: List, + ) + + private data class Wrap(val open: String, val close: String, val shape: Shape) + + private companion object { + val SKIP = setOf("classdef", "class", "click", "style", "linkstyle", "direction", "acctitle", "accdescr") + + val RAIL = Regex("""[-.=]{2,}[>ox]?""") + + val WRAPS = listOf( + Wrap("[[", "]]", Shape.Subroutine), + Wrap("[(", ")]", Shape.Cylinder), + Wrap("([", "])", Shape.Stadium), + Wrap("(((", ")))", Shape.Doubled), + Wrap("((", "))", Shape.Circle), + Wrap("[/", "/]", Shape.Skew), + Wrap("[\\", "\\]", Shape.SkewAlt), + Wrap("[/", "\\]", Shape.Trapezoid), + Wrap("[\\", "/]", Shape.TrapezoidAlt), + Wrap("{{", "}}", Shape.Hexagon), + Wrap("{", "}", Shape.Rhombus), + Wrap("(", ")", Shape.Round), + Wrap("[", "]", Shape.Rect), + Wrap(">", "]", Shape.Flag), + ) + + fun dirOf(text: String) = when (text.trim().uppercase()) { + "BT" -> Dir.Up + "LR" -> Dir.Right + "RL" -> Dir.Left + else -> Dir.Down + } + + fun headOf(char: Char?) = when (char) { + '>' -> Head.Arrow + 'o' -> Head.Dot + 'x' -> Head.Cross + else -> Head.None + } + + fun linkOf(rail: String) = when { + rail.contains('.') -> Link.Dotted + rail.contains('=') -> Link.Thick + else -> Link.Solid + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt new file mode 100644 index 00000000000..c44fc958eb2 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt @@ -0,0 +1,408 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Measure +import ai.kilocode.client.ui.diagram.Pt +import ai.kilocode.client.ui.diagram.Rect +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Spec +import kotlin.coroutines.coroutineContext +import kotlin.math.abs +import kotlin.math.max +import kotlinx.coroutines.ensureActive + +internal data class Slot(val id: String, val rect: Rect, val node: FlowNode?) + +internal data class Route(val edge: FlowEdge, val points: List) + +internal data class Placed(val graph: Graph, val slots: Map, val routes: List, val size: Size) + +/** + * Layered flowchart layout. All phases use fixed iteration counts and insertion-ordered maps so the + * result is byte-stable for a given input and [Measure]. + * + * Layout always runs in top-down space; `LR`/`RL` swap node extents up front and the finished + * geometry is transposed once at the end. + */ +internal class FlowLayout(private val measure: Measure, private val spec: Spec) { + private val gap get() = spec.metrics.gap + private val step get() = spec.metrics.rank + + suspend fun run(graph: Graph): Placed { + val flip = graph.dir == Dir.Left || graph.dir == Dir.Right + val sizes = sizes(graph, flip) + val links = graph.edges.filter { it.from != it.to && graph.nodes.containsKey(it.to) } + val backs = backs(graph, links) + val rank = ranks(graph.nodes.keys.toList(), links, backs) + coroutineContext.ensureActive() + + val paths = paths(links, rank, sizes) + val order = order(graph, rank, paths) + coroutineContext.ensureActive() + + val x = place(order, sizes, pairs(paths)) + val rows = rows(order, sizes) + val boxes = boxes(graph, order, sizes, x, rows) + coroutineContext.ensureActive() + + val routes = routes(graph, boxes, paths) + return finish(graph, boxes, routes, flip) + } + + private suspend fun sizes(graph: Graph, flip: Boolean): MutableMap { + val out = linkedMapOf() + for (node in graph.nodes.values) { + coroutineContext.ensureActive() + val size = size(node) + out[node.id] = if (flip) Size(size.h, size.w) else size + } + return out + } + + private fun size(node: FlowNode): Size { + val pad = spec.metrics.pad + val text = node.label.maxOf { measure.width(it, spec.font) } + val high = measure.height(spec.font) * node.label.size + return when (node.shape) { + Shape.Rhombus -> Size(text + pad * 4, high + pad * 4) + Shape.Circle, Shape.Doubled -> { + val side = max(text, high) + pad * 3 + Size(side, side) + } + Shape.Hexagon, Shape.Skew, Shape.SkewAlt, Shape.Trapezoid, Shape.TrapezoidAlt -> + Size(text + pad * 4, high + pad * 2) + else -> Size(text + pad * 2, high + pad * 2) + } + } + + /** Edge indices that close a cycle, found by DFS colouring in declaration order. */ + private fun backs(graph: Graph, links: List): Set { + val adj = linkedMapOf>() + for (id in graph.nodes.keys) adj[id] = mutableListOf() + links.forEachIndexed { idx, edge -> adj[edge.from]?.add(idx) } + val state = linkedMapOf() + val out = linkedSetOf() + + fun visit(id: String) { + state[id] = GRAY + for (idx in adj[id] ?: mutableListOf()) { + val to = links[idx].to + when (state[to] ?: WHITE) { + GRAY -> out.add(idx) + WHITE -> visit(to) + else -> Unit + } + } + state[id] = BLACK + } + + for (id in graph.nodes.keys) { + if ((state[id] ?: WHITE) == WHITE) visit(id) + } + return out + } + + /** Longest-path ranking by bounded relaxation over the acyclic orientation. */ + private fun ranks(ids: List, links: List, backs: Set): Map { + val out = linkedMapOf() + for (id in ids) out[id] = 0 + val dag = links.mapIndexed { idx, edge -> + if (idx in backs) edge.to to edge.from else edge.from to edge.to + } + repeat(ids.size) { + var moved = false + for ((from, to) in dag) { + val next = (out[from] ?: 0) + 1 + if ((out[to] ?: 0) >= next) continue + out[to] = next + moved = true + } + if (!moved) return out + } + return out + } + + /** Expands each edge into the chain of ranks it crosses, adding a virtual slot per crossed rank. */ + private fun paths(links: List, rank: Map, sizes: MutableMap): List { + val out = mutableListOf() + for (edge in links) { + val from = rank[edge.from] ?: 0 + val to = rank[edge.to] ?: 0 + val ids = mutableListOf(edge.from) + val dir = if (to >= from) 1 else -1 + var at = from + dir + while (at != to && from != to) { + val id = "~${edge.index}@$at" + sizes[id] = Size(spec.metrics.line, 0.0) + ids.add(id) + at += dir + } + ids.add(edge.to) + out.add(Path(edge, ids, from, to)) + } + return out + } + + private fun order( + graph: Graph, + rank: Map, + paths: List, + ): List> { + val depth = (rank.values.maxOrNull() ?: 0) + 1 + val out = List(depth) { mutableListOf() } + for (node in graph.nodes.values) out[rank[node.id] ?: 0].add(node.id) + for (path in paths) { + val dir = if (path.to >= path.from) 1 else -1 + path.ids.drop(1).dropLast(1).forEachIndexed { idx, id -> + val at = path.from + dir * (idx + 1) + if (at in out.indices) out[at].add(id) + } + } + sweep(graph, out, pairs(paths)) + return out + } + + private fun sweep( + graph: Graph, + order: List>, + pairs: List>, + ) { + val adj = adjacency(pairs) + val index = linkedMapOf() + order.forEach { ids -> ids.forEach { index[it] = index.size } } + repeat(SWEEPS) { pass -> + val down = pass % 2 == 0 + val ranks = if (down) order.indices.drop(1) else order.indices.reversed().drop(1) + for (at in ranks) { + val other = order[at + if (down) -1 else 1] + val slot = linkedMapOf() + other.forEachIndexed { idx, id -> slot[id] = idx.toDouble() } + val group = order[at].associateWith { key(graph, it) } + order[at].sortWith( + compareBy( + { group[it] }, + { median(adj[it]?.mapNotNull { peer -> slot[peer] } ?: emptyList()) ?: Double.MAX_VALUE }, + { index[it] ?: 0 }, + ), + ) + } + } + } + + /** Keeps subgraph members adjacent inside a rank; ungrouped nodes sort first. */ + private fun key(graph: Graph, id: String): Int { + val node = graph.nodes[id] ?: return -1 + val cluster = node.cluster ?: return -1 + return graph.clusters[cluster]?.index ?: -1 + } + + private fun pairs(paths: List): List> { + val out = mutableListOf>() + for (path in paths) { + path.ids.zipWithNext().forEach { out.add(it) } + } + return out + } + + private fun adjacency(pairs: List>): Map> { + val out = linkedMapOf>() + for ((from, to) in pairs) { + out.getOrPut(from) { mutableListOf() }.add(to) + out.getOrPut(to) { mutableListOf() }.add(from) + } + return out + } + + private fun place( + order: List>, + sizes: Map, + pairs: List>, + ): MutableMap { + val x = linkedMapOf() + for (ids in order) { + var cursor = 0.0 + for (id in ids) { + x[id] = cursor + cursor += width(sizes, id) + gap + } + } + val adj = adjacency(pairs) + repeat(PASSES) { pass -> align(order, sizes, adj, x, pass % 2 == 0) } + return x + } + + private fun align( + order: List>, + sizes: Map, + adj: Map>, + x: MutableMap, + down: Boolean, + ) { + val ranks = if (down) order.indices.drop(1) else order.indices.reversed().drop(1) + for (at in ranks) { + val other = order[at + if (down) -1 else 1] + val centers = linkedMapOf() + for (id in other) centers[id] = (x[id] ?: 0.0) + width(sizes, id) / 2 + var min = 0.0 + for (id in order[at]) { + val want = median(adj[id]?.mapNotNull { centers[it] } ?: emptyList()) + val wide = width(sizes, id) + val left = if (want == null) x.getValue(id) else want - wide / 2 + val at2 = max(min, left) + x[id] = at2 + min = at2 + wide + gap + } + } + } + + private fun rows(order: List>, sizes: Map): List { + var cursor = 0.0 + return order.map { ids -> + val top = cursor + cursor += (ids.maxOfOrNull { height(sizes, it) } ?: 0.0) + step + top + } + } + + private fun boxes( + graph: Graph, + order: List>, + sizes: Map, + x: Map, + rows: List, + ): MutableMap { + val out = linkedMapOf() + order.forEachIndexed { at, ids -> + val tall = ids.maxOfOrNull { height(sizes, it) } ?: 0.0 + for (id in ids) { + val wide = width(sizes, id) + val high = height(sizes, id) + val top = rows[at] + (tall - high) / 2 + out[id] = Slot(id, Rect(x[id] ?: 0.0, top, wide, high), graph.nodes[id]) + } + } + return out + } + + private fun routes(graph: Graph, boxes: Map, paths: List): List { + val out = mutableListOf() + val seen = linkedMapOf() + for (edge in graph.edges) { + val from = boxes[edge.from] ?: continue + if (edge.from == edge.to) { + out.add(Route(edge, loop(from.rect))) + continue + } + val to = boxes[edge.to] ?: continue + val path = paths.firstOrNull { it.edge.index == edge.index } ?: continue + val lane = seen.getOrDefault(lane(edge), 0) + seen[lane(edge)] = lane + 1 + out.add(Route(edge, trace(path, boxes, from.rect, to.rect, lane))) + } + return out + } + + private fun lane(edge: FlowEdge) = if (edge.from <= edge.to) "${edge.from}>${edge.to}" else "${edge.to}>${edge.from}" + + private fun trace(path: Path, boxes: Map, from: Rect, to: Rect, lane: Int): List { + val mid = path.ids.drop(1).dropLast(1).mapNotNull { boxes[it]?.rect }.map { Pt(it.x + it.w / 2, it.y) } + val bend = if (mid.isNotEmpty() || lane == 0) mid else listOf(bend(from, to, lane)) + val head = bend.firstOrNull() ?: Pt(to.x + to.w / 2, to.y + to.h / 2) + val tail = bend.lastOrNull() ?: Pt(from.x + from.w / 2, from.y + from.h / 2) + return listOf(exit(from, head)) + bend + listOf(exit(to, tail)) + } + + private fun bend(from: Rect, to: Rect, lane: Int): Pt { + val cx = (from.x + from.w / 2 + to.x + to.w / 2) / 2 + val cy = (from.y + from.h / 2 + to.y + to.h / 2) / 2 + return Pt(cx + lane * gap / 2, cy) + } + + private fun loop(rect: Rect): List { + val right = rect.x + rect.w + val out = right + gap / 2 + val top = rect.y + rect.h / 4 + val low = rect.y + rect.h * 3 / 4 + return listOf(Pt(right, top), Pt(out, top), Pt(out, low), Pt(right, low)) + } + + /** Point where the straight line from the rect centre towards [to] leaves the rect. */ + private fun exit(rect: Rect, to: Pt): Pt { + val cx = rect.x + rect.w / 2 + val cy = rect.y + rect.h / 2 + val dx = to.x - cx + val dy = to.y - cy + if (dx == 0.0 && dy == 0.0) return Pt(cx, cy) + val tx = if (dx == 0.0) Double.MAX_VALUE else rect.w / 2 / abs(dx) + val ty = if (dy == 0.0) Double.MAX_VALUE else rect.h / 2 / abs(dy) + val t = minOf(tx, ty) + return Pt(cx + dx * t, cy + dy * t) + } + + private fun finish(graph: Graph, boxes: Map, routes: List, flip: Boolean): Placed { + val real = boxes.values.filter { it.node != null } + val minX = real.minOfOrNull { it.rect.x } ?: 0.0 + val minY = real.minOfOrNull { it.rect.y } ?: 0.0 + val pad = spec.metrics.pad + val shifted = boxes.values.map { slot -> + slot.copy(rect = Rect(slot.rect.x - minX + pad, slot.rect.y - minY + pad, slot.rect.w, slot.rect.h)) + } + val moved = routes.map { route -> + route.copy(points = route.points.map { Pt(it.x - minX + pad, it.y - minY + pad) }) + } + val turned = Turn(graph.dir, flip, span(shifted, moved, pad)) + val slots = linkedMapOf() + for (slot in shifted) slots[slot.id] = slot.copy(rect = turned.rect(slot.rect)) + val out = moved.map { route -> route.copy(points = route.points.map { turned.point(it) }) } + return Placed(graph, slots, out, turned.size()) + } + + private fun span(slots: List, routes: List, pad: Double): Size { + val xs = slots.map { it.rect.x + it.rect.w } + routes.flatMap { route -> route.points.map { it.x } } + val ys = slots.map { it.rect.y + it.rect.h } + routes.flatMap { route -> route.points.map { it.y } } + return Size((xs.maxOrNull() ?: 0.0) + pad, (ys.maxOrNull() ?: 0.0) + pad) + } + + private fun width(sizes: Map, id: String) = sizes[id]?.w ?: 0.0 + + private fun height(sizes: Map, id: String) = sizes[id]?.h ?: 0.0 + + private fun median(values: List): Double? { + if (values.isEmpty()) return null + val sorted = values.sorted() + val mid = sorted.size / 2 + if (sorted.size % 2 == 1) return sorted[mid] + return (sorted[mid - 1] + sorted[mid]) / 2 + } + + private data class Path(val edge: FlowEdge, val ids: List, val from: Int, val to: Int) + + /** Single geometry transform applied after layout, so only top-down space is ever computed. */ + private class Turn(private val dir: Dir, private val flip: Boolean, private val bounds: Size) { + fun rect(rect: Rect): Rect = when (dir) { + Dir.Down -> rect + Dir.Up -> Rect(rect.x, bounds.h - rect.y - rect.h, rect.w, rect.h) + Dir.Right -> Rect(rect.y, rect.x, rect.h, rect.w) + Dir.Left -> Rect(bounds.h - rect.y - rect.h, rect.x, rect.h, rect.w) + } + + fun point(pt: Pt): Pt = when (dir) { + Dir.Down -> pt + Dir.Up -> Pt(pt.x, bounds.h - pt.y) + Dir.Right -> Pt(pt.y, pt.x) + Dir.Left -> Pt(bounds.h - pt.y, pt.x) + } + + fun size(): Size { + if (!flip) return bounds + return Size(bounds.h, bounds.w) + } + } + + private companion object { + const val WHITE = 0 + const val GRAY = 1 + const val BLACK = 2 + const val SWEEPS = 4 + const val PASSES = 2 + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt new file mode 100644 index 00000000000..69076212f9e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt @@ -0,0 +1,229 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Anchor +import ai.kilocode.client.ui.diagram.Mark +import ai.kilocode.client.ui.diagram.Measure +import ai.kilocode.client.ui.diagram.Pt +import ai.kilocode.client.ui.diagram.Rect +import ai.kilocode.client.ui.diagram.Role +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.Type +import kotlin.math.max + +/** Turns laid-out flowchart geometry into marks. Cluster frames are emitted first so they paint behind. */ +internal class FlowMarks(private val measure: Measure, private val spec: Spec) { + private val pad get() = spec.metrics.pad + + fun run(placed: Placed): Scene { + val frames = frames(placed) + val dx = -minOf(0.0, frames.minOfOrNull { it.rect.x } ?: 0.0) + val dy = -minOf(0.0, frames.minOfOrNull { it.rect.y } ?: 0.0) + val marks = mutableListOf() + for (frame in frames) marks.add(group(frame, dx, dy)) + for (route in placed.routes) { + marks.add(line(route, dx, dy)) + marks.addAll(tag(route, dx, dy)) + } + for (slot in placed.slots.values) { + val node = slot.node ?: continue + val rect = move(slot.rect, dx, dy) + marks.addAll(shape(node, rect)) + marks.addAll(label(node.label, rect, Role.Text)) + } + return Scene(Type.Flowchart, marks, size(placed, frames, dx, dy)) + } + + private fun size(placed: Placed, frames: List, dx: Double, dy: Double): Size { + val wide = max(placed.size.w + dx, frames.maxOfOrNull { it.rect.x + it.rect.w + dx } ?: 0.0) + val high = max(placed.size.h + dy, frames.maxOfOrNull { it.rect.y + it.rect.h + dy } ?: 0.0) + return Size(wide, high) + } + + private fun frames(placed: Placed): List { + val out = mutableListOf() + for (cluster in placed.graph.clusters.values) { + val rects = members(placed.graph, cluster.id).mapNotNull { placed.slots[it]?.rect } + if (rects.isEmpty()) continue + val room = pad * 2 * (1 + deep(placed.graph, cluster.id)) + val title = measure.height(spec.font) * cluster.label.size + val x = rects.minOf { it.x } - room + val y = rects.minOf { it.y } - room - title + val wide = rects.maxOf { it.x + it.w } + room - x + val high = rects.maxOf { it.y + it.h } + room - y + out.add(Frame(cluster, Rect(x, y, wide, high))) + } + return out + } + + private fun members(graph: Graph, id: String): List { + val out = mutableListOf() + for (node in graph.nodes.values) { + if (node.cluster == id) out.add(node.id) + } + for (cluster in graph.clusters.values) { + if (cluster.parent == id) out.addAll(members(graph, cluster.id)) + } + return out + } + + /** Nesting depth below [id]; used so an outer frame reserves room for the frames inside it. */ + private fun deep(graph: Graph, id: String): Int { + val kids = graph.clusters.values.filter { it.parent == id } + if (kids.isEmpty()) return 0 + return 1 + (kids.maxOfOrNull { deep(graph, it.id) } ?: 0) + } + + private fun group(frame: Frame, dx: Double, dy: Double): Mark { + val rect = move(frame.rect, dx, dy) + val box = Mark.Box(rect, spec.metrics.arc, null, Role.Cluster, dash = true) + val high = measure.height(spec.font) + val title = frame.cluster.label.mapIndexed { idx, text -> + Mark.Text(text, Pt(rect.x + rect.w / 2, rect.y + pad + high * (idx + HALF)), Anchor.Center, Role.Muted, true) + } + return Mark.Group(frame.cluster.id, listOf(box) + title) + } + + private fun line(route: Route, dx: Double, dy: Double): Mark { + val points = route.points.map { Pt(it.x + dx, it.y + dy) } + return Mark.Edge( + points, + Role.Line, + dash = route.edge.link == Link.Dotted, + thick = route.edge.link == Link.Thick, + head = route.edge.head, + tail = route.edge.tail, + ) + } + + private fun tag(route: Route, dx: Double, dy: Double): List { + if (route.edge.label.isEmpty()) return emptyList() + val at = mid(route.points) + val high = measure.height(spec.font) + val top = at.y + dy - high * route.edge.label.size / 2 + return route.edge.label.mapIndexed { idx, text -> + Mark.Text(text, Pt(at.x + dx, top + high * (idx + HALF)), Anchor.Center, Role.Muted) + } + } + + private fun mid(points: List): Pt { + if (points.isEmpty()) return Pt(0.0, 0.0) + if (points.size % 2 == 1) return points[points.size / 2] + val left = points[points.size / 2 - 1] + val right = points[points.size / 2] + return Pt((left.x + right.x) / 2, (left.y + right.y) / 2) + } + + private fun label(lines: List, rect: Rect, role: Role): List { + val high = measure.height(spec.font) + val top = rect.y + (rect.h - high * lines.size) / 2 + return lines.mapIndexed { idx, text -> + Mark.Text(text, Pt(rect.x + rect.w / 2, top + high * (idx + HALF)), Anchor.Center, role) + } + } + + private fun shape(node: FlowNode, rect: Rect): List = when (node.shape) { + Shape.Rect -> listOf(box(rect, 0.0)) + Shape.Round, Shape.Cylinder -> listOf(box(rect, spec.metrics.arc * 2)) + Shape.Stadium -> listOf(box(rect, rect.h / 2)) + Shape.Subroutine -> listOf(box(rect, 0.0)) + bars(rect) + Shape.Circle -> listOf(Mark.Oval(rect, Role.Surface, Role.Border)) + Shape.Doubled -> listOf(Mark.Oval(rect, Role.Surface, Role.Border), Mark.Oval(inset(rect), null, Role.Border)) + Shape.Rhombus -> listOf(poly(diamond(rect))) + Shape.Hexagon -> listOf(poly(hexagon(rect))) + Shape.Skew -> listOf(poly(skew(rect, false))) + Shape.SkewAlt -> listOf(poly(skew(rect, true))) + Shape.Trapezoid -> listOf(poly(trapezoid(rect, false))) + Shape.TrapezoidAlt -> listOf(poly(trapezoid(rect, true))) + Shape.Flag -> listOf(poly(flag(rect))) + } + + private fun box(rect: Rect, arc: Double) = Mark.Box(rect, arc, Role.Surface, Role.Border) + + private fun poly(points: List) = Mark.Poly(points, Role.Surface, Role.Border) + + private fun bars(rect: Rect): List { + val left = rect.x + pad + val right = rect.x + rect.w - pad + val top = rect.y + val low = rect.y + rect.h + return listOf( + Mark.Edge(listOf(Pt(left, top), Pt(left, low)), Role.Border), + Mark.Edge(listOf(Pt(right, top), Pt(right, low)), Role.Border), + ) + } + + private fun inset(rect: Rect): Rect { + val room = pad / 2 + return Rect(rect.x + room, rect.y + room, rect.w - room * 2, rect.h - room * 2) + } + + private fun diamond(rect: Rect) = listOf( + Pt(rect.x + rect.w / 2, rect.y), + Pt(rect.x + rect.w, rect.y + rect.h / 2), + Pt(rect.x + rect.w / 2, rect.y + rect.h), + Pt(rect.x, rect.y + rect.h / 2), + ) + + private fun hexagon(rect: Rect): List { + val cut = minOf(pad * 2, rect.w / 3) + return listOf( + Pt(rect.x + cut, rect.y), + Pt(rect.x + rect.w - cut, rect.y), + Pt(rect.x + rect.w, rect.y + rect.h / 2), + Pt(rect.x + rect.w - cut, rect.y + rect.h), + Pt(rect.x + cut, rect.y + rect.h), + Pt(rect.x, rect.y + rect.h / 2), + ) + } + + private fun skew(rect: Rect, back: Boolean): List { + val cut = minOf(pad * 2, rect.w / 4) + val lean = if (back) -cut else cut + return listOf( + Pt(rect.x + max(0.0, lean), rect.y), + Pt(rect.x + rect.w + minOf(0.0, lean), rect.y), + Pt(rect.x + rect.w - max(0.0, lean), rect.y + rect.h), + Pt(rect.x - minOf(0.0, lean), rect.y + rect.h), + ) + } + + private fun trapezoid(rect: Rect, flip: Boolean): List { + val cut = minOf(pad * 2, rect.w / 4) + if (flip) { + return listOf( + Pt(rect.x, rect.y), + Pt(rect.x + rect.w, rect.y), + Pt(rect.x + rect.w - cut, rect.y + rect.h), + Pt(rect.x + cut, rect.y + rect.h), + ) + } + return listOf( + Pt(rect.x + cut, rect.y), + Pt(rect.x + rect.w - cut, rect.y), + Pt(rect.x + rect.w, rect.y + rect.h), + Pt(rect.x, rect.y + rect.h), + ) + } + + private fun flag(rect: Rect): List { + val cut = minOf(pad * 2, rect.w / 5) + return listOf( + Pt(rect.x, rect.y), + Pt(rect.x + rect.w - cut, rect.y), + Pt(rect.x + rect.w, rect.y + rect.h / 2), + Pt(rect.x + rect.w - cut, rect.y + rect.h), + Pt(rect.x, rect.y + rect.h), + Pt(rect.x + cut, rect.y + rect.h / 2), + ) + } + + private fun move(rect: Rect, dx: Double, dy: Double) = Rect(rect.x + dx, rect.y + dy, rect.w, rect.h) + + private data class Frame(val cluster: Cluster, val rect: Rect) + + private companion object { + const val HALF = 0.5 + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt new file mode 100644 index 00000000000..3db8af82e44 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt @@ -0,0 +1,58 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.Fault +import ai.kilocode.client.ui.diagram.Measure +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.Type + +/** + * In-process mermaid engine covering flowcharts and sequence diagrams. + * + * [Measure] is a construction-time capability rather than part of [Spec] so an out-of-process engine + * can implement the same interface while doing its own text measurement. + */ +internal class Mermaid(private val measure: Measure) : Engine { + override fun accepts(type: Type) = type == Type.Flowchart || type == Type.Sequence + + override suspend fun draw(source: String, spec: Spec): Out { + val clean = Source.clean(source) + if (clean.lines.size > spec.limits.lines) { + return Out.Err(Fault.Limit, "source exceeds ${spec.limits.lines} lines") + } + val type = Type.of(clean) + if (!accepts(type)) return Out.Err(Fault.Unsupported, "unsupported diagram type: $type") + if (type == Type.Flowchart) return flow(clean, spec) + return seq(clean, spec) + } + + private suspend fun flow(clean: Clean, spec: Spec): Out { + val parsed = Flow().parse(clean) + if (parsed is FlowOut.Err) return Out.Err(Fault.Syntax, parsed.message, parsed.line) + val graph = (parsed as FlowOut.Ok).graph + if (graph.nodes.isEmpty()) return Out.Err(Fault.Syntax, "flowchart has no nodes") + if (graph.nodes.size > spec.limits.nodes) { + return Out.Err(Fault.Limit, "flowchart exceeds ${spec.limits.nodes} nodes") + } + if (graph.edges.size > spec.limits.edges) { + return Out.Err(Fault.Limit, "flowchart exceeds ${spec.limits.edges} links") + } + val placed = FlowLayout(measure, spec).run(graph) + return Out.Ok(FlowMarks(measure, spec).run(placed)) + } + + private suspend fun seq(clean: Clean, spec: Spec): Out { + val parsed = Seq().parse(clean) + if (parsed is SeqOut.Err) return Out.Err(Fault.Syntax, parsed.message, parsed.line) + val script = (parsed as SeqOut.Ok).script + if (script.actors.isEmpty()) return Out.Err(Fault.Syntax, "sequence diagram has no participants") + if (script.actors.size > spec.limits.nodes) { + return Out.Err(Fault.Limit, "sequence diagram exceeds ${spec.limits.nodes} participants") + } + if (script.steps.size > spec.limits.edges) { + return Out.Err(Fault.Limit, "sequence diagram exceeds ${spec.limits.edges} steps") + } + return Out.Ok(SeqLayout(measure, spec).run(script)) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt new file mode 100644 index 00000000000..ce24ea24c2e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt @@ -0,0 +1,182 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Head + +internal enum class NoteAt { Left, Right, Over } + +internal enum class BlockKind { Loop, Alt, Opt, Par, Critical, Break } + +internal data class Actor(val id: String, val label: List, val index: Int) + +internal sealed interface Step { + data class Msg( + val from: String, + val to: String, + val label: List, + val link: Link, + val head: Head, + ) : Step + + data class Note(val at: NoteAt, val actors: List, val label: List) : Step + data class Open(val kind: BlockKind, val label: List) : Step + data class Split(val label: List) : Step + data class Toggle(val actor: String, val on: Boolean) : Step + data object Close : Step +} + +internal data class Script( + val actors: Map, + val steps: List, + val title: List, + val numbered: Boolean, +) + +internal sealed interface SeqOut { + data class Ok(val script: Script) : SeqOut + data class Err(val message: String, val line: Int) : SeqOut +} + +/** Line-oriented sequence diagram parser. Unknown statements are skipped rather than failing. */ +internal class Seq { + private val actors = linkedMapOf() + private val steps = mutableListOf() + private var title = emptyList() + private var numbered = false + private var depth = 0 + + fun parse(clean: Clean): SeqOut { + var first = true + for (line in clean.lines) { + val text = line.text.trim() + if (text.isEmpty()) continue + if (first) { + first = false + if (text.substringBefore(' ').lowercase() == "sequencediagram") continue + } + val err = stmt(text, line.at) ?: continue + return SeqOut.Err(err, line.at) + } + if (depth > 0) return SeqOut.Err("block is missing a matching end", clean.lines.lastOrNull()?.at ?: 1) + return SeqOut.Ok(Script(actors, steps, title, numbered)) + } + + private fun stmt(text: String, at: Int): String? { + val token = text.substringBefore(' ').lowercase() + when (token) { + "autonumber" -> { + numbered = true + return null + } + "title" -> { + title = Source.label(text.substringAfter(' ', "").removePrefix(":").trim()) + return null + } + "participant", "actor" -> return actor(text) + "activate", "deactivate" -> { + val id = text.substringAfter(' ', "").trim() + if (id.isEmpty()) return "$token needs a participant" + steps.add(Step.Toggle(name(id), token == "activate")) + return null + } + "end" -> { + if (depth == 0) return "end without a matching block" + depth-- + steps.add(Step.Close) + return null + } + "else", "and" -> { + if (depth == 0) return "$token outside a block" + steps.add(Step.Split(Source.label(text.substringAfter(' ', "").trim()))) + return null + } + in BLOCKS.keys -> { + depth++ + steps.add(Step.Open(BLOCKS.getValue(token), Source.label(text.substringAfter(' ', "").trim()))) + return null + } + else -> Unit + } + if (token == "note") return note(text, at) + if (token in SKIP) return null + return message(text, at) + } + + private fun actor(text: String): String? { + val rest = text.substringAfter(' ', "").trim() + if (rest.isEmpty()) return "participant needs a name" + val cut = rest.indexOf(" as ") + val id = if (cut < 0) rest else rest.substring(0, cut).trim() + val label = if (cut < 0) rest else rest.substring(cut + 4).trim() + add(id, Source.label(label)) + return null + } + + private fun note(text: String, at: Int): String? { + val match = NOTE.find(text) ?: return null + val where = match.groupValues[1].lowercase() + val kind = when { + where.startsWith("left") -> NoteAt.Left + where.startsWith("right") -> NoteAt.Right + else -> NoteAt.Over + } + val targets = match.groupValues[2].split(',').map { name(it.trim()) }.filter { it.isNotEmpty() } + if (targets.isEmpty()) return "note on line $at needs a participant" + targets.forEach { add(it, listOf(it)) } + steps.add(Step.Note(kind, targets, Source.label(match.groupValues[3]))) + return null + } + + private fun message(text: String, at: Int): String? { + val match = MSG.find(text) ?: return null + val from = name(match.groupValues[1]) + val arrow = match.groupValues[2] + val sign = match.groupValues[3] + val to = name(match.groupValues[4]) + if (from.isEmpty() || to.isEmpty()) return "message on line $at needs both participants" + add(from, listOf(from)) + add(to, listOf(to)) + if (sign == "+") steps.add(Step.Toggle(to, true)) + steps.add(Step.Msg(from, to, Source.label(match.groupValues[5]), linkOf(arrow), headOf(arrow))) + if (sign == "-") steps.add(Step.Toggle(from, false)) + return null + } + + private fun add(id: String, label: List) { + if (id.isEmpty()) return + val prior = actors[id] + if (prior == null) { + actors[id] = Actor(id, label, actors.size) + return + } + if (label == listOf(id) || prior.label != listOf(id)) return + actors[id] = prior.copy(label = label) + } + + private fun name(text: String) = Source.unquote(text.trim()).trim() + + private companion object { + val BLOCKS = mapOf( + "loop" to BlockKind.Loop, + "alt" to BlockKind.Alt, + "opt" to BlockKind.Opt, + "par" to BlockKind.Par, + "critical" to BlockKind.Critical, + "break" to BlockKind.Break, + ) + + val SKIP = setOf("box", "rect", "link", "links", "accdescr", "acctitle", "create", "destroy", "option") + + val NOTE = Regex("""^[Nn]ote\s+(left of|right of|over)\s+([^:]+):\s*(.*)$""") + + val MSG = Regex("""^(.+?)\s*(--?>>|--?>|--?[x)])\s*([+-]?)\s*(.+?)\s*:\s*(.*)$""") + + fun linkOf(arrow: String) = if (arrow.startsWith("--")) Link.Dotted else Link.Solid + + fun headOf(arrow: String) = when { + arrow.endsWith(">>") -> Head.Arrow + arrow.endsWith("x") -> Head.Cross + arrow.endsWith(")") -> Head.Dot + else -> Head.Open + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt new file mode 100644 index 00000000000..93925903673 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt @@ -0,0 +1,264 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Anchor +import ai.kilocode.client.ui.diagram.Mark +import ai.kilocode.client.ui.diagram.Measure +import ai.kilocode.client.ui.diagram.Pt +import ai.kilocode.client.ui.diagram.Rect +import ai.kilocode.client.ui.diagram.Role +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.Type +import kotlin.coroutines.coroutineContext +import kotlin.math.max +import kotlinx.coroutines.ensureActive + +/** + * Sequence diagram geometry. Participants form columns, steps advance a single cursor downwards, so + * the layout is deterministic without any graph work. + */ +internal class SeqLayout(private val measure: Measure, private val spec: Spec) { + private val pad get() = spec.metrics.pad + private val gap get() = spec.metrics.gap + private val step get() = spec.metrics.rank + + private val marks = mutableListOf() + private val heads = linkedMapOf() + private val live = linkedMapOf>() + private val blocks = ArrayDeque() + private var cursor = 0.0 + private var count = 0 + private var wide = 0.0 + + suspend fun run(script: Script): Scene { + val high = measure.height(spec.font) + cursor = pad + title(script, high) + columns(script, high) + cursor += pad + step + for (item in script.steps) { + coroutineContext.ensureActive() + when (item) { + is Step.Msg -> message(item, script.numbered, high) + is Step.Note -> note(item, high) + is Step.Open -> open(item, high) + is Step.Split -> split(item, high) + is Step.Toggle -> toggle(item) + Step.Close -> close() + } + } + return scene(lines() + marks) + } + + private fun title(script: Script, high: Double): Double { + if (script.title.isEmpty()) return 0.0 + script.title.forEachIndexed { idx, text -> + marks.add(Mark.Text(text, Pt(pad, pad + high * (idx + HALF)), Anchor.Left, Role.Text, true)) + } + return high * script.title.size + pad + } + + private suspend fun columns(script: Script, high: Double) { + var cursorX = pad + for (actor in script.actors.values) { + coroutineContext.ensureActive() + val text = actor.label.maxOf { measure.width(it, spec.font) } + val box = Rect(cursorX, cursor, text + pad * 4, high * actor.label.size + pad * 2) + heads[actor.id] = box + marks.add(Mark.Box(box, spec.metrics.arc, Role.Surface, Role.Border)) + actor.label.forEachIndexed { idx, label -> + val at = Pt(box.x + box.w / 2, box.y + pad + high * (idx + HALF)) + marks.add(Mark.Text(label, at, Anchor.Center, Role.Text, true)) + } + cursorX = box.x + box.w + gap * 2 + wide = max(wide, box.x + box.w) + } + cursor += heads.values.maxOfOrNull { it.h } ?: 0.0 + } + + private fun message(item: Step.Msg, numbered: Boolean, high: Double) { + val from = center(item.from) ?: return + val to = center(item.to) ?: return + count++ + val label = if (numbered) prefix(item.label) else item.label + if (item.from == item.to) { + self(item, from, label, high) + return + } + cursor += high * label.size + pad + val fromX = edge(item.from, from, to > from) + val toX = edge(item.to, to, from > to) + marks.add(line(listOf(Pt(fromX, cursor), Pt(toX, cursor)), item)) + val top = cursor - high * label.size - pad / 2 + label.forEachIndexed { idx, text -> + val at = Pt((fromX + toX) / 2, top + high * (idx + HALF)) + marks.add(Mark.Text(text, at, Anchor.Center, Role.Muted)) + } + cursor += step + } + + private fun self(item: Step.Msg, at: Double, label: List, high: Double) { + val out = at + gap * 2 + val top = cursor + pad + val low = top + max(step, high * label.size + pad) + val from = edge(item.from, at, true) + marks.add(line(listOf(Pt(from, top), Pt(out, top), Pt(out, low), Pt(from, low)), item)) + label.forEachIndexed { idx, text -> + marks.add(Mark.Text(text, Pt(out + pad, top + high * (idx + HALF)), Anchor.Left, Role.Muted)) + } + wide = max(wide, out + pad + label.maxOf { measure.width(it, spec.font) }) + cursor = low + step + } + + private fun note(item: Step.Note, high: Double) { + val rects = item.actors.mapNotNull { heads[it] } + if (rects.isEmpty()) return + val text = item.label.maxOfOrNull { measure.width(it, spec.font) } ?: 0.0 + val body = text + pad * 4 + val tall = high * item.label.size + pad * 2 + val anchor = rects.first() + val rect = when (item.at) { + NoteAt.Left -> Rect(anchor.x + anchor.w / 2 - gap - body, cursor, body, tall) + NoteAt.Right -> Rect(anchor.x + anchor.w / 2 + gap, cursor, body, tall) + NoteAt.Over -> over(rects, body, tall) + } + marks.add(Mark.Box(rect, spec.metrics.arc, Role.Note, Role.Border)) + item.label.forEachIndexed { idx, label -> + val at = Pt(rect.x + rect.w / 2, rect.y + pad + high * (idx + HALF)) + marks.add(Mark.Text(label, at, Anchor.Center, Role.Text)) + } + wide = max(wide, rect.x + rect.w) + cursor = rect.y + rect.h + step + } + + private fun over(rects: List, body: Double, tall: Double): Rect { + val left = rects.minOf { it.x + it.w / 2 } + val right = rects.maxOf { it.x + it.w / 2 } + val span = max(body, right - left + body / 2) + return Rect((left + right) / 2 - span / 2, cursor, span, tall) + } + + private fun open(item: Step.Open, high: Double) { + val label = kind(item.kind) + item.label.firstOrNull().orEmpty() + blocks.addLast(Frame(cursor, label)) + cursor += high + pad * 2 + } + + private fun split(item: Step.Split, high: Double) { + if (blocks.isEmpty()) return + val inset = pad * blocks.size + marks.add( + Mark.Edge(listOf(Pt(inset, cursor), Pt(max(inset, wide - inset), cursor)), Role.Cluster, dash = true), + ) + val text = item.label.firstOrNull().orEmpty() + if (text.isNotEmpty()) marks.add(Mark.Text(text, Pt(inset + pad, cursor + high * HALF), Anchor.Left, Role.Muted)) + cursor += high + pad + } + + private fun close() { + val frame = blocks.removeLastOrNull() ?: return + val inset = pad * (blocks.size + 1) + val rect = Rect(inset, frame.top, max(pad, wide - inset * 2), cursor - frame.top) + val high = measure.height(spec.font) + val tab = Rect(rect.x, rect.y, measure.width(frame.text, spec.font) + pad * 2, high + pad) + marks.add(Mark.Box(rect, spec.metrics.arc, null, Role.Cluster, dash = true)) + marks.add(Mark.Box(tab, spec.metrics.arc, Role.Note, Role.Cluster)) + marks.add(Mark.Text(frame.text, Pt(tab.x + pad, tab.y + tab.h / 2), Anchor.Left, Role.Muted, true)) + cursor += pad + } + + private fun toggle(item: Step.Toggle) { + val stack = live.getOrPut(item.actor) { mutableListOf() } + if (item.on) { + stack.add(cursor) + return + } + val start = stack.removeLastOrNull() ?: return + val at = center(item.actor) ?: return + val width = pad + val rect = Rect(at - width / 2 + stack.size * width, start, width, max(step / 2, cursor - start)) + marks.add(Mark.Box(rect, 0.0, Role.Accent, Role.Border)) + } + + /** Lifelines are emitted first so messages and boxes paint on top of them. */ + private fun lines(): List = heads.values.map { rect -> + val at = rect.x + rect.w / 2 + Mark.Edge(listOf(Pt(at, rect.y + rect.h), Pt(at, cursor + step / 2)), Role.Muted, dash = true) + } + + private fun line(points: List, item: Step.Msg) = Mark.Edge( + points, + Role.Line, + dash = item.link == Link.Dotted, + head = item.head, + ) + + /** Shifts an endpoint clear of any activation bar currently open on that participant. */ + private fun edge(actor: String, at: Double, rightward: Boolean): Double { + val open = live[actor]?.size ?: 0 + if (open == 0) return at + val shift = pad / 2 + (open - 1) * pad + return if (rightward) at + shift else at - shift + } + + private fun center(actor: String): Double? { + val rect = heads[actor] ?: return null + return rect.x + rect.w / 2 + } + + private fun prefix(label: List): List { + if (label.isEmpty()) return listOf("$count") + return listOf("$count. ${label.first()}") + label.drop(1) + } + + private fun scene(source: List): Scene { + val pts = source.flatMap(::pts) + val minX = pts.minOfOrNull { it.x } ?: 0.0 + val minY = pts.minOfOrNull { it.y } ?: 0.0 + val dx = -minOf(0.0, minX) + val dy = -minOf(0.0, minY) + val marks = source.map { move(it, dx, dy) } + val moved = pts.map { Pt(it.x + dx, it.y + dy) } + val size = Size((moved.maxOfOrNull { it.x } ?: 0.0) + pad, (moved.maxOfOrNull { it.y } ?: 0.0) + pad) + return Scene(Type.Sequence, marks, size) + } + + private fun pts(mark: Mark): List = when (mark) { + is Mark.Box -> corners(mark.rect) + is Mark.Oval -> corners(mark.rect) + is Mark.Poly -> mark.points + is Mark.Edge -> mark.points + is Mark.Text -> listOf(mark.at) + is Mark.Group -> mark.marks.flatMap(::pts) + } + + private fun corners(rect: Rect) = listOf(Pt(rect.x, rect.y), Pt(rect.x + rect.w, rect.y + rect.h)) + + private fun move(mark: Mark, dx: Double, dy: Double): Mark = when (mark) { + is Mark.Box -> mark.copy(rect = move(mark.rect, dx, dy)) + is Mark.Oval -> mark.copy(rect = move(mark.rect, dx, dy)) + is Mark.Poly -> mark.copy(points = mark.points.map { move(it, dx, dy) }) + is Mark.Edge -> mark.copy(points = mark.points.map { move(it, dx, dy) }) + is Mark.Text -> mark.copy(at = move(mark.at, dx, dy)) + is Mark.Group -> mark.copy(marks = mark.marks.map { move(it, dx, dy) }) + } + + private fun move(rect: Rect, dx: Double, dy: Double) = Rect(rect.x + dx, rect.y + dy, rect.w, rect.h) + + private fun move(pt: Pt, dx: Double, dy: Double) = Pt(pt.x + dx, pt.y + dy) + + private data class Frame(val top: Double, val text: String) + + private companion object { + const val HALF = 0.5 + + fun kind(kind: BlockKind) = when (kind) { + BlockKind.Loop -> "loop " + BlockKind.Alt -> "alt " + BlockKind.Opt -> "opt " + BlockKind.Par -> "par " + BlockKind.Critical -> "critical " + BlockKind.Break -> "break " + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt new file mode 100644 index 00000000000..f41172cf537 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt @@ -0,0 +1,93 @@ +package ai.kilocode.client.ui.diagram.mermaid + +/** One preprocessed line plus the 1-based line number it came from in the user's text. */ +internal data class Line(val text: String, val at: Int) + +internal data class Clean(val lines: List) + +/** + * Mermaid source preprocessing. + * + * Normalizes line endings and tabs, drops leading YAML frontmatter, blanks `%%{ ... }%%` init + * directives without shifting line numbers, and cuts `%%` comments outside quoted strings. Every + * surviving line keeps its original line number so parse failures can point at the user's text. + */ +internal object Source { + private const val RAILS = "-.=" + private val DIRECTIVE = Regex("""%%\{[\s\S]*?}%%""") + + fun clean(text: String): Clean { + val raw = mask(normalize(text)).split("\n") + val start = front(raw) + val out = ArrayList(raw.size - start) + for (idx in start until raw.size) out.add(Line(cut(raw[idx]), idx + 1)) + return Clean(out) + } + + /** Splits label text on the break forms mermaid accepts. */ + fun label(text: String): List { + val body = unquote(text.trim()) + val parts = body.split("
", "
", "
", "\\n", "\n").map { it.trim() } + val kept = parts.filter { it.isNotEmpty() } + if (kept.isEmpty()) return listOf("") + return kept + } + + fun unquote(text: String): String { + if (text.length < 2) return text + if (text.first() == '"' && text.last() == '"') return text.substring(1, text.length - 1) + return text + } + + /** True when [index] sits outside quotes and outside any bracket group. */ + fun open(text: String, index: Int): Boolean { + var quote = false + var depth = 0 + for (idx in 0 until index) { + val char = text[idx] + if (char == '"') quote = !quote + if (quote) continue + if (char == '[' || char == '(' || char == '{') depth++ + if (char == ']' || char == ')' || char == '}') depth-- + } + return !quote && depth <= 0 + } + + fun rail(char: Char) = RAILS.indexOf(char) >= 0 + + private fun normalize(text: String) = text + .replace("\r\n", "\n") + .replace("\r", "\n") + .replace("\t", " ") + + private fun mask(text: String) = DIRECTIVE.replace(text) { match -> + match.value.map { if (it == '\n') '\n' else ' ' }.joinToString("") + } + + /** Returns the index of the first content line, skipping terminated frontmatter. */ + private fun front(raw: List): Int { + var idx = 0 + while (idx < raw.size && raw[idx].isBlank()) idx++ + if (idx >= raw.size || raw[idx].trim() != "---") return 0 + var scan = idx + 1 + while (scan < raw.size) { + if (raw[scan].trim() == "---") return scan + 1 + scan++ + } + return 0 + } + + private fun cut(line: String): String { + var quote = false + var idx = 0 + while (idx < line.length) { + val char = line[idx] + if (char == '"') quote = !quote + if (!quote && char == '%' && idx + 1 < line.length && line[idx + 1] == '%') { + return line.substring(0, idx).trimEnd() + } + idx++ + } + return line + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt new file mode 100644 index 00000000000..31687620a71 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt @@ -0,0 +1,50 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job +import kotlinx.coroutines.async +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Cancellation is driven from the measurement hook rather than a timeout so the test proves the + * cooperative checks exist without depending on timing. + */ +class CancelTest { + @Test + fun `flowchart layout stops when the job is cancelled`() { + val source = "flowchart TD\n" + (1..40).joinToString("\n") { " n$it --> n${it + 1}" } + + assertFailsWith { cancel(source) } + } + + @Test + fun `sequence layout stops when the job is cancelled`() { + val source = "sequenceDiagram\n" + (1..40).joinToString("\n") { " p$it->>p${it + 1}: step $it" } + + assertFailsWith { cancel(source) } + } + + @Test + fun `uncancelled work completes`() { + val out = runBlocking { Mermaid(FakeMeasure()).draw("flowchart TD\n A --> B", spec()) } + + assertTrue(scene(out).marks.isNotEmpty()) + } + + private fun cancel(source: String) = runBlocking { + val job = Job() + val scope = CoroutineScope(job + Dispatchers.Unconfined) + val measure = FakeMeasure { calls -> if (calls >= CUT) job.cancel() } + scope.async { Mermaid(measure).draw(source, spec()) }.await() + } + + private companion object { + const val CUT = 3 + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ConformanceTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ConformanceTest.kt new file mode 100644 index 00000000000..6f6af6c29bc --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ConformanceTest.kt @@ -0,0 +1,78 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * The contract any diagram engine must satisfy, not a test of one implementation. A replacement + * engine should be pointed at this corpus first. + */ +class ConformanceTest { + private val engine = Mermaid(FakeMeasure()) + + @Test + fun `every corpus diagram produces a finite scene`() { + for (name in CORPUS) { + val out = runBlocking { engine.draw(read(name), spec()) } + val scene = scene(out) + + assertTrue(scene.marks.isNotEmpty(), "$name produced no marks") + assertTrue(scene.size.w > 0 && scene.size.h > 0, "$name has an empty size ${scene.size}") + assertTrue(scene.size.w.isFinite() && scene.size.h.isFinite(), "$name has a non-finite size") + } + } + + @Test + fun `corpus diagrams report the detected type`() { + for (name in CORPUS) { + val out = runBlocking { engine.draw(read(name), spec()) } + val expected = if (name.startsWith("flow")) Type.Flowchart else Type.Sequence + + assertEquals(expected, scene(out).type, "$name resolved the wrong type") + } + } + + @Test + fun `rendering is deterministic across runs`() { + for (name in CORPUS) { + val first = runBlocking { engine.draw(read(name), spec()) } + val second = runBlocking { Mermaid(FakeMeasure()).draw(read(name), spec()) } + + assertEquals(scene(first).toString(), scene(second).toString(), "$name is not deterministic") + } + } + + @Test + fun `text marks never lose their content`() { + for (name in CORPUS) { + val out = runBlocking { engine.draw(read(name), spec()) } + val texts = flatten(scene(out).marks).filterIsInstance() + + assertTrue(texts.isNotEmpty(), "$name produced no labels") + assertTrue(texts.none { it.text.isEmpty() }, "$name produced an empty label") + } + } + + private fun read(name: String): String { + val stream = javaClass.getResourceAsStream("/diagram/$name.mmd") + assertNotNull(stream, "missing corpus file $name.mmd") + return stream.bufferedReader().use { it.readText() } + } + + internal companion object { + val CORPUS = listOf( + "flow-basic", + "flow-shapes", + "flow-subgraph", + "flow-cycle", + "flow-long", + "seq-basic", + "seq-blocks", + "seq-notes", + ) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt new file mode 100644 index 00000000000..3de00e53d8f --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt @@ -0,0 +1,115 @@ +package ai.kilocode.client.ui.diagram + +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +private const val EPS = 1e-6 + +internal fun scene(out: Out): Scene { + assertTrue(out is Out.Ok, "expected Out.Ok but was $out") + val art = (out as Out.Ok).art + assertTrue(art is Scene, "expected a Scene but was $art") + return art as Scene +} + +/** Compares the whole scene against a snapshot, mirroring the `assertModel` idiom in session tests. */ +internal fun assertScene(expected: String, out: Out) { + assertEquals(expected.trimIndent().trim(), scene(out).toString().trim()) +} + +internal fun err(out: Out): Out.Err { + assertTrue(out is Out.Err, "expected Out.Err but was $out") + return out as Out.Err +} + +/** Every mark must sit inside the reported scene size; a renderer relies on that for scroll bounds. */ +internal fun assertInBounds(scene: Scene) { + for (mark in flatten(scene.marks)) { + for (pt in points(mark)) { + assertTrue(pt.x >= -EPS, "mark left of origin: $mark") + assertTrue(pt.y >= -EPS, "mark above origin: $mark") + assertTrue(pt.x <= scene.size.w + EPS, "mark past width ${scene.size.w}: $mark") + assertTrue(pt.y <= scene.size.h + EPS, "mark past height ${scene.size.h}: $mark") + } + } +} + +/** Node surfaces must not overlap; this is font-independent so it also holds under real metrics. */ +internal fun assertNoOverlap(scene: Scene) { + val rects = surfaces(scene) + for (left in rects.indices) { + for (right in left + 1 until rects.size) { + assertTrue(apart(rects[left], rects[right]), "overlapping nodes ${rects[left]} ${rects[right]}") + } + } +} + +/** Flowchart links must start and end on a node outline rather than floating in space. */ +internal fun assertEdgesTouchNodes(scene: Scene) { + if (scene.type != Type.Flowchart) return + val rects = surfaces(scene) + for (mark in flatten(scene.marks)) { + if (mark !is Mark.Edge || mark.role != Role.Line) continue + val ends = listOf(mark.points.first(), mark.points.last()) + for (pt in ends) { + assertTrue(rects.any { edgeOf(it, pt) }, "link endpoint $pt is not on a node outline") + } + } +} + +private fun surfaces(scene: Scene): List { + val out = mutableListOf() + for (mark in flatten(scene.marks)) { + when (mark) { + is Mark.Box -> if (mark.fill == Role.Surface) out.add(mark.rect) + is Mark.Oval -> if (mark.fill == Role.Surface) out.add(mark.rect) + is Mark.Poly -> if (mark.fill == Role.Surface) out.add(bounds(mark.points)) + else -> Unit + } + } + return out +} + +private fun bounds(points: List): Rect { + val x = points.minOf { it.x } + val y = points.minOf { it.y } + return Rect(x, y, points.maxOf { it.x } - x, points.maxOf { it.y } - y) +} + +private fun apart(left: Rect, right: Rect): Boolean { + if (left.x + left.w <= right.x + EPS || right.x + right.w <= left.x + EPS) return true + return left.y + left.h <= right.y + EPS || right.y + right.h <= left.y + EPS +} + +private fun edgeOf(rect: Rect, pt: Pt): Boolean { + val insideX = pt.x >= rect.x - EPS && pt.x <= rect.x + rect.w + EPS + val insideY = pt.y >= rect.y - EPS && pt.y <= rect.y + rect.h + EPS + val onVertical = near(pt.x, rect.x) || near(pt.x, rect.x + rect.w) + val onHorizontal = near(pt.y, rect.y) || near(pt.y, rect.y + rect.h) + return (onVertical && insideY) || (onHorizontal && insideX) +} + +private fun near(left: Double, right: Double) = kotlin.math.abs(left - right) < 1e-3 + +internal fun flatten(marks: List): List { + val out = mutableListOf() + for (mark in marks) { + if (mark is Mark.Group) { + out.addAll(flatten(mark.marks)) + continue + } + out.add(mark) + } + return out +} + +private fun points(mark: Mark): List = when (mark) { + is Mark.Box -> corners(mark.rect) + is Mark.Oval -> corners(mark.rect) + is Mark.Poly -> mark.points + is Mark.Edge -> mark.points + is Mark.Text -> listOf(mark.at) + is Mark.Group -> emptyList() +} + +private fun corners(rect: Rect) = listOf(Pt(rect.x, rect.y), Pt(rect.x + rect.w, rect.y + rect.h)) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ErrorTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ErrorTest.kt new file mode 100644 index 00000000000..65f3ac9571d --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ErrorTest.kt @@ -0,0 +1,54 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertEquals + +class ErrorTest { + private val engine = Mermaid(FakeMeasure()) + + @Test + fun `unsupported diagram types are rejected without parsing`() { + val out = draw("pie title Pets\n \"Dogs\" : 40") + + assertEquals(Fault.Unsupported, err(out).fault) + } + + @Test + fun `unknown keywords are unsupported`() { + assertEquals(Fault.Unsupported, err(draw("hello world")).fault) + assertEquals(Fault.Unsupported, err(draw("")).fault) + } + + @Test + fun `syntax errors report the original line number`() { + val out = draw("flowchart TD\n A --> B\n end") + + assertEquals(Fault.Syntax, err(out).fault) + assertEquals(3, err(out).line) + } + + @Test + fun `line numbers survive frontmatter and comments`() { + val out = draw("---\ntitle: Demo\n---\n%% a note\nflowchart TD\n A --> B\n end") + + assertEquals(7, err(out).line) + } + + @Test + fun `empty diagrams are a syntax error rather than an empty scene`() { + assertEquals(Fault.Syntax, err(draw("flowchart TD")).fault) + assertEquals(Fault.Syntax, err(draw("sequenceDiagram")).fault) + } + + @Test + fun `sequence blocks report unbalanced ends`() { + val out = draw("sequenceDiagram\n A->>B: hi\n end") + + assertEquals(Fault.Syntax, err(out).fault) + assertEquals(3, err(out).line) + } + + private fun draw(source: String) = runBlocking { engine.draw(source, spec()) } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FakeMeasure.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FakeMeasure.kt new file mode 100644 index 00000000000..2e2714c0e71 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FakeMeasure.kt @@ -0,0 +1,31 @@ +package ai.kilocode.client.ui.diagram + +/** + * Deterministic text measurement for engine tests. + * + * Geometry snapshots must not depend on which fonts a machine or CI image happens to have, so tests + * measure with fixed per-character widths instead of AWT metrics. [onCall] receives the running call + * count and is used by the cancellation test to cancel mid-layout. + */ +internal class FakeMeasure(private val onCall: (Int) -> Unit = {}) : Measure { + private var calls = 0 + + override fun width(text: String, font: FontSpec): Double { + calls++ + onCall(calls) + val bold = if (font.bold) BOLD else 1.0 + return text.length * UNIT * bold + } + + override fun height(font: FontSpec) = font.size * LINE + + override fun ascent(font: FontSpec) = font.size.toDouble() + + private companion object { + const val UNIT = 7.0 + const val BOLD = 1.1 + const val LINE = 1.4 + } +} + +internal fun spec(size: Int = 10) = Spec(FontSpec("Test", size)) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt new file mode 100644 index 00000000000..03132d440e2 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt @@ -0,0 +1,88 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test + +class FlowLayoutTest { + private val engine = Mermaid(FakeMeasure()) + + @Test + fun `two nodes stack vertically`() = assertScene( + """ + scene Flowchart 39x124 + edge 20,38 20,86 role=Line dash=false thick=false head=Arrow tail=None + box 8,8 23x30 arc=0 fill=Surface line=Border dash=false + text "A" at=20,23 anchor=Center role=Text bold=false + box 8,86 23x30 arc=0 fill=Surface line=Border dash=false + text "B" at=20,101 anchor=Center role=Text bold=false + """, + draw("flowchart TD\n A --> B"), + ) + + @Test + fun `branches route labelled links`() = assertScene( + """ + scene Flowchart 114x218 + edge 54,38 54,86 role=Line dash=false thick=false head=Arrow tail=None + edge 45,132 28,180 role=Line dash=false thick=false head=Arrow tail=None + text "yes" at=37,156 anchor=Center role=Muted bold=false + edge 62,132 79,180 role=Line dash=false thick=false head=Arrow tail=None + text "no" at=70,156 anchor=Center role=Muted bold=false + box 28,8 51x30 arc=0 fill=Surface line=Border dash=false + text "Start" at=54,23 anchor=Center role=Text bold=false + poly 54,86 80,109 54,132 27,109 fill=Surface line=Border + text "Ok?" at=54,109 anchor=Center role=Text bold=false + box 8,180 30x30 arc=0 fill=Surface line=Border dash=false + text "Go" at=23,195 anchor=Center role=Text bold=false + box 62,180 44x30 arc=0 fill=Surface line=Border dash=false + text "Stop" at=84,195 anchor=Center role=Text bold=false + """, + draw("flowchart TD\n A[Start] --> B{Ok?}\n B -->|yes| C[Go]\n B -->|no| D[Stop]"), + ) + + @Test + fun `self links render as loops`() = assertScene( + """ + scene Flowchart 51x46 + edge 31,16 43,16 43,31 31,31 role=Line dash=false thick=false head=Arrow tail=None + box 8,8 23x30 arc=0 fill=Surface line=Border dash=false + text "A" at=20,23 anchor=Center role=Text bold=false + """, + draw("flowchart TD\n A --> A"), + ) + + @Test + fun `direction transforms are applied after layout`() = assertScene( + """ + scene Flowchart 110x46 + edge 31,23 79,23 role=Line dash=false thick=false head=Arrow tail=None + box 8,8 23x30 arc=0 fill=Surface line=Border dash=false + text "A" at=20,23 anchor=Center role=Text bold=false + box 79,8 23x30 arc=0 fill=Surface line=Border dash=false + text "B" at=91,23 anchor=Center role=Text bold=false + """, + draw("flowchart LR\n A --> B"), + ) + + @Test + fun `clusters emit grouped dashed frames`() = assertScene( + """ + scene Flowchart 55x210 + group s + box 0,134 55x76 arc=4 fill=- line=Cluster dash=true + text "Group" at=28,149 anchor=Center role=Muted bold=true + edge 28,38 28,86 role=Line dash=false thick=false head=Arrow tail=None + edge 28,116 28,164 role=Line dash=false thick=false head=Arrow tail=None + box 16,8 23x30 arc=0 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=false + box 16,86 23x30 arc=0 fill=Surface line=Border dash=false + text "B" at=28,101 anchor=Center role=Text bold=false + box 16,164 23x30 arc=0 fill=Surface line=Border dash=false + text "C" at=28,179 anchor=Center role=Text bold=false + """, + draw("flowchart TD\n A --> B\n subgraph s [Group]\n B --> C\n end"), + ) + + private fun draw(source: String) = runBlocking { engine.draw(source, spec()) } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt new file mode 100644 index 00000000000..0779e8a7f9e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt @@ -0,0 +1,42 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertNotNull + +/** + * Font-independent invariants, checked under both the deterministic fake and real AWT metrics. + * + * Snapshots pin exact geometry for one measurement model; these assertions catch layout bugs that a + * different font would expose, which is the failure mode snapshots cannot see. + */ +class InvariantTest { + @Test + fun `fake metrics keep nodes separated and inside bounds`() { + check(FakeMeasure()) + } + + @Test + fun `real font metrics keep nodes separated and inside bounds`() { + check(AwtMeasure()) + } + + private fun check(measure: Measure) { + val engine = Mermaid(measure) + for (name in ConformanceTest.CORPUS) { + val out = runBlocking { engine.draw(read(name), spec(size = 12)) } + val scene = scene(out) + + assertInBounds(scene) + assertNoOverlap(scene) + assertEdgesTouchNodes(scene) + } + } + + private fun read(name: String): String { + val stream = javaClass.getResourceAsStream("/diagram/$name.mmd") + assertNotNull(stream, "missing corpus file $name.mmd") + return stream.bufferedReader().use { it.readText() } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt new file mode 100644 index 00000000000..1acf3137a0d --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt @@ -0,0 +1,54 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertEquals + +/** Model output can be pathological; the engine must refuse rather than hang or exhaust memory. */ +class LimitsTest { + private val engine = Mermaid(FakeMeasure()) + + @Test + fun `line cap is enforced before parsing`() { + val source = "flowchart TD\n" + (1..50).joinToString("\n") { " n$it --> n${it + 1}" } + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(lines = 10))) } + + assertEquals(Fault.Limit, err(out).fault) + } + + @Test + fun `node cap is enforced`() { + val source = "flowchart TD\n" + (1..30).joinToString("\n") { " n$it --> n${it + 1}" } + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(nodes = 5))) } + + assertEquals(Fault.Limit, err(out).fault) + } + + @Test + fun `link cap is enforced`() { + val source = "flowchart TD\n" + (1..30).joinToString("\n") { " a --> n$it" } + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(edges = 5))) } + + assertEquals(Fault.Limit, err(out).fault) + } + + @Test + fun `sequence caps are enforced`() { + val source = "sequenceDiagram\n" + (1..30).joinToString("\n") { " a->>b: step $it" } + val steps = runBlocking { engine.draw(source, spec().copy(limits = Limits(edges = 5))) } + val actors = "sequenceDiagram\n" + (1..30).joinToString("\n") { " a->>p$it: step" } + val people = runBlocking { engine.draw(actors, spec().copy(limits = Limits(nodes = 5))) } + + assertEquals(Fault.Limit, err(steps).fault) + assertEquals(Fault.Limit, err(people).fault) + } + + @Test + fun `a graph at the cap still renders`() { + val source = "flowchart TD\n" + (1..9).joinToString("\n") { " n$it --> n${it + 1}" } + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(nodes = 10, edges = 9))) } + + assertEquals(10, scene(out).marks.count { it is Mark.Box }) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt new file mode 100644 index 00000000000..030e1bceb30 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt @@ -0,0 +1,102 @@ +package ai.kilocode.client.ui.diagram + +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import kotlinx.coroutines.runBlocking +import kotlin.test.Test + +class SeqLayoutTest { + private val engine = Mermaid(FakeMeasure()) + + @Test + fun `self messages render loops`() = assertScene( + """ + scene Sequence 92x230 + edge 28,38 28,222 role=Muted dash=true thick=false head=None tail=None + box 8,8 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=true + edge 28,102 76,102 76,150 28,150 role=Line dash=false thick=false head=Arrow tail=None + text "retry" at=84,109 anchor=Left role=Muted bold=false + """, + draw("sequenceDiagram\n A->>A: retry"), + ) + + @Test + fun `notes can span participants`() = assertScene( + """ + scene Sequence 142x274 + edge 28,38 28,266 role=Muted dash=true thick=false head=None tail=None + edge 115,38 115,266 role=Muted dash=true thick=false head=None tail=None + box 8,8 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=true + box 95,8 39x30 arc=4 fill=Surface line=Border dash=false + text "B" at=115,23 anchor=Center role=Text bold=true + box 9,94 124x30 arc=4 fill=Note line=Border dash=false + text "shared" at=71,109 anchor=Center role=Text bold=false + edge 28,194 115,194 role=Line dash=false thick=false head=Arrow tail=None + text "go" at=71,183 anchor=Center role=Muted bold=false + """, + draw("sequenceDiagram\n participant A\n participant B\n Note over A,B: shared\n A->>B: go"), + ) + + @Test + fun `blocks render dashed frames with split labels`() = assertScene( + """ + scene Sequence 142x396 + edge 28,38 28,388 role=Muted dash=true thick=false head=None tail=None + edge 115,38 115,388 role=Muted dash=true thick=false head=None tail=None + box 8,8 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=true + box 95,8 39x30 arc=4 fill=Surface line=Border dash=false + text "B" at=115,23 anchor=Center role=Text bold=true + edge 28,116 115,116 role=Line dash=false thick=false head=Arrow tail=None + text "go" at=71,105 anchor=Center role=Muted bold=false + edge 28,216 115,216 role=Line dash=false thick=false head=Arrow tail=None + text "one" at=71,205 anchor=Center role=Muted bold=false + edge 8,264 126,264 role=Cluster dash=true thick=false head=None tail=None + text "no" at=16,271 anchor=Left role=Muted bold=false + edge 28,308 115,308 role=Line dash=false thick=false head=Arrow tail=None + text "two" at=71,297 anchor=Center role=Muted bold=false + box 8,164 118x192 arc=4 fill=- line=Cluster dash=true + box 8,164 65x22 arc=4 fill=Note line=Cluster dash=false + text "alt yes" at=16,175 anchor=Left role=Muted bold=true + """, + draw("sequenceDiagram\n A->>B: go\n alt yes\n A->>B: one\n else no\n A->>B: two\n end"), + ) + + @Test + fun `autonumber prefixes message labels`() = assertScene( + """ + scene Sequence 142x266 + edge 28,38 28,258 role=Muted dash=true thick=false head=None tail=None + edge 115,38 115,258 role=Muted dash=true thick=false head=None tail=None + box 8,8 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=true + box 95,8 39x30 arc=4 fill=Surface line=Border dash=false + text "B" at=115,23 anchor=Center role=Text bold=true + edge 28,116 115,116 role=Line dash=false thick=false head=Arrow tail=None + text "1. one" at=71,105 anchor=Center role=Muted bold=false + edge 115,186 28,186 role=Line dash=false thick=false head=Arrow tail=None + text "2. two" at=71,175 anchor=Center role=Muted bold=false + """, + draw("sequenceDiagram\n autonumber\n A->>B: one\n B->>A: two"), + ) + + @Test + fun `titles reserve space before participants`() = assertScene( + """ + scene Sequence 142x218 + edge 28,60 28,210 role=Muted dash=true thick=false head=None tail=None + edge 115,60 115,210 role=Muted dash=true thick=false head=None tail=None + text "Flow" at=8,15 anchor=Left role=Text bold=true + box 8,30 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,45 anchor=Center role=Text bold=true + box 95,30 39x30 arc=4 fill=Surface line=Border dash=false + text "B" at=115,45 anchor=Center role=Text bold=true + edge 28,138 115,138 role=Line dash=false thick=false head=Arrow tail=None + text "go" at=71,127 anchor=Center role=Muted bold=false + """, + draw("sequenceDiagram\n title Flow\n A->>B: go"), + ) + + private fun draw(source: String) = runBlocking { engine.draw(source, spec()) } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt new file mode 100644 index 00000000000..9721efc71d4 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt @@ -0,0 +1,57 @@ +package ai.kilocode.client.ui.diagram + +import kotlinx.serialization.json.Json +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * The IR must survive a process boundary so a future out-of-process engine is a drop-in. This test is + * the mechanical guarantee: it fails the moment an AWT or otherwise unserializable type leaks in. + */ +class SerializeTest { + private val json = Json + + @Test + fun `every mark variant round trips`() { + val scene = sample() + + assertEquals(scene, json.decodeFromString(json.encodeToString(scene))) + } + + @Test + fun `art round trips polymorphically`() { + val scene: Art = sample() + val text = json.encodeToString(scene) + + assertTrue(text.contains("Scene"), "expected a discriminator in $text") + assertEquals(scene, json.decodeFromString(text)) + } + + @Test + fun `spec round trips`() { + val value = Spec(FontSpec("Inter", 13, bold = true), Metrics(pad = 3.0), Limits(nodes = 7)) + + assertEquals(value, json.decodeFromString(json.encodeToString(value))) + } + + private fun sample() = Scene( + Type.Flowchart, + listOf( + Mark.Box(Rect(1.0, 2.0, 30.0, 40.0), 4.0, Role.Surface, Role.Border, dash = true), + Mark.Oval(Rect(5.0, 6.0, 10.0, 10.0), Role.Note, null), + Mark.Poly(listOf(Pt(0.0, 0.0), Pt(4.0, 0.0), Pt(2.0, 6.0)), Role.Surface, Role.Border), + Mark.Edge( + listOf(Pt(0.0, 0.0), Pt(9.0, 9.0)), + Role.Line, + dash = true, + thick = true, + head = Head.Arrow, + tail = Head.Dot, + ), + Mark.Text("hello", Pt(3.0, 4.0), Anchor.Center, Role.Text, bold = true), + Mark.Group("cluster", listOf(Mark.Box(Rect(0.0, 0.0, 2.0, 2.0), 0.0, null, Role.Cluster))), + ), + Size(80.0, 90.0), + ) +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/TypeTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/TypeTest.kt new file mode 100644 index 00000000000..a0a1c16af0b --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/TypeTest.kt @@ -0,0 +1,43 @@ +package ai.kilocode.client.ui.diagram + +import kotlin.test.Test +import kotlin.test.assertEquals + +class TypeTest { + @Test + fun `detects flowchart aliases and directions`() { + assertEquals(Type.Flowchart, Type.of("graph TD\n A --> B")) + assertEquals(Type.Flowchart, Type.of("flowchart LR\n A --> B")) + assertEquals(Type.Flowchart, Type.of("flowchart\n A --> B")) + } + + @Test + fun `detects sequence diagrams regardless of case`() { + assertEquals(Type.Sequence, Type.of("sequenceDiagram\n A->>B: hi")) + assertEquals(Type.Sequence, Type.of("SEQUENCEDIAGRAM\n A->>B: hi")) + } + + @Test + fun `ignores blank lines, comments, frontmatter and directives`() { + assertEquals(Type.Flowchart, Type.of("\n\n \ngraph TD\n A --> B")) + assertEquals(Type.Flowchart, Type.of("%% a comment\ngraph TD\n A --> B")) + assertEquals(Type.Sequence, Type.of("---\ntitle: x\n---\nsequenceDiagram\n A->>B: hi")) + assertEquals(Type.Sequence, Type.of("%%{init: {'theme':'dark'}}%%\nsequenceDiagram\n A->>B: hi")) + } + + @Test + fun `maps other known diagram keywords`() { + assertEquals(Type.Class, Type.of("classDiagram\n class A")) + assertEquals(Type.State, Type.of("stateDiagram-v2\n [*] --> A")) + assertEquals(Type.Er, Type.of("erDiagram\n A ||--o{ B : has")) + assertEquals(Type.Gantt, Type.of("gantt\n title x")) + assertEquals(Type.Pie, Type.of("pie title Pets")) + } + + @Test + fun `unknown and empty sources fall through`() { + assertEquals(Type.Unknown, Type.of("hello world")) + assertEquals(Type.Unknown, Type.of("")) + assertEquals(Type.Unknown, Type.of("%% only a comment")) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt new file mode 100644 index 00000000000..3b7fda319c5 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt @@ -0,0 +1,181 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Head +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.test.assertTrue + +class FlowParseTest { + @Test + fun `header sets direction`() { + assertEquals(Dir.Down, graph("graph TD\n A --> B").dir) + assertEquals(Dir.Down, graph("graph TB\n A --> B").dir) + assertEquals(Dir.Up, graph("graph BT\n A --> B").dir) + assertEquals(Dir.Right, graph("flowchart LR\n A --> B").dir) + assertEquals(Dir.Left, graph("flowchart RL\n A --> B").dir) + } + + @Test + fun `node shapes are recognised`() { + val source = """ + flowchart TD + a[Rect] --> b(Round) + c([Stadium]) --> d[[Sub]] + e[(Cyl)] --> f((Circle)) + g{Rhombus} --> h{{Hex}} + i[/Skew/] --> j[\SkewAlt\] + k[/Trap\] --> l[\TrapAlt/] + m>Flag] --> n(((Doubled))) + """ + val nodes = graph(source).nodes + + assertEquals(Shape.Rect, nodes.getValue("a").shape) + assertEquals(Shape.Round, nodes.getValue("b").shape) + assertEquals(Shape.Stadium, nodes.getValue("c").shape) + assertEquals(Shape.Subroutine, nodes.getValue("d").shape) + assertEquals(Shape.Cylinder, nodes.getValue("e").shape) + assertEquals(Shape.Circle, nodes.getValue("f").shape) + assertEquals(Shape.Rhombus, nodes.getValue("g").shape) + assertEquals(Shape.Hexagon, nodes.getValue("h").shape) + assertEquals(Shape.Skew, nodes.getValue("i").shape) + assertEquals(Shape.SkewAlt, nodes.getValue("j").shape) + assertEquals(Shape.Trapezoid, nodes.getValue("k").shape) + assertEquals(Shape.TrapezoidAlt, nodes.getValue("l").shape) + assertEquals(Shape.Flag, nodes.getValue("m").shape) + assertEquals(Shape.Doubled, nodes.getValue("n").shape) + } + + @Test + fun `labels keep declaration order and break into lines`() { + val graph = graph("flowchart TD\n A --> B\n B[\"Second
line\"]") + + assertEquals(listOf("A", "B"), graph.nodes.keys.toList()) + assertEquals(listOf("A"), graph.nodes.getValue("A").label) + assertEquals(listOf("Second", "line"), graph.nodes.getValue("B").label) + } + + @Test + fun `link styles and heads are classified`() { + val edges = graph( + """ + flowchart TD + A --> B + A --- C + A -.-> D + A ==> E + A --o F + A --x G + A <--> H + """, + ).edges + + assertEquals(Link.Solid to Head.Arrow, edges[0].link to edges[0].head) + assertEquals(Link.Solid to Head.None, edges[1].link to edges[1].head) + assertEquals(Link.Dotted to Head.Arrow, edges[2].link to edges[2].head) + assertEquals(Link.Thick to Head.Arrow, edges[3].link to edges[3].head) + assertEquals(Link.Solid to Head.Dot, edges[4].link to edges[4].head) + assertEquals(Link.Solid to Head.Cross, edges[5].link to edges[5].head) + assertEquals(Head.Arrow, edges[6].tail) + } + + @Test + fun `edge labels come from pipes and inline text`() { + val edges = graph( + """ + flowchart TD + A -->|yes| B + A -- maybe --> C + A -. later .-> D + A == fast ==> E + """, + ).edges + + assertEquals(listOf("yes"), edges[0].label) + assertEquals(listOf("maybe"), edges[1].label) + assertEquals(listOf("later"), edges[2].label) + assertEquals(listOf("fast"), edges[3].label) + assertEquals(Link.Dotted, edges[2].link) + assertEquals(Link.Thick, edges[3].link) + } + + @Test + fun `chains and ampersand groups expand into edges`() { + val edges = graph("flowchart TD\n A --> B --> C\n X --> Y & Z").edges + + assertEquals(listOf("A" to "B", "B" to "C", "X" to "Y", "X" to "Z"), edges.map { it.from to it.to }) + } + + @Test + fun `dashes inside labels do not split statements`() { + val graph = graph("flowchart TD\n A[\"a --> b\"] --> B") + + assertEquals(1, graph.edges.size) + assertEquals(listOf("a --> b"), graph.nodes.getValue("A").label) + } + + @Test + fun `subgraphs nest and assign membership`() { + val graph = graph( + """ + flowchart TD + Client --> Gate + subgraph core [Core] + Gate --> Auth + subgraph store [Store] + Auth --> Db + end + end + """, + ) + + assertEquals(listOf("core", "store"), graph.clusters.keys.toList()) + assertEquals(listOf("Core"), graph.clusters.getValue("core").label) + assertNull(graph.clusters.getValue("core").parent) + assertEquals("core", graph.clusters.getValue("store").parent) + assertNull(graph.nodes.getValue("Client").cluster) + assertEquals("core", graph.nodes.getValue("Auth").cluster) + assertEquals("store", graph.nodes.getValue("Db").cluster) + } + + @Test + fun `styling statements are skipped and class suffixes dropped`() { + val graph = graph( + """ + flowchart TD + classDef hot fill:#f00 + A:::hot --> B + class B hot + style A stroke:#000 + click A "https://example.com" + linkStyle 0 stroke:#0f0 + """, + ) + + assertEquals(listOf("A", "B"), graph.nodes.keys.toList()) + assertEquals(1, graph.edges.size) + } + + @Test + fun `self links are preserved`() { + val edges = graph("flowchart TD\n A --> A").edges + + assertEquals(1, edges.size) + assertTrue(edges.single().from == edges.single().to) + } + + @Test + fun `dangling subgraph and stray end are reported with line numbers`() { + val open = Flow().parse(Source.clean("flowchart TD\n subgraph s\n A --> B")) + val stray = Flow().parse(Source.clean("flowchart TD\n A --> B\n end")) + + assertEquals(3, (open as FlowOut.Err).line) + assertEquals(3, (stray as FlowOut.Err).line) + } + + private fun graph(source: String): Graph { + val out = Flow().parse(Source.clean(source.trimIndent())) + assertTrue(out is FlowOut.Ok, "expected a parsed graph but was $out") + return (out as FlowOut.Ok).graph + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt new file mode 100644 index 00000000000..ec188c5af0e --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt @@ -0,0 +1,140 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import ai.kilocode.client.ui.diagram.Head +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +class SeqParseTest { + @Test + fun `declared participants keep order and aliases`() { + val script = script("sequenceDiagram\n participant C as Client\n actor S as Server\n C->>S: hi") + + assertEquals(listOf("C", "S"), script.actors.keys.toList()) + assertEquals(listOf("Client"), script.actors.getValue("C").label) + assertEquals(listOf("Server"), script.actors.getValue("S").label) + } + + @Test + fun `undeclared participants appear in first use order`() { + val script = script("sequenceDiagram\n B->>A: first\n A->>C: second") + + assertEquals(listOf("B", "A", "C"), script.actors.keys.toList()) + } + + @Test + fun `arrow forms map to link styles and heads`() { + val steps = script( + """ + sequenceDiagram + A->>B: solid arrow + A-->>B: dotted arrow + A->B: solid open + A-->B: dotted open + A-xB: solid cross + A--xB: dotted cross + A-)B: solid dot + """, + ).steps.filterIsInstance() + + assertEquals(Link.Solid to Head.Arrow, steps[0].link to steps[0].head) + assertEquals(Link.Dotted to Head.Arrow, steps[1].link to steps[1].head) + assertEquals(Link.Solid to Head.Open, steps[2].link to steps[2].head) + assertEquals(Link.Dotted to Head.Open, steps[3].link to steps[3].head) + assertEquals(Link.Solid to Head.Cross, steps[4].link to steps[4].head) + assertEquals(Link.Dotted to Head.Cross, steps[5].link to steps[5].head) + assertEquals(Link.Solid to Head.Dot, steps[6].link to steps[6].head) + } + + @Test + fun `participant names may contain dashes`() { + val steps = script("sequenceDiagram\n web-app->>db-main: query").steps.filterIsInstance() + + assertEquals("web-app", steps.single().from) + assertEquals("db-main", steps.single().to) + } + + @Test + fun `activation shorthand wraps the message`() { + val steps = script("sequenceDiagram\n A->>+B: open\n B-->>-A: close").steps + + assertEquals(Step.Toggle("B", true), steps[0]) + assertTrue(steps[1] is Step.Msg) + assertTrue(steps[2] is Step.Msg) + assertEquals(Step.Toggle("B", false), steps[3]) + } + + @Test + fun `explicit activate and deactivate are recorded`() { + val steps = script("sequenceDiagram\n activate A\n A->>B: work\n deactivate A").steps + + assertEquals(Step.Toggle("A", true), steps[0]) + assertEquals(Step.Toggle("A", false), steps[2]) + } + + @Test + fun `notes carry placement and targets`() { + val notes = script( + """ + sequenceDiagram + participant A + participant B + Note left of A: left side + Note right of B: right side + Note over A,B: spanning + """, + ).steps.filterIsInstance() + + assertEquals(NoteAt.Left, notes[0].at) + assertEquals(listOf("A"), notes[0].actors) + assertEquals(listOf("left side"), notes[0].label) + assertEquals(NoteAt.Right, notes[1].at) + assertEquals(NoteAt.Over, notes[2].at) + assertEquals(listOf("A", "B"), notes[2].actors) + } + + @Test + fun `blocks open split and close`() { + val steps = script( + """ + sequenceDiagram + alt in stock + A->>B: reserve + else sold out + A->>B: refuse + end + loop twice + A->>B: retry + end + """, + ).steps + + assertEquals(Step.Open(BlockKind.Alt, listOf("in stock")), steps[0]) + assertEquals(Step.Split(listOf("sold out")), steps[2]) + assertEquals(Step.Close, steps[4]) + assertEquals(Step.Open(BlockKind.Loop, listOf("twice")), steps[5]) + } + + @Test + fun `title and autonumber are captured`() { + val script = script("sequenceDiagram\n title Checkout\n autonumber\n A->>B: go") + + assertEquals(listOf("Checkout"), script.title) + assertTrue(script.numbered) + } + + @Test + fun `unbalanced blocks are reported with line numbers`() { + val open = Seq().parse(Source.clean("sequenceDiagram\n loop forever\n A->>B: x")) + val stray = Seq().parse(Source.clean("sequenceDiagram\n A->>B: x\n end")) + + assertEquals(3, (open as SeqOut.Err).line) + assertEquals(3, (stray as SeqOut.Err).line) + } + + private fun script(source: String): Script { + val out = Seq().parse(Source.clean(source.trimIndent())) + assertTrue(out is SeqOut.Ok, "expected a parsed script but was $out") + return (out as SeqOut.Ok).script + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt new file mode 100644 index 00000000000..340c822b180 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt @@ -0,0 +1,74 @@ +package ai.kilocode.client.ui.diagram.mermaid + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +class SourceTest { + @Test + fun `normalizes line endings and tabs`() { + val clean = Source.clean("graph TD\r\n\tA --> B\r") + + assertEquals(listOf("graph TD", " A --> B", ""), clean.lines.map { it.text }) + assertEquals(listOf(1, 2, 3), clean.lines.map { it.at }) + } + + @Test + fun `skips terminated frontmatter and keeps original line numbers`() { + val clean = Source.clean("---\ntitle: Demo\n---\ngraph TD\n A --> B") + + assertEquals(listOf("graph TD", " A --> B"), clean.lines.map { it.text }) + assertEquals(listOf(4, 5), clean.lines.map { it.at }) + } + + @Test + fun `unterminated frontmatter is treated as content`() { + val clean = Source.clean("---\ntitle: Demo\ngraph TD") + + assertEquals(3, clean.lines.size) + assertEquals(1, clean.lines.first().at) + } + + @Test + fun `line comments are cut but quoted percent signs survive`() { + val clean = Source.clean("graph TD %% direction note\n A[\"50%% done\"] --> B %% trailing") + + assertEquals("graph TD", clean.lines[0].text) + assertEquals(" A[\"50%% done\"] --> B", clean.lines[1].text) + } + + @Test + fun `init directives are blanked without shifting line numbers`() { + val clean = Source.clean("%%{init: {'theme':'dark'}}%%\ngraph TD\n A --> B") + + assertTrue(clean.lines[0].text.isBlank()) + assertEquals("graph TD", clean.lines[1].text) + assertEquals(2, clean.lines[1].at) + } + + @Test + fun `multi line directives keep the line map aligned`() { + val clean = Source.clean("%%{init: {\n 'theme':'dark'\n}}%%\ngraph TD\n end") + + assertEquals("graph TD", clean.lines[3].text) + assertEquals(4, clean.lines[3].at) + assertEquals(5, clean.lines[4].at) + } + + @Test + fun `labels split on break forms and drop quotes`() { + assertEquals(listOf("one", "two"), Source.label("\"one
two\"")) + assertEquals(listOf("a", "b"), Source.label("a
b")) + assertEquals(listOf("a", "b"), Source.label("a
b")) + assertEquals(listOf("a", "b"), Source.label("a\\nb")) + assertEquals(listOf(""), Source.label(" ")) + } + + @Test + fun `open reports bracket and quote nesting`() { + val text = "A[x --> y] --> B" + + assertTrue(Source.open(text, text.lastIndexOf("-->"))) + assertTrue(!Source.open(text, text.indexOf("-->"))) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-basic.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-basic.mmd new file mode 100644 index 00000000000..f042dfc27f3 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-basic.mmd @@ -0,0 +1,6 @@ +flowchart TD + A[Start] --> B{Is valid?} + B -->|Yes| C[Process] + B -->|No| D[Reject] + C --> E[Done] + D --> E diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-cycle.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-cycle.mmd new file mode 100644 index 00000000000..3f593341658 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-cycle.mmd @@ -0,0 +1,9 @@ +graph TD + Retry --> Fetch + Fetch --> Parse + Parse -.-> Retry + Parse ==> Store + Store --- Audit + Fetch --> Fetch + Store --> Audit + Orphan[Detached node] diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-long.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-long.mmd new file mode 100644 index 00000000000..94652bba9c7 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-long.mmd @@ -0,0 +1,7 @@ +flowchart LR + Start --> Mid + Start --> Skip + Mid --> Late + Skip --> Late + Start --> Late + Late --> End["Wrap up
and finish"] diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-shapes.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-shapes.mmd new file mode 100644 index 00000000000..3952c8e6551 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-shapes.mmd @@ -0,0 +1,17 @@ +--- +title: Shapes +--- +%%{init: {'theme':'dark'}}%% +graph LR + a[Rect] --> b(Round) + b --> c([Stadium]) + c --> d[[Subroutine]] + d --> e[(Cylinder)] + e --> f((Circle)) + f --> g{Rhombus} + g --> h{{Hexagon}} + h --> i[/Skew/] + i --> j[\SkewAlt\] + j --> k[/Trapezoid\] + k --> l[\TrapezoidAlt/] + l --> m>Flag] diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-subgraph.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-subgraph.mmd new file mode 100644 index 00000000000..8df3064c987 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/flow-subgraph.mmd @@ -0,0 +1,10 @@ +flowchart TD + Client --> Gateway + subgraph core [Core Services] + Gateway --> Auth + subgraph store [Storage] + Auth --> Cache + Cache --> Db[(Database)] + end + end + Db --> Report diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-basic.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-basic.mmd new file mode 100644 index 00000000000..a1df2a61311 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-basic.mmd @@ -0,0 +1,7 @@ +sequenceDiagram + autonumber + participant C as Client + participant S as Server + C->>S: GET /users + S-->>C: 200 OK + C->>C: cache result diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-blocks.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-blocks.mmd new file mode 100644 index 00000000000..f85bab28b41 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-blocks.mmd @@ -0,0 +1,16 @@ +sequenceDiagram + title Checkout flow + participant U as User + participant A as Api + participant D as Db + U->>+A: POST /order + alt in stock + A->>D: reserve + D-->>A: ok + else sold out + A-->>U: 409 Conflict + end + loop retry twice + A->>D: charge + end + A-->>-U: 201 Created diff --git a/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-notes.mmd b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-notes.mmd new file mode 100644 index 00000000000..724892e4f18 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/resources/diagram/seq-notes.mmd @@ -0,0 +1,8 @@ +sequenceDiagram + participant A + participant B + Note left of A: starts here + A->>B: ping + Note over A,B: handshake done + B-)A: async pong + Note right of B: fire and forget From b9d4d874f76c034626e628a29f50a82b028901d0 Mon Sep 17 00:00:00 2001 From: kirillk Date: Mon, 24 Aug 2026 15:07:04 -0400 Subject: [PATCH 02/11] feat(jetbrains): render mermaid diagrams inline in markdown Add a scene painter, palette, async render service, and Swing panel so closed mermaid fences show a real diagram in the transcript instead of source text. Streaming fences and render errors keep the code block visible, and a source toggle reveals the fence text on demand. Ship behind the kilo.diagram.inline.enabled registry key since this changes default rendering. --- .changeset/mermaid-diagrams-jetbrains.md | 5 + .../client/session/ui/style/SessionUiStyle.kt | 6 + .../ai/kilocode/client/ui/diagram/Metrics.kt | 18 +- .../ai/kilocode/client/ui/diagram/Painter.kt | 15 ++ .../ai/kilocode/client/ui/diagram/Palette.kt | 28 +++ .../client/ui/diagram/ScenePainter.kt | 170 +++++++++++++++++ .../client/ui/diagram/ui/DiagramPanel.kt | 111 +++++++++++ .../kilocode/client/ui/diagram/ui/Diagrams.kt | 78 ++++++++ .../client/ui/md/hybrid/MdLanguage.kt | 2 + .../client/ui/md/hybrid/MdProjector.kt | 6 +- .../client/ui/md/hybrid/MdViewHybrid.kt | 179 ++++++++++++++++++ .../resources/kilo.jetbrains.frontend.xml | 7 +- .../resources/messages/KiloBundle.properties | 4 + .../client/ui/diagram/ScenePainterTest.kt | 66 +++++++ .../client/ui/diagram/ui/DiagramPanelTest.kt | 59 ++++++ .../client/ui/diagram/ui/DiagramsTest.kt | 91 +++++++++ .../kilocode/client/ui/md/MdLanguageTest.kt | 5 + .../kilocode/client/ui/md/MdProjectorTest.kt | 5 +- .../client/ui/md/MdViewDiagramTest.kt | 140 ++++++++++++++ 19 files changed, 980 insertions(+), 15 deletions(-) create mode 100644 .changeset/mermaid-diagrams-jetbrains.md create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Palette.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ScenePainter.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ScenePainterTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt diff --git a/.changeset/mermaid-diagrams-jetbrains.md b/.changeset/mermaid-diagrams-jetbrains.md new file mode 100644 index 00000000000..82e47dfe679 --- /dev/null +++ b/.changeset/mermaid-diagrams-jetbrains.md @@ -0,0 +1,5 @@ +--- +"@kilocode/kilo-jetbrains": minor +--- + +Render Mermaid code fences as inline diagrams in JetBrains chat markdown, with a source toggle and fallback errors. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/style/SessionUiStyle.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/style/SessionUiStyle.kt index c17ce03ea94..ef13a3df680 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/style/SessionUiStyle.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/style/SessionUiStyle.kt @@ -281,6 +281,12 @@ object SessionUiStyle { fun topPadding(): Int = VIEWPORT_TOP_PADDING + UiStyle.Gap.lg() } + object Diagram { + const val MAX_HEIGHT = 480 + const val PADDING = 16 + const val EMPTY_HEIGHT = 96 + } + /** Permission session-view command preview limits. */ object Permission { const val COMMAND_LINES = 3 diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt index 0489211e69e..de0b8c7a8bc 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt @@ -1,22 +1,20 @@ package ai.kilocode.client.ui.diagram import java.awt.Font -import java.awt.FontMetrics -import java.awt.image.BufferedImage +import java.awt.font.FontRenderContext internal class AwtMeasure : Measure { - private val img = BufferedImage(1, 1, BufferedImage.TYPE_INT_ARGB) - private val g = img.createGraphics() - private val cache = linkedMapOf() + private val ctx = FontRenderContext(null, true, true) + private val cache = linkedMapOf() - override fun width(text: String, font: FontSpec) = metrics(font).stringWidth(text).toDouble() - override fun height(font: FontSpec) = metrics(font).height.toDouble() - override fun ascent(font: FontSpec) = metrics(font).ascent.toDouble() + override fun width(text: String, font: FontSpec) = font(font).getStringBounds(text, ctx).width + override fun height(font: FontSpec) = font(font).getLineMetrics("Ag", ctx).height.toDouble() + override fun ascent(font: FontSpec) = font(font).getLineMetrics("Ag", ctx).ascent.toDouble() - private fun metrics(font: FontSpec): FontMetrics { + private fun font(font: FontSpec): Font { cache[font]?.let { return it } val style = if (font.bold) Font.BOLD else Font.PLAIN - val value = g.getFontMetrics(Font(font.family, style, font.size)) + val value = Font(font.family, style, font.size) cache[font] = value return value } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt new file mode 100644 index 00000000000..da4bd51cb4c --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt @@ -0,0 +1,15 @@ +package ai.kilocode.client.ui.diagram + +import java.awt.Graphics2D + +internal interface Painter { + fun accepts(art: Art): Boolean + fun size(art: Art): Size + fun paint(g: Graphics2D, art: Art, palette: Palette) +} + +internal object Painters { + private val all = listOf(ScenePainter) + + fun of(art: Art) = all.first { it.accepts(art) } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Palette.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Palette.kt new file mode 100644 index 00000000000..9bef0f4dac9 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Palette.kt @@ -0,0 +1,28 @@ +package ai.kilocode.client.ui.diagram + +import java.awt.Color +import java.awt.Font + +internal data class Palette( + val surface: Color, + val border: Color, + val text: Color, + val muted: Color, + val accent: Color, + val note: Color, + val cluster: Color, + val line: Color, + val font: Font, + val bold: Font, +) { + fun color(role: Role): Color = when (role) { + Role.Surface -> surface + Role.Border -> border + Role.Text -> text + Role.Muted -> muted + Role.Accent -> accent + Role.Note -> note + Role.Cluster -> cluster + Role.Line -> line + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ScenePainter.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ScenePainter.kt new file mode 100644 index 00000000000..8f7e484c9a9 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ScenePainter.kt @@ -0,0 +1,170 @@ +package ai.kilocode.client.ui.diagram + +import java.awt.BasicStroke +import java.awt.Graphics2D +import java.awt.RenderingHints +import java.awt.geom.Ellipse2D +import java.awt.geom.Line2D +import java.awt.geom.Path2D +import java.awt.geom.RoundRectangle2D +import kotlin.math.PI +import kotlin.math.atan2 +import kotlin.math.cos +import kotlin.math.max +import kotlin.math.sin + +internal object ScenePainter : Painter { + private const val THIN = 1.5f + private const val THICK = 3.0f + private const val DASH = 6.0f + private const val HEAD = 10.0 + private const val DOT = 4.0 + private const val CROSS = 5.0 + + override fun accepts(art: Art) = art is Scene + + override fun size(art: Art) = (art as Scene).size + + override fun paint(g: Graphics2D, art: Art, palette: Palette) { + val old = g.renderingHints.clone() as RenderingHints + g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON) + g.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON) + try { + for (mark in (art as Scene).marks) draw(g, mark, palette) + } finally { + g.setRenderingHints(old) + } + } + + private fun draw(g: Graphics2D, mark: Mark, palette: Palette) { + when (mark) { + is Mark.Box -> box(g, mark, palette) + is Mark.Oval -> oval(g, mark, palette) + is Mark.Poly -> poly(g, mark, palette) + is Mark.Edge -> edge(g, mark, palette) + is Mark.Text -> text(g, mark, palette) + is Mark.Group -> mark.marks.forEach { draw(g, it, palette) } + } + } + + private fun box(g: Graphics2D, mark: Mark.Box, palette: Palette) { + val rect = mark.rect + val shape = RoundRectangle2D.Double(rect.x, rect.y, rect.w, rect.h, mark.arc, mark.arc) + mark.fill?.let { + g.color = palette.color(it) + g.fill(shape) + } + mark.line?.let { + g.color = palette.color(it) + g.stroke = stroke(mark.dash) + g.draw(shape) + } + } + + private fun oval(g: Graphics2D, mark: Mark.Oval, palette: Palette) { + val rect = mark.rect + val shape = Ellipse2D.Double(rect.x, rect.y, rect.w, rect.h) + mark.fill?.let { + g.color = palette.color(it) + g.fill(shape) + } + mark.line?.let { + g.color = palette.color(it) + g.stroke = stroke() + g.draw(shape) + } + } + + private fun poly(g: Graphics2D, mark: Mark.Poly, palette: Palette) { + val shape = path(mark.points, true) + mark.fill?.let { + g.color = palette.color(it) + g.fill(shape) + } + mark.line?.let { + g.color = palette.color(it) + g.stroke = stroke() + g.draw(shape) + } + } + + private fun edge(g: Graphics2D, mark: Mark.Edge, palette: Palette) { + if (mark.points.size < 2) return + g.color = palette.color(mark.role) + g.stroke = stroke(mark.dash, mark.thick) + g.draw(path(mark.points, false)) + head(g, mark.points[mark.points.lastIndex - 1], mark.points.last(), mark.head) + head(g, mark.points[1], mark.points.first(), mark.tail) + } + + private fun text(g: Graphics2D, mark: Mark.Text, palette: Palette) { + g.font = if (mark.bold) palette.bold else palette.font + g.color = palette.color(mark.role) + val fm = g.fontMetrics + val width = fm.stringWidth(mark.text).toDouble() + val height = fm.height.toDouble() + val x = when (mark.anchor) { + Anchor.TopLeft, Anchor.Left, Anchor.BottomLeft -> mark.at.x + Anchor.Top, Anchor.Center, Anchor.Bottom -> mark.at.x - width / 2.0 + Anchor.TopRight, Anchor.Right, Anchor.BottomRight -> mark.at.x - width + } + val y = when (mark.anchor) { + Anchor.TopLeft, Anchor.Top, Anchor.TopRight -> mark.at.y + fm.ascent + Anchor.Left, Anchor.Center, Anchor.Right -> mark.at.y - height / 2.0 + fm.ascent + Anchor.BottomLeft, Anchor.Bottom, Anchor.BottomRight -> mark.at.y - fm.descent + } + g.drawString(mark.text, x.toFloat(), y.toFloat()) + } + + private fun head(g: Graphics2D, from: Pt, to: Pt, head: Head) { + if (head == Head.None) return + val angle = atan2(to.y - from.y, to.x - from.x) + when (head) { + Head.Arrow -> { + val p = arrow(to, angle) + g.fill(p) + } + Head.Open -> g.draw(arrow(to, angle)) + Head.Cross -> cross(g, to, angle) + Head.Dot -> g.fill(Ellipse2D.Double(to.x - DOT, to.y - DOT, DOT * 2, DOT * 2)) + Head.None -> Unit + } + } + + private fun arrow(to: Pt, angle: Double): Path2D { + val left = point(to, angle + PI * 0.82, HEAD) + val right = point(to, angle - PI * 0.82, HEAD) + return Path2D.Double().apply { + moveTo(to.x, to.y) + lineTo(left.x, left.y) + lineTo(right.x, right.y) + closePath() + } + } + + private fun cross(g: Graphics2D, to: Pt, angle: Double) { + val a = point(to, angle + PI / 4.0, CROSS) + val b = point(to, angle + PI + PI / 4.0, CROSS) + val c = point(to, angle - PI / 4.0, CROSS) + val d = point(to, angle + PI - PI / 4.0, CROSS) + g.draw(Line2D.Double(a.x, a.y, b.x, b.y)) + g.draw(Line2D.Double(c.x, c.y, d.x, d.y)) + } + + private fun point(pt: Pt, angle: Double, len: Double) = Pt(pt.x + cos(angle) * len, pt.y + sin(angle) * len) + + private fun path(points: List, close: Boolean): Path2D { + val path = Path2D.Double() + val first = points.firstOrNull() ?: return path + path.moveTo(first.x, first.y) + points.drop(1).forEach { path.lineTo(it.x, it.y) } + if (close) path.closePath() + return path + } + + private fun stroke(dash: Boolean = false, thick: Boolean = false): BasicStroke { + val width = if (thick) THICK else THIN + if (!dash) return BasicStroke(width, BasicStroke.CAP_ROUND, BasicStroke.JOIN_ROUND) + return BasicStroke(width, BasicStroke.CAP_ROUND, BasicStroke.JOIN_ROUND, max(DASH, width), floatArrayOf(DASH, DASH), 0f) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt new file mode 100644 index 00000000000..bd0b032ac06 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -0,0 +1,111 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.session.ui.SessionSurface +import ai.kilocode.client.session.ui.selection.SessionCopyTarget +import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.ui.diagram.Art +import ai.kilocode.client.ui.diagram.Painters +import ai.kilocode.client.ui.diagram.Palette +import com.intellij.util.concurrency.annotations.RequiresEdt +import com.intellij.util.ui.JBUI +import java.awt.Dimension +import java.awt.Graphics +import java.awt.Graphics2D +import java.awt.RenderingHints +import javax.swing.JComponent +import kotlin.math.roundToInt + +internal class DiagramPanel( + private var source: String, + private var palette: Palette, +) : JComponent(), SessionCopyTarget { + private var art: Art? = null + private var last = Dimension(0, 0) + + override val copyAnchor: JComponent get() = this + + @RequiresEdt + override fun copyText() = source + + @RequiresEdt + fun source(value: String) { + source = value + } + + @RequiresEdt + fun art(value: Art) { + art = value + resize() + repaint() + } + + @RequiresEdt + fun palette(value: Palette) { + palette = value + repaint() + } + + override fun getPreferredSize() = fitSize() + + override fun getMinimumSize() = fitSize() + + override fun getMaximumSize() = Dimension(Int.MAX_VALUE, fitSize().height) + + override fun setBounds(x: Int, y: Int, width: Int, height: Int) { + val before = fitSize() + super.setBounds(x, y, width, height) + if (before.height != fitSize().height) resize() + } + + override fun paintComponent(g: Graphics) { + val g2 = g.create() as Graphics2D + try { + g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON) + g2.color = background + val arc = JBUI.scale(SessionUiStyle.View.BLOCK_ARC) + g2.fillRoundRect(0, 0, width, height, arc, arc) + } finally { + g2.dispose() + } + val value = art ?: return + val painter = Painters.of(value) + val scale = scale(value) + SessionSurface.clipped(g, width, height) { clipped -> + val inner = clipped.create() as Graphics2D + try { + inner.translate(pad(), pad()) + inner.scale(scale, scale) + painter.paint(inner, value, palette) + } finally { + inner.dispose() + } + } + } + + private fun resize() { + val next = fitSize() + if (last == next) return + last = next + revalidate() + } + + private fun fitSize(): Dimension { + val value = art ?: return Dimension(0, emptyHeight()) + val size = Painters.of(value).size(value) + val height = (size.h * scale(value)).roundToInt() + pad() * 2 + return Dimension(0, height.coerceAtLeast(emptyHeight())) + } + + private fun scale(value: Art): Double { + val size = Painters.of(value).size(value) + val avail = (width.takeIf { it > 0 } ?: parent?.width ?: 0) - pad() * 2 + val byWidth = if (avail > 0) minOf(1.0, avail / size.w) else 1.0 + val max = JBUI.scale(SessionUiStyle.View.Diagram.MAX_HEIGHT) - pad() * 2 + val byHeight = if (size.h > 0.0) minOf(1.0, max / size.h) else 1.0 + return minOf(byWidth, byHeight).coerceAtLeast(0.1) + } + + private fun pad() = JBUI.scale(SessionUiStyle.View.Diagram.PADDING) + + private fun emptyHeight() = JBUI.scale(SessionUiStyle.View.Diagram.EMPTY_HEIGHT) +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt new file mode 100644 index 00000000000..0eb361ff0a7 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt @@ -0,0 +1,78 @@ +@file:Suppress("UnstableApiUsage", "DEPRECATION") + +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.ui.diagram.AwtMeasure +import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.Fault +import ai.kilocode.client.ui.diagram.FontSpec +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.EDT +import com.intellij.openapi.application.ModalityState +import com.intellij.openapi.application.asContextElement +import com.intellij.openapi.components.Service +import com.intellij.openapi.util.Disposer +import com.intellij.util.concurrency.annotations.RequiresEdt +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext + +@Service(Service.Level.APP) +internal class Diagrams internal constructor( + private val cs: CoroutineScope, + private val engine: Engine?, +) { + constructor(cs: CoroutineScope) : this(cs, null) + + private val measure = AwtMeasure() + private val cache = object : LinkedHashMap(CACHE, 0.75f, true) { + override fun removeEldestEntry(eldest: MutableMap.MutableEntry?) = size > CACHE + } + + @RequiresEdt + fun render(source: String, spec: Spec, owner: Disposable, done: (Out) -> Unit) { + val key = Key(hash(source), spec.font) + cache[key]?.let { + done(it) + return + } + val job = cs.launch { + val out = try { + impl().draw(source, spec) + } catch (err: CancellationException) { + throw err + } catch (err: Exception) { + Out.Err(Fault.Internal, err.message ?: err.javaClass.simpleName) + } + withContext(edt) { + if (Disposer.isDisposed(owner)) return@withContext + cache[key] = out + done(out) + } + } + Disposer.register(owner) { job.cancel() } + } + + private fun impl() = engine ?: Mermaid(measure) + + private data class Key(val hash: Long, val font: FontSpec) + + private companion object { + const val CACHE = 64 + val edt = Dispatchers.EDT + ModalityState.any().asContextElement() + + fun hash(text: String): Long { + var value = -3750763034362895579L + for (char in text) { + value = value xor char.code.toLong() + value *= 1099511628211L + } + return value + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt index 3be503139e0..bf5215d9769 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt @@ -8,6 +8,7 @@ import com.intellij.openapi.fileTypes.UnknownFileType internal sealed class Kind { data class Source(val file: FileType, val highlight: Highlight = Highlight.None) : Kind() data class Terminal(val stream: Stream, val mode: Mode) : Kind() + data class Diagram(val file: FileType) : Kind() } internal enum class Stream { Stdout, Stderr } @@ -69,6 +70,7 @@ internal object MdLanguage { terms[key]?.let { return it } if (key == "shell script") return Kind.Source(type("sh")) val single = key.substringBefore(' ') + if (key == "mermaid" || key == "mmd" || single == "mermaid" || single == "mmd") return Kind.Diagram(type("mmd")) if (key in pure || single in pure) return Kind.Source(PlainTextFileType.INSTANCE, Highlight.DiffPure) if (key in diffs || single in diffs) return Kind.Source(PlainTextFileType.INSTANCE, Highlight.Diff) terms[single]?.let { return it } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdProjector.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdProjector.kt index 1d3e077951b..1e9929e7cbe 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdProjector.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdProjector.kt @@ -60,7 +60,7 @@ internal class MdProjector { val pending = idx == lines.lastIndex && pendingOpener(line.text) if (pending) { flush() - blocks.add(Desc.Code("", Kind.Source(PlainTextFileType.INSTANCE))) + blocks.add(Desc.Code("", Kind.Source(PlainTextFileType.INSTANCE), open = true)) html.append(codeHtml("")) } else { md.append(line.text).append(line.end) @@ -87,7 +87,7 @@ internal class MdProjector { if (!partial) code.append(item.text).append(item.end) idx++ } - val desc = Desc.Code(code.toString(), MdLanguage.kind(open.info)) + val desc = Desc.Code(code.toString(), MdLanguage.kind(open.info), open = !closed || trimmed) blocks.add(desc) html.append(codeHtml(desc.text)) trailing = if (!closed && !trimmed) open else null @@ -215,7 +215,7 @@ internal class MdProjector { internal sealed class Desc { data class Html(val body: String) : Desc() - data class Code(val text: String, val kind: Kind) : Desc() + data class Code(val text: String, val kind: Kind, val open: Boolean = false) : Desc() data class Table(val body: String) : Desc() } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index f4fbe03a7af..281e1204ce8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -5,6 +5,15 @@ import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.ui.selection.SessionSelection import ai.kilocode.client.session.ui.style.SessionEditorStyle import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.FontSpec +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.ui.DiagramPanel +import ai.kilocode.client.ui.diagram.ui.Diagrams +import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.md.MdCodeBlockBorder import ai.kilocode.client.ui.md.MdCodeBlockFactory import ai.kilocode.client.ui.md.MdCommon @@ -15,6 +24,7 @@ import com.intellij.execution.ui.ConsoleViewContentType import com.intellij.openapi.Disposable import com.intellij.openapi.actionSystem.DataSink import com.intellij.openapi.actionSystem.UiDataProvider +import com.intellij.openapi.components.service import com.intellij.openapi.editor.EditorFactory import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.editor.markup.HighlighterLayer @@ -23,10 +33,13 @@ import com.intellij.openapi.fileTypes.FileType import com.intellij.openapi.fileTypes.PlainTextFileType import com.intellij.openapi.project.ProjectManager import com.intellij.openapi.util.Disposer +import com.intellij.openapi.util.registry.Registry +import com.intellij.ui.HyperlinkLabel import com.intellij.ui.components.JBTextArea import com.intellij.ui.components.JBHtmlPane import com.intellij.ui.components.JBHtmlPaneConfiguration import com.intellij.ui.components.JBHtmlPaneStyleConfiguration +import com.intellij.ui.components.JBLabel import com.intellij.ui.components.JBScrollPane import com.intellij.util.ui.JBUI import java.awt.Color @@ -324,6 +337,7 @@ internal open class MdViewHybrid( is Desc.Code -> when (val kind = desc.kind) { is Kind.Source -> CodeView(desc, codeBlock(desc.text, kind, disposable), disposable) is Kind.Terminal -> TermView(desc, terminalBlock(desc.text, kind, disposable), disposable) + is Kind.Diagram -> DiagramView(desc, kind, disposable) } } } @@ -501,6 +515,24 @@ internal open class MdViewHybrid( return pane } + private fun palette(opts: MdStyle): Palette { + val font = style.editorFont + return Palette( + surface = UiStyle.Colors.contrast(opts.preBg, 8), + border = opts.codeBorder, + text = opts.foreground, + muted = opts.quoteFg, + accent = opts.linkColor, + note = opts.quoteBg, + cluster = opts.codeBorder, + line = opts.quoteFg, + font = font, + bold = style.boldEditorFont, + ) + } + + private fun spec() = Spec(FontSpec(style.editorFamily, style.editorSize)) + private fun styleCodePane(pane: JBScrollPane, opts: MdStyle) { pane.apply { val width = SessionUiStyle.View.Code.BORDER_WIDTH @@ -944,6 +976,153 @@ internal open class MdViewHybrid( } } + private inner class DiagramView(desc: Desc.Code, kind: Kind.Diagram, disposable: Disposable) : + View(desc, Stack.vertical(gap = UiStyle.Gap.sm()), disposable) { + private val root = component as Stack + private val codePane = codeBlock(desc.text, Kind.Source(kind.file), disposable) + private val toggle = HyperlinkLabel(KiloBundle.message("diagram.diagram")) + private val label = JBLabel(KiloBundle.message("diagram.rendering")).apply { + foreground = SessionUiStyle.Text.Secondary.foreground() + } + private val row = Stack.horizontal(gap = UiStyle.Gap.md()).apply { + next(toggle) + next(label) + } + private var panel: DiagramPanel? = null + private var hash = 0 + private var gen = 0 + private var font = spec().font + + init { + root.next(codePane).next(row) + toggle.addHyperlinkListener { toggle() } + kick() + } + + override fun compatible(desc: Desc) = desc is Desc.Code && desc.kind is Kind.Diagram + + override fun update(desc: Desc) { + if (this.desc == desc) return + this.desc = desc + val item = desc as Desc.Code + panel?.source(item.text) + updateCode(item.text) + kick() + } + + override fun grow(delta: String) { + val item = desc as Desc.Code + update(item.copy(text = item.text + delta, open = true)) + } + + override fun style(opts: MdStyle) { + styleCodePane(codePane, opts) + val view = codePane.viewport.view + when (view) { + is CodeField -> { + view.font = style.editorFont + view.background = opts.preBg + view.getEditor(false)?.let { ed -> applyEditorChrome(ed, opts, view.soft) } + } + is JBTextArea -> styleTextArea(view, opts) + } + if (view is JComponent) { + sizeCodeField(view, fieldText(view)) + sizeCodePane(codePane, view) + } + panel?.background = opts.preBg + panel?.palette(palette(opts)) + val next = spec().font + if (font == next) return + font = next + hash = 0 + kick() + } + + private fun kick() { + val item = desc as Desc.Code + if (!Registry.`is`("kilo.diagram.inline.enabled", true)) { + label.text = "" + showSource() + return + } + if (item.open) { + showSource() + label.text = KiloBundle.message("diagram.rendering") + label.foreground = SessionUiStyle.Text.Secondary.foreground() + return + } + val code = item.text.hashCode() + if (hash == code) return + hash = code + label.text = KiloBundle.message("diagram.rendering") + label.foreground = SessionUiStyle.Text.Secondary.foreground() + val seq = ++gen + service().render(item.text, spec(), disposable) { out -> + if (seq != gen) return@render + when (out) { + is Out.Ok -> ok(out) + is Out.Err -> fail(out.message) + } + } + } + + private fun ok(out: Out.Ok) { + val pane = panel ?: DiagramPanel((desc as Desc.Code).text, palette(opts())).also { + it.background = opts().preBg + panel = it + } + pane.art(out.art) + if (pane.parent == null) root.add(pane, 0) + showDiagram() + label.text = "" + root.revalidate() + root.repaint() + } + + private fun fail(message: String) { + val text = message.ifBlank { KiloBundle.message("diagram.rendering") } + label.text = KiloBundle.message("diagram.error", text) + label.foreground = UiStyle.Colors.errorLabelForeground() + showSource() + root.revalidate() + root.repaint() + } + + private fun toggle() { + if (panel?.parent === root && codePane.parent == null) showSource() else showDiagram() + root.revalidate() + root.repaint() + } + + private fun showDiagram() { + val pane = panel ?: return + if (pane.parent == null) root.add(pane, 0) + if (codePane.parent === root) root.remove(codePane) + toggle.setHyperlinkText(KiloBundle.message("diagram.source")) + } + + private fun showSource() { + val pane = panel + if (pane?.parent === root) root.remove(pane) + if (codePane.parent == null) root.add(codePane, 0) + toggle.setHyperlinkText(KiloBundle.message("diagram.diagram")) + } + + private fun updateCode(text: String) { + val value = text.trimEnd('\n') + val view = codePane.viewport.view + when (view) { + is CodeField -> view.text = value + is JBTextArea -> view.text = value + } + if (view is JComponent) { + sizeCodeField(view, value) + sizeCodePane(codePane, view) + } + } + } + private inner class TermView(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).kind == desc.kind diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml index aeecfb90e2b..340505af2b7 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml +++ b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml @@ -142,7 +142,12 @@ restartRequired="false" overrides="false"/> + diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index 0f0bd29d267..0f9fc4489ca 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -5,6 +5,10 @@ common.rename=Rename common.rename.help=Use a custom name that describes your task. common.save=Save common.dont.show.again=Don''t show again +diagram.source=Source +diagram.diagram=Diagram +diagram.error=Couldn''t render diagram: {0} +diagram.rendering=Rendering diagram... session.action.cancel=Cancel session.connection.connecting=Loading... diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ScenePainterTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ScenePainterTest.kt new file mode 100644 index 00000000000..77189e9aff3 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ScenePainterTest.kt @@ -0,0 +1,66 @@ +package ai.kilocode.client.ui.diagram + +import java.awt.Color +import java.awt.Font +import java.awt.image.BufferedImage +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotEquals +import kotlin.test.assertTrue + +class ScenePainterTest { + @Test + fun `test painter renders all mark variants`() { + val scene = Scene( + Type.Flowchart, + listOf( + Mark.Box(Rect(10.0, 10.0, 30.0, 20.0), 6.0, Role.Surface, Role.Border), + Mark.Oval(Rect(50.0, 10.0, 24.0, 20.0), Role.Note, Role.Line), + Mark.Poly(listOf(Pt(15.0, 50.0), Pt(35.0, 45.0), Pt(45.0, 65.0)), Role.Cluster, Role.Border), + Mark.Edge(listOf(Pt(70.0, 50.0), Pt(110.0, 50.0)), Role.Accent, dash = true, thick = true, head = Head.Arrow), + Mark.Group("g", listOf(Mark.Text("T", Pt(90.0, 25.0), Anchor.Center, Role.Text, bold = true))), + ), + Size(130.0, 80.0), + ) + val img = BufferedImage(140, 90, BufferedImage.TYPE_INT_ARGB) + + ScenePainter.paint(img.createGraphics(), scene, palette()) + + assertNotEquals(0, img.rgb(20, 20)) + assertNotEquals(0, img.rgb(60, 20)) + assertNotEquals(0, img.rgb(25, 55)) + assertNotEquals(0, img.rgb(105, 50)) + assertTrue(nonEmpty(img) > 300) + } + + @Test + fun `test registry chooses scene painter`() { + val scene = Scene(Type.Sequence, emptyList(), Size(1.0, 2.0)) + + assertEquals(ScenePainter, Painters.of(scene)) + assertEquals(Size(1.0, 2.0), Painters.of(scene).size(scene)) + } + + private fun palette() = Palette( + surface = Color(0xEE, 0xEE, 0xEE), + border = Color(0x11, 0x11, 0x11), + text = Color(0x22, 0x22, 0x22), + muted = Color(0x77, 0x77, 0x77), + accent = Color(0x00, 0x66, 0xCC), + note = Color(0xFF, 0xF5, 0xCC), + cluster = Color(0xDD, 0xEE, 0xFF), + line = Color(0x33, 0x33, 0x33), + font = Font(Font.SANS_SERIF, Font.PLAIN, 12), + bold = Font(Font.SANS_SERIF, Font.BOLD, 12), + ) + + private fun BufferedImage.rgb(x: Int, y: Int) = getRGB(x, y) ushr 24 + + private fun nonEmpty(img: BufferedImage): Int { + var count = 0 + for (x in 0 until img.width) { + for (y in 0 until img.height) if (img.rgb(x, y) != 0) count++ + } + return count + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt new file mode 100644 index 00000000000..e6ad1216cbf --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -0,0 +1,59 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.ui.diagram.Mark +import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.diagram.Rect +import ai.kilocode.client.ui.diagram.Role +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Type +import java.awt.Color +import java.awt.Font +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +class DiagramPanelTest { + @Test + fun `test panel fits width and caps height`() { + val panel = DiagramPanel("source", palette()) + panel.setSize(100, 1) + panel.art(scene(300.0, 100.0)) + + assertTrue(panel.preferredSize.height < 100) + assertEquals(0, panel.preferredSize.width) + + panel.setSize(2_000, 1) + panel.art(scene(100.0, 2_000.0)) + + assertTrue(panel.preferredSize.height <= 520) + } + + @Test + fun `test copy returns source`() { + val panel = DiagramPanel("flowchart TD", palette()) + + assertEquals("flowchart TD", panel.copyText()) + panel.source("sequenceDiagram") + assertEquals("sequenceDiagram", panel.copyText()) + } + + private fun scene(w: Double, h: Double) = Scene( + Type.Flowchart, + listOf(Mark.Box(Rect(0.0, 0.0, w, h), 4.0, Role.Surface, Role.Border)), + Size(w, h), + ) + + private fun palette() = Palette( + surface = Color.WHITE, + border = Color.BLACK, + text = Color.BLACK, + muted = Color.GRAY, + accent = Color.BLUE, + note = Color.YELLOW, + cluster = Color.LIGHT_GRAY, + line = Color.DARK_GRAY, + font = Font(Font.SANS_SERIF, Font.PLAIN, 12), + bold = Font(Font.SANS_SERIF, Font.BOLD, 12), + ) +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt new file mode 100644 index 00000000000..d572fd2e636 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt @@ -0,0 +1,91 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.testing.TestCoroutines +import ai.kilocode.client.testing.pumpEdt +import ai.kilocode.client.ui.diagram.Art +import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.FontSpec +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.Type +import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import kotlinx.coroutines.awaitCancellation + +class DiagramsTest : BasePlatformTestCase() { + private lateinit var coroutines: TestCoroutines + private lateinit var engine: FakeEngine + private lateinit var service: Diagrams + + override fun setUp() { + super.setUp() + coroutines = TestCoroutines() + engine = FakeEngine() + service = Diagrams(coroutines.scope, engine) + } + + override fun tearDown() { + try { + coroutines.close() + } finally { + super.tearDown() + } + } + + fun `test miss resolves then identical request is synchronous cache hit`() { + val owner = Disposer.newDisposable("diagram") + val calls = mutableListOf() + + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + assertTrue(calls.isEmpty()) + coroutines.drain() + assertEquals(1, calls.size) + assertEquals(1, engine.calls) + + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + assertEquals(2, calls.size) + assertEquals(1, engine.calls) + Disposer.dispose(owner) + } + + fun `test different font misses cache`() { + val owner = Disposer.newDisposable("diagram") + + service.render("flowchart TD\nA-->B", spec(12), owner) {} + coroutines.drain() + service.render("flowchart TD\nA-->B", spec(13), owner) {} + coroutines.drain() + + assertEquals(2, engine.calls) + Disposer.dispose(owner) + } + + fun `test owner dispose cancels render callback`() { + val owner = Disposer.newDisposable("diagram") + engine.pause = true + var called = false + + service.render("flowchart TD\nA-->B", spec(), owner) { called = true } + Disposer.dispose(owner) + coroutines.drain(::pumpEdt) + + assertFalse(called) + } + + private fun spec(size: Int = 12) = Spec(FontSpec("Test", size)) + + private class FakeEngine : Engine { + var calls = 0 + var pause = false + + override fun accepts(type: Type) = true + + override suspend fun draw(source: String, spec: Spec): Out { + calls++ + if (pause) awaitCancellation() + return Out.Ok(Scene(Type.Flowchart, emptyList(), Size(20.0, 10.0)) as Art) + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdLanguageTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdLanguageTest.kt index 06793dd446f..878013e79ee 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdLanguageTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdLanguageTest.kt @@ -56,6 +56,11 @@ class MdLanguageTest : BasePlatformTestCase() { assertKind(" ansi-stdout ignored metadata ", Stream.Stdout, Mode.Ansi) } + fun `test mermaid resolves to diagram kind`() { + assertSame(type("mmd"), (MdLanguage.kind("mermaid") as Kind.Diagram).file) + assertSame(type("mmd"), (MdLanguage.kind("mmd title") as Kind.Diagram).file) + } + private fun assertKind(lang: String, stream: Stream, mode: Mode) { val kind = MdLanguage.kind(lang) as Kind.Terminal diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdProjectorTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdProjectorTest.kt index 2e34f69dc37..61c9c26b5e6 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdProjectorTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdProjectorTest.kt @@ -45,7 +45,7 @@ class MdProjectorTest : BasePlatformTestCase() { fun `test partial opener renders empty code block`() { val out = projector.project("``") - assertEquals(listOf(Desc.Code("", Kind.Source(PlainTextFileType.INSTANCE))), out.blocks) + assertEquals(listOf(Desc.Code("", Kind.Source(PlainTextFileType.INSTANCE), open = true)), out.blocks) assertEquals("
\n", out.html) assertNull(out.open) } @@ -55,6 +55,7 @@ class MdProjectorTest : BasePlatformTestCase() { val code = out.blocks.single() as Desc.Code assertEquals("print(1)\n", code.text) + assertTrue(code.open) assertFalse(out.html.contains("python")) assertEquals('`', out.open!!.char) } @@ -64,9 +65,11 @@ class MdProjectorTest : BasePlatformTestCase() { val complete = projector.project("```python\nprint(1)\n```\n\nafter") assertEquals("print(1)\n", (partial.blocks.single() as Desc.Code).text) + assertTrue((partial.blocks.single() as Desc.Code).open) assertNull(partial.open) assertEquals(2, complete.blocks.size) assertEquals("print(1)\n", (complete.blocks[0] as Desc.Code).text) + assertFalse((complete.blocks[0] as Desc.Code).open) assertTrue((complete.blocks[1] as Desc.Html).body.contains("after")) } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt new file mode 100644 index 00000000000..d71c2130e1d --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -0,0 +1,140 @@ +package ai.kilocode.client.ui.md + +import ai.kilocode.client.testing.TestCoroutines +import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.Mark +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.diagram.Pt +import ai.kilocode.client.ui.diagram.Role +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.Type +import ai.kilocode.client.ui.diagram.ui.DiagramPanel +import ai.kilocode.client.ui.diagram.ui.Diagrams +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import com.intellij.testFramework.replaceService +import com.intellij.ui.EditorTextField +import com.intellij.util.ui.UIUtil +import javax.swing.JPanel + +@Suppress("UnstableApiUsage") +class MdViewDiagramTest : BasePlatformTestCase() { + private lateinit var coroutines: TestCoroutines + private lateinit var engine: FakeEngine + private lateinit var view: MdView + + override fun setUp() { + super.setUp() + coroutines = TestCoroutines() + engine = FakeEngine() + ApplicationManager.getApplication().replaceService(Diagrams::class.java, Diagrams(coroutines.scope, engine), testRootDisposable) + view = MdViewFactory.hybrid() + } + + override fun tearDown() { + try { + if (this::view.isInitialized) Disposer.dispose(view) + coroutines.close() + } finally { + super.tearDown() + } + } + + fun `test mermaid fence renders and hides source`() { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + + assertEquals(1, diagrams().size) + assertEquals(0, editors().size) + assertEquals(1, engine.calls) + } + + fun `test engine error keeps source visible`() { + engine.out = Out.Err(ai.kilocode.client.ui.diagram.Fault.Syntax, "bad syntax") + + view.set("```mermaid\nflowchart TD\nA-->\n```") + drain() + + assertEquals(0, diagrams().size) + assertEquals(1, editors().size) + assertTrue(labels().contains("bad syntax")) + } + + fun `test streaming waits for closed fence`() { + view.append("```mermaid\n") + view.append("flowchart TD\n") + view.append("A-->B\n") + drain() + + assertEquals(0, engine.calls) + assertEquals(1, editors().size) + + view.append("```") + drain() + + assertEquals(1, engine.calls) + assertEquals(1, diagrams().size) + } + + fun `test repeated set retains diagram view and does not leak editors`() { + val base = EditorFactory.getInstance().allEditors.size + + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + val panel = diagramContainers().single() + + repeat(50) { i -> + view.set("```mermaid\nflowchart TD\nA-->B$i\n```") + drain() + assertSame(panel, diagramContainers().single()) + } + + view.clear() + UIUtil.dispatchAllInvocationEvents() + + assertEquals(base, EditorFactory.getInstance().allEditors.size) + } + + private fun drain() = coroutines.drain() + + private fun root() = view.component as JPanel + + private fun diagramContainers() = root().components.filterIsInstance() + + private fun diagrams() = descendants(root()).filterIsInstance() + + private fun editors() = descendants(root()).filterIsInstance() + + private fun labels() = descendants(root()).joinToString("\n") { (it as? javax.swing.JLabel)?.text.orEmpty() } + + private fun descendants(root: java.awt.Container): List { + val out = mutableListOf() + for (comp in root.components) { + out.add(comp) + if (comp is java.awt.Container) out.addAll(descendants(comp)) + } + return out + } + + private class FakeEngine : Engine { + var calls = 0 + var out: Out? = null + + override fun accepts(type: Type) = true + + override suspend fun draw(source: String, spec: Spec): Out { + calls++ + return out ?: Out.Ok( + Scene( + Type.Flowchart, + listOf(Mark.Edge(listOf(Pt(10.0, 10.0), Pt(80.0, 10.0)), Role.Line, head = ai.kilocode.client.ui.diagram.Head.Arrow)), + Size(100.0, 30.0), + ), + ) + } + } +} From f557ba61c524b573b60722fe5570431907b43bc1 Mon Sep 17 00:00:00 2001 From: kirillk Date: Mon, 24 Aug 2026 15:11:48 -0400 Subject: [PATCH 03/11] feat(jetbrains): add open-in-editor action to rendered diagrams The hover overlay on a rendered mermaid diagram only offered copy. Expose the shared copy toolbar with an extra standard open action that puts the mermaid source in a scratch editor tab, so the diagram keeps the affordances of the code block it replaces and the source can be edited or previewed with IDE tooling. --- .../client/ui/diagram/ui/DiagramPanel.kt | 17 +++++++++++++ .../client/ui/diagram/ui/DiagramSource.kt | 25 +++++++++++++++++++ .../resources/messages/KiloBundle.properties | 1 + .../client/ui/diagram/ui/DiagramPanelTest.kt | 22 ++++++++++++++++ 4 files changed, 65 insertions(+) create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index bd0b032ac06..1708ad93a01 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -1,11 +1,15 @@ package ai.kilocode.client.ui.diagram.ui +import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.SessionSurface import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.session.views.MessageToolbar +import ai.kilocode.client.ui.ToolbarButtonAction import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Painters import ai.kilocode.client.ui.diagram.Palette +import com.intellij.icons.AllIcons import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import java.awt.Dimension @@ -22,8 +26,21 @@ internal class DiagramPanel( private var art: Art? = null private var last = Dimension(0, 0) + // Copy plus the standard open-in-editor action, so a rendered diagram keeps every hover + // affordance the code block it replaced had. + private val toolbar = MessageToolbar( + text = { source }, + actions = listOf( + ToolbarButtonAction(AllIcons.Actions.EditSource, KiloBundle.message("diagram.open")) { + openDiagram(this, source) + }, + ), + ) + override val copyAnchor: JComponent get() = this + override val copyToolbar: JComponent get() = toolbar + @RequiresEdt override fun copyText() = source diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt new file mode 100644 index 00000000000..703e80e7a6b --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt @@ -0,0 +1,25 @@ +package ai.kilocode.client.ui.diagram.ui + +import com.intellij.ide.DataManager +import com.intellij.ide.scratch.ScratchRootType +import com.intellij.openapi.actionSystem.CommonDataKeys +import com.intellij.openapi.fileEditor.FileEditorManager +import com.intellij.util.concurrency.annotations.RequiresEdt +import javax.swing.JComponent + +private const val NAME = "diagram.mmd" + +/** + * Opens the mermaid source of a rendered diagram in a real editor tab. + * + * A scratch file rather than an in-memory light file, so the text is editable, savable, and picked up + * by whatever mermaid tooling the IDE has installed for the `.mmd` file type. + */ +@RequiresEdt +internal fun openDiagram(anchor: JComponent, source: String): Boolean { + val ctx = DataManager.getInstance().getDataContext(anchor) + val project = CommonDataKeys.PROJECT.getData(ctx) ?: return false + val file = ScratchRootType.getInstance().createScratchFile(project, NAME, null, source) ?: return false + FileEditorManager.getInstance(project).openFile(file, true) + return true +} diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index 0f9fc4489ca..3341e6ec873 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -8,6 +8,7 @@ common.dont.show.again=Don''t show again diagram.source=Source diagram.diagram=Diagram diagram.error=Couldn''t render diagram: {0} +diagram.open=Open in Editor diagram.rendering=Rendering diagram... session.action.cancel=Cancel diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt index e6ad1216cbf..749c34dd24a 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -1,5 +1,6 @@ package ai.kilocode.client.ui.diagram.ui +import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.ui.diagram.Mark import ai.kilocode.client.ui.diagram.Palette import ai.kilocode.client.ui.diagram.Rect @@ -9,6 +10,7 @@ import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Type import java.awt.Color import java.awt.Font +import javax.swing.AbstractButton import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertTrue @@ -38,6 +40,26 @@ class DiagramPanelTest { assertEquals("sequenceDiagram", panel.copyText()) } + @Test + fun `test hover toolbar offers copy and open in editor`() { + val panel = DiagramPanel("flowchart TD", palette()) + + val buttons = buttons(panel.copyToolbar) + + assertEquals(2, buttons.size) + assertTrue(buttons.any { it.toolTipText == KiloBundle.message("diagram.open") }) + assertTrue(buttons.any { it.toolTipText == KiloBundle.message("session.copy.hover") }) + } + + private fun buttons(root: java.awt.Container): List { + val out = mutableListOf() + for (comp in root.components) { + if (comp is AbstractButton) out.add(comp) + if (comp is java.awt.Container) out.addAll(buttons(comp)) + } + return out + } + private fun scene(w: Double, h: Double) = Scene( Type.Flowchart, listOf(Mark.Box(Rect(0.0, 0.0, w, h), 4.0, Role.Surface, Role.Border)), From c1ded75ba82f7287a5992ac825ac32f7f535eb86 Mon Sep 17 00:00:00 2001 From: kirillk Date: Mon, 24 Aug 2026 15:35:25 -0400 Subject: [PATCH 04/11] fix(jetbrains): show diagram hover actions and order block children Two bugs hid the diagram affordances. The hover resolver always handed the message-level toolbar to nested targets, so the diagram's copy and open buttons never appeared, and Stack ignores add() indexes, so the toggle row rendered above the diagram instead of below it. Let the deepest target win when it brings its own toolbar, and make the diagram block a stable Stack whose children keep a fixed order and switch between diagram and source by visibility. The block, not the painted panel, is now the hover target so the toolbar stays put while the pointer moves within the block. --- .../ui/selection/SessionTargetResolver.kt | 8 ++- .../client/ui/diagram/ui/DiagramBlock.kt | 40 +++++++++++ .../client/ui/diagram/ui/DiagramPanel.kt | 33 +-------- .../client/ui/md/hybrid/MdViewHybrid.kt | 42 ++++++----- .../session/ui/SessionSelectionCopyTest.kt | 15 ++++ .../client/ui/diagram/ui/DiagramPanelTest.kt | 19 ++--- .../client/ui/md/MdViewDiagramTest.kt | 70 +++++++++++++++---- 7 files changed, 145 insertions(+), 82 deletions(-) create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionTargetResolver.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionTargetResolver.kt index 78a05e8b9a2..754d36b3f68 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionTargetResolver.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionTargetResolver.kt @@ -69,8 +69,12 @@ internal object SessionTargetResolver { if (current is SessionCopyTarget && current.copyEligible) targets.add(current) current = current.parent } - val toolbar = targets.indexOfFirst { it.copyToolbar != null } - if (toolbar > 0) return targets.take(toolbar).firstOrNull { it.copyToolbar == null } + val own = targets.indexOfFirst { it.copyToolbar != null } + // The deepest target under the pointer wins when it brings its own toolbar (rendered + // diagram); a toolbar-owning ancestor instead yields to a plain inner target (code block), + // and plain targets anchor on the outermost one for a stable hover position. + if (own == 0) return targets.first() + if (own > 0) return targets.take(own).firstOrNull { it.copyToolbar == null } return targets.lastOrNull() } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt new file mode 100644 index 00000000000..bb770883ba5 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt @@ -0,0 +1,40 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.ui.selection.SessionCopyTarget +import ai.kilocode.client.session.views.MessageToolbar +import ai.kilocode.client.ui.ToolbarButtonAction +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.layout.Stack +import ai.kilocode.client.ui.layout.StackAxis +import com.intellij.icons.AllIcons +import com.intellij.util.concurrency.annotations.RequiresEdt +import javax.swing.JComponent + +/** + * Container for one rendered diagram plus its source fallback. + * + * The block, not the painted [DiagramPanel], is the hover target so the floating toolbar stays put + * while the pointer travels between the diagram and the toggle row underneath it. Copy is paired with + * the standard open-in-editor action, matching the affordances of the code block it replaces. + */ +internal class DiagramBlock : Stack(StackAxis.VERTICAL, UiStyle.Gap.sm()), SessionCopyTarget { + /** Source of the fence text; the owning view rebinds it as the diagram streams or updates. */ + var text: () -> String = { "" } + + private val bar = MessageToolbar( + text = { text() }, + actions = listOf( + ToolbarButtonAction(AllIcons.Actions.EditSource, KiloBundle.message("diagram.open")) { + openDiagram(this, text()) + }, + ), + ) + + override val copyAnchor: JComponent get() = this + + override val copyToolbar: JComponent get() = bar + + @RequiresEdt + override fun copyText() = text() +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index 1708ad93a01..b3ebb337f97 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -1,15 +1,10 @@ package ai.kilocode.client.ui.diagram.ui -import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.SessionSurface -import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.ui.style.SessionUiStyle -import ai.kilocode.client.session.views.MessageToolbar -import ai.kilocode.client.ui.ToolbarButtonAction import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Painters import ai.kilocode.client.ui.diagram.Palette -import com.intellij.icons.AllIcons import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import java.awt.Dimension @@ -19,36 +14,10 @@ import java.awt.RenderingHints import javax.swing.JComponent import kotlin.math.roundToInt -internal class DiagramPanel( - private var source: String, - private var palette: Palette, -) : JComponent(), SessionCopyTarget { +internal class DiagramPanel(private var palette: Palette) : JComponent() { private var art: Art? = null private var last = Dimension(0, 0) - // Copy plus the standard open-in-editor action, so a rendered diagram keeps every hover - // affordance the code block it replaced had. - private val toolbar = MessageToolbar( - text = { source }, - actions = listOf( - ToolbarButtonAction(AllIcons.Actions.EditSource, KiloBundle.message("diagram.open")) { - openDiagram(this, source) - }, - ), - ) - - override val copyAnchor: JComponent get() = this - - override val copyToolbar: JComponent get() = toolbar - - @RequiresEdt - override fun copyText() = source - - @RequiresEdt - fun source(value: String) { - source = value - } - @RequiresEdt fun art(value: Art) { art = value diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index 281e1204ce8..76a2776a480 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -11,6 +11,7 @@ import ai.kilocode.client.ui.diagram.FontSpec import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.Palette import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.diagram.ui.DiagramBlock import ai.kilocode.client.ui.diagram.ui.DiagramPanel import ai.kilocode.client.ui.diagram.ui.Diagrams import ai.kilocode.client.ui.layout.Stack @@ -977,9 +978,10 @@ internal open class MdViewHybrid( } private inner class DiagramView(desc: Desc.Code, kind: Kind.Diagram, disposable: Disposable) : - View(desc, Stack.vertical(gap = UiStyle.Gap.sm()), disposable) { - private val root = component as Stack + View(desc, DiagramBlock(), disposable) { + private val root = component as DiagramBlock private val codePane = codeBlock(desc.text, Kind.Source(kind.file), disposable) + private val panel = DiagramPanel(palette(opts())) private val toggle = HyperlinkLabel(KiloBundle.message("diagram.diagram")) private val label = JBLabel(KiloBundle.message("diagram.rendering")).apply { foreground = SessionUiStyle.Text.Secondary.foreground() @@ -988,13 +990,18 @@ internal open class MdViewHybrid( next(toggle) next(label) } - private var panel: DiagramPanel? = null private var hash = 0 private var gen = 0 private var font = spec().font init { - root.next(codePane).next(row) + // Children keep a fixed order — the diagram above its source, both above the toggle row — + // because Stack lays children out in insertion order and ignores add() indexes. Switching + // between diagram and source flips visibility instead of re-adding components. + panel.background = opts().preBg + panel.isVisible = false + root.next(panel).next(codePane).next(row) + root.text = { (this.desc as Desc.Code).text } toggle.addHyperlinkListener { toggle() } kick() } @@ -1004,9 +1011,7 @@ internal open class MdViewHybrid( override fun update(desc: Desc) { if (this.desc == desc) return this.desc = desc - val item = desc as Desc.Code - panel?.source(item.text) - updateCode(item.text) + updateCode((desc as Desc.Code).text) kick() } @@ -1030,8 +1035,8 @@ internal open class MdViewHybrid( sizeCodeField(view, fieldText(view)) sizeCodePane(codePane, view) } - panel?.background = opts.preBg - panel?.palette(palette(opts)) + panel.background = opts.preBg + panel.palette(palette(opts)) val next = spec().font if (font == next) return font = next @@ -1068,12 +1073,7 @@ internal open class MdViewHybrid( } private fun ok(out: Out.Ok) { - val pane = panel ?: DiagramPanel((desc as Desc.Code).text, palette(opts())).also { - it.background = opts().preBg - panel = it - } - pane.art(out.art) - if (pane.parent == null) root.add(pane, 0) + panel.art(out.art) showDiagram() label.text = "" root.revalidate() @@ -1090,22 +1090,20 @@ internal open class MdViewHybrid( } private fun toggle() { - if (panel?.parent === root && codePane.parent == null) showSource() else showDiagram() + if (panel.isVisible) showSource() else showDiagram() root.revalidate() root.repaint() } private fun showDiagram() { - val pane = panel ?: return - if (pane.parent == null) root.add(pane, 0) - if (codePane.parent === root) root.remove(codePane) + panel.isVisible = true + codePane.isVisible = false toggle.setHyperlinkText(KiloBundle.message("diagram.source")) } private fun showSource() { - val pane = panel - if (pane?.parent === root) root.remove(pane) - if (codePane.parent == null) root.add(codePane, 0) + panel.isVisible = false + codePane.isVisible = true toggle.setHyperlinkText(KiloBundle.message("diagram.diagram")) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt index e8d4d59ea74..45acd5c1f56 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt @@ -143,6 +143,21 @@ class SessionSelectionCopyTest : SessionUiTestBase() { assertSame(target, item) } + fun `test hover copy resolver prefers innermost target that owns a toolbar`() { + val root = JPanel(null) + val outer = InlineTarget(JPanel(), JPanel()) + val inner = InlineTarget(JPanel(), JPanel()) + root.setBounds(0, 0, 100, 100) + outer.setBounds(10, 10, 80, 80) + inner.setBounds(5, 5, 20, 20) + root.add(outer) + outer.add(inner) + + val item = SessionTargetResolver.copy(root, root, Point(20, 20)) + + assertSame(inner, item) + } + fun `test code block hover copy target copies full content despite selection`() { showText("```text\nalpha code\n```") val field = textEditors(ui).first { it.text.contains("alpha code") } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt index 749c34dd24a..c1549d023bd 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -18,7 +18,7 @@ import kotlin.test.assertTrue class DiagramPanelTest { @Test fun `test panel fits width and caps height`() { - val panel = DiagramPanel("source", palette()) + val panel = DiagramPanel(palette()) panel.setSize(100, 1) panel.art(scene(300.0, 100.0)) @@ -32,20 +32,13 @@ class DiagramPanelTest { } @Test - fun `test copy returns source`() { - val panel = DiagramPanel("flowchart TD", palette()) + fun `test block copies fence text and offers copy plus open in editor`() { + val block = DiagramBlock() + block.text = { "flowchart TD" } - assertEquals("flowchart TD", panel.copyText()) - panel.source("sequenceDiagram") - assertEquals("sequenceDiagram", panel.copyText()) - } - - @Test - fun `test hover toolbar offers copy and open in editor`() { - val panel = DiagramPanel("flowchart TD", palette()) - - val buttons = buttons(panel.copyToolbar) + val buttons = buttons(block.copyToolbar) + assertEquals("flowchart TD", block.copyText()) assertEquals(2, buttons.size) assertTrue(buttons.any { it.toolTipText == KiloBundle.message("diagram.open") }) assertTrue(buttons.any { it.toolTipText == KiloBundle.message("session.copy.hover") }) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt index d71c2130e1d..a72acc3b94f 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -1,5 +1,7 @@ package ai.kilocode.client.ui.md +import ai.kilocode.client.session.ui.selection.SessionCopyTarget +import ai.kilocode.client.session.ui.selection.SessionTargetResolver import ai.kilocode.client.testing.TestCoroutines import ai.kilocode.client.ui.diagram.Engine import ai.kilocode.client.ui.diagram.Mark @@ -10,6 +12,7 @@ import ai.kilocode.client.ui.diagram.Scene import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.Type +import ai.kilocode.client.ui.diagram.ui.DiagramBlock import ai.kilocode.client.ui.diagram.ui.DiagramPanel import ai.kilocode.client.ui.diagram.ui.Diagrams import com.intellij.openapi.application.ApplicationManager @@ -17,8 +20,9 @@ import com.intellij.openapi.editor.EditorFactory import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.testFramework.replaceService -import com.intellij.ui.EditorTextField +import com.intellij.ui.HyperlinkLabel import com.intellij.util.ui.UIUtil +import java.awt.Point import javax.swing.JPanel @Suppress("UnstableApiUsage") @@ -44,13 +48,45 @@ class MdViewDiagramTest : BasePlatformTestCase() { } } - fun `test mermaid fence renders and hides source`() { + fun `test mermaid fence renders above toggle row and hides source`() { view.set("```mermaid\nflowchart TD\nA-->B\n```") drain() - assertEquals(1, diagrams().size) - assertEquals(0, editors().size) + val children = block().components.toList() + assertEquals(1, engine.calls) + assertSame(diagram(), children.first()) + assertSame(row(), children.last()) + assertTrue(diagram().isVisible) + assertFalse(codePane().isVisible) + } + + fun `test block is the hover target and copies the fence text`() { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + block().setSize(400, 200) + block().doLayout() + + val target = SessionTargetResolver.copy(block(), diagram(), Point(1, 1)) + + assertSame(block(), target) + assertEquals("flowchart TD\nA-->B\n", block().copyText()) + assertSame(block().copyToolbar, (target as SessionCopyTarget).copyToolbar) + } + + fun `test toggle switches between diagram and source`() { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + + toggle().doClick() + + assertFalse(diagram().isVisible) + assertTrue(codePane().isVisible) + + toggle().doClick() + + assertTrue(diagram().isVisible) + assertFalse(codePane().isVisible) } fun `test engine error keeps source visible`() { @@ -59,8 +95,8 @@ class MdViewDiagramTest : BasePlatformTestCase() { view.set("```mermaid\nflowchart TD\nA-->\n```") drain() - assertEquals(0, diagrams().size) - assertEquals(1, editors().size) + assertFalse(diagram().isVisible) + assertTrue(codePane().isVisible) assertTrue(labels().contains("bad syntax")) } @@ -71,13 +107,14 @@ class MdViewDiagramTest : BasePlatformTestCase() { drain() assertEquals(0, engine.calls) - assertEquals(1, editors().size) + assertTrue(codePane().isVisible) + assertFalse(diagram().isVisible) view.append("```") drain() assertEquals(1, engine.calls) - assertEquals(1, diagrams().size) + assertTrue(diagram().isVisible) } fun `test repeated set retains diagram view and does not leak editors`() { @@ -85,12 +122,15 @@ class MdViewDiagramTest : BasePlatformTestCase() { view.set("```mermaid\nflowchart TD\nA-->B\n```") drain() - val panel = diagramContainers().single() + val block = block() + val panel = diagram() repeat(50) { i -> view.set("```mermaid\nflowchart TD\nA-->B$i\n```") drain() - assertSame(panel, diagramContainers().single()) + assertSame(block, block()) + assertSame(panel, diagram()) + assertEquals(3, block().components.size) } view.clear() @@ -103,11 +143,15 @@ class MdViewDiagramTest : BasePlatformTestCase() { private fun root() = view.component as JPanel - private fun diagramContainers() = root().components.filterIsInstance() + private fun block() = descendants(root()).filterIsInstance().single() + + private fun diagram() = descendants(root()).filterIsInstance().single() + + private fun codePane() = block().components[1] - private fun diagrams() = descendants(root()).filterIsInstance() + private fun row() = block().components.last() - private fun editors() = descendants(root()).filterIsInstance() + private fun toggle() = descendants(root()).filterIsInstance().single() private fun labels() = descendants(root()).joinToString("\n") { (it as? javax.swing.JLabel)?.text.orEmpty() } From 83ec1b33df4920241bf826b556a83bb0491f9629 Mon Sep 17 00:00:00 2001 From: kirillk Date: Wed, 26 Aug 2026 13:01:26 -0400 Subject: [PATCH 05/11] fix(jetbrains): harden mermaid parsing, limits and sequence layout Addresses PR review on the mermaid diagram engine: - Share AwtMeasure safely. The engine service holds one instance and draws off the EDT on the default dispatcher, so the font cache is now a ConcurrentHashMap. The Graphics2D that had no dispose() is already gone. - Refuse pathological sources earlier. A character cap is checked before any preprocessing, node and link caps are enforced while the model is built so `A & B & ... --> ...` cannot expand first, and both parsers check cancellation per line. The `%%{...}%%` mask is a linear scan instead of a lazy regex that rescanned to end-of-text per unterminated opener, and the message regex can no longer backtrack over a colon-free line. - Assign subgraph membership on re-mention. `Client --> Gateway` followed by `subgraph core` / `Gateway --> Auth` now puts Gateway inside the frame, which is how mermaid reads it. - Report sequence scene size from glyph extents, not anchors, so a left-anchored self-message label is inside the bounds a renderer clips to. - Close activations left open at the end of a script, so `A->>+B: hi` with no matching deactivate still draws its bar. - Parse participant ids through unquote and match ` as ` case-insensitively outside quotes, so `participant "Alice"` and `PARTICIPANT C AS Client` resolve to one column. assertInBounds now checks text extents rather than anchors, so this class of overflow fails in tests instead of clipping at paint time. --- .../ai/kilocode/client/ui/diagram/Engine.kt | 12 ++++- .../ai/kilocode/client/ui/diagram/Metrics.kt | 18 +++++--- .../client/ui/diagram/mermaid/Flow.kt | 38 +++++++++++++--- .../client/ui/diagram/mermaid/Mermaid.kt | 24 +++++----- .../kilocode/client/ui/diagram/mermaid/Seq.kt | 39 ++++++++++++---- .../client/ui/diagram/mermaid/SeqLayout.kt | 27 +++++++++++- .../client/ui/diagram/mermaid/Source.kt | 25 +++++++++-- .../kilocode/client/ui/diagram/CancelTest.kt | 25 +++++++++++ .../client/ui/diagram/DiagramAsserts.kt | 24 +++++++--- .../client/ui/diagram/FlowLayoutTest.kt | 4 +- .../client/ui/diagram/InvariantTest.kt | 5 ++- .../kilocode/client/ui/diagram/LimitsTest.kt | 18 ++++++++ .../kilocode/client/ui/diagram/MetricsTest.kt | 30 +++++++++++++ .../client/ui/diagram/SeqLayoutTest.kt | 21 ++++++++- .../ui/diagram/mermaid/FlowParseTest.kt | 44 +++++++++++++++++-- .../client/ui/diagram/mermaid/SeqParseTest.kt | 30 +++++++++++-- .../client/ui/diagram/mermaid/SourceTest.kt | 17 +++++++ 17 files changed, 345 insertions(+), 56 deletions(-) create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/MetricsTest.kt diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt index 06f205d595b..ba1ef1dbe4e 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt @@ -36,5 +36,15 @@ internal data class Metrics( val wrap: Double = 0.0, ) +/** + * Guards against pathological model output. [chars] is checked before any preprocessing so a single + * enormous line cannot reach the parsers; [nodes] and [edges] are enforced while the model is built + * rather than after, so `A & B & … --> …` cannot expand into a huge edge list first. + */ @Serializable -internal data class Limits(val nodes: Int = 400, val edges: Int = 800, val lines: Int = 2_000) +internal data class Limits( + val nodes: Int = 400, + val edges: Int = 800, + val lines: Int = 2_000, + val chars: Int = 100_000, +) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt index de0b8c7a8bc..fbc63c8d4f0 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Metrics.kt @@ -2,20 +2,24 @@ package ai.kilocode.client.ui.diagram import java.awt.Font import java.awt.font.FontRenderContext +import java.util.concurrent.ConcurrentHashMap +/** + * AWT text measurement. Holds no native Java2D state, so there is nothing to dispose. + * + * A single instance is shared by every [Engine.draw] call, and those run off the EDT on the default + * dispatcher, so the font cache must tolerate concurrent access. [FontRenderContext] and [Font] are + * both immutable for these queries. + */ internal class AwtMeasure : Measure { private val ctx = FontRenderContext(null, true, true) - private val cache = linkedMapOf() + private val cache = ConcurrentHashMap() override fun width(text: String, font: FontSpec) = font(font).getStringBounds(text, ctx).width override fun height(font: FontSpec) = font(font).getLineMetrics("Ag", ctx).height.toDouble() override fun ascent(font: FontSpec) = font(font).getLineMetrics("Ag", ctx).ascent.toDouble() - private fun font(font: FontSpec): Font { - cache[font]?.let { return it } - val style = if (font.bold) Font.BOLD else Font.PLAIN - val value = Font(font.family, style, font.size) - cache[font] = value - return value + private fun font(spec: FontSpec): Font = cache.computeIfAbsent(spec) { + Font(it.family, if (it.bold) Font.BOLD else Font.PLAIN, it.size) } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt index c5c19536568..f9aab06adf0 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt @@ -1,6 +1,9 @@ package ai.kilocode.client.ui.diagram.mermaid import ai.kilocode.client.ui.diagram.Head +import ai.kilocode.client.ui.diagram.Limits +import kotlin.coroutines.coroutineContext +import kotlinx.coroutines.ensureActive internal enum class Dir { Down, Up, Left, Right } @@ -53,27 +56,30 @@ internal data class Graph( internal sealed interface FlowOut { data class Ok(val graph: Graph) : FlowOut data class Err(val message: String, val line: Int) : FlowOut + data class Over(val message: String) : FlowOut } /** Line-oriented flowchart parser. Unknown statements are skipped rather than failing the diagram. */ -internal class Flow { +internal class Flow(private val limits: Limits = Limits()) { private val nodes = linkedMapOf() private val edges = mutableListOf() private val clusters = linkedMapOf() private val stack = ArrayDeque() private var dir = Dir.Down - fun parse(clean: Clean): FlowOut { + suspend fun parse(clean: Clean): FlowOut { var first = true for (line in clean.lines) { + coroutineContext.ensureActive() val text = line.text.trim() if (text.isEmpty()) continue if (first) { first = false if (header(text)) continue } - val err = stmt(text, line.at) ?: continue - return FlowOut.Err(err, line.at) + val err = stmt(text, line.at) + if (err != null) return FlowOut.Err(err, line.at) + over()?.let { return it } } if (stack.isNotEmpty()) { return FlowOut.Err("subgraph is missing a matching end", clean.lines.lastOrNull()?.at ?: 1) @@ -81,6 +87,16 @@ internal class Flow { return FlowOut.Ok(Graph(dir, nodes, edges, clusters)) } + /** + * Caps are checked per statement, and [add] / [chain] stop one item past the cap, so a single + * pathological line cannot build an unbounded model before the refusal is reported. + */ + private fun over(): FlowOut.Over? { + if (nodes.size > limits.nodes) return FlowOut.Over("flowchart exceeds ${limits.nodes} nodes") + if (edges.size > limits.edges) return FlowOut.Over("flowchart exceeds ${limits.edges} links") + return null + } + private fun header(text: String): Boolean { val token = text.substringBefore(' ').lowercase() if (token != "graph" && token != "flowchart") return false @@ -133,6 +149,7 @@ internal class Flow { val label = labels[idx] ?: hit.label for (from in groups[idx]) { for (to in groups[idx + 1]) { + if (edges.size > limits.edges) return null edges.add(FlowEdge(from, to, hit.link, hit.head, hit.tail, label, edges.size)) } } @@ -219,15 +236,24 @@ internal class Flow { return text.substring(0, cut) } + /** + * A node first mentioned outside a subgraph still joins the first subgraph that mentions it, + * which is how mermaid reads `Client --> Gateway` followed by `subgraph core` / `Gateway --> Auth`. + */ private fun add(id: String, label: List, shape: Shape) { val prior = nodes[id] if (prior == null) { + if (nodes.size > limits.nodes) return nodes[id] = FlowNode(id, label, shape, nodes.size, stack.lastOrNull()) return } + val cluster = prior.cluster ?: stack.lastOrNull() val implicit = prior.label == listOf(prior.id) && prior.shape == Shape.Rect - if (!implicit || label == listOf(id)) return - nodes[id] = prior.copy(label = label, shape = shape) + if (!implicit || label == listOf(id)) { + if (cluster != prior.cluster) nodes[id] = prior.copy(cluster = cluster) + return + } + nodes[id] = prior.copy(label = label, shape = shape, cluster = cluster) } private fun hits(text: String): List { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt index 3db8af82e44..0d8831a4c36 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt @@ -6,6 +6,8 @@ import ai.kilocode.client.ui.diagram.Measure import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.Type +import kotlin.coroutines.coroutineContext +import kotlinx.coroutines.ensureActive /** * In-process mermaid engine covering flowcharts and sequence diagrams. @@ -17,6 +19,10 @@ internal class Mermaid(private val measure: Measure) : Engine { override fun accepts(type: Type) = type == Type.Flowchart || type == Type.Sequence override suspend fun draw(source: String, spec: Spec): Out { + if (source.length > spec.limits.chars) { + return Out.Err(Fault.Limit, "source exceeds ${spec.limits.chars} characters") + } + coroutineContext.ensureActive() val clean = Source.clean(source) if (clean.lines.size > spec.limits.lines) { return Out.Err(Fault.Limit, "source exceeds ${spec.limits.lines} lines") @@ -28,31 +34,21 @@ internal class Mermaid(private val measure: Measure) : Engine { } private suspend fun flow(clean: Clean, spec: Spec): Out { - val parsed = Flow().parse(clean) + val parsed = Flow(spec.limits).parse(clean) + if (parsed is FlowOut.Over) return Out.Err(Fault.Limit, parsed.message) if (parsed is FlowOut.Err) return Out.Err(Fault.Syntax, parsed.message, parsed.line) val graph = (parsed as FlowOut.Ok).graph if (graph.nodes.isEmpty()) return Out.Err(Fault.Syntax, "flowchart has no nodes") - if (graph.nodes.size > spec.limits.nodes) { - return Out.Err(Fault.Limit, "flowchart exceeds ${spec.limits.nodes} nodes") - } - if (graph.edges.size > spec.limits.edges) { - return Out.Err(Fault.Limit, "flowchart exceeds ${spec.limits.edges} links") - } val placed = FlowLayout(measure, spec).run(graph) return Out.Ok(FlowMarks(measure, spec).run(placed)) } private suspend fun seq(clean: Clean, spec: Spec): Out { - val parsed = Seq().parse(clean) + val parsed = Seq(spec.limits).parse(clean) + if (parsed is SeqOut.Over) return Out.Err(Fault.Limit, parsed.message) if (parsed is SeqOut.Err) return Out.Err(Fault.Syntax, parsed.message, parsed.line) val script = (parsed as SeqOut.Ok).script if (script.actors.isEmpty()) return Out.Err(Fault.Syntax, "sequence diagram has no participants") - if (script.actors.size > spec.limits.nodes) { - return Out.Err(Fault.Limit, "sequence diagram exceeds ${spec.limits.nodes} participants") - } - if (script.steps.size > spec.limits.edges) { - return Out.Err(Fault.Limit, "sequence diagram exceeds ${spec.limits.edges} steps") - } return Out.Ok(SeqLayout(measure, spec).run(script)) } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt index ce24ea24c2e..297a20b6a2e 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt @@ -1,6 +1,9 @@ package ai.kilocode.client.ui.diagram.mermaid import ai.kilocode.client.ui.diagram.Head +import ai.kilocode.client.ui.diagram.Limits +import kotlin.coroutines.coroutineContext +import kotlinx.coroutines.ensureActive internal enum class NoteAt { Left, Right, Over } @@ -34,32 +37,42 @@ internal data class Script( internal sealed interface SeqOut { data class Ok(val script: Script) : SeqOut data class Err(val message: String, val line: Int) : SeqOut + data class Over(val message: String) : SeqOut } /** Line-oriented sequence diagram parser. Unknown statements are skipped rather than failing. */ -internal class Seq { +internal class Seq(private val limits: Limits = Limits()) { private val actors = linkedMapOf() private val steps = mutableListOf() private var title = emptyList() private var numbered = false private var depth = 0 - fun parse(clean: Clean): SeqOut { + suspend fun parse(clean: Clean): SeqOut { var first = true for (line in clean.lines) { + coroutineContext.ensureActive() val text = line.text.trim() if (text.isEmpty()) continue if (first) { first = false if (text.substringBefore(' ').lowercase() == "sequencediagram") continue } - val err = stmt(text, line.at) ?: continue - return SeqOut.Err(err, line.at) + val err = stmt(text, line.at) + if (err != null) return SeqOut.Err(err, line.at) + over()?.let { return it } } if (depth > 0) return SeqOut.Err("block is missing a matching end", clean.lines.lastOrNull()?.at ?: 1) return SeqOut.Ok(Script(actors, steps, title, numbered)) } + /** Caps are checked per statement so a refusal never waits for the whole script to be built. */ + private fun over(): SeqOut.Over? { + if (actors.size > limits.nodes) return SeqOut.Over("sequence diagram exceeds ${limits.nodes} participants") + if (steps.size > limits.edges) return SeqOut.Over("sequence diagram exceeds ${limits.edges} steps") + return null + } + private fun stmt(text: String, at: Int): String? { val token = text.substringBefore(' ').lowercase() when (token) { @@ -101,12 +114,17 @@ internal class Seq { return message(text, at) } + /** + * `participant "Alice"` must land on the same column as a later `Alice->>Bob`, so the id goes + * through [name] the way message endpoints do. The ` as ` separator is matched case-insensitively + * and only outside quotes, so `participant "Bob as builder"` stays a single quoted name. + */ private fun actor(text: String): String? { val rest = text.substringAfter(' ', "").trim() if (rest.isEmpty()) return "participant needs a name" - val cut = rest.indexOf(" as ") - val id = if (cut < 0) rest else rest.substring(0, cut).trim() - val label = if (cut < 0) rest else rest.substring(cut + 4).trim() + val cut = AS.findAll(rest).firstOrNull { Source.open(rest, it.range.first) } + val id = name(if (cut == null) rest else rest.substring(0, cut.range.first)) + val label = if (cut == null) rest else rest.substring(cut.range.last + 1) add(id, Source.label(label)) return null } @@ -127,6 +145,7 @@ internal class Seq { } private fun message(text: String, at: Int): String? { + if (!text.contains(':')) return null val match = MSG.find(text) ?: return null val from = name(match.groupValues[1]) val arrow = match.groupValues[2] @@ -145,6 +164,7 @@ internal class Seq { if (id.isEmpty()) return val prior = actors[id] if (prior == null) { + if (actors.size > limits.nodes) return actors[id] = Actor(id, label, actors.size) return } @@ -168,7 +188,10 @@ internal class Seq { val NOTE = Regex("""^[Nn]ote\s+(left of|right of|over)\s+([^:]+):\s*(.*)$""") - val MSG = Regex("""^(.+?)\s*(--?>>|--?>|--?[x)])\s*([+-]?)\s*(.+?)\s*:\s*(.*)$""") + val AS = Regex("""\s+as\s+""", RegexOption.IGNORE_CASE) + + /** The receiver is `[^:]+?` rather than `.+?` so a colon-free line cannot backtrack quadratically. */ + val MSG = Regex("""^(.+?)\s*(--?>>|--?>|--?[x)])\s*([+-]?)\s*([^:]+?)\s*:\s*(.*)$""") fun linkOf(arrow: String) = if (arrow.startsWith("--")) Link.Dotted else Link.Solid diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt index 93925903673..dccb1f076c9 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqLayout.kt @@ -22,6 +22,7 @@ internal class SeqLayout(private val measure: Measure, private val spec: Spec) { private val pad get() = spec.metrics.pad private val gap get() = spec.metrics.gap private val step get() = spec.metrics.rank + private val bold get() = spec.font.copy(bold = true) private val marks = mutableListOf() private val heads = linkedMapOf() @@ -47,9 +48,20 @@ internal class SeqLayout(private val measure: Measure, private val spec: Spec) { Step.Close -> close() } } + drain() return scene(lines() + marks) } + /** + * `A->>+B: hi` without a matching deactivate is normal mermaid, so any activation still open at + * the end of the script is closed at the cursor instead of being dropped without a bar. + */ + private fun drain() { + for (entry in live.entries.toList()) { + while (entry.value.isNotEmpty()) toggle(Step.Toggle(entry.key, false)) + } + } + private fun title(script: Script, high: Double): Double { if (script.title.isEmpty()) return 0.0 script.title.forEachIndexed { idx, text -> @@ -228,10 +240,23 @@ internal class SeqLayout(private val measure: Measure, private val spec: Spec) { is Mark.Oval -> corners(mark.rect) is Mark.Poly -> mark.points is Mark.Edge -> mark.points - is Mark.Text -> listOf(mark.at) + is Mark.Text -> span(mark) is Mark.Group -> mark.marks.flatMap(::pts) } + /** + * A text mark contributes its glyph extent, not just its anchor. Renderers use [Scene.size] for + * scroll and clip bounds, so a left-anchored self-message label would otherwise be clipped. + */ + private fun span(mark: Mark.Text): List { + val room = measure.width(mark.text, if (mark.bold) bold else spec.font) + return when (mark.anchor) { + Anchor.Left, Anchor.TopLeft, Anchor.BottomLeft -> listOf(mark.at, Pt(mark.at.x + room, mark.at.y)) + Anchor.Right, Anchor.TopRight, Anchor.BottomRight -> listOf(Pt(mark.at.x - room, mark.at.y), mark.at) + else -> listOf(Pt(mark.at.x - room / 2, mark.at.y), Pt(mark.at.x + room / 2, mark.at.y)) + } + } + private fun corners(rect: Rect) = listOf(Pt(rect.x, rect.y), Pt(rect.x + rect.w, rect.y + rect.h)) private fun move(mark: Mark, dx: Double, dy: Double): Mark = when (mark) { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt index f41172cf537..12760821b5c 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt @@ -14,7 +14,8 @@ internal data class Clean(val lines: List) */ internal object Source { private const val RAILS = "-.=" - private val DIRECTIVE = Regex("""%%\{[\s\S]*?}%%""") + private const val OPEN = "%%{" + private const val CLOSE = "}%%" fun clean(text: String): Clean { val raw = mask(normalize(text)).split("\n") @@ -60,8 +61,26 @@ internal object Source { .replace("\r", "\n") .replace("\t", " ") - private fun mask(text: String) = DIRECTIVE.replace(text) { match -> - match.value.map { if (it == '\n') '\n' else ' ' }.joinToString("") + /** + * Blanks `%%{ ... }%%` directives in place. Scanned by hand rather than with a lazy regex + * because `%%\{[\s\S]*?}%%` rescans to the end of the text for every unterminated `%%{`, which + * is quadratic on pathological input. + */ + private fun mask(text: String): String { + if (!text.contains(OPEN)) return text + val out = StringBuilder(text) + var idx = 0 + while (idx < out.length) { + val open = out.indexOf(OPEN, idx) + if (open < 0) return out.toString() + val close = out.indexOf(CLOSE, open + OPEN.length) + if (close < 0) return out.toString() + for (at in open until close + CLOSE.length) { + if (out[at] != '\n') out[at] = ' ' + } + idx = close + CLOSE.length + } + return out.toString() } /** Returns the index of the first content line, skipping terminated frontmatter. */ diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt index 31687620a71..ba57232abe7 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt @@ -1,11 +1,15 @@ package ai.kilocode.client.ui.diagram +import ai.kilocode.client.ui.diagram.mermaid.Flow import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import ai.kilocode.client.ui.diagram.mermaid.Seq +import ai.kilocode.client.ui.diagram.mermaid.Source import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.async +import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlin.test.Test import kotlin.test.assertFailsWith @@ -30,6 +34,27 @@ class CancelTest { assertFailsWith { cancel(source) } } + /** Parsing runs before any measurement, so its cancellation checks need their own proof. */ + @Test + fun `parsing stops before layout when the job is cancelled`() { + val flow = sink { Flow().parse(Source.clean("flowchart TD\n A --> B")) } + val seq = sink { Seq().parse(Source.clean("sequenceDiagram\n A->>B: hi")) } + + assertTrue(flow.isEmpty(), "flowchart parsing ignored cancellation") + assertTrue(seq.isEmpty(), "sequence parsing ignored cancellation") + } + + /** Runs [body] in a coroutine that cancels itself first; a result only lands if that was ignored. */ + private fun sink(body: suspend () -> Any): List = runBlocking { + val out = mutableListOf() + val job = Job() + CoroutineScope(job + Dispatchers.Unconfined).launch { + job.cancel() + out.add(body()) + }.join() + out + } + @Test fun `uncancelled work completes`() { val out = runBlocking { Mermaid(FakeMeasure()).draw("flowchart TD\n A --> B", spec()) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt index 3de00e53d8f..449969bed0a 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/DiagramAsserts.kt @@ -22,10 +22,15 @@ internal fun err(out: Out): Out.Err { return out as Out.Err } -/** Every mark must sit inside the reported scene size; a renderer relies on that for scroll bounds. */ -internal fun assertInBounds(scene: Scene) { +/** + * Every mark must sit inside the reported scene size; a renderer relies on that for scroll bounds. + * + * Text marks are checked by glyph extent rather than by anchor, so a label that overflows the + * reported size fails here instead of silently clipping at paint time. + */ +internal fun assertInBounds(scene: Scene, measure: Measure, spec: Spec) { for (mark in flatten(scene.marks)) { - for (pt in points(mark)) { + for (pt in points(mark, measure, spec)) { assertTrue(pt.x >= -EPS, "mark left of origin: $mark") assertTrue(pt.y >= -EPS, "mark above origin: $mark") assertTrue(pt.x <= scene.size.w + EPS, "mark past width ${scene.size.w}: $mark") @@ -103,13 +108,22 @@ internal fun flatten(marks: List): List { return out } -private fun points(mark: Mark): List = when (mark) { +private fun points(mark: Mark, measure: Measure, spec: Spec): List = when (mark) { is Mark.Box -> corners(mark.rect) is Mark.Oval -> corners(mark.rect) is Mark.Poly -> mark.points is Mark.Edge -> mark.points - is Mark.Text -> listOf(mark.at) + is Mark.Text -> span(mark, measure, spec) is Mark.Group -> emptyList() } +private fun span(mark: Mark.Text, measure: Measure, spec: Spec): List { + val room = measure.width(mark.text, spec.font.copy(bold = mark.bold)) + return when (mark.anchor) { + Anchor.Left, Anchor.TopLeft, Anchor.BottomLeft -> listOf(mark.at, Pt(mark.at.x + room, mark.at.y)) + Anchor.Right, Anchor.TopRight, Anchor.BottomRight -> listOf(Pt(mark.at.x - room, mark.at.y), mark.at) + else -> listOf(Pt(mark.at.x - room / 2, mark.at.y), Pt(mark.at.x + room / 2, mark.at.y)) + } +} + private fun corners(rect: Rect) = listOf(Pt(rect.x, rect.y), Pt(rect.x + rect.w, rect.y + rect.h)) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt index 03132d440e2..94acb3eac88 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/FlowLayoutTest.kt @@ -70,8 +70,8 @@ class FlowLayoutTest { """ scene Flowchart 55x210 group s - box 0,134 55x76 arc=4 fill=- line=Cluster dash=true - text "Group" at=28,149 anchor=Center role=Muted bold=true + box 0,56 55x154 arc=4 fill=- line=Cluster dash=true + text "Group" at=28,71 anchor=Center role=Muted bold=true edge 28,38 28,86 role=Line dash=false thick=false head=Arrow tail=None edge 28,116 28,164 role=Line dash=false thick=false head=Arrow tail=None box 16,8 23x30 arc=0 fill=Surface line=Border dash=false diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt index 0779e8a7f9e..4980ae5a389 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/InvariantTest.kt @@ -24,11 +24,12 @@ class InvariantTest { private fun check(measure: Measure) { val engine = Mermaid(measure) + val spec = spec(size = 12) for (name in ConformanceTest.CORPUS) { - val out = runBlocking { engine.draw(read(name), spec(size = 12)) } + val out = runBlocking { engine.draw(read(name), spec) } val scene = scene(out) - assertInBounds(scene) + assertInBounds(scene, measure, spec) assertNoOverlap(scene) assertEdgesTouchNodes(scene) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt index 1acf3137a0d..20aaca87f17 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt @@ -9,6 +9,24 @@ import kotlin.test.assertEquals class LimitsTest { private val engine = Mermaid(FakeMeasure()) + @Test + fun `character cap is enforced before preprocessing`() { + val source = "flowchart TD\n A --> B" + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(chars = 10))) } + + assertEquals(Fault.Limit, err(out).fault) + } + + /** A single line can fan out to n*m links, so the cap has to bite while the edges are built. */ + @Test + fun `ampersand fan out is capped as it expands`() { + val left = (1..40).joinToString(" & ") { "a$it" } + val right = (1..40).joinToString(" & ") { "b$it" } + val out = runBlocking { engine.draw("flowchart TD\n $left --> $right", spec().copy(limits = Limits(edges = 5))) } + + assertEquals(Fault.Limit, err(out).fault) + } + @Test fun `line cap is enforced before parsing`() { val source = "flowchart TD\n" + (1..50).joinToString("\n") { " n$it --> n${it + 1}" } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/MetricsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/MetricsTest.kt new file mode 100644 index 00000000000..cf8594427cb --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/MetricsTest.kt @@ -0,0 +1,30 @@ +package ai.kilocode.client.ui.diagram + +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +import kotlinx.coroutines.awaitAll +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertEquals + +/** One [AwtMeasure] is shared by every off-EDT draw, so its font cache must tolerate concurrency. */ +class MetricsTest { + @Test + fun `a shared instance measures consistently from many coroutines`() { + val measure = AwtMeasure() + val fonts = (8..24).map { FontSpec("Dialog", it, bold = it % 2 == 0) } + val want = fonts.associateWith { AwtMeasure().width(TEXT, it) } + + val got = runBlocking(Dispatchers.Default) { + List(64) { async { fonts.map { it to measure.width(TEXT, it) } } }.awaitAll() + } + + for (batch in got) { + for ((font, width) in batch) assertEquals(want.getValue(font), width, "$font measured differently") + } + } + + private companion object { + const val TEXT = "Ag quick brown fox" + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt index 030e1bceb30..1efb8fd8c04 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SeqLayoutTest.kt @@ -7,10 +7,11 @@ import kotlin.test.Test class SeqLayoutTest { private val engine = Mermaid(FakeMeasure()) + /** The reported width must cover the left-anchored loop label, not just its anchor. */ @Test fun `self messages render loops`() = assertScene( """ - scene Sequence 92x230 + scene Sequence 127x230 edge 28,38 28,222 role=Muted dash=true thick=false head=None tail=None box 8,8 39x30 arc=4 fill=Surface line=Border dash=false text "A" at=28,23 anchor=Center role=Text bold=true @@ -81,6 +82,24 @@ class SeqLayoutTest { draw("sequenceDiagram\n autonumber\n A->>B: one\n B->>A: two"), ) + /** `A->>+B` with no matching deactivate is normal mermaid and must still draw the bar. */ + @Test + fun `unmatched activations still draw a bar`() = assertScene( + """ + scene Sequence 142x196 + edge 28,38 28,188 role=Muted dash=true thick=false head=None tail=None + edge 115,38 115,188 role=Muted dash=true thick=false head=None tail=None + box 8,8 39x30 arc=4 fill=Surface line=Border dash=false + text "A" at=28,23 anchor=Center role=Text bold=true + box 95,8 39x30 arc=4 fill=Surface line=Border dash=false + text "B" at=115,23 anchor=Center role=Text bold=true + edge 28,116 111,116 role=Line dash=false thick=false head=Arrow tail=None + text "open" at=69,105 anchor=Center role=Muted bold=false + box 111,94 8x70 arc=0 fill=Accent line=Border dash=false + """, + draw("sequenceDiagram\n A->>+B: open"), + ) + @Test fun `titles reserve space before participants`() = assertScene( """ diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt index 3b7fda319c5..4b38c170cf8 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowParseTest.kt @@ -1,6 +1,7 @@ package ai.kilocode.client.ui.diagram.mermaid import ai.kilocode.client.ui.diagram.Head +import kotlinx.coroutines.runBlocking import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertNull @@ -134,10 +135,47 @@ class FlowParseTest { assertNull(graph.clusters.getValue("core").parent) assertEquals("core", graph.clusters.getValue("store").parent) assertNull(graph.nodes.getValue("Client").cluster) + assertEquals("core", graph.nodes.getValue("Gate").cluster) assertEquals("core", graph.nodes.getValue("Auth").cluster) assertEquals("store", graph.nodes.getValue("Db").cluster) } + @Test + fun `a node mentioned before a subgraph still joins it`() { + val graph = graph( + """ + flowchart TD + Client --> Gateway + subgraph core [Core] + Gateway --> Auth + end + Gateway[API Gateway] --> Report + """, + ) + + assertEquals("core", graph.nodes.getValue("Gateway").cluster) + assertEquals(listOf("API Gateway"), graph.nodes.getValue("Gateway").label) + assertNull(graph.nodes.getValue("Report").cluster) + } + + @Test + fun `the first subgraph to mention a node wins`() { + val graph = graph( + """ + flowchart TD + subgraph one [One] + A --> B + end + subgraph two [Two] + B --> C + end + """, + ) + + assertEquals("one", graph.nodes.getValue("B").cluster) + assertEquals("two", graph.nodes.getValue("C").cluster) + } + @Test fun `styling statements are skipped and class suffixes dropped`() { val graph = graph( @@ -166,15 +204,15 @@ class FlowParseTest { @Test fun `dangling subgraph and stray end are reported with line numbers`() { - val open = Flow().parse(Source.clean("flowchart TD\n subgraph s\n A --> B")) - val stray = Flow().parse(Source.clean("flowchart TD\n A --> B\n end")) + val open = runBlocking { Flow().parse(Source.clean("flowchart TD\n subgraph s\n A --> B")) } + val stray = runBlocking { Flow().parse(Source.clean("flowchart TD\n A --> B\n end")) } assertEquals(3, (open as FlowOut.Err).line) assertEquals(3, (stray as FlowOut.Err).line) } private fun graph(source: String): Graph { - val out = Flow().parse(Source.clean(source.trimIndent())) + val out = runBlocking { Flow().parse(Source.clean(source.trimIndent())) } assertTrue(out is FlowOut.Ok, "expected a parsed graph but was $out") return (out as FlowOut.Ok).graph } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt index ec188c5af0e..2865d94da69 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SeqParseTest.kt @@ -1,6 +1,7 @@ package ai.kilocode.client.ui.diagram.mermaid import ai.kilocode.client.ui.diagram.Head +import kotlinx.coroutines.runBlocking import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertTrue @@ -15,6 +16,29 @@ class SeqParseTest { assertEquals(listOf("Server"), script.actors.getValue("S").label) } + @Test + fun `the as separator is case insensitive`() { + val script = script("sequenceDiagram\n PARTICIPANT C AS Client\n C->>S: hi") + + assertEquals(listOf("C", "S"), script.actors.keys.toList()) + assertEquals(listOf("Client"), script.actors.getValue("C").label) + } + + @Test + fun `quoted participant ids match unquoted usages`() { + val script = script("sequenceDiagram\n participant \"Alice\"\n Alice->>Bob: hi") + + assertEquals(listOf("Alice", "Bob"), script.actors.keys.toList()) + assertEquals(listOf("Alice"), script.actors.getValue("Alice").label) + } + + @Test + fun `as inside a quoted name is not an alias`() { + val script = script("sequenceDiagram\n participant \"Bob as builder\"\n") + + assertEquals(listOf("Bob as builder"), script.actors.keys.toList()) + } + @Test fun `undeclared participants appear in first use order`() { val script = script("sequenceDiagram\n B->>A: first\n A->>C: second") @@ -125,15 +149,15 @@ class SeqParseTest { @Test fun `unbalanced blocks are reported with line numbers`() { - val open = Seq().parse(Source.clean("sequenceDiagram\n loop forever\n A->>B: x")) - val stray = Seq().parse(Source.clean("sequenceDiagram\n A->>B: x\n end")) + val open = runBlocking { Seq().parse(Source.clean("sequenceDiagram\n loop forever\n A->>B: x")) } + val stray = runBlocking { Seq().parse(Source.clean("sequenceDiagram\n A->>B: x\n end")) } assertEquals(3, (open as SeqOut.Err).line) assertEquals(3, (stray as SeqOut.Err).line) } private fun script(source: String): Script { - val out = Seq().parse(Source.clean(source.trimIndent())) + val out = runBlocking { Seq().parse(Source.clean(source.trimIndent())) } assertTrue(out is SeqOut.Ok, "expected a parsed script but was $out") return (out as SeqOut.Ok).script } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt index 340c822b180..8ca6d266d72 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt @@ -55,6 +55,23 @@ class SourceTest { assertEquals(5, clean.lines[4].at) } + @Test + fun `an unterminated directive leaves the rest of the source intact`() { + val clean = Source.clean("%%{init: {'theme':'dark'}\ngraph TD\n A --> B") + + assertEquals(3, clean.lines.size) + assertEquals("graph TD", clean.lines[1].text) + assertEquals(2, clean.lines[1].at) + } + + /** A lazy `%%\{[\s\S]*?}%%` regex rescans to the end of the text per opener; this would stall. */ + @Test + fun `many unterminated directives do not stall preprocessing`() { + val clean = Source.clean("%%{".repeat(50_000) + "\ngraph TD\n A --> B") + + assertEquals("graph TD", clean.lines[1].text) + } + @Test fun `labels split on break forms and drop quotes`() { assertEquals(listOf("one", "two"), Source.label("\"one
two\"")) From c0773eb77d7268f6ac6447f24fd31b763a53122d Mon Sep 17 00:00:00 2001 From: kirillk Date: Wed, 26 Aug 2026 14:29:58 -0400 Subject: [PATCH 06/11] fix(jetbrains): detect corrupt IDE extractions --- packages/kilo-jetbrains/AGENTS.md | 1 + packages/kilo-jetbrains/build.gradle.kts | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/packages/kilo-jetbrains/AGENTS.md b/packages/kilo-jetbrains/AGENTS.md index c570086d3d0..7f3ee9c665c 100644 --- a/packages/kilo-jetbrains/AGENTS.md +++ b/packages/kilo-jetbrains/AGENTS.md @@ -235,6 +235,7 @@ For the full release process (resolve version, pin verification, prepare, change - **Via Turbo**: `bun turbo build --filter=@kilocode/kilo-jetbrains` from repo root. - **Run split mode**: `./gradlew --no-configuration-cache runIdeSplitMode` or the checked-in `Run IDE (Split Mode)` configuration — launches backend and frontend locally. Emulate latency via the Split Mode widget (requires internal mode: `-Didea.is.internal=true`). - **Run split backend**: `./gradlew --no-configuration-cache runIdeBackend` — if it exits shortly after startup, check for an orphaned Java process from a previous backend run and kill it before restarting. +- **Corrupt IDE extraction**: if `runIdeBackend` or `runIdeSplitMode` fails before startup with `coroutinesJavaAgentFile` / `Collection contains no element matching the predicate`, the extracted IDE under `.intellijPlatform/ides/` is likely incomplete. Health check: `ls .intellijPlatform/ides/*/lib/*.jar | wc -l` should be in the hundreds. Repair by removing `.intellijPlatform/ides`, `.intellijPlatform/localPlatformArtifacts`, `.intellijPlatform/layoutIndex`, and `.intellijPlatform/coroutines-javaagent.jar`, then rerun the Gradle task. - **Run in monolithic sandbox**: `./gradlew runIde` — launches sandboxed IntelliJ with the plugin. Does not build or bundle CLI binaries; the backend downloads the pinned release at connect time. ### CLI/SDK Change Awareness diff --git a/packages/kilo-jetbrains/build.gradle.kts b/packages/kilo-jetbrains/build.gradle.kts index 20bb5da9c8f..21fdb5af949 100644 --- a/packages/kilo-jetbrains/build.gradle.kts +++ b/packages/kilo-jetbrains/build.gradle.kts @@ -102,6 +102,29 @@ val worktreeRoot = providers.gradleProperty("kilo.dev.worktree.root").orElse( providers.provider { rootProject.layout.projectDirectory.asFile.parentFile.parentFile.canonicalPath } ) +val ides = file(".intellijPlatform/ides") +val corrupt = ides.listFiles() + ?.filter { ide -> + ide.isDirectory && ( + ide.walkTopDown().none { it.name == "product-info.json" } || + ide.resolve("lib").listFiles()?.any { jar -> jar.isFile && jar.extension == "jar" } != true + ) + } + .orEmpty() + +if (corrupt.isNotEmpty()) { + val paths = corrupt.joinToString("\n") { ide -> "- ${ide.absolutePath}" } + error( + """ + Incomplete IntelliJ Platform extraction detected: + $paths + + Remove .intellijPlatform/ides, .intellijPlatform/localPlatformArtifacts, .intellijPlatform/layoutIndex, + and .intellijPlatform/coroutines-javaagent.jar, then rerun the Gradle task. + """.trimIndent(), + ) +} + version = ver plugins { From dedd439befc5fa00364e768eda8dbec0c5b1e241 Mon Sep 17 00:00:00 2001 From: kirillk Date: Thu, 27 Aug 2026 09:38:58 -0400 Subject: [PATCH 07/11] feat(jetbrains): open diagrams in editor tabs --- .changeset/mermaid-diagrams-jetbrains.md | 2 +- .../KiloFrontendDynamicPluginListener.kt | 2 + .../session/ui/selection/SessionCopyTarget.kt | 2 + .../ui/selection/SessionHoverCopyOverlay.kt | 3 +- .../settings/agents/SkillsConfigurable.kt | 4 +- .../settings/agents/WorkflowsConfigurable.kt | 4 +- .../settings/base/SettingsContentEditor.kt | 51 +---- .../client/settings/rules/RulesSettingsUi.kt | 5 +- .../ai/kilocode/client/ui/CodeViewField.kt | 48 +++++ .../client/ui/diagram/ui/DiagramBlock.kt | 6 +- .../client/ui/diagram/ui/DiagramEditorKind.kt | 186 ++++++++++++++++++ .../client/ui/diagram/ui/DiagramPanel.kt | 14 +- .../client/ui/diagram/ui/DiagramSource.kt | 25 --- .../client/ui/diagram/ui/DiagramTheme.kt | 24 +++ .../client/ui/md/hybrid/MdLanguage.kt | 2 +- .../client/ui/md/hybrid/MdViewHybrid.kt | 62 ++---- .../ai/kilocode/client/vfs/KiloEditorKind.kt | 8 +- .../ai/kilocode/client/vfs/KiloFileEditor.kt | 8 +- .../client/vfs/KiloFileEditorProvider.kt | 42 ++-- .../client/vfs/KiloSourceEditorProvider.kt | 37 ++++ .../client/vfs/KiloVirtualFileKind.kt | 5 + .../resources/kilo.jetbrains.frontend.xml | 3 + .../resources/messages/KiloBundle.properties | 4 +- .../session/ui/SessionSelectionCopyTest.kt | 29 +++ .../settings/rules/RulesSettingsUiTest.kt | 2 +- .../ui/diagram/ui/DiagramEditorKindTest.kt | 134 +++++++++++++ .../client/ui/diagram/ui/DiagramPanelTest.kt | 14 ++ .../client/ui/md/MdViewDiagramTest.kt | 33 ++-- 28 files changed, 577 insertions(+), 182 deletions(-) create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/CodeViewField.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt delete mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramTheme.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloSourceEditorProvider.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt diff --git a/.changeset/mermaid-diagrams-jetbrains.md b/.changeset/mermaid-diagrams-jetbrains.md index 82e47dfe679..85a78c021d3 100644 --- a/.changeset/mermaid-diagrams-jetbrains.md +++ b/.changeset/mermaid-diagrams-jetbrains.md @@ -2,4 +2,4 @@ "@kilocode/kilo-jetbrains": minor --- -Render Mermaid code fences as inline diagrams in JetBrains chat markdown, with a source toggle and fallback errors. +Render Mermaid code fences as inline diagrams in JetBrains chat markdown, and open any diagram in its own editor tab with Diagram and Source views. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt index 1afa5e3f521..ec2b13326ee 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt @@ -3,6 +3,7 @@ package ai.kilocode.client.plugin import ai.kilocode.KiloPlugin import ai.kilocode.client.agentManager.worktree.unregisterWorktreeSessionEditorKind import ai.kilocode.client.session.ui.attachment.unregisterAttachmentEditorKind +import ai.kilocode.client.ui.diagram.ui.unregisterDiagramEditorKind import ai.kilocode.client.vfs.KiloEditorKindRegistry import ai.kilocode.client.vfs.KiloVirtualFileSystem import ai.kilocode.log.KiloLog @@ -39,6 +40,7 @@ object KiloFrontendUnloadCleanup { } unregisterAttachmentEditorKind() unregisterWorktreeSessionEditorKind() + unregisterDiagramEditorKind() service().clear() KiloVirtualFileSystem.getInstance().clear() } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyTarget.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyTarget.kt index 650fc442de5..eeaba8401aa 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyTarget.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyTarget.kt @@ -12,6 +12,8 @@ internal interface SessionCopyTarget { val copyToolbar: JComponent? get() = null + val copyCorner: Boolean get() = false + @RequiresEdt fun copyText(): String? } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionHoverCopyOverlay.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionHoverCopyOverlay.kt index e25b033939a..c69c6d56123 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionHoverCopyOverlay.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionHoverCopyOverlay.kt @@ -60,11 +60,12 @@ internal class SessionHoverCopyOverlay( val gap = JBUI.scale(4) val limit = limit(pane) if (limit.isEmpty) return Rectangle() - if (item.copyToolbar != null) { + if (item.copyToolbar != null && !item.copyCorner) { val pt = SwingUtilities.convertPoint(anchor, Point(visible.x, visible.y), pane) // A zero-height anchor is an inline header placeholder (edit/modified open-diff): center // the floating button on the header row so it lines up with the change badge. A real-height // anchor is a footer row (message/text copy): keep the button bottom-aligned inside it. + // Targets that opt into corner placement fall through to the code-block positioning below. val inline = anchor.preferredSize.height == 0 val offset = if (inline) (visible.height - size.height) / 2 else visible.height - size.height val x = clamp(pt.x + visible.width - size.width, limit.x, limit.x + limit.width - size.width) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/SkillsConfigurable.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/SkillsConfigurable.kt index 7a2694bb259..d2423a4109d 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/SkillsConfigurable.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/SkillsConfigurable.kt @@ -5,7 +5,6 @@ import ai.kilocode.client.app.KiloAgentBehaviorService import ai.kilocode.client.app.KiloAppService import ai.kilocode.client.app.KiloWorkspaceService import ai.kilocode.client.plugin.KiloBundle -import ai.kilocode.client.settings.base.SettingsContentField import ai.kilocode.client.settings.base.SettingsDraftPage import ai.kilocode.client.settings.base.SettingsDraftState import ai.kilocode.client.settings.base.SettingsListPanel @@ -15,6 +14,7 @@ import ai.kilocode.client.settings.base.SettingsPathDialogHandle import ai.kilocode.client.settings.base.settingsChoosePath import ai.kilocode.client.settings.base.settingsContentScroll import ai.kilocode.client.settings.base.settingsEditorFileType +import ai.kilocode.client.ui.CodeViewField import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.list.ActiveListBadge @@ -320,7 +320,7 @@ private fun saved(base: SkillsDraft, draft: SkillsDraft): Boolean = base == draf internal class SkillEditDialog(private val skill: SkillDto, private val savable: Boolean) : DialogWrapper(true), SkillEditDialogHandle { private val base = initial() - private val editor = SettingsContentField(base, skillFileType(skill.location, base), savable) + private val editor = CodeViewField(base, skillFileType(skill.location, base), savable) init { title = skill.name diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/WorkflowsConfigurable.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/WorkflowsConfigurable.kt index 38ca2fa250e..429f7a54513 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/WorkflowsConfigurable.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/agents/WorkflowsConfigurable.kt @@ -4,13 +4,13 @@ import ai.kilocode.client.KiloNotifications import ai.kilocode.client.app.KiloAgentBehaviorService import ai.kilocode.client.app.KiloWorkspaceService import ai.kilocode.client.plugin.KiloBundle -import ai.kilocode.client.settings.base.SettingsContentField import ai.kilocode.client.settings.base.SettingsDraftPage import ai.kilocode.client.settings.base.SettingsDraftState import ai.kilocode.client.settings.base.SettingsListPanel import ai.kilocode.client.settings.base.SettingsMessageException import ai.kilocode.client.settings.base.settingsContentScroll import ai.kilocode.client.settings.base.settingsEditorFileType +import ai.kilocode.client.ui.CodeViewField import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.list.ActiveListBadge import ai.kilocode.client.ui.list.ActiveListCell @@ -283,7 +283,7 @@ private fun saved(base: WorkflowsDraft, draft: WorkflowsDraft): Boolean = base = internal class WorkflowEditDialog(private val flow: CommandFileDto, private val savable: Boolean) : DialogWrapper(true), WorkflowEditDialogHandle { private val base = initial() - private val editor = SettingsContentField(base, workflowFileType(flow.location, base), savable) + private val editor = CodeViewField(base, workflowFileType(flow.location, base), savable) init { title = "/${flow.name}" diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/SettingsContentEditor.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/SettingsContentEditor.kt index 2ce0f6a0f32..c842eb99787 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/SettingsContentEditor.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/SettingsContentEditor.kt @@ -1,59 +1,14 @@ package ai.kilocode.client.settings.base -import ai.kilocode.client.session.ui.style.SessionUiStyle -import com.intellij.openapi.editor.EditorFactory +import ai.kilocode.client.ui.CodeViewField +import ai.kilocode.client.ui.codeViewScroll import com.intellij.openapi.fileTypes.FileType import com.intellij.openapi.fileTypes.FileTypeManager import com.intellij.openapi.fileTypes.PlainTextFileType import com.intellij.openapi.fileTypes.UnknownFileType -import com.intellij.openapi.project.ProjectManager -import com.intellij.ui.EditorTextField -import com.intellij.ui.components.JBScrollPane import com.intellij.util.ui.JBUI -import javax.swing.ScrollPaneConstants -/** - * Shared code-editor primitives for settings dialogs (skill content, instruction files). - * - * Keeps the tuned [EditorTextField] configuration, scroll chrome, and content-aware file-type - * detection in one place so pages don't each hand-roll their own editor. - */ -internal class SettingsContentField( - content: String, - fileType: FileType, - editable: Boolean, -) : EditorTextField( - EditorFactory.getInstance().createDocument(content), - ProjectManager.getInstance().defaultProject, - fileType, - !editable, - false, -) { - init { - border = JBUI.Borders.empty() - setOneLineMode(false) - addSettingsProvider { ed -> - ed.setBorder(JBUI.Borders.empty()) - ed.scrollPane.border = JBUI.Borders.empty() - ed.scrollPane.viewportBorder = JBUI.Borders.empty() - ed.settings.isUseSoftWraps = true - ed.settings.isPaintSoftWraps = false - ed.settings.isAdditionalPageAtBottom = false - ed.scrollPane.horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER - ed.scrollPane.verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED - } - } -} - -internal fun settingsContentScroll(field: SettingsContentField) = JBScrollPane(field).apply { - viewportBorder = JBUI.Borders.empty( - JBUI.scale(SessionUiStyle.View.Prompt.SHELL_VERTICAL_PADDING), - JBUI.scale(SessionUiStyle.View.Prompt.SHELL_HORIZONTAL_PADDING), - JBUI.scale(SessionUiStyle.View.Prompt.SHELL_VERTICAL_PADDING), - JBUI.scale(SessionUiStyle.View.Prompt.SHELL_HORIZONTAL_PADDING), - ) - horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER - verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED +internal fun settingsContentScroll(field: CodeViewField) = codeViewScroll(field).apply { preferredSize = JBUI.size(720, 520) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUi.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUi.kt index 71483f226b6..fc1a7132c28 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUi.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUi.kt @@ -5,7 +5,6 @@ import ai.kilocode.client.app.KiloAgentBehaviorService import ai.kilocode.client.app.KiloAppService import ai.kilocode.client.app.KiloWorkspaceService import ai.kilocode.client.plugin.KiloBundle -import ai.kilocode.client.settings.base.SettingsContentField import ai.kilocode.client.settings.base.SettingsDraftPage import ai.kilocode.client.settings.base.SettingsDraftState import ai.kilocode.client.settings.base.SettingsListPanel @@ -16,6 +15,7 @@ import ai.kilocode.client.settings.base.SettingsToolbarAction import ai.kilocode.client.settings.base.settingsChoosePath import ai.kilocode.client.settings.base.settingsContentScroll import ai.kilocode.client.settings.base.settingsEditorFileType +import ai.kilocode.client.ui.CodeViewField import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.layout.StackAxis @@ -294,7 +294,7 @@ internal class InstructionEditDialog( content: String, ) : DialogWrapper(true), RuleContentDialogHandle { private val base = content - private val field = SettingsContentField(base, settingsEditorFileType(heading, base), true) + private val field = CodeViewField(base, settingsEditorFileType(heading, base), true) init { title = heading @@ -363,4 +363,3 @@ private fun writeInstruction(root: String?, path: String, text: String): Boolean } return ok } - diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/CodeViewField.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/CodeViewField.kt new file mode 100644 index 00000000000..bece24e8af3 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/CodeViewField.kt @@ -0,0 +1,48 @@ +package ai.kilocode.client.ui + +import ai.kilocode.client.session.ui.style.SessionUiStyle +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.fileTypes.FileType +import com.intellij.openapi.project.ProjectManager +import com.intellij.ui.EditorTextField +import com.intellij.ui.components.JBScrollPane +import com.intellij.util.ui.JBUI +import javax.swing.ScrollPaneConstants + +internal class CodeViewField( + content: String, + fileType: FileType, + editable: Boolean, +) : EditorTextField( + EditorFactory.getInstance().createDocument(content), + ProjectManager.getInstance().defaultProject, + fileType, + !editable, + false, +) { + init { + border = JBUI.Borders.empty() + setOneLineMode(false) + addSettingsProvider { ed -> + ed.setBorder(JBUI.Borders.empty()) + ed.scrollPane.border = JBUI.Borders.empty() + ed.scrollPane.viewportBorder = JBUI.Borders.empty() + ed.settings.isUseSoftWraps = true + ed.settings.isPaintSoftWraps = false + ed.settings.isAdditionalPageAtBottom = false + ed.scrollPane.horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER + ed.scrollPane.verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED + } + } +} + +internal fun codeViewScroll(field: CodeViewField) = JBScrollPane(field).apply { + viewportBorder = JBUI.Borders.empty( + JBUI.scale(SessionUiStyle.View.Prompt.SHELL_VERTICAL_PADDING), + JBUI.scale(SessionUiStyle.View.Prompt.SHELL_HORIZONTAL_PADDING), + JBUI.scale(SessionUiStyle.View.Prompt.SHELL_VERTICAL_PADDING), + JBUI.scale(SessionUiStyle.View.Prompt.SHELL_HORIZONTAL_PADDING), + ) + horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER + verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt index bb770883ba5..eed83dabb5f 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt @@ -3,11 +3,11 @@ package ai.kilocode.client.ui.diagram.ui import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.views.MessageToolbar +import ai.kilocode.client.session.views.SessionViewIcons import ai.kilocode.client.ui.ToolbarButtonAction import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.layout.StackAxis -import com.intellij.icons.AllIcons import com.intellij.util.concurrency.annotations.RequiresEdt import javax.swing.JComponent @@ -25,7 +25,7 @@ internal class DiagramBlock : Stack(StackAxis.VERTICAL, UiStyle.Gap.sm()), Sessi private val bar = MessageToolbar( text = { text() }, actions = listOf( - ToolbarButtonAction(AllIcons.Actions.EditSource, KiloBundle.message("diagram.open")) { + ToolbarButtonAction(SessionViewIcons.openDiff, KiloBundle.message("diagram.open")) { openDiagram(this, text()) }, ), @@ -35,6 +35,8 @@ internal class DiagramBlock : Stack(StackAxis.VERTICAL, UiStyle.Gap.sm()), Sessi override val copyToolbar: JComponent get() = bar + override val copyCorner: Boolean get() = true + @RequiresEdt override fun copyText() = text() } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt new file mode 100644 index 00000000000..a2d802fb064 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt @@ -0,0 +1,186 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.ui.style.SessionEditorStyle +import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.ui.CodeViewField +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.codeViewScroll +import ai.kilocode.client.ui.diagram.Out +import ai.kilocode.client.ui.md.hybrid.MdLanguage +import ai.kilocode.client.vfs.KiloEditorKind +import ai.kilocode.client.vfs.KiloEditorKindRegistry +import ai.kilocode.client.vfs.KiloEditorView +import ai.kilocode.client.vfs.KiloVfsManager +import ai.kilocode.client.vfs.KiloVirtualFile +import com.intellij.ide.DataManager +import com.intellij.openapi.Disposable +import com.intellij.openapi.actionSystem.CommonDataKeys +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.editor.colors.EditorColorsListener +import com.intellij.openapi.editor.colors.EditorColorsManager +import com.intellij.openapi.fileTypes.FileType +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.Disposer +import com.intellij.ui.components.JBLabel +import com.intellij.ui.components.JBScrollPane +import com.intellij.util.concurrency.annotations.RequiresEdt +import com.intellij.util.ui.Centerizer +import com.intellij.util.ui.JBUI +import java.awt.BorderLayout +import java.security.MessageDigest +import java.util.Collections +import javax.swing.JComponent +import javax.swing.JPanel + +private const val TOKEN = "token" + +/** + * Hands the mermaid source of an "Open in Editor" click to the editor that opens for it. + * + * Application level because [KiloEditorKind.isValid] has no project and the fence text is not + * project scoped. Bounded by a small access-ordered LRU so the sources of long-closed tabs are not + * retained for the IDE session's lifetime. The token is a content hash, so reopening the same + * diagram resolves to the same virtual file and therefore reuses its tab. + */ +@Service(Service.Level.APP) +internal class DiagramStore { + private val items = Collections.synchronizedMap( + object : LinkedHashMap(16, 0.75f, true) { + override fun removeEldestEntry(eldest: Map.Entry): Boolean = size > MAX + }, + ) + + fun put(source: String): String = token(source).also { items[it] = source } + + fun get(token: String): String? = items[token] + + private companion object { + const val MAX = 32 + } +} + +private fun token(source: String): String { + val bytes = MessageDigest.getInstance("SHA-256").digest(source.toByteArray(Charsets.UTF_8)) + return bytes.take(16).joinToString("") { "%02x".format(it) } +} + +private fun source(params: Map): String? { + val token = params[TOKEN]?.takeIf { it.isNotBlank() } ?: return null + return service().get(token) +} + +private fun mmd(): FileType = MdLanguage.type("mmd") + +/** + * Diagram editor tab: the rendered diagram plus a read-only view of its mermaid source. + * + * The source view is a second [KiloEditorView], so the platform composes both into one tab with a + * bottom tab strip (see `EditorComposite`) instead of us hand-rolling a toggle. + */ +internal object DiagramEditorKind : KiloEditorKind { + const val ID = "diagram" + + override val id: String = ID + + override fun title(params: Map): String = KiloBundle.message("diagram.title") + + // No fileType override on purpose: the tab file must stay binary (see KiloVirtualFileKind.fileType). + // The mermaid type is only used for highlighting the source view below. + + override fun presentablePath(params: Map): String = + KiloBundle.message("diagram.path", params[TOKEN].orEmpty()) + + override fun isValid(params: Map): Boolean = source(params) != null + + override val source: KiloEditorView get() = DiagramSourceView + + @RequiresEdt + override fun createContent(project: Project, file: KiloVirtualFile, parent: Disposable): JComponent { + val text = source(file.path.params) ?: return center(KiloBundle.message("diagram.missing")) + val root = JPanel(BorderLayout()) + val label = JBLabel().apply { + border = JBUI.Borders.empty(UiStyle.Gap.sm(), UiStyle.Gap.pad()) + isVisible = false + } + val panel = DiagramPanel(diagramPalette(SessionEditorStyle.current()), fit = true) + root.add(panel, BorderLayout.CENTER) + root.add(label, BorderLayout.SOUTH) + + fun render() { + val style = SessionEditorStyle.current() + panel.background = SessionUiStyle.Colors.codeBlockBackground() + panel.palette(diagramPalette(style)) + label.text = KiloBundle.message("diagram.rendering") + label.foreground = SessionUiStyle.Text.Secondary.foreground() + label.isVisible = true + service().render(text, diagramSpec(style), parent) { out -> + when (out) { + is Out.Ok -> { + panel.art(out.art) + label.isVisible = false + } + + is Out.Err -> { + label.text = KiloBundle.message("diagram.error", out.message) + label.foreground = UiStyle.Colors.errorLabelForeground() + label.isVisible = true + } + } + root.revalidate() + root.repaint() + } + } + + render() + ApplicationManager.getApplication().messageBus.connect(parent) + .subscribe(EditorColorsManager.TOPIC, EditorColorsListener { ApplicationManager.getApplication().invokeLater(::render) }) + return root + } +} + +private object DiagramSourceView : KiloEditorView { + override fun title(params: Map): String = KiloBundle.message("diagram.source") + + @RequiresEdt + override fun createContent(project: Project, file: KiloVirtualFile, parent: Disposable): JComponent { + val text = source(file.path.params) ?: return center(KiloBundle.message("diagram.missing")) + val field = CodeViewField(text, mmd(), editable = false) + Disposer.register(parent) { + field.editor?.let(EditorFactory.getInstance()::releaseEditor) + } + return codeViewScroll(field) + } + + // Stateless on purpose: this view is a singleton shared by every open diagram tab. + override fun preferredFocus(component: JComponent): JComponent? = + (component as? JBScrollPane)?.viewport?.view as? JComponent +} + +private fun center(text: String): JComponent = Centerizer(JBLabel(text)) + +fun ensureDiagramEditorKind() { + service().register(DiagramEditorKind) +} + +internal fun unregisterDiagramEditorKind() { + service().unregister(DiagramEditorKind.ID) +} + +/** + * Opens the mermaid source of a rendered diagram as an editor tab. + * + * Routed through the Kilo virtual file system, the same way session attachments open, so the tab is + * identified by content and reuses the shared editor-kind plumbing. + */ +@RequiresEdt +internal fun openDiagram(anchor: JComponent, source: String): Boolean { + val ctx = DataManager.getInstance().getDataContext(anchor) + val project = CommonDataKeys.PROJECT.getData(ctx) ?: return false + ensureDiagramEditorKind() + val token = service().put(source) + return project.service().open(DiagramEditorKind.ID, mapOf(TOKEN to token)) +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index b3ebb337f97..a5f73c47d78 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -14,7 +14,7 @@ import java.awt.RenderingHints import javax.swing.JComponent import kotlin.math.roundToInt -internal class DiagramPanel(private var palette: Palette) : JComponent() { +internal class DiagramPanel(private var palette: Palette, private val fit: Boolean = false) : JComponent() { private var art: Art? = null private var last = Dimension(0, 0) @@ -31,15 +31,19 @@ internal class DiagramPanel(private var palette: Palette) : JComponent() { repaint() } - override fun getPreferredSize() = fitSize() + override fun getPreferredSize() = if (fit) Dimension(0, 0) else fitSize() - override fun getMinimumSize() = fitSize() + override fun getMinimumSize() = if (fit) Dimension(0, 0) else fitSize() - override fun getMaximumSize() = Dimension(Int.MAX_VALUE, fitSize().height) + override fun getMaximumSize() = if (fit) Dimension(Int.MAX_VALUE, Int.MAX_VALUE) else Dimension(Int.MAX_VALUE, fitSize().height) override fun setBounds(x: Int, y: Int, width: Int, height: Int) { val before = fitSize() super.setBounds(x, y, width, height) + if (fit) { + repaint() + return + } if (before.height != fitSize().height) resize() } @@ -86,7 +90,7 @@ internal class DiagramPanel(private var palette: Palette) : JComponent() { val size = Painters.of(value).size(value) val avail = (width.takeIf { it > 0 } ?: parent?.width ?: 0) - pad() * 2 val byWidth = if (avail > 0) minOf(1.0, avail / size.w) else 1.0 - val max = JBUI.scale(SessionUiStyle.View.Diagram.MAX_HEIGHT) - pad() * 2 + val max = if (fit) height - pad() * 2 else JBUI.scale(SessionUiStyle.View.Diagram.MAX_HEIGHT) - pad() * 2 val byHeight = if (size.h > 0.0) minOf(1.0, max / size.h) else 1.0 return minOf(byWidth, byHeight).coerceAtLeast(0.1) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt deleted file mode 100644 index 703e80e7a6b..00000000000 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramSource.kt +++ /dev/null @@ -1,25 +0,0 @@ -package ai.kilocode.client.ui.diagram.ui - -import com.intellij.ide.DataManager -import com.intellij.ide.scratch.ScratchRootType -import com.intellij.openapi.actionSystem.CommonDataKeys -import com.intellij.openapi.fileEditor.FileEditorManager -import com.intellij.util.concurrency.annotations.RequiresEdt -import javax.swing.JComponent - -private const val NAME = "diagram.mmd" - -/** - * Opens the mermaid source of a rendered diagram in a real editor tab. - * - * A scratch file rather than an in-memory light file, so the text is editable, savable, and picked up - * by whatever mermaid tooling the IDE has installed for the `.mmd` file type. - */ -@RequiresEdt -internal fun openDiagram(anchor: JComponent, source: String): Boolean { - val ctx = DataManager.getInstance().getDataContext(anchor) - val project = CommonDataKeys.PROJECT.getData(ctx) ?: return false - val file = ScratchRootType.getInstance().createScratchFile(project, NAME, null, source) ?: return false - FileEditorManager.getInstance(project).openFile(file, true) - return true -} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramTheme.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramTheme.kt new file mode 100644 index 00000000000..d5ff58b2fa9 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramTheme.kt @@ -0,0 +1,24 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.session.ui.style.SessionEditorStyle +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.FontSpec +import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.diagram.Spec +import ai.kilocode.client.ui.md.MdCommon +import ai.kilocode.client.ui.md.MdStyle + +internal fun diagramPalette(style: SessionEditorStyle, opts: MdStyle = MdCommon.defaults(style)) = Palette( + surface = UiStyle.Colors.contrast(opts.preBg, 8), + border = opts.codeBorder, + text = opts.foreground, + muted = opts.quoteFg, + accent = opts.linkColor, + note = opts.quoteBg, + cluster = opts.codeBorder, + line = opts.quoteFg, + font = style.editorFont, + bold = style.boldEditorFont, +) + +internal fun diagramSpec(style: SessionEditorStyle) = Spec(FontSpec(style.editorFamily, style.editorSize)) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt index bf5215d9769..d61d1de9e97 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdLanguage.kt @@ -80,7 +80,7 @@ internal object MdLanguage { return Kind.Source(type(single)) } - private fun type(ext: String): FileType { + internal fun type(ext: String): FileType { val type = FileTypeRegistry.getInstance().getFileTypeByExtension(ext) if (type == UnknownFileType.INSTANCE) return PlainTextFileType.INSTANCE return type diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index 76a2776a480..a39a634e966 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -7,13 +7,12 @@ import ai.kilocode.client.session.ui.style.SessionEditorStyle import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.ui.UiStyle -import ai.kilocode.client.ui.diagram.FontSpec import ai.kilocode.client.ui.diagram.Out -import ai.kilocode.client.ui.diagram.Palette -import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.ui.DiagramBlock import ai.kilocode.client.ui.diagram.ui.DiagramPanel import ai.kilocode.client.ui.diagram.ui.Diagrams +import ai.kilocode.client.ui.diagram.ui.diagramPalette +import ai.kilocode.client.ui.diagram.ui.diagramSpec import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.md.MdCodeBlockBorder import ai.kilocode.client.ui.md.MdCodeBlockFactory @@ -35,7 +34,6 @@ import com.intellij.openapi.fileTypes.PlainTextFileType import com.intellij.openapi.project.ProjectManager import com.intellij.openapi.util.Disposer import com.intellij.openapi.util.registry.Registry -import com.intellij.ui.HyperlinkLabel import com.intellij.ui.components.JBTextArea import com.intellij.ui.components.JBHtmlPane import com.intellij.ui.components.JBHtmlPaneConfiguration @@ -516,23 +514,9 @@ internal open class MdViewHybrid( return pane } - private fun palette(opts: MdStyle): Palette { - val font = style.editorFont - return Palette( - surface = UiStyle.Colors.contrast(opts.preBg, 8), - border = opts.codeBorder, - text = opts.foreground, - muted = opts.quoteFg, - accent = opts.linkColor, - note = opts.quoteBg, - cluster = opts.codeBorder, - line = opts.quoteFg, - font = font, - bold = style.boldEditorFont, - ) - } + private fun palette(opts: MdStyle) = diagramPalette(style, opts) - private fun spec() = Spec(FontSpec(style.editorFamily, style.editorSize)) + private fun spec() = diagramSpec(style) private fun styleCodePane(pane: JBScrollPane, opts: MdStyle) { pane.apply { @@ -982,27 +966,18 @@ internal open class MdViewHybrid( private val root = component as DiagramBlock private val codePane = codeBlock(desc.text, Kind.Source(kind.file), disposable) private val panel = DiagramPanel(palette(opts())) - private val toggle = HyperlinkLabel(KiloBundle.message("diagram.diagram")) private val label = JBLabel(KiloBundle.message("diagram.rendering")).apply { foreground = SessionUiStyle.Text.Secondary.foreground() } - private val row = Stack.horizontal(gap = UiStyle.Gap.md()).apply { - next(toggle) - next(label) - } private var hash = 0 private var gen = 0 private var font = spec().font init { - // Children keep a fixed order — the diagram above its source, both above the toggle row — - // because Stack lays children out in insertion order and ignores add() indexes. Switching - // between diagram and source flips visibility instead of re-adding components. panel.background = opts().preBg panel.isVisible = false - root.next(panel).next(codePane).next(row) + root.next(panel).next(codePane).next(label) root.text = { (this.desc as Desc.Code).text } - toggle.addHyperlinkListener { toggle() } kick() } @@ -1047,21 +1022,19 @@ internal open class MdViewHybrid( private fun kick() { val item = desc as Desc.Code if (!Registry.`is`("kilo.diagram.inline.enabled", true)) { - label.text = "" + status("") showSource() return } if (item.open) { showSource() - label.text = KiloBundle.message("diagram.rendering") - label.foreground = SessionUiStyle.Text.Secondary.foreground() + status(KiloBundle.message("diagram.rendering")) return } val code = item.text.hashCode() if (hash == code) return hash = code - label.text = KiloBundle.message("diagram.rendering") - label.foreground = SessionUiStyle.Text.Secondary.foreground() + status(KiloBundle.message("diagram.rendering")) val seq = ++gen service().render(item.text, spec(), disposable) { out -> if (seq != gen) return@render @@ -1075,36 +1048,33 @@ internal open class MdViewHybrid( private fun ok(out: Out.Ok) { panel.art(out.art) showDiagram() - label.text = "" + status("") root.revalidate() root.repaint() } private fun fail(message: String) { val text = message.ifBlank { KiloBundle.message("diagram.rendering") } - label.text = KiloBundle.message("diagram.error", text) - label.foreground = UiStyle.Colors.errorLabelForeground() + status(KiloBundle.message("diagram.error", text), true) showSource() root.revalidate() root.repaint() } - private fun toggle() { - if (panel.isVisible) showSource() else showDiagram() - root.revalidate() - root.repaint() - } - private fun showDiagram() { panel.isVisible = true codePane.isVisible = false - toggle.setHyperlinkText(KiloBundle.message("diagram.source")) } private fun showSource() { panel.isVisible = false codePane.isVisible = true - toggle.setHyperlinkText(KiloBundle.message("diagram.diagram")) + } + + private fun status(text: String, error: Boolean = false) { + label.text = text + label.foreground = if (error) UiStyle.Colors.errorLabelForeground() else SessionUiStyle.Text.Secondary.foreground() + label.isVisible = text.isNotEmpty() } private fun updateCode(text: String) { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKind.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKind.kt index e3276f68529..fa6e954a836 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKind.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKind.kt @@ -5,9 +5,15 @@ import com.intellij.openapi.project.Project import com.intellij.util.concurrency.annotations.RequiresEdt import javax.swing.JComponent -interface KiloEditorKind : KiloVirtualFileKind { +interface KiloEditorView { + fun title(params: Map): String + @RequiresEdt fun createContent(project: Project, file: KiloVirtualFile, parent: Disposable): JComponent fun preferredFocus(component: JComponent): JComponent? = null } + +interface KiloEditorKind : KiloVirtualFileKind, KiloEditorView { + val source: KiloEditorView? get() = null +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditor.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditor.kt index 1b0dc0f7605..ad25d9f86c9 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditor.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditor.kt @@ -9,15 +9,15 @@ class KiloFileEditor( private val project: Project, private val file: VirtualFile, private val kilo: KiloVirtualFile, - private val kind: KiloEditorKind, + private val view: KiloEditorView, ) : KiloFileEditorBase() { - private val ui: JComponent by lazy { kind.createContent(project, kilo, this) } + private val ui: JComponent by lazy { view.createContent(project, kilo, this) } @RequiresEdt override fun getComponent(): JComponent = ui - override fun getPreferredFocusedComponent(): JComponent? = kind.preferredFocus(ui) - override fun getName(): String = kind.title(kilo.path.params) + override fun getPreferredFocusedComponent(): JComponent? = view.preferredFocus(ui) + override fun getName(): String = view.title(kilo.path.params) override fun getFile(): VirtualFile = file override fun isValid(): Boolean = super.isValid() && kilo.isValid diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditorProvider.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditorProvider.kt index 1fd68ea6d72..5bd630c458f 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditorProvider.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloFileEditorProvider.kt @@ -4,6 +4,7 @@ import ai.kilocode.client.agentManager.worktree.ensureWorktreeSessionEditorKind import ai.kilocode.client.diff.ensureDiffEditorKind import ai.kilocode.client.session.subagent.ensureSubagentSessionEditorKind import ai.kilocode.client.session.ui.attachment.ensureAttachmentEditorKind +import ai.kilocode.client.ui.diagram.ui.ensureDiagramEditorKind import com.intellij.openapi.components.service import com.intellij.openapi.fileEditor.FileEditor import com.intellij.openapi.fileEditor.FileEditorPolicy @@ -15,16 +16,14 @@ import com.intellij.openapi.vfs.VirtualFile class KiloFileEditorProvider : FileEditorProvider, DumbAware { override fun accept(project: Project, file: VirtualFile): Boolean { - ensureKinds() - val path = path(file) ?: return false - return service().get(path.kind) != null + return kiloKind(file) != null } override fun acceptRequiresReadAction(): Boolean = false override fun createEditor(project: Project, file: VirtualFile): FileEditor { - ensureKinds() - val path = path(file) ?: error("Invalid Kilo virtual file: ${file.path}") + ensureKiloKinds() + val path = kiloPath(file) ?: error("Invalid Kilo virtual file: ${file.path}") val kilo = file as? KiloVirtualFile ?: KiloVirtualFile(path) val kind = service().get(kilo.path.kind) ?: error("Unknown Kilo editor kind: ${kilo.path.kind}") return KiloFileEditor(project, file, kilo, kind) @@ -39,18 +38,25 @@ class KiloFileEditorProvider : FileEditorProvider, DumbAware { companion object { const val EDITOR_TYPE_ID = "KiloVfsEditor" - - private fun ensureKinds() { - ensureAttachmentEditorKind() - ensureDiffEditorKind() - ensureSubagentSessionEditorKind() - ensureWorktreeSessionEditorKind() - } - - private fun path(file: VirtualFile): KiloPath? { - if (file is KiloVirtualFile) return file.path - if (file.fileSystem.protocol != KiloVirtualFileSystem.PROTOCOL && !file.url.startsWith("${KiloVirtualFileSystem.PROTOCOL}://")) return null - return KiloVirtualFileSystem.decode(file.path) ?: KiloVirtualFileSystem.decode(file.url) - } } } + +internal fun kiloKind(file: VirtualFile): KiloEditorKind? { + ensureKiloKinds() + val path = kiloPath(file) ?: return null + return service().get(path.kind) +} + +internal fun kiloPath(file: VirtualFile): KiloPath? { + if (file is KiloVirtualFile) return file.path + if (file.fileSystem.protocol != KiloVirtualFileSystem.PROTOCOL && !file.url.startsWith("${KiloVirtualFileSystem.PROTOCOL}://")) return null + return KiloVirtualFileSystem.decode(file.path) ?: KiloVirtualFileSystem.decode(file.url) +} + +private fun ensureKiloKinds() { + ensureAttachmentEditorKind() + ensureDiffEditorKind() + ensureSubagentSessionEditorKind() + ensureWorktreeSessionEditorKind() + ensureDiagramEditorKind() +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloSourceEditorProvider.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloSourceEditorProvider.kt new file mode 100644 index 00000000000..d4d7e028c75 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloSourceEditorProvider.kt @@ -0,0 +1,37 @@ +package ai.kilocode.client.vfs + +import com.intellij.openapi.components.service +import com.intellij.openapi.fileEditor.FileEditor +import com.intellij.openapi.fileEditor.FileEditorPolicy +import com.intellij.openapi.fileEditor.FileEditorProvider +import com.intellij.openapi.project.DumbAware +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.Disposer +import com.intellij.openapi.vfs.VirtualFile + +class KiloSourceEditorProvider : FileEditorProvider, DumbAware { + override fun accept(project: Project, file: VirtualFile): Boolean { + return kiloKind(file)?.source != null + } + + override fun acceptRequiresReadAction(): Boolean = false + + override fun createEditor(project: Project, file: VirtualFile): FileEditor { + val path = kiloPath(file) ?: error("Invalid Kilo virtual file: ${file.path}") + val kilo = file as? KiloVirtualFile ?: KiloVirtualFile(path) + val kind = service().get(kilo.path.kind) ?: error("Unknown Kilo editor kind: ${kilo.path.kind}") + val view = kind.source ?: error("Kilo editor kind has no source view: ${kilo.path.kind}") + return KiloFileEditor(project, file, kilo, view) + } + + override fun disposeEditor(editor: FileEditor) { + Disposer.dispose(editor) + } + + override fun getEditorTypeId(): String = EDITOR_TYPE_ID + override fun getPolicy(): FileEditorPolicy = FileEditorPolicy.HIDE_OTHER_EDITORS + + companion object { + const val EDITOR_TYPE_ID = "KiloVfsSourceEditor" + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKind.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKind.kt index e86f47fd891..6e2d31f1ceb 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKind.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKind.kt @@ -11,6 +11,11 @@ interface KiloVirtualFileKind { fun icon(params: Map): Icon? = null + /** + * Must be a binary [FileType]. Kilo virtual files carry no content, so a text file type makes + * `FileDocumentManagerBase.getDocument` load text while the editor composite is built and + * [KiloVirtualFile.contentsToByteArray] throws, which cancels the tab. + */ fun fileType(params: Map): FileType = FileTypes.UNKNOWN fun presentablePath(params: Map): String = title(params) diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml index dbf8a7908f3..39435984bc0 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml +++ b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml @@ -24,6 +24,9 @@ factoryClass="ai.kilocode.client.KiloToolWindowFactory"/> + diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index afb5d05235a..98f86d7773f 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -6,9 +6,11 @@ common.rename.help=Use a custom name that describes your task. common.save=Save common.dont.show.again=Don''t show again diagram.source=Source -diagram.diagram=Diagram +diagram.title=Diagram diagram.error=Couldn''t render diagram: {0} +diagram.missing=Diagram source is no longer available. diagram.open=Open in Editor +diagram.path=Kilo / Diagrams / {0} diagram.rendering=Rendering diagram... session.action.cancel=Cancel diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt index 45acd5c1f56..844e30f00ba 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionSelectionCopyTest.kt @@ -21,6 +21,7 @@ import com.intellij.openapi.util.Disposer import com.intellij.openapi.ide.CopyPasteManager import com.intellij.ui.EditorTextField import com.intellij.ui.components.JBScrollPane +import com.intellij.util.ui.JBUI import java.awt.Component import java.awt.Container import java.awt.Cursor @@ -312,6 +313,32 @@ class SessionSelectionCopyTest : SessionUiTestBase() { } } + fun `test hover toolbar opting into corner placement matches the code block copy button`() { + val root = ShowingPanel().also { it.setBounds(0, 0, 300, 300) } + val area = ShowingPanel().also { it.setBounds(0, 0, 300, 300) } + val toolbar = JPanel().also { it.preferredSize = Dimension(48, 24) } + root.add(area) + + val corner = InlineTarget(fixed(24, 40), toolbar, corner = true) + corner.copyAnchor.setBounds(100, 50, 120, 40) + area.add(corner.copyAnchor) + val parent = Disposer.newDisposable("overlay-corner") + val overlay = SessionHoverCopyOverlay(root, area, parent) + root.add(overlay) + + try { + show(overlay, corner) + val bounds = overlay.bounds(root, toolbar) + val gap = JBUI.scale(4) + + // Same formula as a plain code block copy button: inset from the anchor's top-right corner. + assertEquals(100 + 120 - 48 - gap, bounds.x) + assertEquals(50 + gap, bounds.y) + } finally { + Disposer.dispose(parent) + } + } + fun `test session context menu can reinstall after parent disposal`() { val root = JPanel(null) val one = Disposer.newDisposable("context-one") @@ -496,9 +523,11 @@ class SessionSelectionCopyTest : SessionUiTestBase() { private class InlineTarget( private val anchor: JComponent, private val toolbar: JComponent, + private val corner: Boolean = false, ) : JPanel(), SessionCopyTarget { override val copyAnchor: JComponent get() = anchor override val copyToolbar: JComponent get() = toolbar + override val copyCorner: Boolean get() = corner override fun copyText(): String? = null } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUiTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUiTest.kt index b72995256ed..65f006aed77 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUiTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/rules/RulesSettingsUiTest.kt @@ -270,7 +270,7 @@ class RulesSettingsUiTest : BasePlatformTestCase() { fun `test content scroll renders an editor field`() { edt { - val field = ai.kilocode.client.settings.base.SettingsContentField( + val field = ai.kilocode.client.ui.CodeViewField( "# Rules", ai.kilocode.client.settings.base.settingsEditorFileType("./RULES.md", "# Rules"), true, diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt new file mode 100644 index 00000000000..bbd0fde01e8 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt @@ -0,0 +1,134 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.ui.attachment.AttachmentEditorKind +import ai.kilocode.client.session.ui.attachment.ensureAttachmentEditorKind +import ai.kilocode.client.ui.CodeViewField +import ai.kilocode.client.util.edtWait +import ai.kilocode.client.vfs.KiloEditorKindRegistry +import ai.kilocode.client.vfs.KiloFileEditorProvider +import ai.kilocode.client.vfs.KiloPath +import ai.kilocode.client.vfs.KiloSourceEditorProvider +import ai.kilocode.client.vfs.KiloVfsManager +import ai.kilocode.client.vfs.KiloVirtualFile +import ai.kilocode.client.vfs.KiloVirtualFileSystem +import com.intellij.openapi.application.runReadActionBlocking +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.fileEditor.FileDocumentManager +import com.intellij.openapi.fileEditor.FileEditorManager +import com.intellij.openapi.fileEditor.FileEditorPolicy +import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import java.awt.Component +import java.awt.Container + +/** + * The bottom Diagram/Source tab strip itself cannot be asserted here: unit tests run against + * `TestEditorManagerImpl`, which builds a single editor per file and honours + * `FileEditorProvider.KEY`. These tests cover the pieces the platform composes instead — which + * providers accept the file, and what each provider builds. + */ +class DiagramEditorKindTest : BasePlatformTestCase() { + private val flow = "flowchart TD\nA-->B\n" + + override fun setUp() { + super.setUp() + ensureDiagramEditorKind() + } + + override fun tearDown() { + try { + unregisterDiagramEditorKind() + KiloVirtualFileSystem.getInstance().clear() + } finally { + super.tearDown() + } + } + + fun `test opening the same source reuses one tab and a different source opens another`() { + val vfs = project.service() + val manager = FileEditorManager.getInstance(project) + + edtWait { assertTrue(vfs.open(DiagramEditorKind.ID, params(flow))) } + edtWait { assertTrue(vfs.open(DiagramEditorKind.ID, params(flow))) } + + val files = manager.openFiles.filterIsInstance() + assertEquals(1, files.size) + assertEquals(DiagramEditorKind.ID, files.single().path.kind) + + edtWait { assertTrue(vfs.open(DiagramEditorKind.ID, params("flowchart TD\nA-->C\n"))) } + + assertEquals(2, manager.openFiles.filterIsInstance().size) + } + + fun `test the platform builds no document for a diagram file`() { + // A text file type would make FileDocumentManager load text while EditorComposite is built, + // and KiloVirtualFile has no content to give. + val diagram = file(params(flow)) + + assertNull(runReadActionBlocking { FileDocumentManager.getInstance().getDocument(diagram) }) + assertTrue(diagram.fileType.isBinary) + } + + fun `test kind is invalid once the source is unknown`() { + assertTrue(DiagramEditorKind.isValid(params(flow))) + assertFalse(DiagramEditorKind.isValid(mapOf("token" to "deadbeef"))) + assertFalse(DiagramEditorKind.isValid(emptyMap())) + } + + fun `test both providers accept a diagram file and only the diagram kind offers a source view`() { + ensureAttachmentEditorKind() + val diagram = file(params(flow)) + val attachment = KiloVirtualFile(KiloPath(AttachmentEditorKind.ID, mapOf("partId" to "prt1"))) + + assertTrue(KiloFileEditorProvider().accept(project, diagram)) + assertTrue(KiloSourceEditorProvider().accept(project, diagram)) + assertTrue(KiloFileEditorProvider().accept(project, attachment)) + assertFalse(KiloSourceEditorProvider().accept(project, attachment)) + } + + fun `test both providers hide other editors so the platform keeps them side by side`() { + // FileEditorProviderManagerImpl drops every provider whose policy is not HIDE_OTHER_EDITORS + // as soon as one provider requests it, so both kilo providers must agree. + assertEquals(FileEditorPolicy.HIDE_OTHER_EDITORS, KiloFileEditorProvider().policy) + assertEquals(FileEditorPolicy.HIDE_OTHER_EDITORS, KiloSourceEditorProvider().policy) + } + + fun `test diagram and source views build named editors and release their editor`() { + val base = EditorFactory.getInstance().allEditors.size + val diagram = file(params(flow)) + + edtWait { + val main = KiloFileEditorProvider().createEditor(project, diagram) + val source = KiloSourceEditorProvider().createEditor(project, diagram) + try { + assertEquals(KiloBundle.message("diagram.title"), main.name) + assertEquals(KiloBundle.message("diagram.source"), source.name) + assertNotNull(main.component) + + val field = descendants(source.component).filterIsInstance().single() + assertEquals(flow.trim(), field.text.trim()) + assertTrue(field.isViewer) + } finally { + Disposer.dispose(main) + Disposer.dispose(source) + } + } + + assertEquals(base, EditorFactory.getInstance().allEditors.size) + } + + private fun params(source: String) = mapOf("token" to service().put(source)) + + private fun file(params: Map) = KiloVirtualFile(KiloPath(DiagramEditorKind.ID, params)) + + private fun descendants(root: Container): List { + val out = mutableListOf() + for (comp in root.components) { + out.add(comp) + if (comp is Container) out.addAll(descendants(comp)) + } + return out + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt index c1549d023bd..13f856b6c4f 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -1,6 +1,7 @@ package ai.kilocode.client.ui.diagram.ui import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.views.SessionViewIcons import ai.kilocode.client.ui.diagram.Mark import ai.kilocode.client.ui.diagram.Palette import ai.kilocode.client.ui.diagram.Rect @@ -31,6 +32,17 @@ class DiagramPanelTest { assertTrue(panel.preferredSize.height <= 520) } + @Test + fun `test fit mode scales to the component bounds instead of the transcript cap`() { + val panel = DiagramPanel(palette(), fit = true) + panel.setSize(1_000, 1_000) + panel.art(scene(100.0, 2_000.0)) + + // The transcript cap (480) no longer applies; the panel fills whatever the tab gives it. + assertEquals(0, panel.preferredSize.height) + assertTrue(panel.maximumSize.height > 520) + } + @Test fun `test block copies fence text and offers copy plus open in editor`() { val block = DiagramBlock() @@ -40,8 +52,10 @@ class DiagramPanelTest { assertEquals("flowchart TD", block.copyText()) assertEquals(2, buttons.size) + assertTrue(block.copyCorner) assertTrue(buttons.any { it.toolTipText == KiloBundle.message("diagram.open") }) assertTrue(buttons.any { it.toolTipText == KiloBundle.message("session.copy.hover") }) + assertTrue(buttons.any { it.icon === SessionViewIcons.openDiff }) } private fun buttons(root: java.awt.Container): List { diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt index a72acc3b94f..0fe4900f9d5 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -20,7 +20,6 @@ import com.intellij.openapi.editor.EditorFactory import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.testFramework.replaceService -import com.intellij.ui.HyperlinkLabel import com.intellij.util.ui.UIUtil import java.awt.Point import javax.swing.JPanel @@ -48,7 +47,7 @@ class MdViewDiagramTest : BasePlatformTestCase() { } } - fun `test mermaid fence renders above toggle row and hides source`() { + fun `test mermaid fence renders above status label and hides source`() { view.set("```mermaid\nflowchart TD\nA-->B\n```") drain() @@ -56,9 +55,17 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertEquals(1, engine.calls) assertSame(diagram(), children.first()) - assertSame(row(), children.last()) + assertSame(label(), children.last()) assertTrue(diagram().isVisible) assertFalse(codePane().isVisible) + assertFalse("a rendered diagram leaves no status row behind", label().isVisible) + } + + fun `test block anchors its toolbar in the corner like a code block`() { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + + assertTrue(block().copyCorner) } fun `test block is the hover target and copies the fence text`() { @@ -74,21 +81,6 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertSame(block().copyToolbar, (target as SessionCopyTarget).copyToolbar) } - fun `test toggle switches between diagram and source`() { - view.set("```mermaid\nflowchart TD\nA-->B\n```") - drain() - - toggle().doClick() - - assertFalse(diagram().isVisible) - assertTrue(codePane().isVisible) - - toggle().doClick() - - assertTrue(diagram().isVisible) - assertFalse(codePane().isVisible) - } - fun `test engine error keeps source visible`() { engine.out = Out.Err(ai.kilocode.client.ui.diagram.Fault.Syntax, "bad syntax") @@ -97,6 +89,7 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertFalse(diagram().isVisible) assertTrue(codePane().isVisible) + assertTrue(label().isVisible) assertTrue(labels().contains("bad syntax")) } @@ -149,9 +142,7 @@ class MdViewDiagramTest : BasePlatformTestCase() { private fun codePane() = block().components[1] - private fun row() = block().components.last() - - private fun toggle() = descendants(root()).filterIsInstance().single() + private fun label() = block().components.last() as javax.swing.JLabel private fun labels() = descendants(root()).joinToString("\n") { (it as? javax.swing.JLabel)?.text.orEmpty() } From 54b7225c4d30775cd5918aa77e410fc55889221f Mon Sep 17 00:00:00 2001 From: kirillk Date: Thu, 27 Aug 2026 13:30:31 -0400 Subject: [PATCH 08/11] fix(jetbrains): repair diagram viewer controls --- .changeset/diagram-viewer-window-jetbrains.md | 5 + .../KiloFrontendDynamicPluginListener.kt | 2 + .../client/ui/diagram/ui/DiagramCanvas.kt | 172 ++++++++++ .../client/ui/diagram/ui/DiagramContent.kt | 60 ++++ .../client/ui/diagram/ui/DiagramEditorKind.kt | 49 +-- .../client/ui/diagram/ui/DiagramPanel.kt | 24 +- .../client/ui/diagram/ui/DiagramViewer.kt | 159 ++++++++++ .../client/ui/diagram/ui/DiagramWindow.kt | 145 +++++++++ .../client/ui/md/hybrid/MdViewHybrid.kt | 14 + .../resources/messages/KiloBundle.properties | 4 + .../ui/diagram/ui/DiagramEditorKindTest.kt | 3 +- .../client/ui/diagram/ui/DiagramPanelTest.kt | 11 - .../client/ui/diagram/ui/DiagramViewerTest.kt | 294 ++++++++++++++++++ .../client/ui/diagram/ui/DiagramWindowTest.kt | 111 +++++++ .../client/ui/md/MdViewDiagramTest.kt | 83 +++++ 15 files changed, 1058 insertions(+), 78 deletions(-) create mode 100644 .changeset/diagram-viewer-window-jetbrains.md create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindow.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindowTest.kt diff --git a/.changeset/diagram-viewer-window-jetbrains.md b/.changeset/diagram-viewer-window-jetbrains.md new file mode 100644 index 00000000000..aa13a9437b0 --- /dev/null +++ b/.changeset/diagram-viewer-window-jetbrains.md @@ -0,0 +1,5 @@ +--- +"@kilocode/kilo-jetbrains": patch +--- + +Click a diagram in chat to open it in a resizable viewer window with zoom controls, trackpad pinch zoom, drag to pan and scrollbars. The diagram editor tab uses the same viewer. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt index ec2b13326ee..3972766cffd 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt @@ -3,6 +3,7 @@ package ai.kilocode.client.plugin import ai.kilocode.KiloPlugin import ai.kilocode.client.agentManager.worktree.unregisterWorktreeSessionEditorKind import ai.kilocode.client.session.ui.attachment.unregisterAttachmentEditorKind +import ai.kilocode.client.ui.diagram.ui.DiagramWindows import ai.kilocode.client.ui.diagram.ui.unregisterDiagramEditorKind import ai.kilocode.client.vfs.KiloEditorKindRegistry import ai.kilocode.client.vfs.KiloVirtualFileSystem @@ -30,6 +31,7 @@ object KiloFrontendUnloadCleanup { runEdt { ProjectManager.getInstance().openProjects.forEach { project -> if (project.isDisposed) return@forEach + project.getServiceIfCreated(DiagramWindows::class.java)?.closeAll() ToolWindowManager.getInstance(project).getToolWindow("Kilo Code") ?.contentManager ?.removeAllContents(true) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt new file mode 100644 index 00000000000..1f60018df50 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt @@ -0,0 +1,172 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.ui.diagram.Art +import ai.kilocode.client.ui.diagram.Painters +import ai.kilocode.client.ui.diagram.Palette +import com.intellij.ui.components.Magnificator +import com.intellij.util.concurrency.annotations.RequiresEdt +import com.intellij.util.ui.JBUI +import java.awt.Dimension +import java.awt.Graphics +import java.awt.Graphics2D +import java.awt.Point +import java.awt.Rectangle +import java.awt.RenderingHints +import javax.swing.JComponent +import javax.swing.JViewport +import javax.swing.Scrollable +import kotlin.math.roundToInt + +/** + * Scrollable diagram surface for the diagram viewer. + * + * Two states: fit (no explicit factor) tracks the viewport on both axes so the whole diagram is + * visible without scrollbars, while an explicit factor reports the scaled art as its preferred size + * so the enclosing scroll pane can scroll it. The fit scale is derived from the **viewport** extent, + * never from this component's own bounds, so sizing cannot feed back into itself. + */ +internal class DiagramCanvas(private var palette: Palette) : JComponent(), Scrollable { + private var art: Art? = null + private var factor: Double? = null + + init { + // Trackpad pinch: JBViewport reads this off its view and drives it through ZoomingDelegate, + // which does the scrolling itself from the returned point, so no anchoring here. + putClientProperty( + Magnificator.CLIENT_PROPERTY_KEY, + Magnificator { scale, at -> + zoom(this.scale() * scale) + Point((at.x * scale).roundToInt(), (at.y * scale).roundToInt()) + }, + ) + } + + @RequiresEdt + fun art(value: Art) { + art = value + revalidate() + repaint() + } + + @RequiresEdt + fun palette(value: Palette) { + palette = value + repaint() + } + + /** + * Sets an explicit scale, or restores fit when [value] is null. + * + * [at] is a point in **viewport** coordinates that should stay put across the zoom. + */ + @RequiresEdt + fun zoom(value: Double?, at: Point? = null) { + val before = scale() + factor = value?.coerceIn(MIN, maxOf(MAX, fitScale() * FIT_ZOOM)) + // Size the view up front so the viewport clamps the anchored position against the new bounds. + if (factor != null) size = preferredSize + revalidate() + repaint() + if (at != null) anchor(at, before, scale()) + } + + @RequiresEdt + fun fit() { + zoom(null) + } + + @RequiresEdt + fun scale(): Double = factor ?: fitScale() + + override fun getPreferredSize(): Dimension { + if (factor == null) return Dimension(0, 0) + val value = art ?: return Dimension(0, 0) + val size = Painters.of(value).size(value) + val scale = scale() + return Dimension( + (size.w * scale).roundToInt() + pad() * 2, + (size.h * scale).roundToInt() + pad() * 2, + ) + } + + override fun paintComponent(g: Graphics) { + background?.let { + g.color = it + g.fillRect(0, 0, width, height) + } + val value = art ?: return + val size = Painters.of(value).size(value) + val scale = scale() + val x = ((width - size.w * scale) / 2).roundToInt().coerceAtLeast(pad()) + val y = ((height - size.h * scale) / 2).roundToInt().coerceAtLeast(pad()) + paintDiagram(g, value, palette, scale, x, y) + } + + override fun getPreferredScrollableViewportSize(): Dimension = preferredSize + + override fun getScrollableUnitIncrement(visibleRect: Rectangle, orientation: Int, direction: Int) = step() + + override fun getScrollableBlockIncrement(visibleRect: Rectangle, orientation: Int, direction: Int) = step() + + override fun getScrollableTracksViewportWidth(): Boolean = tracks { extent -> preferredSize.width <= extent.width } + + override fun getScrollableTracksViewportHeight(): Boolean = tracks { extent -> preferredSize.height <= extent.height } + + private fun tracks(fits: (Dimension) -> Boolean): Boolean { + if (factor == null) return true + val viewport = parent as? JViewport ?: return false + return fits(viewport.extentSize) + } + + private fun anchor(at: Point, before: Double, after: Double) { + if (before <= 0.0) return + val viewport = parent as? JViewport ?: return + val ratio = after / before + val pos = viewport.viewPosition + val x = ((pos.x + at.x) * ratio - at.x).roundToInt() + val y = ((pos.y + at.y) * ratio - at.y).roundToInt() + viewport.viewPosition = clamped(viewport, Point(x, y)) + } + + private fun fitScale(): Double { + val value = art ?: return 1.0 + val size = Painters.of(value).size(value) + if (size.w <= 0.0 || size.h <= 0.0) return 1.0 + val extent = (parent as? JViewport)?.extentSize ?: Dimension(width, height) + val w = (extent.width - pad() * 2).coerceAtLeast(1) + val h = (extent.height - pad() * 2).coerceAtLeast(1) + return minOf(w / size.w, h / size.h).coerceAtLeast(MIN) + } + + private fun pad() = JBUI.scale(SessionUiStyle.View.Diagram.PADDING) + + private fun step() = JBUI.scale(SessionUiStyle.SessionLayout.SCROLL_INCREMENT) + + private companion object { + const val MIN = 0.1 + const val MAX = 4.0 + const val FIT_ZOOM = 4.0 + } +} + +/** Keeps a viewport position inside the scrollable range of its view. */ +internal fun clamped(viewport: JViewport, at: Point): Point { + val view = viewport.view ?: return at + val x = (view.width - viewport.extentSize.width).coerceAtLeast(0) + val y = (view.height - viewport.extentSize.height).coerceAtLeast(0) + return Point(at.x.coerceIn(0, x), at.y.coerceIn(0, y)) +} + +/** Paints [art] scaled by [scale] with its top-left corner at ([x], [y]). */ +internal fun paintDiagram(g: Graphics, art: Art, palette: Palette, scale: Double, x: Int, y: Int) { + val g2 = g.create() as Graphics2D + try { + g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON) + g2.translate(x, y) + g2.scale(scale, scale) + Painters.of(art).paint(g2, art, palette) + } finally { + g2.dispose() + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt new file mode 100644 index 00000000000..91d2bc99ea9 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt @@ -0,0 +1,60 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.ui.style.SessionEditorStyle +import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.Out +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.colors.EditorColorsListener +import com.intellij.openapi.editor.colors.EditorColorsManager +import com.intellij.ui.components.JBLabel +import com.intellij.util.concurrency.annotations.RequiresEdt +import com.intellij.util.ui.JBUI +import java.awt.BorderLayout +import javax.swing.JComponent +import javax.swing.JPanel + +@RequiresEdt +internal fun diagramContent(source: String, parent: Disposable): JComponent { + val root = JPanel(BorderLayout()) + val label = JBLabel().apply { + border = JBUI.Borders.empty(UiStyle.Gap.sm(), UiStyle.Gap.pad()) + isVisible = false + } + val viewer = DiagramViewer(diagramPalette(SessionEditorStyle.current())) + root.add(viewer, BorderLayout.CENTER) + root.add(label, BorderLayout.SOUTH) + + fun render() { + val style = SessionEditorStyle.current() + viewer.surface(SessionUiStyle.Colors.codeBlockBackground()) + viewer.palette(diagramPalette(style)) + label.text = KiloBundle.message("diagram.rendering") + label.foreground = SessionUiStyle.Text.Secondary.foreground() + label.isVisible = true + service().render(source, diagramSpec(style), parent) { out -> + when (out) { + is Out.Ok -> { + viewer.art(out.art) + label.isVisible = false + } + + is Out.Err -> { + label.text = KiloBundle.message("diagram.error", out.message) + label.foreground = UiStyle.Colors.errorLabelForeground() + label.isVisible = true + } + } + root.revalidate() + root.repaint() + } + } + + render() + ApplicationManager.getApplication().messageBus.connect(parent) + .subscribe(EditorColorsManager.TOPIC, EditorColorsListener { ApplicationManager.getApplication().invokeLater(::render) }) + return root +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt index a2d802fb064..b885202aa3c 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKind.kt @@ -1,12 +1,8 @@ package ai.kilocode.client.ui.diagram.ui import ai.kilocode.client.plugin.KiloBundle -import ai.kilocode.client.session.ui.style.SessionEditorStyle -import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.ui.CodeViewField -import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.codeViewScroll -import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.md.hybrid.MdLanguage import ai.kilocode.client.vfs.KiloEditorKind import ai.kilocode.client.vfs.KiloEditorKindRegistry @@ -16,12 +12,9 @@ import ai.kilocode.client.vfs.KiloVirtualFile import com.intellij.ide.DataManager import com.intellij.openapi.Disposable import com.intellij.openapi.actionSystem.CommonDataKeys -import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.editor.EditorFactory -import com.intellij.openapi.editor.colors.EditorColorsListener -import com.intellij.openapi.editor.colors.EditorColorsManager import com.intellij.openapi.fileTypes.FileType import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer @@ -29,12 +22,9 @@ import com.intellij.ui.components.JBLabel import com.intellij.ui.components.JBScrollPane import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.Centerizer -import com.intellij.util.ui.JBUI -import java.awt.BorderLayout import java.security.MessageDigest import java.util.Collections import javax.swing.JComponent -import javax.swing.JPanel private const val TOKEN = "token" @@ -101,44 +91,7 @@ internal object DiagramEditorKind : KiloEditorKind { @RequiresEdt override fun createContent(project: Project, file: KiloVirtualFile, parent: Disposable): JComponent { val text = source(file.path.params) ?: return center(KiloBundle.message("diagram.missing")) - val root = JPanel(BorderLayout()) - val label = JBLabel().apply { - border = JBUI.Borders.empty(UiStyle.Gap.sm(), UiStyle.Gap.pad()) - isVisible = false - } - val panel = DiagramPanel(diagramPalette(SessionEditorStyle.current()), fit = true) - root.add(panel, BorderLayout.CENTER) - root.add(label, BorderLayout.SOUTH) - - fun render() { - val style = SessionEditorStyle.current() - panel.background = SessionUiStyle.Colors.codeBlockBackground() - panel.palette(diagramPalette(style)) - label.text = KiloBundle.message("diagram.rendering") - label.foreground = SessionUiStyle.Text.Secondary.foreground() - label.isVisible = true - service().render(text, diagramSpec(style), parent) { out -> - when (out) { - is Out.Ok -> { - panel.art(out.art) - label.isVisible = false - } - - is Out.Err -> { - label.text = KiloBundle.message("diagram.error", out.message) - label.foreground = UiStyle.Colors.errorLabelForeground() - label.isVisible = true - } - } - root.revalidate() - root.repaint() - } - } - - render() - ApplicationManager.getApplication().messageBus.connect(parent) - .subscribe(EditorColorsManager.TOPIC, EditorColorsListener { ApplicationManager.getApplication().invokeLater(::render) }) - return root + return diagramContent(text, parent) } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index a5f73c47d78..f8f6baeacca 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -14,7 +14,7 @@ import java.awt.RenderingHints import javax.swing.JComponent import kotlin.math.roundToInt -internal class DiagramPanel(private var palette: Palette, private val fit: Boolean = false) : JComponent() { +internal class DiagramPanel(private var palette: Palette) : JComponent() { private var art: Art? = null private var last = Dimension(0, 0) @@ -31,19 +31,15 @@ internal class DiagramPanel(private var palette: Palette, private val fit: Boole repaint() } - override fun getPreferredSize() = if (fit) Dimension(0, 0) else fitSize() + override fun getPreferredSize() = fitSize() - override fun getMinimumSize() = if (fit) Dimension(0, 0) else fitSize() + override fun getMinimumSize() = fitSize() - override fun getMaximumSize() = if (fit) Dimension(Int.MAX_VALUE, Int.MAX_VALUE) else Dimension(Int.MAX_VALUE, fitSize().height) + override fun getMaximumSize() = Dimension(Int.MAX_VALUE, fitSize().height) override fun setBounds(x: Int, y: Int, width: Int, height: Int) { val before = fitSize() super.setBounds(x, y, width, height) - if (fit) { - repaint() - return - } if (before.height != fitSize().height) resize() } @@ -58,17 +54,9 @@ internal class DiagramPanel(private var palette: Palette, private val fit: Boole g2.dispose() } val value = art ?: return - val painter = Painters.of(value) val scale = scale(value) SessionSurface.clipped(g, width, height) { clipped -> - val inner = clipped.create() as Graphics2D - try { - inner.translate(pad(), pad()) - inner.scale(scale, scale) - painter.paint(inner, value, palette) - } finally { - inner.dispose() - } + paintDiagram(clipped, value, palette, scale, pad(), pad()) } } @@ -90,7 +78,7 @@ internal class DiagramPanel(private var palette: Palette, private val fit: Boole val size = Painters.of(value).size(value) val avail = (width.takeIf { it > 0 } ?: parent?.width ?: 0) - pad() * 2 val byWidth = if (avail > 0) minOf(1.0, avail / size.w) else 1.0 - val max = if (fit) height - pad() * 2 else JBUI.scale(SessionUiStyle.View.Diagram.MAX_HEIGHT) - pad() * 2 + val max = JBUI.scale(SessionUiStyle.View.Diagram.MAX_HEIGHT) - pad() * 2 val byHeight = if (size.h > 0.0) minOf(1.0, max / size.h) else 1.0 return minOf(byWidth, byHeight).coerceAtLeast(0.1) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt new file mode 100644 index 00000000000..5b120b45305 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt @@ -0,0 +1,159 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.ui.ToolbarButtonAction +import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.Art +import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.layout.Stack +import ai.kilocode.client.ui.toolbarButton +import com.intellij.icons.AllIcons +import com.intellij.ui.components.JBLayeredPane +import com.intellij.ui.components.JBScrollPane +import com.intellij.util.concurrency.annotations.RequiresEdt +import com.intellij.util.ui.JBUI +import java.awt.Color +import java.awt.Cursor +import java.awt.Point +import java.awt.event.MouseAdapter +import java.awt.event.MouseEvent +import java.awt.event.MouseWheelEvent +import java.awt.event.MouseWheelListener +import javax.swing.Icon +import javax.swing.ScrollPaneConstants +import javax.swing.SwingUtilities + +/** + * Reusable zoomable diagram surface: a scrollable [DiagramCanvas] with floating zoom controls. + * + * Shared by the diagram editor tab and the detached diagram window. Zoom comes from three sources: + * trackpad pinch (via the canvas [com.intellij.ui.components.Magnificator]), Ctrl/Cmd + wheel, and + * the overlay buttons. Dragging pans whenever the scaled diagram overflows the viewport. + */ +internal class DiagramViewer(palette: Palette) : JBLayeredPane() { + private val canvas = DiagramCanvas(palette) + private val scroll = JBScrollPane(canvas).apply { + border = JBUI.Borders.empty() + viewportBorder = JBUI.Borders.empty() + horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_AS_NEEDED + verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED + } + // Built by chaining rather than `apply`, so these lambdas cannot bind to Stack's own fit(). + private val controls = Stack.horizontal(UiStyle.Gap.xs()) + .next(control(AllIcons.General.ZoomIn, "diagram.zoom.in") { zoomIn() }) + .next(control(AllIcons.General.ZoomOut, "diagram.zoom.out") { zoomOut() }) + .next(control(AllIcons.General.FitContent, "diagram.zoom.fit") { fit() }) + private val wheel = Wheel() + private val drag = Drag() + + init { + // Layer first, then add: add(Component, Int) binds to Container.add(comp, index) from Kotlin, + // so the layer would be taken as an insertion index and both children would end up in the + // default layer, with the scroll pane painting over the controls and swallowing their clicks. + setLayer(scroll, DEFAULT_LAYER) + setLayer(controls, PALETTE_LAYER) + add(scroll) + add(controls) + scroll.addMouseWheelListener(wheel) + canvas.addMouseListener(drag) + canvas.addMouseMotionListener(drag) + } + + @RequiresEdt + fun art(value: Art) { + canvas.art(value) + } + + @RequiresEdt + fun palette(value: Palette) { + canvas.palette(value) + } + + /** Paints the diagram surface (canvas and viewport) with [color]. */ + @RequiresEdt + fun surface(color: Color) { + background = color + scroll.background = color + scroll.viewport.background = color + canvas.background = color + } + + @RequiresEdt + fun zoomIn(at: Point? = null) { + canvas.zoom(canvas.scale() * STEP, at) + } + + @RequiresEdt + fun zoomOut(at: Point? = null) { + canvas.zoom(canvas.scale() / STEP, at) + } + + @RequiresEdt + fun fit() { + canvas.fit() + } + + override fun doLayout() { + scroll.setBounds(0, 0, width, height) + val size = controls.preferredSize + controls.setBounds(width - size.width - UiStyle.Gap.pad(), UiStyle.Gap.pad(), size.width, size.height) + controls.doLayout() + } + + private inner class Wheel : MouseWheelListener { + override fun mouseWheelMoved(e: MouseWheelEvent) { + if (!e.isControlDown && !e.isMetaDown) return + val at = SwingUtilities.convertPoint(e.component, e.point, scroll.viewport) + if (e.wheelRotation < 0) zoomIn(at) + if (e.wheelRotation > 0) zoomOut(at) + e.consume() + } + } + + private inner class Drag : MouseAdapter() { + private var from: Point? = null + private var origin: Point? = null + + override fun mousePressed(e: MouseEvent) { + if (e.button != MouseEvent.BUTTON1 || !overflows()) return + from = e.point + origin = scroll.viewport.viewPosition + canvas.cursor = Cursor.getPredefinedCursor(Cursor.MOVE_CURSOR) + } + + override fun mouseDragged(e: MouseEvent) { + val start = from ?: return + val base = origin ?: return + val at = Point(base.x + start.x - e.x, base.y + start.y - e.y) + scroll.viewport.viewPosition = clamped(scroll.viewport, at) + } + + override fun mouseReleased(e: MouseEvent) { + release() + } + + override fun mouseExited(e: MouseEvent) { + release() + } + + private fun release() { + if (from == null) return + from = null + origin = null + canvas.cursor = Cursor.getDefaultCursor() + } + + private fun overflows(): Boolean { + val viewport = scroll.viewport + val view = viewport.view ?: return false + return view.width > viewport.extentSize.width || view.height > viewport.extentSize.height + } + } + + private companion object { + const val STEP = 1.25 + + fun control(icon: Icon, key: String, handler: () -> Unit) = + toolbarButton(ToolbarButtonAction(icon, KiloBundle.message(key), handler), fill = true) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindow.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindow.kt new file mode 100644 index 00000000000..34434ab5ef5 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindow.kt @@ -0,0 +1,145 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.telemetry.Telemetry +import com.intellij.ide.DataManager +import com.intellij.openapi.Disposable +import com.intellij.openapi.actionSystem.CommonDataKeys +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.openapi.project.Project +import com.intellij.openapi.ui.FrameWrapper +import com.intellij.openapi.util.Disposer +import com.intellij.openapi.util.WindowState +import com.intellij.openapi.wm.WindowManager +import com.intellij.util.concurrency.annotations.RequiresEdt +import java.awt.Rectangle +import java.util.function.BooleanSupplier +import javax.swing.JComponent +import javax.swing.RootPaneContainer + +private const val DIMENSION_KEY = "ai.kilocode.DiagramViewer" +private const val SHARE = 0.75 + +internal fun diagramWindowBounds(frame: Rectangle): Rectangle { + val w = (frame.width * SHARE).toInt().coerceAtLeast(1) + val h = (frame.height * SHARE).toInt().coerceAtLeast(1) + val x = frame.x + (frame.width - w) / 2 + val y = frame.y + (frame.height - h) / 2 + return Rectangle(x, y, w, h) +} + +internal interface DiagramHandle : Disposable { + fun show() + fun focus() +} + +@Service(Service.Level.PROJECT) +internal class DiagramWindows internal constructor( + project: Project, + private val factory: (String) -> DiagramHandle, + private val send: (String, Map) -> Unit, +) { + constructor(project: Project) : this(project, { source -> FrameHandle(project, source) }, Telemetry::send) + + private val windows = mutableMapOf() + + @RequiresEdt + fun open(source: String): Boolean { + val token = service().put(source) + val handle = windows[token] + if (handle != null) { + handle.focus() + track(true) + return true + } + val next = factory(source) + windows[token] = next + Disposer.register(next) { + if (windows[token] === next) windows.remove(token) + } + next.show() + track(false) + return true + } + + @RequiresEdt + fun closeAll() { + val all = windows.values.toList() + windows.clear() + all.forEach(Disposer::dispose) + } + + private fun track(reused: Boolean) { + send( + "Diagram Viewer Opened", + mapOf( + "surface" to "session", + "reused" to reused.toString(), + ), + ) + } +} + +private class FrameHandle(project: Project, source: String) : DiagramHandle { + private val frame = DiagramFrame(project).apply { + component = diagramContent(source, this) + preferredFocusedComponent = component + closeOnEsc() + setOnCloseHandler(BooleanSupplier { + Disposer.dispose(this@FrameHandle) + false + }) + } + + override fun show() { + frame.show(true) + } + + override fun focus() { + val window = frame.getFrame() + window.toFront() + window.requestFocus() + } + + override fun dispose() { + if (!frame.isDisposed) Disposer.dispose(frame) + } +} + +/** + * A frame rather than a dialog on purpose. + * + * The viewer is a document surface, so it belongs in the Window menu and should live on its own + * instead of floating over the IDE frame. It also keeps the window closer to the editor tab, which + * matters for trackpad zoom: magnification is routed per window by the platform's + * [com.intellij.openapi.actionSystem.impl.MouseGestureManager], and the editor tab (a plain IDE + * frame) is the surface where that routing is known to reach our canvas. + */ +private class DiagramFrame(private val project: Project) : FrameWrapper( + project, + DIMENSION_KEY, + false, + KiloBundle.message("diagram.title"), +) { + override fun loadFrameState(state: WindowState?) { + if (state != null) { + super.loadFrameState(state) + return + } + val base = WindowManager.getInstance().getFrame(project)?.bounds + if (base == null) { + super.loadFrameState(null) + return + } + getFrame().bounds = diagramWindowBounds(base) + (getFrame() as RootPaneContainer).rootPane.revalidate() + } +} + +@RequiresEdt +internal fun openDiagramWindow(anchor: JComponent, source: String): Boolean { + val ctx = DataManager.getInstance().getDataContext(anchor) + val project = CommonDataKeys.PROJECT.getData(ctx) ?: return false + return project.service().open(source) +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index a39a634e966..6b77d23a81c 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -13,6 +13,7 @@ import ai.kilocode.client.ui.diagram.ui.DiagramPanel import ai.kilocode.client.ui.diagram.ui.Diagrams import ai.kilocode.client.ui.diagram.ui.diagramPalette import ai.kilocode.client.ui.diagram.ui.diagramSpec +import ai.kilocode.client.ui.diagram.ui.openDiagramWindow import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.md.MdCodeBlockBorder import ai.kilocode.client.ui.md.MdCodeBlockFactory @@ -43,6 +44,7 @@ import com.intellij.ui.components.JBScrollPane import com.intellij.util.ui.JBUI import java.awt.Color import java.awt.Component +import java.awt.Cursor import java.awt.Dimension import java.awt.Font import java.awt.Graphics @@ -50,6 +52,7 @@ import java.awt.Graphics2D import java.awt.Point import java.awt.RenderingHints import java.awt.event.HierarchyEvent +import java.awt.event.MouseAdapter import java.awt.event.MouseEvent import javax.swing.Box import javax.swing.BoxLayout @@ -972,10 +975,21 @@ internal open class MdViewHybrid( private var hash = 0 private var gen = 0 private var font = spec().font + private val click = object : MouseAdapter() { + override fun mouseClicked(e: MouseEvent) { + if (e.button != MouseEvent.BUTTON1 || e.clickCount != 1) return + if (!panel.isVisible) return + openDiagramWindow(panel, (this@DiagramView.desc as Desc.Code).text) + } + } init { panel.background = opts().preBg panel.isVisible = false + panel.cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) + panel.toolTipText = KiloBundle.message("diagram.viewer.hint") + panel.addMouseListener(click) + Disposer.register(disposable) { panel.removeMouseListener(click) } root.next(panel).next(codePane).next(label) root.text = { (this.desc as Desc.Code).text } kick() diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index 98f86d7773f..eeb37fb8872 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -12,6 +12,10 @@ diagram.missing=Diagram source is no longer available. diagram.open=Open in Editor diagram.path=Kilo / Diagrams / {0} diagram.rendering=Rendering diagram... +diagram.viewer.hint=Open diagram viewer +diagram.zoom.fit=Fit to Window +diagram.zoom.in=Zoom In +diagram.zoom.out=Zoom Out session.action.cancel=Cancel session.connection.connecting=Loading... diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt index bbd0fde01e8..df327b8410b 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramEditorKindTest.kt @@ -105,7 +105,8 @@ class DiagramEditorKindTest : BasePlatformTestCase() { try { assertEquals(KiloBundle.message("diagram.title"), main.name) assertEquals(KiloBundle.message("diagram.source"), source.name) - assertNotNull(main.component) + // The tab hosts the same zoomable viewer the diagram window uses. + assertEquals(1, descendants(main.component).filterIsInstance().size) val field = descendants(source.component).filterIsInstance().single() assertEquals(flow.trim(), field.text.trim()) diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt index 13f856b6c4f..ef4e0974651 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -32,17 +32,6 @@ class DiagramPanelTest { assertTrue(panel.preferredSize.height <= 520) } - @Test - fun `test fit mode scales to the component bounds instead of the transcript cap`() { - val panel = DiagramPanel(palette(), fit = true) - panel.setSize(1_000, 1_000) - panel.art(scene(100.0, 2_000.0)) - - // The transcript cap (480) no longer applies; the panel fills whatever the tab gives it. - assertEquals(0, panel.preferredSize.height) - assertTrue(panel.maximumSize.height > 520) - } - @Test fun `test block copies fence text and offers copy plus open in editor`() { val block = DiagramBlock() diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt new file mode 100644 index 00000000000..bff4a11cdc2 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt @@ -0,0 +1,294 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.ui.diagram.Mark +import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.diagram.Rect +import ai.kilocode.client.ui.diagram.Role +import ai.kilocode.client.ui.diagram.Scene +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.client.ui.diagram.Type +import ai.kilocode.client.util.edtWait +import com.intellij.icons.AllIcons +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import com.intellij.ui.components.JBScrollPane +import com.intellij.ui.components.Magnificator +import java.awt.Color +import java.awt.Container +import java.awt.Font +import java.awt.Point +import java.awt.event.MouseEvent +import java.awt.event.MouseWheelEvent +import javax.swing.AbstractButton +import javax.swing.JComponent +import javax.swing.JViewport +import javax.swing.SwingUtilities + +class DiagramViewerTest : BasePlatformTestCase() { + fun `test fit tracks the viewport and needs no scrolling`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + + assertTrue(scale(viewer) < 1.0) + val canvas = canvas(viewer) + assertTrue(canvas.getScrollableTracksViewportWidth()) + assertTrue(canvas.getScrollableTracksViewportHeight()) + } + + fun `test fit upscales a diagram smaller than the viewport`() = edtWait { + val viewer = viewer(800, 600) + viewer.art(scene(100.0, 100.0)) + layout(viewer) + + assertTrue("fit should fill the window, not cap at native size", scale(viewer) > 1.0) + } + + fun `test zooming in leaves fit and lets the scroll pane scroll`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val fit = scale(viewer) + + viewer.zoomIn() + layout(viewer) + + assertEquals(fit * 1.25, scale(viewer), 1e-6) + val canvas = canvas(viewer) + assertTrue(canvas.preferredSize.width > 0) + assertFalse(canvas.getScrollableTracksViewportWidth()) + } + + fun `test zoom out then fit restores viewport tracking`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val fit = scale(viewer) + + viewer.zoomOut() + assertEquals(fit / 1.25, scale(viewer), 1e-6) + + viewer.fit() + layout(viewer) + + assertEquals(fit, scale(viewer), 1e-6) + assertEquals(0, canvas(viewer).preferredSize.width) + } + + fun `test fit refits after the viewport is resized`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val narrow = scale(viewer) + + viewer.setSize(800, 600) + layout(viewer) + + assertTrue(scale(viewer) > narrow) + } + + fun `test zoom clamps to the supported range`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(400.0, 300.0)) + layout(viewer) + + repeat(30) { viewer.zoomIn() } + assertEquals(4.0, scale(viewer), 1e-6) + + repeat(60) { viewer.zoomOut() } + assertEquals(0.1, scale(viewer), 1e-6) + } + + fun `test canvas exposes a magnificator that scales and reports the anchor`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(400.0, 300.0)) + layout(viewer) + val canvas = canvas(viewer) + val magnificator = canvas.getClientProperty(Magnificator.CLIENT_PROPERTY_KEY) as Magnificator + val before = scale(viewer) + + val at = magnificator.magnify(2.0, Point(30, 40)) + + assertEquals(before * 2.0, scale(viewer), 1e-6) + assertEquals(Point(60, 80), at) + } + + fun `test control wheel zooms and consumes while a plain wheel scrolls`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val fit = scale(viewer) + + val plain = wheel(viewer, control = false, rotation = -1) + viewport(viewer).parent.dispatchEvent(plain) + + assertFalse(plain.isConsumed) + assertEquals(fit, scale(viewer), 1e-6) + + val zoom = wheel(viewer, control = true, rotation = -1) + viewport(viewer).parent.dispatchEvent(zoom) + + assertTrue(zoom.isConsumed) + assertEquals(fit * 1.25, scale(viewer), 1e-6) + } + + fun `test dragging pans the viewport and clamps at the edges`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + repeat(4) { viewer.zoomIn() } + layout(viewer) + val canvas = canvas(viewer) + + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_PRESSED, 200, 150)) + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_DRAGGED, 150, 120)) + + assertEquals(Point(50, 30), viewport(viewer).viewPosition) + + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_DRAGGED, 400, 350)) + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_RELEASED, 400, 350)) + + assertEquals(Point(0, 0), viewport(viewer).viewPosition) + } + + fun `test overlay offers zoom in zoom out and fit`() = edtWait { + val viewer = viewer(400, 300) + + val buttons = buttons(viewer) + + assertEquals(3, buttons.size) + assertEquals( + listOf(AllIcons.General.ZoomIn, AllIcons.General.ZoomOut, AllIcons.General.FitContent), + buttons.map { it.icon }, + ) + assertEquals( + listOf( + KiloBundle.message("diagram.zoom.in"), + KiloBundle.message("diagram.zoom.out"), + KiloBundle.message("diagram.zoom.fit"), + ), + buttons.map { it.toolTipText }, + ) + } + + fun `test overlay floats above the scroll pane and receives its own clicks`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val zoom = buttons(viewer).first() + + val at = SwingUtilities.convertPoint(zoom, zoom.width / 2, zoom.height / 2, viewer) + + assertSame("the scroll pane must not cover the controls", zoom, SwingUtilities.getDeepestComponentAt(viewer, at.x, at.y)) + assertFalse("overlapping layers cannot use optimized drawing", viewer.isOptimizedDrawingEnabled) + } + + fun `test overlay buttons drive the zoom`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val fit = scale(viewer) + val buttons = buttons(viewer) + + buttons[0].doClick() + assertEquals(fit * 1.25, scale(viewer), 1e-6) + + buttons[1].doClick() + assertEquals(fit, scale(viewer), 1e-6) + + buttons[0].doClick() + buttons[2].doClick() + assertEquals(fit, scale(viewer), 1e-6) + assertEquals(0, canvas(viewer).preferredSize.width) + } + + private fun viewer(width: Int, height: Int) = DiagramViewer(palette()).apply { + setSize(width, height) + surface(Color.WHITE) + } + + private fun layout(viewer: DiagramViewer) { + viewer.doLayout() + layout(viewer as Container) + } + + private fun layout(root: Container) { + root.doLayout() + root.components.filterIsInstance().forEach(::layout) + } + + private fun viewport(viewer: DiagramViewer): JViewport = descendants(viewer) + .filterIsInstance() + .single() + .viewport + + private fun canvas(viewer: DiagramViewer) = viewport(viewer).view as DiagramCanvas + + private fun scale(viewer: DiagramViewer) = canvas(viewer).scale() + + private fun wheel(viewer: DiagramViewer, control: Boolean, rotation: Int): MouseWheelEvent { + val scroll = viewport(viewer).parent + return MouseWheelEvent( + scroll, + MouseEvent.MOUSE_WHEEL, + System.currentTimeMillis(), + if (control) MouseEvent.CTRL_DOWN_MASK else 0, + 10, + 10, + 0, + false, + MouseWheelEvent.WHEEL_UNIT_SCROLL, + 1, + rotation, + ) + } + + private fun mouse(target: JComponent, id: Int, x: Int, y: Int) = MouseEvent( + target, + id, + System.currentTimeMillis(), + MouseEvent.BUTTON1_DOWN_MASK, + x, + y, + 1, + false, + MouseEvent.BUTTON1, + ) + + private fun buttons(root: Container): List { + val out = mutableListOf() + for (comp in root.components) { + if (comp is AbstractButton) out.add(comp) + if (comp is Container) out.addAll(buttons(comp)) + } + return out + } + + private fun descendants(root: Container): List { + val out = mutableListOf() + for (comp in root.components) { + out.add(comp) + if (comp is Container) out.addAll(descendants(comp)) + } + return out + } + + private fun scene(w: Double, h: Double) = Scene( + Type.Flowchart, + listOf(Mark.Box(Rect(0.0, 0.0, w, h), 4.0, Role.Surface, Role.Border)), + Size(w, h), + ) + + private fun palette() = Palette( + surface = Color.WHITE, + border = Color.BLACK, + text = Color.BLACK, + muted = Color.GRAY, + accent = Color.BLUE, + note = Color.YELLOW, + cluster = Color.LIGHT_GRAY, + line = Color.DARK_GRAY, + font = Font(Font.SANS_SERIF, Font.PLAIN, 12), + bold = Font(Font.SANS_SERIF, Font.BOLD, 12), + ) +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindowTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindowTest.kt new file mode 100644 index 00000000000..16e168fda31 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramWindowTest.kt @@ -0,0 +1,111 @@ +package ai.kilocode.client.ui.diagram.ui + +import ai.kilocode.client.util.edtWait +import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import java.awt.Rectangle + +/** + * Covers the window bookkeeping without opening a real window: [DiagramWindows] takes its handle + * factory as a dependency, so the reuse, dispose and telemetry paths are exercised against fakes + * while the [com.intellij.openapi.ui.FrameWrapper] wiring stays out of the test. + */ +class DiagramWindowTest : BasePlatformTestCase() { + private val events = mutableListOf>>() + private val handles = mutableListOf() + + fun `test bounds take three quarters of the frame and stay centred`() { + val bounds = diagramWindowBounds(Rectangle(100, 50, 1000, 800)) + + assertEquals(Rectangle(225, 150, 750, 600), bounds) + } + + fun `test bounds survive a degenerate frame`() { + val bounds = diagramWindowBounds(Rectangle(0, 0, 1, 1)) + + assertEquals(Rectangle(0, 0, 1, 1), bounds) + } + + fun `test the same source reuses one window and a different source opens another`() = edtWait { + val windows = windows() + + assertTrue(windows.open("flowchart TD\nA-->B")) + + assertEquals(1, handles.size) + assertEquals(1, handles.single().shown) + assertEquals(0, handles.single().focused) + + assertTrue(windows.open("flowchart TD\nA-->B")) + + assertEquals(1, handles.size) + assertEquals(1, handles.single().shown) + assertEquals(1, handles.single().focused) + + assertTrue(windows.open("flowchart TD\nA-->C")) + + assertEquals(2, handles.size) + assertEquals(listOf(1, 1), handles.map { it.shown }) + } + + fun `test disposing a window drops it so the next click opens a fresh one`() = edtWait { + val windows = windows() + windows.open("flowchart TD\nA-->B") + + Disposer.dispose(handles.single()) + windows.open("flowchart TD\nA-->B") + + assertEquals(2, handles.size) + assertEquals(0, handles.last().focused) + assertEquals(1, handles.last().shown) + } + + fun `test closeAll disposes every open window`() = edtWait { + val windows = windows() + windows.open("flowchart TD\nA-->B") + windows.open("flowchart TD\nA-->C") + + windows.closeAll() + windows.open("flowchart TD\nA-->B") + + assertEquals(listOf(1, 1, 0), handles.map { it.disposed }) + assertEquals(3, handles.size) + } + + fun `test opening reports whether the window was reused`() = edtWait { + val windows = windows() + + windows.open("flowchart TD\nA-->B") + windows.open("flowchart TD\nA-->B") + + assertEquals(listOf("Diagram Viewer Opened", "Diagram Viewer Opened"), events.map { it.first }) + assertEquals(listOf("false", "true"), events.map { it.second["reused"] }) + assertEquals(listOf("session", "session"), events.map { it.second["surface"] }) + } + + private fun windows() = DiagramWindows( + project, + { FakeHandle().also(handles::add) }, + { event, props -> events.add(event to props) }, + ) + + private class FakeHandle : DiagramHandle { + var shown = 0 + private set + var focused = 0 + private set + var disposed = 0 + private set + + override fun show() { + shown++ + } + + override fun focus() { + focused++ + } + + override fun dispose() { + disposed++ + } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt index 0fe4900f9d5..5fb234a0213 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -13,15 +13,24 @@ import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.Type import ai.kilocode.client.ui.diagram.ui.DiagramBlock +import ai.kilocode.client.ui.diagram.ui.DiagramHandle import ai.kilocode.client.ui.diagram.ui.DiagramPanel +import ai.kilocode.client.ui.diagram.ui.DiagramWindows import ai.kilocode.client.ui.diagram.ui.Diagrams +import com.intellij.openapi.actionSystem.CommonDataKeys +import com.intellij.openapi.actionSystem.DataSink +import com.intellij.openapi.actionSystem.UiDataProvider import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.testFramework.replaceService import com.intellij.util.ui.UIUtil +import java.awt.Cursor import java.awt.Point +import java.awt.event.MouseEvent +import javax.swing.JComponent import javax.swing.JPanel @Suppress("UnstableApiUsage") @@ -81,6 +90,34 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertSame(block().copyToolbar, (target as SessionCopyTarget).copyToolbar) } + fun `test clicking a rendered diagram opens the viewer window`() { + val opened = windows() + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + attach() + + click(diagram()) + + assertEquals(listOf("flowchart TD\nA-->B\n"), opened) + assertEquals(Cursor.HAND_CURSOR, diagram().cursor.type) + } + + fun `test the streaming source fallback is not a viewer trigger`() { + // Only the rendered diagram opens the window, so the source pane shown while a fence streams + // (and after an engine error) keeps its plain text behaviour. + val opened = windows() + view.append("```mermaid\nflowchart TD\n") + drain() + attach() + + click(codePane() as JComponent) + click(diagram()) + + assertTrue(codePane().isVisible) + assertFalse(diagram().isVisible) + assertTrue(opened.isEmpty()) + } + fun `test engine error keeps source visible`() { engine.out = Out.Err(ai.kilocode.client.ui.diagram.Fault.Syntax, "bad syntax") @@ -134,6 +171,38 @@ class MdViewDiagramTest : BasePlatformTestCase() { private fun drain() = coroutines.drain() + /** Records the sources the transcript hands to the viewer window instead of opening one. */ + private fun windows(): List { + val opened = mutableListOf() + val service = DiagramWindows(project, { source -> opened.add(source); NoopHandle() }, { _, _ -> }) + project.replaceService(DiagramWindows::class.java, service, testRootDisposable) + return opened + } + + /** Puts the transcript under a project data provider so the click can resolve the project. */ + private fun attach() { + val panel = DataPanel(project) + panel.add(root()) + panel.setSize(400, 400) + panel.doLayout() + } + + private fun click(target: JComponent) { + target.dispatchEvent( + MouseEvent( + target, + MouseEvent.MOUSE_CLICKED, + System.currentTimeMillis(), + 0, + 1, + 1, + 1, + false, + MouseEvent.BUTTON1, + ), + ) + } + private fun root() = view.component as JPanel private fun block() = descendants(root()).filterIsInstance().single() @@ -155,6 +224,20 @@ class MdViewDiagramTest : BasePlatformTestCase() { return out } + private class DataPanel(private val project: Project) : JPanel(), UiDataProvider { + override fun uiDataSnapshot(sink: DataSink) { + sink[CommonDataKeys.PROJECT] = project + } + } + + private class NoopHandle : DiagramHandle { + override fun show() = Unit + + override fun focus() = Unit + + override fun dispose() = Unit + } + private class FakeEngine : Engine { var calls = 0 var out: Out? = null From 0ee530d85a1a1e87ece6cc5dab8354e908631c70 Mon Sep 17 00:00:00 2001 From: kirillk Date: Thu, 27 Aug 2026 14:04:08 -0400 Subject: [PATCH 09/11] fix(jetbrains): refine diagram viewer interactions --- .changeset/diagram-viewer-window-jetbrains.md | 2 +- .../session/ui/selection/SessionCopyButton.kt | 27 +++++- .../client/session/views/MessageToolbar.kt | 8 +- .../kotlin/ai/kilocode/client/ui/Clipboard.kt | 30 +++++++ .../kotlin/ai/kilocode/client/ui/HoverIcon.kt | 18 +++- .../client/ui/diagram/ui/DiagramBlock.kt | 5 ++ .../client/ui/diagram/ui/DiagramCanvas.kt | 38 +++++++++ .../client/ui/diagram/ui/DiagramPanel.kt | 5 ++ .../client/ui/diagram/ui/DiagramViewer.kt | 41 ++++++++-- .../client/ui/md/hybrid/MdViewHybrid.kt | 3 + .../resources/messages/KiloBundle.properties | 1 + .../client/ui/diagram/ui/DiagramViewerTest.kt | 82 +++++++++++++++++-- .../client/ui/md/MdViewDiagramTest.kt | 32 ++++++++ 13 files changed, 267 insertions(+), 25 deletions(-) create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/Clipboard.kt diff --git a/.changeset/diagram-viewer-window-jetbrains.md b/.changeset/diagram-viewer-window-jetbrains.md index aa13a9437b0..1eb11678589 100644 --- a/.changeset/diagram-viewer-window-jetbrains.md +++ b/.changeset/diagram-viewer-window-jetbrains.md @@ -2,4 +2,4 @@ "@kilocode/kilo-jetbrains": patch --- -Click a diagram in chat to open it in a resizable viewer window with zoom controls, trackpad pinch zoom, drag to pan and scrollbars. The diagram editor tab uses the same viewer. +Click a diagram in chat to open it in a resizable viewer window with zoom controls, trackpad pinch zoom, drag to pan, double click to fit and scrollbars. The diagram editor tab uses the same viewer. Copying a rendered diagram, from the viewer or from chat, now puts the picture on the clipboard. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyButton.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyButton.kt index 2b7b8b61cb4..6830c065350 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyButton.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/selection/SessionCopyButton.kt @@ -2,6 +2,7 @@ package ai.kilocode.client.session.ui.selection import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.ui.ToolbarButtonAction +import ai.kilocode.client.ui.copyImage import ai.kilocode.client.ui.toolbarButton import com.intellij.openapi.ide.CopyPasteManager import com.intellij.openapi.util.IconLoader @@ -13,17 +14,20 @@ import java.awt.Point import java.awt.datatransfer.StringSelection import java.awt.event.MouseAdapter import java.awt.event.MouseEvent +import java.awt.image.BufferedImage import javax.swing.Icon internal class SessionCopyButton( fill: Boolean = false, tooltip: String = KiloBundle.message("session.copy.hover"), + icon: Icon = COPY_ICON, + private val image: () -> BufferedImage? = { null }, private val text: () -> String?, ) { private var balloon: Balloon? = null val button = toolbarButton( ToolbarButtonAction( - COPY_ICON, + icon, tooltip, ) { copy() }, fill, @@ -45,8 +49,7 @@ internal class SessionCopyButton( @RequiresEdt fun copy() { - val value = text()?.takeIf { it.isNotEmpty() } ?: return - CopyPasteManager.getInstance().setContents(StringSelection(value)) + if (!put()) return dismiss() balloon = JBPopupFactory.getInstance() .createHtmlTextBalloonBuilder(KiloBundle.message("session.copy.copied"), null, null, null) @@ -57,6 +60,24 @@ internal class SessionCopyButton( } } + /** + * Writes the clipboard and reports whether anything was put there. + * + * A picture wins over text, so a rendered diagram is pasted as an image while everything else (and + * a diagram that is still streaming or failed to render) keeps copying its text. + */ + @RequiresEdt + private fun put(): Boolean { + val picture = image() + if (picture != null) { + copyImage(picture) + return true + } + val value = text()?.takeIf { it.isNotEmpty() } ?: return false + CopyPasteManager.getInstance().setContents(StringSelection(value)) + return true + } + companion object { private val COPY_ICON: Icon = IconLoader.getIcon("/icons/copy.svg", SessionCopyButton::class.java) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageToolbar.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageToolbar.kt index 4ffeee14468..102b5acc9a8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageToolbar.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageToolbar.kt @@ -10,23 +10,25 @@ import ai.kilocode.client.plugin.KiloBundle import com.intellij.util.concurrency.annotations.RequiresEdt import java.awt.Dimension import java.awt.FlowLayout +import java.awt.image.BufferedImage import javax.swing.JComponent import javax.swing.JPanel internal class MessageToolbar( text: () -> String?, + image: () -> BufferedImage? = { null }, actions: List = emptyList(), tooltip: String = KiloBundle.message("session.copy.hover"), ) : JPanel(FlowLayout(FlowLayout.LEFT, 0, 0)) { constructor(text: () -> String?, revert: (() -> Unit)?) : this( text, - revert?.let { + actions = revert?.let { listOf(ToolbarButtonAction(AllIcons.Actions.Rollback, KiloBundle.message("revert.message.rollback"), it)) }.orEmpty(), - KiloBundle.message("session.copy.prompt"), + tooltip = KiloBundle.message("session.copy.prompt"), ) - private val copy = SessionCopyButton(text = text, tooltip = tooltip) + private val copy = SessionCopyButton(text = text, image = image, tooltip = tooltip) private val button = copy.button private val buttons = actions.map(::toolbarButton) private val row = Stack.horizontal(UiStyle.Gap.xs()).apply { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/Clipboard.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/Clipboard.kt new file mode 100644 index 00000000000..3670bb7e634 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/Clipboard.kt @@ -0,0 +1,30 @@ +package ai.kilocode.client.ui + +import com.intellij.openapi.ide.CopyPasteManager +import com.intellij.util.concurrency.annotations.RequiresEdt +import java.awt.datatransfer.DataFlavor +import java.awt.datatransfer.Transferable +import java.awt.datatransfer.UnsupportedFlavorException +import java.awt.image.BufferedImage + +/** + * Puts [image] on the clipboard, so pasting lands a picture rather than text. + * + * Only [DataFlavor.imageFlavor] is offered, matching the platform's own image editor: adding a text + * flavor would let apps that prefer text paste that instead of the picture. + */ +@RequiresEdt +internal fun copyImage(image: BufferedImage) { + CopyPasteManager.getInstance().setContents(Picture(image)) +} + +private class Picture(private val image: BufferedImage) : Transferable { + override fun getTransferDataFlavors(): Array = arrayOf(DataFlavor.imageFlavor) + + override fun isDataFlavorSupported(flavor: DataFlavor): Boolean = DataFlavor.imageFlavor == flavor + + override fun getTransferData(flavor: DataFlavor): Any { + if (DataFlavor.imageFlavor != flavor) throw UnsupportedFlavorException(flavor) + return image + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/HoverIcon.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/HoverIcon.kt index 59c340d7881..d453208bf79 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/HoverIcon.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/HoverIcon.kt @@ -25,10 +25,15 @@ class HoverIcon(private val fill: Boolean = false) : JButton() { }) } - // Icon-only buttons keep a fixed 24x24 hit target; labelled buttons size to their content plus - // their (symmetric) border so the hover pill has equal padding on every side. - override fun getPreferredSize(): Dimension = - if (text.isNullOrEmpty()) JBUI.size(24, 24) else super.getPreferredSize() + // Icon-only buttons are square: a 24x24 hit target for the usual 16px icon, growing with the icon + // so a larger one keeps the same padding inside the hover pill. Labelled buttons size to their + // content plus their (symmetric) border, which already gives equal padding on every side. + override fun getPreferredSize(): Dimension { + if (!text.isNullOrEmpty()) return super.getPreferredSize() + val icon = icon ?: return JBUI.size(MIN, MIN) + val side = maxOf(JBUI.scale(MIN), icon.iconWidth + JBUI.scale(PAD), icon.iconHeight + JBUI.scale(PAD)) + return Dimension(side, side) + } override fun getMinimumSize(): Dimension = preferredSize @@ -66,6 +71,11 @@ class HoverIcon(private val fill: Boolean = false) : JButton() { over = value repaint() } + + private companion object { + const val MIN = 24 + const val PAD = 8 + } } fun iconButton(button: JButton) { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt index eed83dabb5f..9cab3795429 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramBlock.kt @@ -9,6 +9,7 @@ import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.layout.Stack import ai.kilocode.client.ui.layout.StackAxis import com.intellij.util.concurrency.annotations.RequiresEdt +import java.awt.image.BufferedImage import javax.swing.JComponent /** @@ -22,8 +23,12 @@ internal class DiagramBlock : Stack(StackAxis.VERTICAL, UiStyle.Gap.sm()), Sessi /** Source of the fence text; the owning view rebinds it as the diagram streams or updates. */ var text: () -> String = { "" } + /** The rendered diagram, when there is one; copy prefers it over the fence text. */ + var image: () -> BufferedImage? = { null } + private val bar = MessageToolbar( text = { text() }, + image = { image() }, actions = listOf( ToolbarButtonAction(SessionViewIcons.openDiff, KiloBundle.message("diagram.open")) { openDiagram(this, text()) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt index 1f60018df50..83b978273cf 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt @@ -1,18 +1,21 @@ package ai.kilocode.client.ui.diagram.ui import ai.kilocode.client.session.ui.style.SessionUiStyle +import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Painters import ai.kilocode.client.ui.diagram.Palette import com.intellij.ui.components.Magnificator import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI +import java.awt.Color import java.awt.Dimension import java.awt.Graphics import java.awt.Graphics2D import java.awt.Point import java.awt.Rectangle import java.awt.RenderingHints +import java.awt.image.BufferedImage import javax.swing.JComponent import javax.swing.JViewport import javax.swing.Scrollable @@ -55,6 +58,10 @@ internal class DiagramCanvas(private var palette: Palette) : JComponent(), Scrol repaint() } + /** The rendered diagram as an image, or null while nothing has been drawn yet. */ + @RequiresEdt + fun image(): BufferedImage? = art?.let { diagramImage(it, palette, background) } + /** * Sets an explicit scale, or restores fit when [value] is null. * @@ -150,6 +157,37 @@ internal class DiagramCanvas(private var palette: Palette) : JComponent(), Scrol } } +/** + * Renders [art] into an image for the clipboard, padded and filled with [background]. + * + * Painted from the scene rather than grabbed off the component, so the result is the whole diagram at + * a fixed [SHOT] scale regardless of the current zoom, scroll position or viewport size. The + * background is filled because the palette follows the IDE theme: a transparent PNG of a dark theme + * diagram would be unreadable once pasted onto white. + */ +internal fun diagramImage(art: Art, palette: Palette, background: Color?): BufferedImage { + val size = Painters.of(art).size(art) + val pad = SessionUiStyle.View.Diagram.PADDING * SHOT + val w = (size.w * SHOT + pad * 2).roundToInt().coerceAtLeast(1) + val h = (size.h * SHOT + pad * 2).roundToInt().coerceAtLeast(1) + // A plain image on purpose: ImageUtil.createImage would add the IDE's HiDPI scale on top of SHOT, + // making the picture depend on the display it was copied from. + val image = BufferedImage(w, h, BufferedImage.TYPE_INT_RGB) + val g = image.createGraphics() + try { + g.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON) + g.color = background ?: UiStyle.Colors.editorBackground() + g.fillRect(0, 0, w, h) + paintDiagram(g, art, palette, SHOT, pad.roundToInt(), pad.roundToInt()) + } finally { + g.dispose() + } + return image +} + +/** Clipboard images render larger than the screen so they stay crisp when pasted. */ +private const val SHOT = 2.0 + /** Keeps a viewport position inside the scrollable range of its view. */ internal fun clamped(viewport: JViewport, at: Point): Point { val view = viewport.view ?: return at diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index f8f6baeacca..9fcfb550259 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -11,6 +11,7 @@ import java.awt.Dimension import java.awt.Graphics import java.awt.Graphics2D import java.awt.RenderingHints +import java.awt.image.BufferedImage import javax.swing.JComponent import kotlin.math.roundToInt @@ -31,6 +32,10 @@ internal class DiagramPanel(private var palette: Palette) : JComponent() { repaint() } + /** The rendered diagram as an image, or null while nothing has been drawn yet. */ + @RequiresEdt + fun image(): BufferedImage? = art?.let { diagramImage(it, palette, background) } + override fun getPreferredSize() = fitSize() override fun getMinimumSize() = fitSize() diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt index 5b120b45305..535da9bdac8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt @@ -1,6 +1,7 @@ package ai.kilocode.client.ui.diagram.ui import ai.kilocode.client.plugin.KiloBundle +import ai.kilocode.client.session.ui.selection.SessionCopyButton import ai.kilocode.client.ui.ToolbarButtonAction import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.diagram.Art @@ -10,6 +11,7 @@ import ai.kilocode.client.ui.toolbarButton import com.intellij.icons.AllIcons import com.intellij.ui.components.JBLayeredPane import com.intellij.ui.components.JBScrollPane +import com.intellij.util.IconUtil import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import java.awt.Color @@ -26,9 +28,10 @@ import javax.swing.SwingUtilities /** * Reusable zoomable diagram surface: a scrollable [DiagramCanvas] with floating zoom controls. * - * Shared by the diagram editor tab and the detached diagram window. Zoom comes from three sources: - * trackpad pinch (via the canvas [com.intellij.ui.components.Magnificator]), Ctrl/Cmd + wheel, and - * the overlay buttons. Dragging pans whenever the scaled diagram overflows the viewport. + * Shared by the diagram editor tab and the detached diagram window. Zoom comes from four sources: + * trackpad pinch (via the canvas [com.intellij.ui.components.Magnificator]), Ctrl/Cmd + wheel, the + * overlay buttons, and a double click to fit again. Dragging pans whenever the scaled diagram + * overflows the viewport, and the overlay can copy the diagram as a picture. */ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { private val canvas = DiagramCanvas(palette) @@ -38,13 +41,21 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_AS_NEEDED verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED } + private val copy = SessionCopyButton( + tooltip = KiloBundle.message("diagram.copy"), + icon = big(AllIcons.Actions.Copy), + image = { canvas.image() }, + ) { null } // Built by chaining rather than `apply`, so these lambdas cannot bind to Stack's own fit(). - private val controls = Stack.horizontal(UiStyle.Gap.xs()) + private val controls = Stack.vertical(UiStyle.Gap.xs()) .next(control(AllIcons.General.ZoomIn, "diagram.zoom.in") { zoomIn() }) .next(control(AllIcons.General.ZoomOut, "diagram.zoom.out") { zoomOut() }) .next(control(AllIcons.General.FitContent, "diagram.zoom.fit") { fit() }) + .gap(UiStyle.Gap.md()) + .next(copy.button) private val wheel = Wheel() private val drag = Drag() + private val click = Click() init { // Layer first, then add: add(Component, Int) binds to Container.add(comp, index) from Kotlin, @@ -57,6 +68,7 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { scroll.addMouseWheelListener(wheel) canvas.addMouseListener(drag) canvas.addMouseMotionListener(drag) + canvas.addMouseListener(click) } @RequiresEdt @@ -93,6 +105,11 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { canvas.fit() } + override fun removeNotify() { + copy.dismiss() + super.removeNotify() + } + override fun doLayout() { scroll.setBounds(0, 0, width, height) val size = controls.preferredSize @@ -100,6 +117,14 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { controls.doLayout() } + /** Double click is the usual "show me all of it again" gesture in image and diagram viewers. */ + private inner class Click : MouseAdapter() { + override fun mouseClicked(e: MouseEvent) { + if (e.button != MouseEvent.BUTTON1 || e.clickCount != 2) return + fit() + } + } + private inner class Wheel : MouseWheelListener { override fun mouseWheelMoved(e: MouseWheelEvent) { if (!e.isControlDown && !e.isMetaDown) return @@ -118,7 +143,7 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { if (e.button != MouseEvent.BUTTON1 || !overflows()) return from = e.point origin = scroll.viewport.viewPosition - canvas.cursor = Cursor.getPredefinedCursor(Cursor.MOVE_CURSOR) + canvas.cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) } override fun mouseDragged(e: MouseEvent) { @@ -152,8 +177,12 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { private companion object { const val STEP = 1.25 + const val SIZE = 2f fun control(icon: Icon, key: String, handler: () -> Unit) = - toolbarButton(ToolbarButtonAction(icon, KiloBundle.message(key), handler), fill = true) + toolbarButton(ToolbarButtonAction(big(icon), KiloBundle.message(key), handler)) + + /** Overlay icons are drawn over the diagram, so they are sized up to stay readable. */ + fun big(icon: Icon): Icon = IconUtil.scale(icon, null, SIZE) } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index 6b77d23a81c..d5923798910 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -992,6 +992,9 @@ internal open class MdViewHybrid( Disposer.register(disposable) { panel.removeMouseListener(click) } root.next(panel).next(codePane).next(label) root.text = { (this.desc as Desc.Code).text } + // Only offer the picture while the rendered diagram is the thing on screen, so copying the + // streaming or failed source still copies that source. + root.image = { panel.takeIf { it.isVisible }?.image() } kick() } diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index eeb37fb8872..ea76c58ba61 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -7,6 +7,7 @@ common.save=Save common.dont.show.again=Don''t show again diagram.source=Source diagram.title=Diagram +diagram.copy=Copy Diagram diagram.error=Couldn''t render diagram: {0} diagram.missing=Diagram source is no longer available. diagram.open=Open in Editor diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt index bff4a11cdc2..3822a99b240 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt @@ -10,15 +10,19 @@ import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Type import ai.kilocode.client.util.edtWait import com.intellij.icons.AllIcons +import com.intellij.openapi.ide.CopyPasteManager import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.ui.components.JBScrollPane import com.intellij.ui.components.Magnificator import java.awt.Color import java.awt.Container +import java.awt.Cursor import java.awt.Font import java.awt.Point +import java.awt.datatransfer.DataFlavor import java.awt.event.MouseEvent import java.awt.event.MouseWheelEvent +import java.awt.image.BufferedImage import javax.swing.AbstractButton import javax.swing.JComponent import javax.swing.JViewport @@ -151,24 +155,81 @@ class DiagramViewerTest : BasePlatformTestCase() { assertEquals(Point(0, 0), viewport(viewer).viewPosition) } - fun `test overlay offers zoom in zoom out and fit`() = edtWait { + fun `test overlay stacks zoom controls and copy in one column`() = edtWait { val viewer = viewer(400, 300) + layout(viewer) val buttons = buttons(viewer) - assertEquals(3, buttons.size) - assertEquals( - listOf(AllIcons.General.ZoomIn, AllIcons.General.ZoomOut, AllIcons.General.FitContent), - buttons.map { it.icon }, - ) assertEquals( listOf( KiloBundle.message("diagram.zoom.in"), KiloBundle.message("diagram.zoom.out"), KiloBundle.message("diagram.zoom.fit"), + KiloBundle.message("diagram.copy"), ), buttons.map { it.toolTipText }, ) + assertEquals("the overlay is a single column", 1, buttons.map { it.x }.distinct().size) + assertTrue( + "each control sits below the previous one", + buttons.zipWithNext().all { (above, below) -> below.y >= above.y + above.height }, + ) + } + + fun `test overlay icons are twice the platform size and size their buttons`() = edtWait { + val viewer = viewer(400, 300) + + val zoom = buttons(viewer).first() + + assertEquals(AllIcons.General.ZoomIn.iconWidth * 2, zoom.icon.iconWidth) + assertTrue("the hit target grows with the icon", zoom.preferredSize.width > zoom.icon.iconWidth) + assertTrue("icon-only controls stay square", zoom.preferredSize.width == zoom.preferredSize.height) + } + + fun `test copy puts the whole diagram on the clipboard as a picture`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(200.0, 100.0)) + layout(viewer) + + buttons(viewer).single { it.toolTipText == KiloBundle.message("diagram.copy") }.doClick() + + val image = CopyPasteManager.getInstance().contents?.getTransferData(DataFlavor.imageFlavor) as BufferedImage + // Rendered from the scene at 2x with padding, so zoom and scroll state cannot crop it. + assertEquals(200 * 2 + PAD * 2, image.width) + assertEquals(100 * 2 + PAD * 2, image.height) + } + + fun `test double clicking restores fit`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + val fit = scale(viewer) + viewer.zoomIn() + layout(viewer) + assertTrue(scale(viewer) > fit) + val canvas = canvas(viewer) + + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_CLICKED, 100, 100, clicks = 2)) + layout(viewer) + + assertEquals(fit, scale(viewer), 1e-6) + assertEquals(0, canvas.preferredSize.width) + } + + fun `test dragging shows the hand cursor until the drag ends`() = edtWait { + val viewer = viewer(400, 300) + viewer.art(scene(2_000.0, 1_000.0)) + layout(viewer) + repeat(4) { viewer.zoomIn() } + layout(viewer) + val canvas = canvas(viewer) + + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_PRESSED, 200, 150)) + assertEquals(Cursor.HAND_CURSOR, canvas.cursor.type) + + canvas.dispatchEvent(mouse(canvas, MouseEvent.MOUSE_RELEASED, 200, 150)) + assertEquals(Cursor.DEFAULT_CURSOR, canvas.cursor.type) } fun `test overlay floats above the scroll pane and receives its own clicks`() = edtWait { @@ -243,14 +304,14 @@ class DiagramViewerTest : BasePlatformTestCase() { ) } - private fun mouse(target: JComponent, id: Int, x: Int, y: Int) = MouseEvent( + private fun mouse(target: JComponent, id: Int, x: Int, y: Int, clicks: Int = 1) = MouseEvent( target, id, System.currentTimeMillis(), MouseEvent.BUTTON1_DOWN_MASK, x, y, - 1, + clicks, false, MouseEvent.BUTTON1, ) @@ -291,4 +352,9 @@ class DiagramViewerTest : BasePlatformTestCase() { font = Font(Font.SANS_SERIF, Font.PLAIN, 12), bold = Font(Font.SANS_SERIF, Font.BOLD, 12), ) + + private companion object { + /** The padding `diagramImage` leaves around a copied diagram, in image pixels. */ + const val PAD = 32 + } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt index 5fb234a0213..4bf4271c074 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -1,5 +1,6 @@ package ai.kilocode.client.ui.md +import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.ui.selection.SessionTargetResolver import ai.kilocode.client.testing.TestCoroutines @@ -22,6 +23,7 @@ import com.intellij.openapi.actionSystem.DataSink import com.intellij.openapi.actionSystem.UiDataProvider import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.editor.EditorFactory +import com.intellij.openapi.ide.CopyPasteManager import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer import com.intellij.testFramework.fixtures.BasePlatformTestCase @@ -29,7 +31,10 @@ import com.intellij.testFramework.replaceService import com.intellij.util.ui.UIUtil import java.awt.Cursor import java.awt.Point +import java.awt.datatransfer.DataFlavor import java.awt.event.MouseEvent +import java.awt.image.BufferedImage +import javax.swing.AbstractButton import javax.swing.JComponent import javax.swing.JPanel @@ -118,6 +123,28 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertTrue(opened.isEmpty()) } + fun `test copying a rendered diagram puts a picture on the clipboard`() { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + + copyButton().doClick() + + val image = CopyPasteManager.getInstance().contents?.getTransferData(DataFlavor.imageFlavor) as BufferedImage + assertTrue(image.width > 0 && image.height > 0) + } + + fun `test copying a streaming fence still copies its source`() { + // The source pane is what the reader sees until the fence closes, so copy follows the text. + view.append("```mermaid\nflowchart TD\n") + drain() + + copyButton().doClick() + + val contents = CopyPasteManager.getInstance().contents!! + assertFalse(contents.isDataFlavorSupported(DataFlavor.imageFlavor)) + assertEquals("flowchart TD", (contents.getTransferData(DataFlavor.stringFlavor) as String).trim()) + } + fun `test engine error keeps source visible`() { engine.out = Out.Err(ai.kilocode.client.ui.diagram.Fault.Syntax, "bad syntax") @@ -203,6 +230,11 @@ class MdViewDiagramTest : BasePlatformTestCase() { ) } + /** The block's toolbar is only parented while the hover overlay shows it, so reach it directly. */ + private fun copyButton() = descendants(block().copyToolbar!!) + .filterIsInstance() + .single { it.toolTipText == KiloBundle.message("session.copy.hover") } + private fun root() = view.component as JPanel private fun block() = descendants(root()).filterIsInstance().single() From 4801e43127bf42d338251c9f44f6a9c0409159b3 Mon Sep 17 00:00:00 2001 From: kirillk Date: Thu, 27 Aug 2026 14:18:00 -0400 Subject: [PATCH 10/11] fix(jetbrains): use standard diagram toolbar icons --- .../ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt | 9 ++------- .../kilocode/client/ui/diagram/ui/DiagramViewerTest.kt | 7 ++++--- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt index 535da9bdac8..191fbc488b6 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewer.kt @@ -11,7 +11,6 @@ import ai.kilocode.client.ui.toolbarButton import com.intellij.icons.AllIcons import com.intellij.ui.components.JBLayeredPane import com.intellij.ui.components.JBScrollPane -import com.intellij.util.IconUtil import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import java.awt.Color @@ -43,7 +42,7 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { } private val copy = SessionCopyButton( tooltip = KiloBundle.message("diagram.copy"), - icon = big(AllIcons.Actions.Copy), + icon = AllIcons.Actions.Copy, image = { canvas.image() }, ) { null } // Built by chaining rather than `apply`, so these lambdas cannot bind to Stack's own fit(). @@ -177,12 +176,8 @@ internal class DiagramViewer(palette: Palette) : JBLayeredPane() { private companion object { const val STEP = 1.25 - const val SIZE = 2f fun control(icon: Icon, key: String, handler: () -> Unit) = - toolbarButton(ToolbarButtonAction(big(icon), KiloBundle.message(key), handler)) - - /** Overlay icons are drawn over the diagram, so they are sized up to stay readable. */ - fun big(icon: Icon): Icon = IconUtil.scale(icon, null, SIZE) + toolbarButton(ToolbarButtonAction(icon, KiloBundle.message(key), handler)) } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt index 3822a99b240..2d0d5ac47c8 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramViewerTest.kt @@ -177,14 +177,15 @@ class DiagramViewerTest : BasePlatformTestCase() { ) } - fun `test overlay icons are twice the platform size and size their buttons`() = edtWait { + fun `test overlay icons keep the standard platform action size`() = edtWait { val viewer = viewer(400, 300) val zoom = buttons(viewer).first() - assertEquals(AllIcons.General.ZoomIn.iconWidth * 2, zoom.icon.iconWidth) - assertTrue("the hit target grows with the icon", zoom.preferredSize.width > zoom.icon.iconWidth) + assertEquals(AllIcons.General.ZoomIn.iconWidth, zoom.icon.iconWidth) + assertEquals(AllIcons.General.ZoomIn.iconHeight, zoom.icon.iconHeight) assertTrue("icon-only controls stay square", zoom.preferredSize.width == zoom.preferredSize.height) + assertTrue("the hit target is larger than the glyph", zoom.preferredSize.width > zoom.icon.iconWidth) } fun `test copy puts the whole diagram on the clipboard as a picture`() = edtWait { From 921bce7cca59d0dccf7bfdf29cd661cdb7bc263b Mon Sep 17 00:00:00 2001 From: kirillk Date: Thu, 27 Aug 2026 16:27:43 -0400 Subject: [PATCH 11/11] fix(jetbrains): harden diagram rendering fallbacks and logging Diagram rendering degraded to the mermaid source correctly for engine refusals, but three paths could not recover and none of them left a trace in the log. Log every failure. An engine crash was converted to a Fault.Internal carrying only err.message, discarding the throwable, so a parser or layout bug reached the user as an unactionable red label and left nothing to report. It now logs with the stack trace, as does a failing completion callback and a failed image allocation. Close the fallback holes. The completion callback ran unguarded inside the render coroutine, so a failure in it was reported as a plugin error and left the block pending forever. Painters.of used first {}, turning a future art type without a painter into an exception in both paint and sizing. Painting itself was unguarded, and by the time it runs the source pane is already hidden, so a failure left a blank surface with no way back; paint failures now hand a fallback to the owner, which restores the source. diagramImage allocated 2x the scene with no ceiling, and Limits caps the model rather than the geometry, so a legal diagram could ask for a multi-gigabyte raster on the EDT the moment someone pressed copy. Bound the work that limits did not cover. Cancellation is only checked between lines, so one line needed its own cap, and per-index open/quote scanning is replaced by a single pass because it made a long line quadratic. FlowMarks was the only phase with no cancellation point at all and resolved frame members by recursing per cluster, rescanning every node each time; it now resolves bounds and depth in one reverse pass. The remaining FlowLayout phases gained per-iteration checks, and a wall clock ceiling ends a render that outlasts its cooperative checks as a limit fault instead of an endless "rendering" state. Unsupported diagram types now read as a muted note rather than an error. classDiagram, stateDiagram and friends are valid mermaid this engine does not draw, and marking each one red reported working markdown as broken. --- .../client/session/ui/SessionSurface.kt | 4 +- .../ai/kilocode/client/ui/diagram/Engine.kt | 7 ++ .../ai/kilocode/client/ui/diagram/Painter.kt | 16 +++- .../client/ui/diagram/mermaid/Flow.kt | 6 +- .../client/ui/diagram/mermaid/FlowLayout.kt | 32 +++++-- .../client/ui/diagram/mermaid/FlowMarks.kt | 81 +++++++++++----- .../client/ui/diagram/mermaid/Mermaid.kt | 5 + .../kilocode/client/ui/diagram/mermaid/Seq.kt | 3 +- .../client/ui/diagram/mermaid/Source.kt | 21 +++- .../client/ui/diagram/ui/DiagramCanvas.kt | 77 +++++++++++++-- .../client/ui/diagram/ui/DiagramContent.kt | 7 +- .../client/ui/diagram/ui/DiagramPanel.kt | 18 +++- .../kilocode/client/ui/diagram/ui/Diagrams.kt | 63 ++++++++++-- .../client/ui/md/hybrid/MdViewHybrid.kt | 20 +++- .../resources/messages/KiloBundle.properties | 2 + .../kilocode/client/ui/diagram/CancelTest.kt | 29 ++++++ .../kilocode/client/ui/diagram/LimitsTest.kt | 34 +++++++ .../client/ui/diagram/SerializeTest.kt | 3 +- .../client/ui/diagram/mermaid/SourceTest.kt | 29 +++++- .../client/ui/diagram/ui/DiagramPanelTest.kt | 38 ++++++++ .../client/ui/diagram/ui/DiagramsTest.kt | 95 ++++++++++++++++++- .../client/ui/md/MdViewDiagramTest.kt | 41 +++++++- 22 files changed, 555 insertions(+), 76 deletions(-) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionSurface.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionSurface.kt index 64fc5fef9b5..eff4c1994c1 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionSurface.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionSurface.kt @@ -31,13 +31,13 @@ internal object SessionSurface { } /** Runs [paint] with the graphics clipped to the rounded block, keeping opaque content rounded. */ - inline fun clipped(g: Graphics, width: Int, height: Int, paint: (Graphics) -> Unit) { + inline fun clipped(g: Graphics, width: Int, height: Int, paint: (Graphics) -> T): T { val g2 = g.create() as Graphics2D try { g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON) val arc = arc().toFloat() g2.clip(RoundRectangle2D.Float(0f, 0f, width.toFloat(), height.toFloat(), arc, arc)) - paint(g2) + return paint(g2) } finally { g2.dispose() } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt index ba1ef1dbe4e..256eaa92c63 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Engine.kt @@ -40,6 +40,11 @@ internal data class Metrics( * Guards against pathological model output. [chars] is checked before any preprocessing so a single * enormous line cannot reach the parsers; [nodes] and [edges] are enforced while the model is built * rather than after, so `A & B & … --> …` cannot expand into a huge edge list first. + * + * [span] caps one line on its own because per-line scanning and the message regex are superlinear in + * line length, and cancellation is only checked between lines. [millis] is the wall clock ceiling the + * caller applies around the whole draw, so a phase that turns out to be slower than its cooperative + * checks can observe still ends as a [Fault.Limit] rather than an endless "rendering" state. */ @Serializable internal data class Limits( @@ -47,4 +52,6 @@ internal data class Limits( val edges: Int = 800, val lines: Int = 2_000, val chars: Int = 100_000, + val span: Int = 4_000, + val millis: Long = 4_000, ) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt index da4bd51cb4c..2a99df74885 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/Painter.kt @@ -11,5 +11,19 @@ internal interface Painter { internal object Painters { private val all = listOf(ScenePainter) - fun of(art: Art) = all.first { it.accepts(art) } + /** + * Never throws: an [Art] no painter accepts draws nothing and measures empty. + * + * `first { }` here would turn a future art type without a painter into an exception on the EDT, in + * paint and in sizing, where there is no way back to the source fallback. + */ + fun of(art: Art): Painter = all.firstOrNull { it.accepts(art) } ?: Blank + + private object Blank : Painter { + override fun accepts(art: Art) = false + + override fun size(art: Art) = Size(0.0, 0.0) + + override fun paint(g: Graphics2D, art: Art, palette: Palette) = Unit + } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt index f9aab06adf0..b21f3bd0861 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Flow.kt @@ -194,10 +194,11 @@ internal class Flow(private val limits: Limits = Limits()) { /** Splits `A & B` groups at bracket depth zero. */ private fun parts(segment: String): List { val out = mutableListOf() + val mask = Source.opens(segment) var start = 0 for (idx in segment.indices) { if (segment[idx] != '&') continue - if (!Source.open(segment, idx)) continue + if (!mask[idx]) continue out.add(segment.substring(start, idx)) start = idx + 1 } @@ -258,9 +259,10 @@ internal class Flow(private val limits: Limits = Limits()) { private fun hits(text: String): List { val out = mutableListOf() + val mask = Source.opens(text) var idx = 0 while (idx < text.length) { - if (!Source.open(text, idx)) { + if (!mask[idx]) { idx++ continue } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt index c44fc958eb2..f489aba2af9 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowLayout.kt @@ -121,10 +121,17 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return out } - /** Expands each edge into the chain of ranks it crosses, adding a virtual slot per crossed rank. */ - private fun paths(links: List, rank: Map, sizes: MutableMap): List { + /** + * Expands each edge into the chain of ranks it crosses, adding a virtual slot per crossed rank. + * + * This is the one phase whose output is not bounded by [Limits] directly: a long edge contributes a + * slot per rank it spans, so within the node and edge caps it can still mint six figures of virtual + * slots. Hence the per-edge cancellation check rather than one at the phase boundary. + */ + private suspend fun paths(links: List, rank: Map, sizes: MutableMap): List { val out = mutableListOf() for (edge in links) { + coroutineContext.ensureActive() val from = rank[edge.from] ?: 0 val to = rank[edge.to] ?: 0 val ids = mutableListOf(edge.from) @@ -142,7 +149,7 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return out } - private fun order( + private suspend fun order( graph: Graph, rank: Map, paths: List, @@ -161,7 +168,7 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return out } - private fun sweep( + private suspend fun sweep( graph: Graph, order: List>, pairs: List>, @@ -173,14 +180,18 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) val down = pass % 2 == 0 val ranks = if (down) order.indices.drop(1) else order.indices.reversed().drop(1) for (at in ranks) { + coroutineContext.ensureActive() val other = order[at + if (down) -1 else 1] val slot = linkedMapOf() other.forEachIndexed { idx, id -> slot[id] = idx.toDouble() } val group = order[at].associateWith { key(graph, it) } + // Both keys are resolved once per id: a comparator that recomputed the median would + // re-sort a node's neighbour positions on every comparison. + val want = order[at].associateWith { median(adj[it]?.mapNotNull { peer -> slot[peer] } ?: emptyList()) } order[at].sortWith( compareBy( { group[it] }, - { median(adj[it]?.mapNotNull { peer -> slot[peer] } ?: emptyList()) ?: Double.MAX_VALUE }, + { want[it] ?: Double.MAX_VALUE }, { index[it] ?: 0 }, ), ) @@ -212,7 +223,7 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return out } - private fun place( + private suspend fun place( order: List>, sizes: Map, pairs: List>, @@ -230,7 +241,7 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return x } - private fun align( + private suspend fun align( order: List>, sizes: Map, adj: Map>, @@ -239,6 +250,7 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) ) { val ranks = if (down) order.indices.drop(1) else order.indices.reversed().drop(1) for (at in ranks) { + coroutineContext.ensureActive() val other = order[at + if (down) -1 else 1] val centers = linkedMapOf() for (id in other) centers[id] = (x[id] ?: 0.0) + width(sizes, id) / 2 @@ -283,17 +295,19 @@ internal class FlowLayout(private val measure: Measure, private val spec: Spec) return out } - private fun routes(graph: Graph, boxes: Map, paths: List): List { + private suspend fun routes(graph: Graph, boxes: Map, paths: List): List { val out = mutableListOf() val seen = linkedMapOf() + val chains = paths.associateBy { it.edge.index } for (edge in graph.edges) { + coroutineContext.ensureActive() val from = boxes[edge.from] ?: continue if (edge.from == edge.to) { out.add(Route(edge, loop(from.rect))) continue } val to = boxes[edge.to] ?: continue - val path = paths.firstOrNull { it.edge.index == edge.index } ?: continue + val path = chains[edge.index] ?: continue val lane = seen.getOrDefault(lane(edge), 0) seen[lane(edge)] = lane + 1 out.add(Route(edge, trace(path, boxes, from.rect, to.rect, lane))) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt index 69076212f9e..fad0823a5b7 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/FlowMarks.kt @@ -10,23 +10,27 @@ import ai.kilocode.client.ui.diagram.Scene import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.Type +import kotlin.coroutines.coroutineContext import kotlin.math.max +import kotlinx.coroutines.ensureActive /** Turns laid-out flowchart geometry into marks. Cluster frames are emitted first so they paint behind. */ internal class FlowMarks(private val measure: Measure, private val spec: Spec) { private val pad get() = spec.metrics.pad - fun run(placed: Placed): Scene { + suspend fun run(placed: Placed): Scene { val frames = frames(placed) val dx = -minOf(0.0, frames.minOfOrNull { it.rect.x } ?: 0.0) val dy = -minOf(0.0, frames.minOfOrNull { it.rect.y } ?: 0.0) val marks = mutableListOf() for (frame in frames) marks.add(group(frame, dx, dy)) for (route in placed.routes) { + coroutineContext.ensureActive() marks.add(line(route, dx, dy)) marks.addAll(tag(route, dx, dy)) } for (slot in placed.slots.values) { + coroutineContext.ensureActive() val node = slot.node ?: continue val rect = move(slot.rect, dx, dy) marks.addAll(shape(node, rect)) @@ -41,39 +45,61 @@ internal class FlowMarks(private val measure: Measure, private val spec: Spec) { return Size(wide, high) } - private fun frames(placed: Placed): List { - val out = mutableListOf() + /** + * A frame spans the member nodes of its subgraph and of every subgraph nested inside it, padded by + * one step per nesting level so an outer frame reserves room for the inner ones. + * + * Both the member set and the nesting depth are resolved in one reverse pass instead of a recursive + * walk per cluster. A cluster's parent is always declared before it, so visiting declaration order + * backwards visits children first, and each cluster only merges the bounds its children already + * resolved. Recursing per cluster instead rescans every node for every cluster, which is quadratic + * on deeply nested subgraphs and, since this phase is not the one holding the layout, was also the + * only phase with no cancellation point at all. + */ + private suspend fun frames(placed: Placed): List { + val kids = linkedMapOf>() for (cluster in placed.graph.clusters.values) { - val rects = members(placed.graph, cluster.id).mapNotNull { placed.slots[it]?.rect } - if (rects.isEmpty()) continue - val room = pad * 2 * (1 + deep(placed.graph, cluster.id)) - val title = measure.height(spec.font) * cluster.label.size - val x = rects.minOf { it.x } - room - val y = rects.minOf { it.y } - room - title - val wide = rects.maxOf { it.x + it.w } + room - x - val high = rects.maxOf { it.y + it.h } + room - y - out.add(Frame(cluster, Rect(x, y, wide, high))) + if (cluster.parent == null) continue + kids.getOrPut(cluster.parent) { mutableListOf() }.add(cluster.id) + } + val owned = linkedMapOf>() + for (node in placed.graph.nodes.values) { + val id = node.cluster ?: continue + val rect = placed.slots[node.id]?.rect ?: continue + owned.getOrPut(id) { mutableListOf() }.add(rect) } - return out - } - private fun members(graph: Graph, id: String): List { - val out = mutableListOf() - for (node in graph.nodes.values) { - if (node.cluster == id) out.add(node.id) + val reach = linkedMapOf() + val deep = linkedMapOf() + for (cluster in placed.graph.clusters.values.reversed()) { + coroutineContext.ensureActive() + val below = kids[cluster.id].orEmpty() + val bounds = owned[cluster.id].orEmpty().map(::span) + below.mapNotNull { reach[it] } + val merged = bounds.reduceOrNull(::merge) + if (merged != null) reach[cluster.id] = merged + deep[cluster.id] = below.maxOfOrNull { 1 + (deep[it] ?: 0) } ?: 0 } - for (cluster in graph.clusters.values) { - if (cluster.parent == id) out.addAll(members(graph, cluster.id)) + + val out = mutableListOf() + for (cluster in placed.graph.clusters.values) { + val bounds = reach[cluster.id] ?: continue + val room = pad * 2 * (1 + (deep[cluster.id] ?: 0)) + val title = measure.height(spec.font) * cluster.label.size + val x = bounds.minX - room + val y = bounds.minY - room - title + out.add(Frame(cluster, Rect(x, y, bounds.maxX + room - x, bounds.maxY + room - y))) } return out } - /** Nesting depth below [id]; used so an outer frame reserves room for the frames inside it. */ - private fun deep(graph: Graph, id: String): Int { - val kids = graph.clusters.values.filter { it.parent == id } - if (kids.isEmpty()) return 0 - return 1 + (kids.maxOfOrNull { deep(graph, it.id) } ?: 0) - } + private fun span(rect: Rect) = Span(rect.x, rect.y, rect.x + rect.w, rect.y + rect.h) + + private fun merge(a: Span, b: Span) = Span( + minOf(a.minX, b.minX), + minOf(a.minY, b.minY), + max(a.maxX, b.maxX), + max(a.maxY, b.maxY), + ) private fun group(frame: Frame, dx: Double, dy: Double): Mark { val rect = move(frame.rect, dx, dy) @@ -223,6 +249,9 @@ internal class FlowMarks(private val measure: Measure, private val spec: Spec) { private data class Frame(val cluster: Cluster, val rect: Rect) + /** Bounding box of the nodes a subgraph reaches, kept separate from the padded [Frame] rect. */ + private data class Span(val minX: Double, val minY: Double, val maxX: Double, val maxY: Double) + private companion object { const val HALF = 0.5 } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt index 0d8831a4c36..bf4c435d429 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Mermaid.kt @@ -27,6 +27,11 @@ internal class Mermaid(private val measure: Measure) : Engine { if (clean.lines.size > spec.limits.lines) { return Out.Err(Fault.Limit, "source exceeds ${spec.limits.lines} lines") } + // Cancellation is only checked between lines, so one line is also the unit of uninterruptible + // work and needs its own cap rather than relying on the whole-source character limit. + if (clean.lines.any { it.text.length > spec.limits.span }) { + return Out.Err(Fault.Limit, "a line exceeds ${spec.limits.span} characters") + } val type = Type.of(clean) if (!accepts(type)) return Out.Err(Fault.Unsupported, "unsupported diagram type: $type") if (type == Type.Flowchart) return flow(clean, spec) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt index 297a20b6a2e..6e4bc7995d7 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Seq.kt @@ -122,7 +122,8 @@ internal class Seq(private val limits: Limits = Limits()) { private fun actor(text: String): String? { val rest = text.substringAfter(' ', "").trim() if (rest.isEmpty()) return "participant needs a name" - val cut = AS.findAll(rest).firstOrNull { Source.open(rest, it.range.first) } + val mask = Source.opens(rest) + val cut = AS.findAll(rest).firstOrNull { mask[it.range.first] } val id = name(if (cut == null) rest else rest.substring(0, cut.range.first)) val label = if (cut == null) rest else rest.substring(cut.range.last + 1) add(id, Source.label(label)) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt index 12760821b5c..422aa280f82 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/mermaid/Source.kt @@ -40,18 +40,29 @@ internal object Source { return text } - /** True when [index] sits outside quotes and outside any bracket group. */ - fun open(text: String, index: Int): Boolean { + /** + * Per-index "sits outside quotes and outside any bracket group" flags for one line. + * + * Computed in a single pass and reused by every scanner on that line. Answering the question per + * index instead (rescanning from 0 each time) is quadratic, which a 100k character line turns into + * seconds of uninterruptible work because cancellation is only checked between lines. + */ + fun opens(text: String): BooleanArray { + val out = BooleanArray(text.length) var quote = false var depth = 0 - for (idx in 0 until index) { + for (idx in text.indices) { + out[idx] = !quote && depth <= 0 val char = text[idx] - if (char == '"') quote = !quote + if (char == '"') { + quote = !quote + continue + } if (quote) continue if (char == '[' || char == '(' || char == '{') depth++ if (char == ']' || char == ')' || char == '}') depth-- } - return !quote && depth <= 0 + return out } fun rail(char: Char) = RAILS.indexOf(char) >= 0 diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt index 83b978273cf..c34166b0d5f 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramCanvas.kt @@ -5,6 +5,9 @@ import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Painters import ai.kilocode.client.ui.diagram.Palette +import ai.kilocode.client.ui.diagram.Size +import ai.kilocode.log.KiloLog +import com.intellij.openapi.progress.ProcessCanceledException import com.intellij.ui.components.Magnificator import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI @@ -16,10 +19,12 @@ import java.awt.Point import java.awt.Rectangle import java.awt.RenderingHints import java.awt.image.BufferedImage +import java.util.concurrent.atomic.AtomicBoolean import javax.swing.JComponent import javax.swing.JViewport import javax.swing.Scrollable import kotlin.math.roundToInt +import kotlin.math.sqrt /** * Scrollable diagram surface for the diagram viewer. @@ -164,30 +169,67 @@ internal class DiagramCanvas(private var palette: Palette) : JComponent(), Scrol * a fixed [SHOT] scale regardless of the current zoom, scroll position or viewport size. The * background is filled because the palette follows the IDE theme: a transparent PNG of a dark theme * diagram would be unreadable once pasted onto white. + * + * Returns null when there is nothing to draw or the allocation fails, and callers fall back to copying + * the diagram source as text. */ -internal fun diagramImage(art: Art, palette: Palette, background: Color?): BufferedImage { +internal fun diagramImage(art: Art, palette: Palette, background: Color?): BufferedImage? { val size = Painters.of(art).size(art) - val pad = SessionUiStyle.View.Diagram.PADDING * SHOT - val w = (size.w * SHOT + pad * 2).roundToInt().coerceAtLeast(1) - val h = (size.h * SHOT + pad * 2).roundToInt().coerceAtLeast(1) - // A plain image on purpose: ImageUtil.createImage would add the IDE's HiDPI scale on top of SHOT, - // making the picture depend on the display it was copied from. - val image = BufferedImage(w, h, BufferedImage.TYPE_INT_RGB) + val scale = shot(size) + val pad = SessionUiStyle.View.Diagram.PADDING * scale + val w = (size.w * scale + pad * 2).roundToInt().coerceAtLeast(1) + val h = (size.h * scale + pad * 2).roundToInt().coerceAtLeast(1) + // A plain image on purpose: ImageUtil.createImage would add the IDE's HiDPI scale on top of the + // shot scale, making the picture depend on the display it was copied from. + val image = try { + BufferedImage(w, h, BufferedImage.TYPE_INT_RGB) + } catch (err: OutOfMemoryError) { + // Bounded by shot() above, so this is a last resort against an IDE that is already short on + // heap rather than the expected path for a large diagram. + LOG.error("kind=diagram image=failed width=$w height=$h", err) + return null + } val g = image.createGraphics() try { g.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON) g.color = background ?: UiStyle.Colors.editorBackground() g.fillRect(0, 0, w, h) - paintDiagram(g, art, palette, SHOT, pad.roundToInt(), pad.roundToInt()) + paintDiagram(g, art, palette, scale, pad.roundToInt(), pad.roundToInt()) } finally { g.dispose() } return image } +/** + * Scale for a clipboard image: [SHOT] unless that would allocate an unreasonable picture. + * + * The engine's [ai.kilocode.client.ui.diagram.Limits] cap the model, not the geometry it lays out, so a + * legal diagram that is both deep and wide can span tens of thousands of units. At [SHOT] that is a + * multi-gigabyte raster, allocated on the EDT the moment someone presses copy, so the scale gives way + * before the allocation does and a huge diagram is copied smaller instead of taking the IDE down. + */ +private fun shot(size: Size): Double { + if (size.w <= 0.0 || size.h <= 0.0) return SHOT + val side = minOf(SIDE / size.w, SIDE / size.h) + val area = sqrt(PIXELS / (size.w * size.h)) + // No lower bound: the area term is scale invariant, so however large the diagram gets the result + // lands on the pixel budget, and the dimensions are floored at one pixel by the caller. A floor here + // would raise the scale back above the budget it was picked to respect. + return minOf(SHOT, side, area) +} + /** Clipboard images render larger than the screen so they stay crisp when pasted. */ private const val SHOT = 2.0 +/** Longest side, in pixels, of a clipboard image. */ +private const val SIDE = 8_000.0 + +/** Pixel budget for a clipboard image; 8M pixels is ~32MB as `TYPE_INT_RGB`. */ +private const val PIXELS = 8_000_000.0 + +private val LOG = KiloLog.create(DiagramCanvas::class.java) + /** Keeps a viewport position inside the scrollable range of its view. */ internal fun clamped(viewport: JViewport, at: Point): Point { val view = viewport.view ?: return at @@ -196,15 +238,30 @@ internal fun clamped(viewport: JViewport, at: Point): Point { return Point(at.x.coerceIn(0, x), at.y.coerceIn(0, y)) } -/** Paints [art] scaled by [scale] with its top-left corner at ([x], [y]). */ -internal fun paintDiagram(g: Graphics, art: Art, palette: Palette, scale: Double, x: Int, y: Int) { +/** + * Paints [art] scaled by [scale] with its top-left corner at ([x], [y]), reporting whether it drew. + * + * A painter that throws must not escape into the enclosing paint pass, where it would take the rest of + * the transcript's painting with it and repeat on every repaint. The failure is logged once and + * reported so the caller can put the diagram source back on screen instead. + */ +internal fun paintDiagram(g: Graphics, art: Art, palette: Palette, scale: Double, x: Int, y: Int): Boolean { val g2 = g.create() as Graphics2D try { g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON) g2.translate(x, y) g2.scale(scale, scale) Painters.of(art).paint(g2, art, palette) + return true + } catch (err: ProcessCanceledException) { + throw err + } catch (err: Throwable) { + if (reported.compareAndSet(false, true)) LOG.error("kind=diagram paint=failed scale=$scale", err) + return false } finally { g2.dispose() } } + +/** Paint runs per frame, so the same broken art must not fill the log. */ +private val reported = AtomicBoolean(false) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt index 91d2bc99ea9..e85c199082a 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramContent.kt @@ -4,6 +4,7 @@ import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.style.SessionEditorStyle import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.Fault import ai.kilocode.client.ui.diagram.Out import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationManager @@ -42,9 +43,11 @@ internal fun diagramContent(source: String, parent: Disposable): JComponent { label.isVisible = false } + // An unsupported diagram type is a note, not a failure; the source tab still has the text. is Out.Err -> { - label.text = KiloBundle.message("diagram.error", out.message) - label.foreground = UiStyle.Colors.errorLabelForeground() + val hint = out.fault == Fault.Unsupported + label.text = if (hint) KiloBundle.message("diagram.unsupported") else KiloBundle.message("diagram.error", out.message) + label.foreground = if (hint) SessionUiStyle.Text.Secondary.foreground() else UiStyle.Colors.errorLabelForeground() label.isVisible = true } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt index 9fcfb550259..b648d64a43b 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanel.kt @@ -5,6 +5,7 @@ import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Painters import ai.kilocode.client.ui.diagram.Palette +import com.intellij.openapi.application.ApplicationManager import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import java.awt.Dimension @@ -18,10 +19,20 @@ import kotlin.math.roundToInt internal class DiagramPanel(private var palette: Palette) : JComponent() { private var art: Art? = null private var last = Dimension(0, 0) + private var faulted = false + + /** + * Called once, off the paint pass, when the diagram could not be drawn. + * + * The owner uses it to go back to showing the source: by the time painting fails the source pane is + * already hidden, so without this the reader is left looking at a blank surface. + */ + var onFault: () -> Unit = {} @RequiresEdt fun art(value: Art) { art = value + faulted = false resize() repaint() } @@ -60,9 +71,14 @@ internal class DiagramPanel(private var palette: Palette) : JComponent() { } val value = art ?: return val scale = scale(value) - SessionSurface.clipped(g, width, height) { clipped -> + val drew = SessionSurface.clipped(g, width, height) { clipped -> paintDiagram(clipped, value, palette, scale, pad(), pad()) } + if (drew || faulted) return + // Swapping the visible component from inside a paint pass is not safe, so hand the fallback to + // the owner on the next event instead. + faulted = true + ApplicationManager.getApplication().invokeLater(onFault) } private fun resize() { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt index 0eb361ff0a7..a716e015277 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/diagram/ui/Diagrams.kt @@ -9,18 +9,23 @@ import ai.kilocode.client.ui.diagram.FontSpec import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.mermaid.Mermaid +import ai.kilocode.log.KiloLog import com.intellij.openapi.Disposable import com.intellij.openapi.application.EDT import com.intellij.openapi.application.ModalityState import com.intellij.openapi.application.asContextElement import com.intellij.openapi.components.Service +import com.intellij.openapi.progress.ProcessCanceledException import com.intellij.openapi.util.Disposer import com.intellij.util.concurrency.annotations.RequiresEdt +import java.util.concurrent.atomic.AtomicBoolean import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.TimeoutCancellationException import kotlinx.coroutines.launch import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout @Service(Service.Level.APP) internal class Diagrams internal constructor( @@ -42,20 +47,57 @@ internal class Diagrams internal constructor( return } val job = cs.launch { - val out = try { - impl().draw(source, spec) - } catch (err: CancellationException) { - throw err - } catch (err: Exception) { - Out.Err(Fault.Internal, err.message ?: err.javaClass.simpleName) - } + val out = draw(source, spec) withContext(edt) { if (Disposer.isDisposed(owner)) return@withContext - cache[key] = out - done(out) + // Everything except an internal failure is a deterministic function of the input, so it + // is worth remembering. A bug is not: caching it would keep a transient failure on + // screen for the rest of the session. + if (out !is Out.Err || out.fault != Fault.Internal) cache[key] = out + deliver(done, out) } } - Disposer.register(owner) { job.cancel() } + // One child disposable per call would accumulate on a long lived owner (every fence re-render and + // every theme change registers one), so the guard is released as soon as the job settles. The + // latch keeps the two directions from re-entering each other: whichever of owner disposal and job + // completion happens first, the other side becomes a no-op. + val once = AtomicBoolean(false) + val guard = Disposable { if (once.compareAndSet(false, true)) job.cancel() } + Disposer.register(owner, guard) + job.invokeOnCompletion { if (once.compareAndSet(false, true)) Disposer.dispose(guard) } + } + + private suspend fun draw(source: String, spec: Spec): Out { + try { + if (spec.limits.millis <= 0) return impl().draw(source, spec) + return withTimeout(spec.limits.millis) { impl().draw(source, spec) } + } catch (err: TimeoutCancellationException) { + LOG.warn("kind=diagram render=timeout millis=${spec.limits.millis} chars=${source.length}", err) + return Out.Err(Fault.Limit, "rendering took longer than ${spec.limits.millis} ms") + } catch (err: CancellationException) { + throw err + } catch (err: Exception) { + // The message reaches the user, but only the log carries the stack trace, so a diagram that + // trips a parser or layout bug is reportable instead of just looking broken. + LOG.error("kind=diagram render=failed chars=${source.length}", err) + return Out.Err(Fault.Internal, err.message ?: err.javaClass.simpleName) + } + } + + /** + * Runs the completion callback so a failure inside it cannot escape into the service scope. + * + * An exception here would be reported as a plugin error and, worse, leave the caller stuck on its + * pending state forever because nothing else is going to answer that render. + */ + private fun deliver(done: (Out) -> Unit, out: Out) { + try { + done(out) + } catch (err: ProcessCanceledException) { + throw err + } catch (err: Exception) { + LOG.error("kind=diagram deliver=failed out=$out", err) + } } private fun impl() = engine ?: Mermaid(measure) @@ -65,6 +107,7 @@ internal class Diagrams internal constructor( private companion object { const val CACHE = 64 val edt = Dispatchers.EDT + ModalityState.any().asContextElement() + val LOG = KiloLog.create(Diagrams::class.java) fun hash(text: String): Long { var value = -3750763034362895579L diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt index d5923798910..1ba0cfeb5e3 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/hybrid/MdViewHybrid.kt @@ -7,6 +7,7 @@ import ai.kilocode.client.session.ui.style.SessionEditorStyle import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.ui.UiStyle +import ai.kilocode.client.ui.diagram.Fault import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.ui.DiagramBlock import ai.kilocode.client.ui.diagram.ui.DiagramPanel @@ -989,6 +990,7 @@ internal open class MdViewHybrid( panel.cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) panel.toolTipText = KiloBundle.message("diagram.viewer.hint") panel.addMouseListener(click) + panel.onFault = { fail(KiloBundle.message("diagram.paint")) } Disposer.register(disposable) { panel.removeMouseListener(click) } root.next(panel).next(codePane).next(label) root.text = { (this.desc as Desc.Code).text } @@ -1057,7 +1059,7 @@ internal open class MdViewHybrid( if (seq != gen) return@render when (out) { is Out.Ok -> ok(out) - is Out.Err -> fail(out.message) + is Out.Err -> fail(out) } } } @@ -1070,6 +1072,22 @@ internal open class MdViewHybrid( root.repaint() } + /** + * A diagram type this engine does not draw is not a broken diagram, so it reads as a note rather + * than an error. `classDiagram`, `stateDiagram` and friends are common in model output and marking + * every one of them red would report working markdown as a failure. + */ + private fun fail(out: Out.Err) { + if (out.fault == Fault.Unsupported) { + status(KiloBundle.message("diagram.unsupported")) + showSource() + root.revalidate() + root.repaint() + return + } + fail(out.message) + } + private fun fail(message: String) { val text = message.ifBlank { KiloBundle.message("diagram.rendering") } status(KiloBundle.message("diagram.error", text), true) diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties index ea76c58ba61..44e12e447a8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties +++ b/packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties @@ -11,8 +11,10 @@ diagram.copy=Copy Diagram diagram.error=Couldn''t render diagram: {0} diagram.missing=Diagram source is no longer available. diagram.open=Open in Editor +diagram.paint=the diagram could not be drawn diagram.path=Kilo / Diagrams / {0} diagram.rendering=Rendering diagram... +diagram.unsupported=Inline preview supports flowcharts and sequence diagrams. diagram.viewer.hint=Open diagram viewer diagram.zoom.fit=Fit to Window diagram.zoom.in=Zoom In diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt index ba57232abe7..e7a35f2de3b 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/CancelTest.kt @@ -1,6 +1,9 @@ package ai.kilocode.client.ui.diagram import ai.kilocode.client.ui.diagram.mermaid.Flow +import ai.kilocode.client.ui.diagram.mermaid.FlowLayout +import ai.kilocode.client.ui.diagram.mermaid.FlowMarks +import ai.kilocode.client.ui.diagram.mermaid.FlowOut import ai.kilocode.client.ui.diagram.mermaid.Mermaid import ai.kilocode.client.ui.diagram.mermaid.Seq import ai.kilocode.client.ui.diagram.mermaid.Source @@ -44,6 +47,23 @@ class CancelTest { assertTrue(seq.isEmpty(), "sequence parsing ignored cancellation") } + /** + * Mark generation measures only line heights, so the measurement hook cannot reach it. It is also the + * phase that used to have no cancellation point at all, which is why it gets its own proof. + */ + @Test + fun `mark generation stops when the job is cancelled`() { + val measure = FakeMeasure() + val placed = runBlocking { + val parsed = Flow().parse(Source.clean(NESTED)) as FlowOut.Ok + FlowLayout(measure, spec()).run(parsed.graph) + } + + val marks = sink { FlowMarks(measure, spec()).run(placed) } + + assertTrue(marks.isEmpty(), "mark generation ignored cancellation") + } + /** Runs [body] in a coroutine that cancels itself first; a result only lands if that was ignored. */ private fun sink(body: suspend () -> Any): List = runBlocking { val out = mutableListOf() @@ -71,5 +91,14 @@ class CancelTest { private companion object { const val CUT = 3 + + val NESTED = """ + flowchart TD + subgraph one + subgraph two + A --> B + end + end + """.trimIndent() } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt index 20aaca87f17..bf5cc2c2721 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/LimitsTest.kt @@ -4,6 +4,7 @@ import ai.kilocode.client.ui.diagram.mermaid.Mermaid import kotlinx.coroutines.runBlocking import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertTrue /** Model output can be pathological; the engine must refuse rather than hang or exhaust memory. */ class LimitsTest { @@ -62,6 +63,24 @@ class LimitsTest { assertEquals(Fault.Limit, err(people).fault) } + /** One line is the unit of uninterruptible work, so it needs a cap of its own. */ + @Test + fun `single line cap is enforced`() { + val source = "flowchart TD\n A[" + "x".repeat(200) + "] --> B" + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(span = 50))) } + + assertEquals(Fault.Limit, err(out).fault) + assertTrue(err(out).message.contains("50")) + } + + @Test + fun `a long source of short lines is not refused by the line span cap`() { + val source = "flowchart TD\n" + (1..20).joinToString("\n") { " n$it --> n${it + 1}" } + val out = runBlocking { engine.draw(source, spec().copy(limits = Limits(span = 40))) } + + assertTrue(scene(out).marks.isNotEmpty()) + } + @Test fun `a graph at the cap still renders`() { val source = "flowchart TD\n" + (1..9).joinToString("\n") { " n$it --> n${it + 1}" } @@ -69,4 +88,19 @@ class LimitsTest { assertEquals(10, scene(out).marks.count { it is Mark.Box }) } + + /** + * Subgraph nesting is capped only by the line limit, so a frame's members and depth have to be + * resolved in one pass. Walking the cluster tree per cluster rescans every node per cluster, which + * turns a few hundred levels into a stall in a phase that cannot be cancelled cheaply. + */ + @Test + fun `deeply nested subgraphs still render`() { + val depth = 200 + val open = (1..depth).joinToString("\n") { " subgraph s$it" } + val close = (1..depth).joinToString("\n") { " end" } + val out = runBlocking { engine.draw("flowchart TD\n$open\n a --> b\n$close", spec()) } + + assertEquals(depth, scene(out).marks.count { it is Mark.Group }) + } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt index 9721efc71d4..325d76a3613 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/SerializeTest.kt @@ -30,7 +30,8 @@ class SerializeTest { @Test fun `spec round trips`() { - val value = Spec(FontSpec("Inter", 13, bold = true), Metrics(pad = 3.0), Limits(nodes = 7)) + val limits = Limits(nodes = 7, span = 11, millis = 13) + val value = Spec(FontSpec("Inter", 13, bold = true), Metrics(pad = 3.0), limits) assertEquals(value, json.decodeFromString(json.encodeToString(value))) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt index 8ca6d266d72..3d8a8bf19c8 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/mermaid/SourceTest.kt @@ -82,10 +82,33 @@ class SourceTest { } @Test - fun `open reports bracket and quote nesting`() { + fun `opens reports bracket and quote nesting`() { val text = "A[x --> y] --> B" + val mask = Source.opens(text) - assertTrue(Source.open(text, text.lastIndexOf("-->"))) - assertTrue(!Source.open(text, text.indexOf("-->"))) + assertTrue(mask[text.lastIndexOf("-->")]) + assertTrue(!mask[text.indexOf("-->")]) + } + + @Test + fun `opens treats bracketed and quoted regions as closed`() { + // The flag is the state *before* each character, so an opening quote or bracket is still open and + // its closing partner is not. + val mask = Source.opens("a\"b\"c(d)e") + + assertEquals( + listOf(true, true, false, false, true, true, false, false, true), + mask.toList(), + ) + } + + /** Answering "is this index open" per index rescans from 0 each time, which stalls on a long line. */ + @Test + fun `opens scans a long line in one pass`() { + val text = "A".repeat(200_000) + "-->B" + val mask = Source.opens(text) + + assertEquals(text.length, mask.size) + assertTrue(mask.all { it }) } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt index ef4e0974651..d8006720221 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramPanelTest.kt @@ -11,9 +11,11 @@ import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Type import java.awt.Color import java.awt.Font +import java.awt.image.BufferedImage import javax.swing.AbstractButton import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertNotNull import kotlin.test.assertTrue class DiagramPanelTest { @@ -47,6 +49,42 @@ class DiagramPanelTest { assertTrue(buttons.any { it.icon === SessionViewIcons.openDiff }) } + @Test + fun `test a normal diagram is copied at the crisp shot scale`() { + val image = diagramImage(scene(100.0, 50.0), palette(), Color.WHITE) + + assertNotNull(image) + assertEquals(100 * 2 + 16 * 2 * 2, image.width) + assertEquals(50 * 2 + 16 * 2 * 2, image.height) + } + + /** + * The engine caps the model, not the geometry, so a legal diagram can still span tens of thousands + * of units. Copying it must downscale rather than ask for a multi-gigabyte raster on the EDT. + */ + @Test + fun `test a huge diagram is copied downscaled instead of allocating gigabytes`() { + val image = diagramImage(scene(120_000.0, 90_000.0), palette(), Color.WHITE) + + assertNotNull(image) + assertTrue(image.width <= 8_000, "width ${image.width}") + assertTrue(image.height <= 8_000, "height ${image.height}") + assertTrue(image.width.toLong() * image.height <= 9_000_000L, "pixels ${image.width * image.height}") + assertTrue(image.width > 0 && image.height > 0) + } + + @Test + fun `test painting a scene reports that it drew`() { + val target = BufferedImage(80, 80, BufferedImage.TYPE_INT_RGB) + val g = target.createGraphics() + + try { + assertTrue(paintDiagram(g, scene(40.0, 20.0), palette(), 1.0, 4, 4)) + } finally { + g.dispose() + } + } + private fun buttons(root: java.awt.Container): List { val out = mutableListOf() for (comp in root.components) { diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt index d572fd2e636..25440c9d813 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/diagram/ui/DiagramsTest.kt @@ -4,13 +4,16 @@ import ai.kilocode.client.testing.TestCoroutines import ai.kilocode.client.testing.pumpEdt import ai.kilocode.client.ui.diagram.Art import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.Fault import ai.kilocode.client.ui.diagram.FontSpec +import ai.kilocode.client.ui.diagram.Limits import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.Scene import ai.kilocode.client.ui.diagram.Size import ai.kilocode.client.ui.diagram.Spec import ai.kilocode.client.ui.diagram.Type import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.LoggedErrorProcessor import com.intellij.testFramework.fixtures.BasePlatformTestCase import kotlinx.coroutines.awaitCancellation @@ -74,18 +77,108 @@ class DiagramsTest : BasePlatformTestCase() { assertFalse(called) } + /** + * A crash in the engine is the caller's cue to keep showing the source, but it is only actionable if + * the stack trace reaches the log: the fault message alone says nothing about where it came from. + */ + fun `test engine crash is logged with its stack trace and reported as an internal fault`() { + val owner = Disposer.newDisposable("diagram") + val calls = mutableListOf() + engine.fail = IllegalStateException("boom") + + val logged = LoggedErrorProcessor.executeAndReturnLoggedError { + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + coroutines.drain() + } + + val out = calls.single() as Out.Err + assertEquals(Fault.Internal, out.fault) + assertEquals("boom", out.message) + assertEquals("boom", logged.message) + assertEquals(IllegalStateException::class.java, logged.javaClass) + Disposer.dispose(owner) + } + + fun `test an internal fault is not cached so the next attempt can recover`() { + val owner = Disposer.newDisposable("diagram") + val calls = mutableListOf() + engine.fail = IllegalStateException("boom") + + LoggedErrorProcessor.executeAndReturnLoggedError { + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + coroutines.drain() + } + engine.fail = null + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + coroutines.drain() + + assertEquals(2, engine.calls) + assertTrue(calls.last() is Out.Ok) + Disposer.dispose(owner) + } + + /** Syntax, limit and unsupported outcomes are deterministic, so they stay cached. */ + fun `test a refusal is cached`() { + val owner = Disposer.newDisposable("diagram") + engine.out = Out.Err(Fault.Syntax, "bad") + + service.render("flowchart TD\nA-->", spec(), owner) {} + coroutines.drain() + service.render("flowchart TD\nA-->", spec(), owner) {} + + assertEquals(1, engine.calls) + Disposer.dispose(owner) + } + + /** Cancellation is cooperative, so a phase that never yields still has to end somewhere. */ + fun `test a hung engine ends as a limit fault instead of rendering forever`() { + val owner = Disposer.newDisposable("diagram") + val calls = mutableListOf() + engine.pause = true + + service.render("flowchart TD\nA-->B", spec().copy(limits = Limits(millis = 1)), owner) { calls.add(it) } + assertTrue("the render never settled", coroutines.pumpUntil { calls.isNotEmpty() }) + + val out = calls.single() as Out.Err + assertEquals(Fault.Limit, out.fault) + assertTrue(out.message, out.message.contains("1 ms")) + Disposer.dispose(owner) + } + + /** + * A callback that throws must not escape into the service scope: it would be reported as a plugin + * error and the result would never be recorded, leaving the caller pending forever. + */ + fun `test a failing callback is contained and the result is still cached`() { + val owner = Disposer.newDisposable("diagram") + val calls = mutableListOf() + + LoggedErrorProcessor.executeAndReturnLoggedError { + service.render("flowchart TD\nA-->B", spec(), owner) { throw IllegalStateException("callback") } + coroutines.drain() + } + service.render("flowchart TD\nA-->B", spec(), owner) { calls.add(it) } + + assertEquals(1, engine.calls) + assertEquals(1, calls.size) + Disposer.dispose(owner) + } + private fun spec(size: Int = 12) = Spec(FontSpec("Test", size)) private class FakeEngine : Engine { var calls = 0 var pause = false + var fail: Exception? = null + var out: Out? = null override fun accepts(type: Type) = true override suspend fun draw(source: String, spec: Spec): Out { calls++ + fail?.let { throw it } if (pause) awaitCancellation() - return Out.Ok(Scene(Type.Flowchart, emptyList(), Size(20.0, 10.0)) as Art) + return out ?: Out.Ok(Scene(Type.Flowchart, emptyList(), Size(20.0, 10.0)) as Art) } } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt index 4bf4271c074..776528aa768 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/md/MdViewDiagramTest.kt @@ -3,8 +3,11 @@ package ai.kilocode.client.ui.md import ai.kilocode.client.plugin.KiloBundle import ai.kilocode.client.session.ui.selection.SessionCopyTarget import ai.kilocode.client.session.ui.selection.SessionTargetResolver +import ai.kilocode.client.session.ui.style.SessionUiStyle import ai.kilocode.client.testing.TestCoroutines +import ai.kilocode.client.ui.UiStyle import ai.kilocode.client.ui.diagram.Engine +import ai.kilocode.client.ui.diagram.Fault import ai.kilocode.client.ui.diagram.Mark import ai.kilocode.client.ui.diagram.Out import ai.kilocode.client.ui.diagram.Pt @@ -26,6 +29,7 @@ import com.intellij.openapi.editor.EditorFactory import com.intellij.openapi.ide.CopyPasteManager import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer +import com.intellij.testFramework.LoggedErrorProcessor import com.intellij.testFramework.fixtures.BasePlatformTestCase import com.intellij.testFramework.replaceService import com.intellij.util.ui.UIUtil @@ -146,7 +150,7 @@ class MdViewDiagramTest : BasePlatformTestCase() { } fun `test engine error keeps source visible`() { - engine.out = Out.Err(ai.kilocode.client.ui.diagram.Fault.Syntax, "bad syntax") + engine.out = Out.Err(Fault.Syntax, "bad syntax") view.set("```mermaid\nflowchart TD\nA-->\n```") drain() @@ -155,6 +159,39 @@ class MdViewDiagramTest : BasePlatformTestCase() { assertTrue(codePane().isVisible) assertTrue(label().isVisible) assertTrue(labels().contains("bad syntax")) + assertEquals(UiStyle.Colors.errorLabelForeground(), label().foreground) + } + + /** + * `classDiagram`, `stateDiagram` and friends are valid mermaid this engine does not draw. Marking them + * red would report working markdown as broken, so they read as a note over the source instead. + */ + fun `test an unsupported diagram type reads as a note rather than an error`() { + engine.out = Out.Err(Fault.Unsupported, "unsupported diagram type: Class") + + view.set("```mermaid\nclassDiagram\nA <|-- B\n```") + drain() + + assertFalse(diagram().isVisible) + assertTrue(codePane().isVisible) + assertTrue(labels().contains(KiloBundle.message("diagram.unsupported"))) + assertFalse("the engine's internal wording should not reach the reader", labels().contains("unsupported diagram type")) + assertEquals(SessionUiStyle.Text.Secondary.foreground(), label().foreground) + } + + /** A crash in the engine has to land on the same source fallback as a refusal, and be logged. */ + fun `test an engine crash keeps source visible and is logged`() { + engine.fail = IllegalStateException("boom") + + val logged = LoggedErrorProcessor.executeAndReturnLoggedError { + view.set("```mermaid\nflowchart TD\nA-->B\n```") + drain() + } + + assertEquals("boom", logged.message) + assertFalse(diagram().isVisible) + assertTrue(codePane().isVisible) + assertTrue(labels().contains("boom")) } fun `test streaming waits for closed fence`() { @@ -273,11 +310,13 @@ class MdViewDiagramTest : BasePlatformTestCase() { private class FakeEngine : Engine { var calls = 0 var out: Out? = null + var fail: Exception? = null override fun accepts(type: Type) = true override suspend fun draw(source: String, spec: Spec): Out { calls++ + fail?.let { throw it } return out ?: Out.Ok( Scene( Type.Flowchart,