Let's Mention Mermaids (#1838)

* updated system prompt

* Update system.ts

---------

Co-authored-by: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>

.changeset/cyan-bags-work.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": minor
+---
+
+Add support for rendering mermaid graphs in the chat.

src/core/prompts/system.ts

@@ -865,9 +865,10 @@ In each user message, the environment_details will specify the current mode. The
 ## What is PLAN MODE?
 
 - While you are usually in ACT MODE, the user may switch to PLAN MODE in order to have a back and forth with you to plan how to best accomplish the task. 
-- When starting in PLAN MODE, depending on the user's request, you may need to do some information gathering e.g. using read_file or search_files to get more context about the task. You may also ask the user clarifying questions to get a better understanding of the task.
-- Once you've gained more context about the user's request, you should architect a detailed plan for how you will accomplish the task.
+- When starting in PLAN MODE, depending on the user's request, you may need to do some information gathering e.g. using read_file or search_files to get more context about the task. You may also ask the user clarifying questions to get a better understanding of the task. You may return mermaid diagrams to visually display your understanding.
+- Once you've gained more context about the user's request, you should architect a detailed plan for how you will accomplish the task. Returning mermaid diagrams may be helpful here as well.
 - Then you might ask the user if they are pleased with this plan, or if they would like to make any changes. Think of this as a brainstorming session where you can discuss the task and plan the best way to accomplish it.
+- If at any point a mermaid diagram would make your plan clearer to help the user quickly see the structure, you are encouraged to include a Mermaid code block in the response.
 - Finally once it seems like you've reached a good plan, ask the user to switch you back to ACT MODE to implement the solution.
 
 ====

webview-ui/package.json

@@ -9,6 +9,7 @@
 		"fast-deep-equal": "^3.1.3",
 		"fuse.js": "^7.0.0",
 		"fzf": "^0.5.2",
+		"mermaid": "^11.4.1",
 		"pretty-bytes": "^6.1.1",
 		"react": "^18.3.1",
 		"react-dom": "^18.3.1",

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

@@ -1,10 +1,11 @@
-import { memo, useEffect } from "react"
+import React, { memo, useEffect } from "react"
 import { useRemark } from "react-remark"
 import rehypeHighlight, { Options } from "rehype-highlight"
 import styled from "styled-components"
 import { visit } from "unist-util-visit"
 import { useExtensionState } from "../../context/ExtensionStateContext"
 import { CODE_BLOCK_BG_COLOR } from "./CodeBlock"
+import MermaidBlock from "./MermaidBlock"
 
 interface MarkdownBlockProps {
 	markdown?: string
@@ -220,7 +221,27 @@ const MarkdownBlock = memo(({ markdown }: MarkdownBlockProps) => {
 		],
 		rehypeReactOptions: {
 			components: {
-				pre: ({ node, ...preProps }: any) => <StyledPre {...preProps} theme={theme} />,
+				pre: ({ node, children, ...preProps }: any) => {
+					if (Array.isArray(children) && children.length === 1 && React.isValidElement(children[0])) {
+						const child = children[0] as React.ReactElement<{ className?: string }>
+						if (child.props?.className?.includes("language-mermaid")) {
+							return child
+						}
+					}
+					return (
+						<StyledPre {...preProps} theme={theme}>
+							{children}
+						</StyledPre>
+					)
+				},
+				code: (props: any) => {
+					const className = props.className || ""
+					if (className.includes("language-mermaid")) {
+						const codeText = String(props.children || "")
+						return <MermaidBlock code={codeText} />
+					}
+					return <code {...props} />
+				},
 			},
 		},
 	})

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

