add a button to fix mermaid syntax errors by calling the LLM

Author: markijbema

.changeset/empty-ghosts-carry.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+add a button to fix mermaid syntax errors by calling the LLM

src/core/prompts/utilities/mermaid.ts

@@ -0,0 +1,31 @@
+/**
+ * Prompts for Mermaid diagram-related tasks
+ */
+
+/**
+ * Generate a prompt for fixing invalid Mermaid diagram syntax
+ * @param error - The error message from Mermaid parser
+ * @param invalidCode - The invalid Mermaid code that needs fixing
+ * @returns The formatted prompt for the AI to fix the Mermaid syntax
+ */
+export const mermaidFixPrompt = (error: string, invalidCode: string): string => {
+	return `You are a Mermaid diagram syntax expert. Fix the following invalid Mermaid diagram syntax and return ONLY the corrected Mermaid code without any explanations or markdown formatting.
+
+Error: ${error}
+
+Invalid Mermaid code:
+\`\`\`
+${invalidCode}
+\`\`\`
+
+Requirements:
+1. Return ONLY the corrected Mermaid syntax
+2. Do not include markdown code blocks or explanations
+3. Ensure the syntax is valid according to Mermaid specifications
+4. Enclose labels and edge label within double quotes even when you do not think it necessary to ensure the syntax is robust
+5. Do not point to multiple nodes with one edge, use multiple edges instead
+6. Preserve the original intent and structure as much as possible
+7. If the diagram type is unclear, default to a flowchart
+
+Corrected Mermaid code:`
+}

src/core/webview/webviewMessageHandler.ts

@@ -39,6 +39,7 @@ import { getModels, flushModels } from "../../api/providers/fetchers/modelCache"
 import { GetModelsOptions } from "../../shared/api"
 import { generateSystemPrompt } from "./generateSystemPrompt"
 import { getCommand } from "../../utils/commands"
+import { mermaidFixPrompt } from "../prompts/utilities/mermaid"
 
 const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"])
 
@@ -1628,6 +1629,34 @@ export const webviewMessageHandler = async (
 			}
 			break
 		}
