Skip to content

Commit ac32040

Browse files
carlrobertohcarlrobertoh
committed
feat: support mermaid diagrams (closes #788)
1 parent 3f32f19 commit ac32040

File tree

11 files changed

+3324
-11
lines changed

11 files changed

+3324
-11
lines changed

‎src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.intellij.openapi.options.ShowSettingsUtil;
1919
import com.intellij.openapi.project.Project;
2020
import com.intellij.openapi.ui.VerticalFlowLayout;
21+
import com.intellij.openapi.util.Disposer;
2122
import com.intellij.openapi.util.io.FileUtil;
2223
import com.intellij.openapi.vfs.LocalFileSystem;
2324
import com.intellij.openapi.vfs.VirtualFile;
@@ -58,6 +59,7 @@
5859
import ee.carlrobert.codegpt.ui.hover.PsiLinkHoverPreview;
5960
import ee.carlrobert.codegpt.util.EditorUtil;
6061
import java.awt.BorderLayout;
62+
import java.util.Locale;
6163
import java.util.Objects;
6264
import java.util.stream.Stream;
6365
import javax.swing.DefaultListModel;
@@ -85,6 +87,7 @@ public class ChatMessageResponseBody extends JPanel {
8587
new JPanel(new VerticalFlowLayout(VerticalFlowLayout.TOP, 0, 4, true, false));
8688

8789
private ResponseEditorPanel currentlyProcessedEditorPanel;
90+
private MermaidResponsePanel currentlyProcessedMermaidPanel;
8891
private JEditorPane currentlyProcessedTextPane;
8992
private JPanel webpageListPanel;
9093

@@ -131,6 +134,7 @@ public ChatMessageResponseBody withResponse(@NotNull String response) {
131134
processResponse(item, false);
132135
currentlyProcessedTextPane = null;
133136
currentlyProcessedEditorPanel = null;
137+
currentlyProcessedMermaidPanel = null;
134138
}
135139
} catch (Exception e) {
136140
LOG.error("Something went wrong while processing input", e);
@@ -141,6 +145,7 @@ public ChatMessageResponseBody withResponse(@NotNull String response) {
141145
public void addToolStatusPanel(JComponent component) {
142146
currentlyProcessedTextPane = null;
143147
currentlyProcessedEditorPanel = null;
148+
currentlyProcessedMermaidPanel = null;
144149
streamOutputParser.clear();
145150
contentPanel.add(component);
146151
}
@@ -317,6 +322,7 @@ private void processResponse(Segment item, boolean caretVisible) {
317322
handleHeaderOnCompletion(currentlyProcessedEditorPanel);
318323
}
319324
currentlyProcessedEditorPanel = null;
325+
currentlyProcessedMermaidPanel = null;
320326
return;
321327
}
322328

@@ -349,6 +355,11 @@ private void processResponse(Segment item, boolean caretVisible) {
349355
}
350356

351357
private void processCode(Segment item) {
358+
if (isMermaidCode(item)) {
359+
processMermaid(item.getContent());
360+
return;
361+
}
362+
352363
var content = item.getContent();
353364
if (currentlyProcessedEditorPanel == null) {
354365
prepareProcessingCode(item);
@@ -376,6 +387,7 @@ private void processText(String markdownText, boolean caretVisible) {
376387
@Synchronized
377388
private void prepareProcessingText(boolean caretVisible) {
378389
currentlyProcessedEditorPanel = null;
390+
currentlyProcessedMermaidPanel = null;
379391
currentlyProcessedTextPane = createTextPane("", caretVisible);
380392
contentPanel.add(currentlyProcessedTextPane);
381393
contentPanel.revalidate();
@@ -386,13 +398,54 @@ private void prepareProcessingText(boolean caretVisible) {
386398
private void prepareProcessingCode(Segment item) {
387399
hideCaret();
388400
currentlyProcessedTextPane = null;
401+
currentlyProcessedMermaidPanel = null;
389402
currentlyProcessedEditorPanel =
390403
new ResponseEditorPanel(project, item, readOnly, compact, parentDisposable);
391404
contentPanel.add(currentlyProcessedEditorPanel);
392405
contentPanel.revalidate();
393406
contentPanel.repaint();
394407
}
395408

409+
@Synchronized
410+
private void prepareProcessingMermaid() {
411+
hideCaret();
412+
currentlyProcessedTextPane = null;
413+
currentlyProcessedEditorPanel = null;
414+
currentlyProcessedMermaidPanel = new MermaidResponsePanel();
415+
Disposer.register(parentDisposable, currentlyProcessedMermaidPanel);
416+
contentPanel.add(currentlyProcessedMermaidPanel);
417+
contentPanel.revalidate();
418+
contentPanel.repaint();
419+
}
420+
421+
private void processMermaid(String source) {
422+
if (currentlyProcessedMermaidPanel == null) {
423+
prepareProcessingMermaid();
424+
}
425+
if (currentlyProcessedMermaidPanel != null) {
426+
currentlyProcessedMermaidPanel.render(source);
427+
}
428+
}
429+
430+
private boolean isMermaidCode(Segment item) {
431+
if (!MermaidResponsePanel.isSupported()) {
432+
return false;
433+
}
434+
if (!(item instanceof Code code)) {
435+
return false;
436+
}
437+
var language = code.getLanguage();
438+
if (language == null) {
439+
return false;
440+
}
441+
var normalized = language.trim().toLowerCase(Locale.ROOT);
442+
if (normalized.isEmpty()) {
443+
return false;
444+
}
445+
var primaryToken = normalized.split("\\s+", 2)[0];
446+
return primaryToken.startsWith("mermaid");
447+
}
448+
396449
private void handleHeaderOnCompletion(ResponseEditorPanel editorPanel) {
397450
var editor = editorPanel.getEditor();
398451
if (editor != null) {

‎src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/CompleteMessageParser.kt‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ class CompleteMessageParser : MessageParser {
77

88
companion object {
99
private val CODE_BLOCK_PATTERN: Pattern =
10-
Pattern.compile("```([a-zA-Z0-9_+-]*)(?::([^\\n]*))?\\n(.*?)```", Pattern.DOTALL)
10+
Pattern.compile("(`{3,})([^\\n`]*)\\n(.*?)\\n\\1", Pattern.DOTALL)
1111
private val SEARCH_REPLACE_PATTERN: Pattern =
1212
Pattern.compile(
1313
"<<<<<<< SEARCH\\n(.*?)\\n=======\\n(.*?)\\n>>>>>>> REPLACE",
@@ -18,8 +18,7 @@ class CompleteMessageParser : MessageParser {
1818

1919
private const val THINK_OPEN_TAG = "<think>"
2020
private const val THINK_CLOSE_TAG = "</think>\n\n"
21-
private const val LANGUAGE_GROUP_INDEX = 1
22-
private const val FILE_PATH_GROUP_INDEX = 2
21+
private const val CODE_HEADER_GROUP_INDEX = 2
2322
private const val CODE_CONTENT_GROUP_INDEX = 3
2423
private const val SEARCH_CONTENT_GROUP_INDEX = 1
2524
private const val REPLACE_CONTENT_GROUP_INDEX = 2
@@ -105,12 +104,11 @@ class CompleteMessageParser : MessageParser {
105104
* Processes a code block and adds all related segments.
106105
*/
107106
private fun MutableList<Segment>.addCodeBlockSegments(codeBlockMatcher: Matcher) {
108-
val language = codeBlockMatcher.group(LANGUAGE_GROUP_INDEX).orEmpty()
109-
val filePath = codeBlockMatcher.group(FILE_PATH_GROUP_INDEX)
107+
val header = parseCodeHeader(codeBlockMatcher.group(CODE_HEADER_GROUP_INDEX).orEmpty())
110108
val codeContent = codeBlockMatcher.group(CODE_CONTENT_GROUP_INDEX).orEmpty()
111109

112-
add(CodeHeader(language, filePath))
113-
processCodeContent(codeContent, language, filePath)
110+
add(CodeHeader(header.language, header.filePath))
111+
processCodeContent(codeContent, header.language, header.filePath)
114112
add(CodeEnd(codeContent))
115113
}
116114

@@ -360,4 +358,18 @@ class CompleteMessageParser : MessageParser {
360358
val segments: List<Segment>,
361359
val lastIndex: Int
362360
)
361+
362+
private fun parseCodeHeader(headerText: String): CodeHeader {
363+
val parts = headerText.trim().split(":", limit = 2)
364+
val rawLanguage = parts.getOrNull(0).orEmpty()
365+
val normalizedLanguage = rawLanguage
366+
.trim()
367+
.trimStart('`')
368+
.takeWhile { !it.isWhitespace() }
369+
val filePath = parts.getOrNull(1)?.trim()?.ifEmpty { null }
370+
return CodeHeader(
371+
language = normalizedLanguage,
372+
filePath = filePath
373+
)
374+
}
363375
}

‎src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/parser/SseMessageParser.kt‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ class SseMessageParser : MessageParser {
121121
consumeFromBuffer(nlIdx + 1)
122122

123123
return when {
124-
line.trimEnd() == state.indentation + CODE_FENCE -> {
124+
isCodeFenceLine(line, state.indentation) -> {
125125
if (state.content.isNotEmpty()) {
126126
segments.add(Code(state.content, state.header.language, state.header.filePath))
127127
}
@@ -204,7 +204,7 @@ class SseMessageParser : MessageParser {
204204
true
205205
}
206206

207-
line.trimEnd() == state.indentation + CODE_FENCE -> {
207+
isCodeFenceLine(line, state.indentation) -> {
208208
segments.add(CodeEnd(""))
209209
parserState = ParserState.Outside
210210
true
@@ -277,7 +277,7 @@ class SseMessageParser : MessageParser {
277277
is ParserState.InCode -> {
278278
val segments = mutableListOf<Segment>()
279279

280-
if (buffer.toString().trimEnd() == state.indentation + CODE_FENCE) {
280+
if (isCodeFenceLine(buffer.toString(), state.indentation)) {
281281
if (state.content.isNotBlank()) {
282282
segments.add(
283283
Code(state.content, state.header.language, state.header.filePath)
@@ -340,13 +340,35 @@ class SseMessageParser : MessageParser {
340340
private fun parseCodeHeader(headerText: String): CodeHeader? {
341341
val parts = headerText.split(HEADER_DELIMITER, limit = HEADER_PARTS_LIMIT)
342342
return if (parts.isNotEmpty()) {
343+
val rawLanguage = parts.getOrNull(0).orEmpty()
344+
val normalizedLanguage = rawLanguage
345+
.trim()
346+
.trimStart('`')
347+
.takeWhile { !it.isWhitespace() }
348+
343349
CodeHeader(
344-
language = parts.getOrNull(0) ?: "",
350+
language = normalizedLanguage,
345351
filePath = parts.getOrNull(1)?.trim()?.ifEmpty { null }
346352
)
347353
} else null
348354
}
349355

356+
private fun isCodeFenceLine(line: String, indentation: String): Boolean {
357+
val trimmedEnd = line.trimEnd()
358+
if (!trimmedEnd.startsWith(indentation)) {
359+
return false
360+
}
361+
val afterIndent = trimmedEnd.substring(indentation.length)
362+
if (afterIndent.isEmpty() || afterIndent.first() != '`') {
363+
return false
364+
}
365+
if (!afterIndent.all { it == '`' || it.isWhitespace() }) {
366+
return false
367+
}
368+
val fenceCandidate = afterIndent.trim()
369+
return fenceCandidate.length >= CODE_FENCE.length && fenceCandidate.all { it == '`' }
370+
}
371+
350372
private fun isSearchStartLine(line: String): Boolean {
351373
val trimmed = line.trim()
352374
return trimmed.startsWith(SEARCH_MARKER) || SEARCH_START_REGEX.matches(trimmed)

‎src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/MermaidResponseHtmlTemplate.kt‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
package ee.carlrobert.codegpt.toolwindow.chat.ui
2+
3+
import ee.carlrobert.codegpt.util.file.FileUtil
4+
import java.nio.file.Files
5+
import java.nio.file.StandardCopyOption
6+
7+
data class MermaidThemePalette(
8+
val foreground: String,
9+
val mutedForeground: String,
10+
val hudBackground: String,
11+
val hudBorder: String,
12+
val buttonBackground: String,
13+
val buttonBackgroundHover: String,
14+
val buttonBorder: String,
15+
val error: String,
16+
val fontFamily: String,
17+
val fontSize: String,
18+
val mermaidTheme: String,
19+
)
20+
21+
object MermaidResponseHtmlTemplate {
22+
23+
private const val DIAGRAM_PLACEHOLDER = "__MERMAID_SOURCE__"
24+
private const val DIAGRAM_TYPE_LABEL_PLACEHOLDER = "__DIAGRAM_TYPE_LABEL__"
25+
private const val NODE_TARGETS_PLACEHOLDER = "__NODE_TARGETS_JSON__"
26+
private const val NODE_CLICK_HANDLER_PLACEHOLDER = "__NODE_CLICK_HANDLER__"
27+
private const val RENDER_ERROR_HANDLER_PLACEHOLDER = "__RENDER_ERROR_HANDLER__"
28+
private const val THEME_FOREGROUND_PLACEHOLDER = "__THEME_FOREGROUND__"
29+
private const val THEME_MUTED_FOREGROUND_PLACEHOLDER = "__THEME_MUTED_FOREGROUND__"
30+
private const val THEME_HUD_BACKGROUND_PLACEHOLDER = "__THEME_HUD_BACKGROUND__"
31+
private const val THEME_HUD_BORDER_PLACEHOLDER = "__THEME_HUD_BORDER__"
32+
private const val THEME_BUTTON_BACKGROUND_PLACEHOLDER = "__THEME_BUTTON_BACKGROUND__"
33+
private const val THEME_BUTTON_BACKGROUND_HOVER_PLACEHOLDER =
34+
"__THEME_BUTTON_BACKGROUND_HOVER__"
35+
private const val THEME_BUTTON_BORDER_PLACEHOLDER = "__THEME_BUTTON_BORDER__"
36+
private const val THEME_ERROR_PLACEHOLDER = "__THEME_ERROR__"
37+
private const val THEME_FONT_FAMILY_PLACEHOLDER = "__THEME_FONT_FAMILY__"
38+
private const val THEME_FONT_SIZE_PLACEHOLDER = "__THEME_FONT_SIZE__"
39+
private const val MERMAID_THEME_PLACEHOLDER = "__MERMAID_THEME__"
40+
private const val MERMAID_SCRIPT_URL_PLACEHOLDER = "__MERMAID_SCRIPT_URL__"
41+
42+
private const val TEMPLATE_RESOURCE_PATH = "/templates/mermaid-response-panel.html"
43+
private const val BUNDLED_MERMAID_SCRIPT_RESOURCE_PATH = "/web/mermaid/mermaid.min.js"
44+
45+
private val placeholderRegex = Regex("__[A-Z0-9_]+__")
46+
47+
private val templateHtml: String by lazy {
48+
FileUtil.getResourceContent(TEMPLATE_RESOURCE_PATH)
49+
.ifBlank { error("Missing Mermaid response template at $TEMPLATE_RESOURCE_PATH") }
50+
}
51+
private val bundledMermaidScriptUrl: String by lazy {
52+
val stream = MermaidResponseHtmlTemplate::class.java
53+
.getResourceAsStream(BUNDLED_MERMAID_SCRIPT_RESOURCE_PATH)
54+
?: error("Missing bundled Mermaid script at $BUNDLED_MERMAID_SCRIPT_RESOURCE_PATH")
55+
56+
stream.use {
57+
val tempFile = Files.createTempFile("proxyai-mermaid-", ".min.js")
58+
tempFile.toFile().deleteOnExit()
59+
Files.copy(it, tempFile, StandardCopyOption.REPLACE_EXISTING)
60+
tempFile.toUri().toString()
61+
}
62+
}
63+
64+
fun render(source: String, diagramTypeLabel: String, theme: MermaidThemePalette): String {
65+
val replacements = mapOf(
66+
DIAGRAM_PLACEHOLDER to toJavaScriptStringLiteral(source),
67+
DIAGRAM_TYPE_LABEL_PLACEHOLDER to toJavaScriptStringLiteral(diagramTypeLabel),
68+
NODE_TARGETS_PLACEHOLDER to "{}",
69+
NODE_CLICK_HANDLER_PLACEHOLDER to "void 0",
70+
RENDER_ERROR_HANDLER_PLACEHOLDER to "void 0",
71+
THEME_FOREGROUND_PLACEHOLDER to theme.foreground,
72+
THEME_MUTED_FOREGROUND_PLACEHOLDER to theme.mutedForeground,
73+
THEME_HUD_BACKGROUND_PLACEHOLDER to theme.hudBackground,
74+
THEME_HUD_BORDER_PLACEHOLDER to theme.hudBorder,
75+
THEME_BUTTON_BACKGROUND_PLACEHOLDER to theme.buttonBackground,
76+
THEME_BUTTON_BACKGROUND_HOVER_PLACEHOLDER to theme.buttonBackgroundHover,
77+
THEME_BUTTON_BORDER_PLACEHOLDER to theme.buttonBorder,
78+
THEME_ERROR_PLACEHOLDER to theme.error,
79+
THEME_FONT_FAMILY_PLACEHOLDER to theme.fontFamily,
80+
THEME_FONT_SIZE_PLACEHOLDER to theme.fontSize,
81+
MERMAID_THEME_PLACEHOLDER to theme.mermaidTheme,
82+
MERMAID_SCRIPT_URL_PLACEHOLDER to toJavaScriptStringLiteral(bundledMermaidScriptUrl),
83+
)
84+
85+
return placeholderRegex.replace(templateHtml) { matchResult ->
86+
replacements[matchResult.value] ?: matchResult.value
87+
}
88+
}
89+
90+
private fun toJavaScriptStringLiteral(value: String): String {
91+
val escaped = value
92+
.replace("\\", "\\\\")
93+
.replace("\"", "\\\"")
94+
.replace("<", "\\x3C")
95+
.replace(">", "\\x3E")
96+
.replace("&", "\\x26")
97+
.replace("\r", "\\r")
98+
.replace("\n", "\\n")
99+
.replace("\u2028", "\\u2028")
100+
.replace("\u2029", "\\u2029")
101+
102+
return "\"$escaped\""
103+
}
104+
}

0 commit comments

Comments
 (0)