@@ -0,0 +1,176 @@
+import { useEffect, useRef, useState } from "react"
+import mermaid from "mermaid"
+import { useDebounceEffect } from "../../utils/useDebounceEffect"
+import styled from "styled-components"
+import { vscode } from "../../utils/vscode"
+
+const MERMAID_THEME = {
+	background: "#1e1e1e",
+	textColor: "#ffffff",
+	mainBkg: "#2d2d2d",
+	lineColor: "#cccccc",
+	primaryColor: "#3c3c3c",
+}
+
+mermaid.initialize({
+	startOnLoad: false,
+	securityLevel: "loose",
+	theme: "dark",
+	themeVariables: {
+		background: MERMAID_THEME.background,
+		textColor: MERMAID_THEME.textColor,
+		mainBkg: MERMAID_THEME.mainBkg,
+		lineColor: MERMAID_THEME.lineColor,
+		fontSize: "16px",
+		primaryColor: MERMAID_THEME.primaryColor,
+	},
+})
+
+interface MermaidBlockProps {
+	code: string
+}
+
+export default function MermaidBlock({ code }: MermaidBlockProps) {
+	const containerRef = useRef<HTMLDivElement>(null)
+	const [isLoading, setIsLoading] = useState(false)
+
+	// 1) Whenever `code` changes, mark that we need to re-render a new chart
+	useEffect(() => {
+		setIsLoading(true)
+	}, [code])
+
+	// 2) Debounce the actual parse/render
+	useDebounceEffect(
+		() => {
+			if (containerRef.current) {
+				containerRef.current.innerHTML = ""
+			}
+			mermaid
+				.parse(code, { suppressErrors: true })
+				.then((isValid) => {
+					if (!isValid) {
+						throw new Error("Invalid or incomplete Mermaid code")
+					}
+					const id = `mermaid-${Math.random().toString(36).substring(2)}`
+					return mermaid.render(id, code)
+				})
+				.then(({ svg }) => {
+					if (containerRef.current) {
+						containerRef.current.innerHTML = svg
+					}
+				})
+				.catch((err) => {
+					console.warn("Mermaid parse/render failed:", err)
+					containerRef.current!.innerHTML = code.replace(/</g, "&lt;").replace(/>/g, "&gt;")
+				})
+				.finally(() => {
+					setIsLoading(false)
+				})
+		},
+		500, // Delay 500ms
+		[code], // Dependencies for scheduling
+	)
+
+	/**
+	 * Called when user clicks the rendered diagram.
+	 * Converts the <svg> to a PNG and sends it to the extension.
+	 */
+	const handleClick = async () => {
+		if (!containerRef.current) return
+		const svgEl = containerRef.current.querySelector("svg")
+		if (!svgEl) return
+
+		try {
+			const pngDataUrl = await svgToPng(svgEl)
+			vscode.postMessage({
+				type: "openImage",
+				text: pngDataUrl,
+			})
+		} catch (err) {
+			console.error("Error converting SVG to PNG:", err)
+		}
+	}
+
+	return (
+		<MermaidBlockContainer>
+			{isLoading && <LoadingMessage>Creating mermaid chart...</LoadingMessage>}
+
+			{/* The container for the final <svg> or raw code. */}
+			<SvgContainer onClick={handleClick} ref={containerRef} $isLoading={isLoading} />
+		</MermaidBlockContainer>
+	)
+}
+
+async function svgToPng(svgEl: SVGElement): Promise<string> {
+	console.log("svgToPng function called")
+	// Clone the SVG to avoid modifying the original
+	const svgClone = svgEl.cloneNode(true) as SVGElement
+
+	// Get the original viewBox
+	const viewBox = svgClone.getAttribute("viewBox")?.split(" ").map(Number) || []
+	const originalWidth = viewBox[2] || svgClone.clientWidth
+	const originalHeight = viewBox[3] || svgClone.clientHeight
+
+	// Calculate the scale factor to fit editor width while maintaining aspect ratio
+
+	// Unless we can find a way to get the actual editor window dimensions through the VS Code API (which might be possible but would require changes to the extension side),
+	// the fixed 1200px width seems like a reliable approach.
+	const editorWidth = 1200
+
+	const scale = editorWidth / originalWidth
+	const scaledHeight = originalHeight * scale
+
+	// Update SVG dimensions
+	svgClone.setAttribute("width", `${editorWidth}`)
+	svgClone.setAttribute("height", `${scaledHeight}`)
+
+	const serializer = new XMLSerializer()
+	const svgString = serializer.serializeToString(svgClone)
+	const svgDataUrl = "data:image/svg+xml;base64," + btoa(decodeURIComponent(encodeURIComponent(svgString)))
+
+	return new Promise((resolve, reject) => {
+		const img = new Image()
+		img.onload = () => {
+			const canvas = document.createElement("canvas")
+			canvas.width = editorWidth
+			canvas.height = scaledHeight
+
+			const ctx = canvas.getContext("2d")
+			if (!ctx) return reject("Canvas context not available")
+
+			// Fill background with Mermaid's dark theme background color
+			ctx.fillStyle = MERMAID_THEME.background
+			ctx.fillRect(0, 0, canvas.width, canvas.height)
+
+			ctx.imageSmoothingEnabled = true
+			ctx.imageSmoothingQuality = "high"
+
+			ctx.drawImage(img, 0, 0, editorWidth, scaledHeight)
+			resolve(canvas.toDataURL("image/png", 1.0))
+		}
+		img.onerror = reject
+		img.src = svgDataUrl
+	})
+}
+
+const MermaidBlockContainer = styled.div`
+	position: relative;
+	margin: 8px 0;
+`
+
+const LoadingMessage = styled.div`
+	padding: 8px 0;
+	color: var(--vscode-descriptionForeground);
+	font-style: italic;
+	font-size: 0.9em;
+`
+
+interface SvgContainerProps {
+	$isLoading: boolean
+}
+
+const SvgContainer = styled.div<SvgContainerProps>`
+	opacity: ${(props) => (props.$isLoading ? 0.3 : 1)};
+	min-height: 20px;
+	transition: opacity 0.2s ease;
+`

webview-ui/src/utils/useDebounceEffect.ts

@@ -0,0 +1,42 @@
+import { useEffect, useRef } from "react"
+
+type VoidFn = () => void
+
+/**
+ * Runs `effectRef.current()` after `delay` ms whenever any of the `deps` change,
+ * but cancels/re-schedules if they change again before the delay.
+ */
+export function useDebounceEffect(effect: VoidFn, delay: number, deps: any[]) {
+	const callbackRef = useRef<VoidFn>(effect)
+	const timeoutRef = useRef<NodeJS.Timeout | null>(null)
+
+	// Keep callbackRef current
+	useEffect(() => {
+		callbackRef.current = effect
+	}, [effect])
+
+	useEffect(() => {
+		// Clear any queued call
+		if (timeoutRef.current) {
+			clearTimeout(timeoutRef.current)
+		}
+
+		// Schedule a new call
+		timeoutRef.current = setTimeout(() => {
+			// always call the *latest* version of effect
+			callbackRef.current()
+		}, delay)
+
+		// Cleanup on unmount or next effect
+		return () => {
+			if (timeoutRef.current) {
+				clearTimeout(timeoutRef.current)
+			}
+		}
+
+		// We want to re‐schedule if any item in `deps` changed,
+		// or if `delay` changed.
+
+		// eslint-disable-next-line react-hooks/exhaustive-deps
+	}, [delay, ...deps])
+}