+		case "fixMermaidSyntax":
+			if (message.text && message.requestId) {
+				try {
+					const { apiConfiguration } = await provider.getState()
+
+					const prompt = mermaidFixPrompt(message.values?.error || "Unknown syntax error", message.text)
+
+					const fixedCode = await singleCompletionHandler(apiConfiguration, prompt)
+
+					provider.postMessageToWebview({
+						type: "mermaidFixResponse",
+						requestId: message.requestId,
+						success: true,
+						fixedCode: fixedCode?.trim() || null,
+					})
+				} catch (error) {
+					const errorMessage = error instanceof Error ? error.message : "Failed to fix Mermaid syntax"
+					provider.log(`Error fixing Mermaid syntax: ${errorMessage}`)
+
+					provider.postMessageToWebview({
+						type: "mermaidFixResponse",
+						requestId: message.requestId,
+						success: false,
+						error: errorMessage,
+					})
+				}
+			}
+			break
 		case "focusPanelRequest": {
 			// Execute the focusPanel command to focus the WebView
 			await vscode.commands.executeCommand(getCommand("focusPanel"))

src/shared/ExtensionMessage.ts

@@ -99,6 +99,7 @@ export interface ExtensionMessage {
 		| "marketplaceInstallResult"
 		| "marketplaceData"
 		| "shareTaskSuccess"
+		| "mermaidFixResponse"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	action?:
@@ -148,6 +149,7 @@ export interface ExtensionMessage {
 	marketplaceItems?: MarketplaceItem[]
 	marketplaceInstalledMetadata?: MarketplaceInstalledMetadata
 	visibility?: ShareVisibility
+	fixedCode?: string | null // For mermaidFixResponse
 }
 
 export type ExtensionState = Pick<

src/shared/WebviewMessage.ts

@@ -163,6 +163,8 @@ export interface WebviewMessage {
 		| "codebaseIndexConfig"
 		| "profileThresholds"
 		| "setHistoryPreviewCollapsed"
+		| "fixMermaidSyntax"
+		| "mermaidFixResponse"
 		| "openExternal"
 		| "filterMarketplaceItems"
 		| "marketplaceButtonClicked"

webview-ui/src/components/common/MermaidBlock.tsx

@@ -5,8 +5,10 @@ import { useDebounceEffect } from "@src/utils/useDebounceEffect"
 import { vscode } from "@src/utils/vscode"
 import { useAppTranslation } from "@src/i18n/TranslationContext"
 import { useCopyToClipboard } from "@src/utils/clipboard"
+import { MermaidSyntaxFixer } from "@src/services/mermaidSyntaxFixer"
 import CodeBlock from "./CodeBlock"
 import { MermaidButton } from "@/components/common/MermaidButton"
+import { MermaidFixButton } from "@/components/common/MermaidFixButton"
 
 // Removed previous attempts at static imports for individual diagram types
 // as the paths were incorrect for Mermaid v11.4.1 and caused errors.
@@ -86,26 +88,50 @@ interface MermaidBlockProps {
 	code: string
 }
 
-export default function MermaidBlock({ code }: MermaidBlockProps) {
+export default function MermaidBlock({ code: originalCode }: MermaidBlockProps) {
 	const containerRef = useRef<HTMLDivElement>(null)
 	const [isLoading, setIsLoading] = useState(false)
 	const [error, setError] = useState<string | null>(null)
 	const [isErrorExpanded, setIsErrorExpanded] = useState(false)
+	const [svgContent, setSvgContent] = useState<string>("")
+	const [isFixing, setIsFixing] = useState(false)
+	const [code, setCode] = useState("")
 	const { showCopyFeedback, copyWithFeedback } = useCopyToClipboard()
 	const { t } = useAppTranslation()
 
 	// 1) Whenever `code` changes, mark that we need to re-render a new chart
 	useEffect(() => {
 		setIsLoading(true)
 		setError(null)
-	}, [code])
+		setCode(originalCode)
+		setIsFixing(false)
+	}, [originalCode])
+
+	const handleSyntaxFix = async () => {
+		if (isFixing) return
+
+		setIsLoading(true)
+		setIsFixing(true)
+		const result = await MermaidSyntaxFixer.autoFixSyntax(code)
+		if (result.fixedCode) {
+			// Use the improved code even if not completely successful
+			setCode(result.fixedCode)
+		}
+
+		if (!result.success) {
+			setError(result.error || t("common:mermaid.errors.fix_failed"))
+		}
+
+		setIsFixing(false)
+		setIsLoading(false)
+	}
 
 	// 2) Debounce the actual parse/render
+	// the LLM is still 'typing', and we do not want to start rendering and/or autofixing before it is fully done.
 	useDebounceEffect(
 		() => {
-			if (containerRef.current) {
-				containerRef.current.innerHTML = ""
-			}
+			if (isFixing) return
+			setIsLoading(true)
 
 			mermaid
 				.parse(code)
@@ -114,20 +140,20 @@ export default function MermaidBlock({ code }: MermaidBlockProps) {
 					return mermaid.render(id, code)
 				})
 				.then(({ svg }) => {
-					if (containerRef.current) {
-						containerRef.current.innerHTML = svg
-					}
+					setError(null)
+					setSvgContent(svg)
 				})
 				.catch((err) => {
 					console.warn("Mermaid parse/render failed:", err)
-					setError(err.message || "Failed to render Mermaid diagram")
+					const errorMessage = err instanceof Error ? err.message : t("common:mermaid.render_error")
+					setError(errorMessage)
 				})
 				.finally(() => {
 					setIsLoading(false)
 				})
 		},
 		500, // Delay 500ms
-		[code], // Dependencies for scheduling
+		[code, isFixing, originalCode, t], // Dependencies for scheduling
 	)
 
 	/**
@@ -154,7 +180,11 @@ export default function MermaidBlock({ code }: MermaidBlockProps) {
 
 	return (
 		<MermaidBlockContainer>
-			{isLoading && <LoadingMessage>{t("common:mermaid.loading")}</LoadingMessage>}
+			{isLoading && (
+				<LoadingMessage>
+					{isFixing ? t("common:mermaid.fixing_syntax") : t("common:mermaid.loading")}
+				</LoadingMessage>
+			)}
 
 			{error ? (
 				<div style={{ marginTop: "0px", overflow: "hidden", marginBottom: "8px" }}>
@@ -188,6 +218,17 @@ export default function MermaidBlock({ code }: MermaidBlockProps) {
 							<span style={{ fontWeight: "bold" }}>{t("common:mermaid.render_error")}</span>
 						</div>
 						<div style={{ display: "flex", alignItems: "center" }}>
+							{!!error && (
+								<MermaidFixButton
+									onClick={(e) => {
+										e.stopPropagation()
+										handleSyntaxFix()
+									}}
+									disabled={isFixing}
+									title={t("common:mermaid.fix_syntax_button")}>
+									<span className={`codicon codicon-${isFixing ? "loading" : "wand"}`}></span>
+								</MermaidFixButton>
+							)}
 							<CopyButton
 								onClick={(e) => {
 									e.stopPropagation()
@@ -210,12 +251,24 @@ export default function MermaidBlock({ code }: MermaidBlockProps) {
 								{error}
 							</div>
 							<CodeBlock language="mermaid" source={code} />
+							{code !== originalCode && (
+								<div style={{ marginTop: "8px" }}>
+									<div style={{ marginBottom: "4px", fontSize: "0.9em", fontWeight: "bold" }}>
+										{t("common:mermaid.original_code")}
+									</div>
+									<CodeBlock language="mermaid" source={originalCode} />
+								</div>
+							)}
 						</div>
 					)}
 				</div>
 			) : (
 				<MermaidButton containerRef={containerRef} code={code} isLoading={isLoading} svgToPng={svgToPng}>
-					<SvgContainer onClick={handleClick} ref={containerRef} $isLoading={isLoading}></SvgContainer>
+					<SvgContainer
+						onClick={handleClick}
+						$isLoading={isLoading}
+						dangerouslySetInnerHTML={{ __html: svgContent }}
+					/>
 				</MermaidButton>
 			)}
 		</MermaidBlockContainer>

webview-ui/src/components/common/MermaidFixButton.tsx

@@ -0,0 +1,38 @@
+import styled from "styled-components"
+
+export const MermaidFixButton = styled.button`
+	padding: 3px;
+	height: 24px;
+	margin-right: 4px;
+	color: var(--vscode-editor-foreground);
+	display: flex;
+	align-items: center;
+	justify-content: center;
+	background: transparent;
+	border: none;
+	cursor: pointer;
+
+	&:hover {
+		opacity: 0.8;
+		color: var(--vscode-button-foreground);
+		background: var(--vscode-button-hoverBackground);
+	}
+
+	&:disabled {
+		opacity: 0.5;
+		cursor: not-allowed;
+	}
+
+	.codicon-loading {
+		animation: spin 1s linear infinite;
+	}
+
+	@keyframes spin {
+		from {
+			transform: rotate(0deg);
+		}
+		to {
+			transform: rotate(360deg);
+		}
+	}
+`

webview-ui/src/i18n/locales/ca/common.json

@@ -17,6 +17,17 @@
 	"mermaid": {
 		"loading": "Generant diagrama mermaid...",
 		"render_error": "No es pot renderitzar el diagrama",
+		"fixing_syntax": "Corregint la sintaxi de Mermaid...",
+		"fix_syntax_button": "Corregir sintaxi amb IA",
+		"original_code": "Codi original:",
+		"errors": {
+			"unknown_syntax": "Error de sintaxi desconegut",
+			"fix_timeout": "La sol·licitud de correcció de l'IA ha esgotat el temps",
+			"fix_failed": "La correcció de l'IA ha fallat",
+			"fix_attempts": "No s'ha pogut corregir la sintaxi després de {{attempts}} intents. Últim error: {{error}}",
+			"no_fix_provided": "L'IA no ha pogut proporcionar una correcció",
+			"fix_request_failed": "La sol·licitud de correcció ha fallat"
+		},
 		"buttons": {
 			"zoom": "Zoom",
 			"zoomIn": "Ampliar",

webview-ui/src/i18n/locales/de/common.json

@@ -17,6 +17,17 @@
 	"mermaid": {
 		"loading": "Mermaid-Diagramm wird generiert...",
 		"render_error": "Diagramm kann nicht gerendert werden",
+		"fixing_syntax": "Mermaid-Syntax wird korrigiert...",
+		"fix_syntax_button": "Syntax mit KI korrigieren",
+		"original_code": "Ursprünglicher Code:",
+		"errors": {
+			"unknown_syntax": "Unbekannter Syntaxfehler",
+			"fix_timeout": "KI-Korrekturanfrage hat das Zeitlimit überschritten",
+			"fix_failed": "KI-Korrektur fehlgeschlagen",
+			"fix_attempts": "Korrektur nach {{attempts}} Versuchen fehlgeschlagen. Letzter Fehler: {{error}}",
+			"no_fix_provided": "KI konnte keine Korrektur bereitstellen",
+			"fix_request_failed": "Korrekturanfrage fehlgeschlagen"
+		},
 		"buttons": {
 			"zoom": "Zoom",
 			"zoomIn": "Vergrößern",

webview-ui/src/i18n/locales/en/common.json

@@ -17,6 +17,17 @@
 	"mermaid": {
 		"loading": "Generating mermaid diagram...",
 		"render_error": "Unable to Render Diagram",
+		"fixing_syntax": "Fixing Mermaid syntax...",
+		"fix_syntax_button": "Fix syntax with AI",
+		"original_code": "Original code:",
+		"errors": {
+			"unknown_syntax": "Unknown syntax error",
+			"fix_timeout": "LLM fix request timed out",
+			"fix_failed": "LLM fix failed",
+			"fix_attempts": "Failed to fix syntax after {{attempts}} attempts. Last error: {{error}}",
+			"no_fix_provided": "LLM failed to provide a fix",
+			"fix_request_failed": "Fix request failed"
+		},
 		"buttons": {
 			"zoom": "Zoom",
 			"zoomIn": "Zoom In",