Add basic mermaid rendering support in core

For #257761

Ports over extension sample + a few improvements to core

.eslint-ignore

@@ -10,6 +10,7 @@
 **/extensions/markdown-language-features/media/**
 **/extensions/markdown-language-features/notebook-out/**
 **/extensions/markdown-math/notebook-out/**
+**/extensions/mermaid-chat-features/chat-webview-out/**
 **/extensions/notebook-renderers/renderer-out/index.js
 **/extensions/simple-browser/media/index.js
 **/extensions/terminal-suggest/src/completions/upstream/**

build/gulpfile.extensions.js

@@ -51,6 +51,7 @@ const compilations = [
 	'extensions/markdown-math/tsconfig.json',
 	'extensions/media-preview/tsconfig.json',
 	'extensions/merge-conflict/tsconfig.json',
+	'extensions/mermaid-chat-features/tsconfig.json',
 	'extensions/terminal-suggest/tsconfig.json',
 	'extensions/microsoft-authentication/tsconfig.json',
 	'extensions/notebook-renderers/tsconfig.json',

build/lib/extensions.ts

@@ -559,11 +559,12 @@ const extensionsPath = path.join(root, 'extensions');
 
 // Additional projects to run esbuild on. These typically build code for webviews
 const esbuildMediaScripts = [
+	'ipynb/esbuild.mjs',
 	'markdown-language-features/esbuild-notebook.mjs',
 	'markdown-language-features/esbuild-preview.mjs',
 	'markdown-math/esbuild.mjs',
+	'mermaid-chat-features/esbuild-chat-webview.mjs',
 	'notebook-renderers/esbuild.mjs',
-	'ipynb/esbuild.mjs',
 	'simple-browser/esbuild-preview.mjs',
 ];
 

build/npm/dirs.js

@@ -33,6 +33,7 @@ const dirs = [
 	'extensions/markdown-math',
 	'extensions/media-preview',
 	'extensions/merge-conflict',
+	'extensions/mermaid-chat-features',
 	'extensions/microsoft-authentication',
 	'extensions/notebook-renderers',
 	'extensions/npm',

eslint.config.js

@@ -1410,6 +1410,7 @@ export default tseslint.config(
 	{
 		files: [
 			'extensions/markdown-language-features/**/*.ts',
+			'extensions/mermaid-chat-features/**/*.ts',
 			'extensions/media-preview/**/*.ts',
 			'extensions/simple-browser/**/*.ts',
 			'extensions/typescript-language-features/**/*.ts',
@@ -1430,6 +1431,10 @@ export default tseslint.config(
 					'extensions/simple-browser/tsconfig.json',
 					'extensions/simple-browser/preview-src/tsconfig.json',
 
+					// Mermaid chat features
+					'extensions/mermaid-chat-features/tsconfig.json',
+					'extensions/mermaid-chat-features/chat-webview-src/tsconfig.json',
+
 					// TypeScript
 					'extensions/typescript-language-features/tsconfig.json',
 					'extensions/typescript-language-features/web/tsconfig.json',

extensions/mermaid-chat-features/.gitignore

@@ -0,0 +1 @@
+chat-webview-out

extensions/mermaid-chat-features/.npmrc

@@ -0,0 +1,2 @@
+legacy-peer-deps="true"
+timeout=180000

extensions/mermaid-chat-features/.vscodeignore

@@ -0,0 +1,8 @@
+src/**
+extension.webpack.config.js
+esbuild.*
+cgmanifest.json
+package-lock.json
+webpack.config.js
+tsconfig*.json
+.gitignore

extensions/mermaid-chat-features/README.md

@@ -0,0 +1,5 @@
+# Markdown Math
+
+**Notice:** This extension is bundled with Visual Studio Code. It can be disabled but not uninstalled.
+
+Adds math rendering using [KaTeX](https://katex.org) to VS Code's built-in markdown preview and markdown cells in notebooks.

extensions/mermaid-chat-features/cgmanifest.json

@@ -0,0 +1,4 @@
+{
+	"registrations": [],
+	"version": 1
+}

extensions/mermaid-chat-features/chat-webview-src/index.ts

@@ -0,0 +1,73 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import mermaid, { MermaidConfig } from 'mermaid';
+
+function getMermaidTheme() {
+	return document.body.classList.contains('vscode-dark') || document.body.classList.contains('vscode-high-contrast')
+		? 'dark'
+		: 'default';
+}
+
+type State = {
+	readonly diagramText: string;
+	readonly theme: 'dark' | 'default';
+};
+
+let state: State | undefined = undefined;
+
+function init() {
+	const diagram = document.querySelector('.mermaid');
+	if (!diagram) {
+		return;
+	}
+
+	const theme = getMermaidTheme();
+	state = {
+		diagramText: diagram.textContent ?? '',
+		theme
+	};
+
+	const config: MermaidConfig = {
+		startOnLoad: true,
+		theme,
+	};
+	mermaid.initialize(config);
+}
+
+function tryUpdate() {
+	const newTheme = getMermaidTheme();
+	if (state?.theme === newTheme) {
+		return;
+	}
+
+	const diagramNode = document.querySelector('.mermaid');
+	if (!diagramNode || !(diagramNode instanceof HTMLElement)) {
+		return;
+	}
+
+	state = {
+		diagramText: state?.diagramText ?? '',
+		theme: newTheme
+	};
+
+	// Re-render
+	diagramNode.textContent = state?.diagramText ?? '';
+	delete diagramNode.dataset.processed;
+
+	mermaid.initialize({
+		theme: newTheme,
+	});
+	mermaid.run({
+		nodes: [diagramNode]
+	});
+}
+
+// Update when theme changes
+new MutationObserver(() => {
+	tryUpdate();
+}).observe(document.body, { attributes: true, attributeFilter: ['class'] });
+
+init();
+

extensions/mermaid-chat-features/chat-webview-src/tsconfig.json

@@ -0,0 +1,12 @@
+{
+	"extends": "../../tsconfig.base.json",
+	"compilerOptions": {
+		"outDir": "./dist/",
+		"jsx": "react",
+		"lib": [
+			"ES2024",
+			"DOM",
+			"DOM.Iterable"
+		]
+	}
+}

extensions/mermaid-chat-features/esbuild-chat-webview.mjs

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+// @ts-check
+import path from 'path';
+import { run } from '../esbuild-webview-common.mjs';
+
+const srcDir = path.join(import.meta.dirname, 'chat-webview-src');
+const outDir = path.join(import.meta.dirname, 'chat-webview-out');
+
+run({
+	entryPoints: [
+		path.join(srcDir, 'index.ts'),
+	],
+	srcDir,
+	outdir: outDir,
+}, process.argv);

extensions/mermaid-chat-features/extension-browser.webpack.config.js

@@ -0,0 +1,13 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+// @ts-check
+import { browser as withBrowserDefaults } from '../shared.webpack.config.mjs';
+
+export default withBrowserDefaults({
+	context: import.meta.dirname,
+	entry: {
+		extension: './src/extension.ts'
+	}
+});

extensions/mermaid-chat-features/extension.webpack.config.js

@@ -0,0 +1,16 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+// @ts-check
+import withDefaults from '../shared.webpack.config.mjs';
+
+export default withDefaults({
+	context: import.meta.dirname,
+	resolve: {
+		mainFields: ['module', 'main']
+	},
+	entry: {
+		extension: './src/extension.ts',
+	}
+});

extensions/mermaid-chat-features/package.json

@@ -0,0 +1,75 @@
+{
+  "name": "marmaid-chat-features",
+  "displayName": "%displayName%",
+  "description": "%description%",
+  "version": "1.0.0",
+  "publisher": "vscode",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/microsoft/vscode.git"
+  },
+  "aiKey": "0c6ae279ed8443289764825290e4f9e2-1a736e7c-1324-4338-be46-fc2a58ae4d14-7255",
+  "engines": {
+    "vscode": "^1.104.0"
+  },
+  "enabledApiProposals": [
+    "chatOutputRenderer"
+  ],
+  "capabilities": {
+    "virtualWorkspaces": true,
+    "untrustedWorkspaces": {
+      "supported": true
+    }
+  },
+  "main": "./out/extension",
+  "browser": "./dist/browser/extension",
+  "activationEvents": [],
+  "contributes": {
+    "chatOutputRenderers": [
+      {
+        "viewType": "vscode.chatMermaidDiagram",
+        "mimeTypes": [
+          "text/vnd.mermaid"
+        ]
+      }
+    ],
+    "languageModelTools": [
+      {
+        "name": "renderMermaidDiagram",
+        "displayName": "Mermaid Renderer",
+        "toolReferenceName": "renderMermaidDiagram",
+        "canBeReferencedInPrompt": true,
+        "modelDescription": "Renders a Mermaid diagram from Mermaid.js markup.",
+        "userDescription": "Render a Mermaid.js diagrams from markup.",
+        "inputSchema": {
+          "type": "object",
+          "properties": {
+            "markup": {
+              "type": "string",
+              "description": "The mermaid diagram markup to render as a Mermaid diagram. This should only be the markup of the diagram. Do not include a wrapping code block."
+            }
+          }
+        }
+      }
+    ]
+  },
+  "scripts": {
+    "compile": "gulp compile-extension:mermaid-chat-features && npm run build-chat-webview",
+    "watch": "npm run build-chat-webview && gulp watch-extension:mermaid-chat-features",
+    "vscode:prepublish": "npm run build-ext && npm run build-chat-webview",
+    "build-ext": "node ../../node_modules/gulp/bin/gulp.js --gulpfile ../../build/gulpfile.extensions.js compile-extension:mermaid-chat-features ./tsconfig.json",
+    "build-chat-webview": "node ./esbuild-chat-webview.mjs",
+    "compile-web": "npx webpack-cli --config extension-browser.webpack.config --mode none",
+    "watch-web": "npx webpack-cli --config extension-browser.webpack.config --mode none --watch --info-verbosity verbose"
+  },
+  "devDependencies": {
+    "@types/jsdom": "^21.1.7",
+    "@types/node": "^22"
+  },
+  "dependencies": {
+    "dompurify": "^3.2.6",
+    "jsdom": "^26.1.0",
+    "mermaid": "^11.11.0"
+  }
+}

extensions/mermaid-chat-features/package.nls.json

@@ -0,0 +1,6 @@
+{
+	"displayName": "Markdown Math",
+	"description": "Adds math support to Markdown in notebooks.",
+	"config.markdown.math.enabled": "Enable/disable rendering math in the built-in Markdown preview.",
+	"config.markdown.math.macros": "A collection of custom macros. Each macro is a key-value pair where the key is a new command name and the value is the expansion of the macro."
+}

extensions/mermaid-chat-features/src/extension.ts

@@ -0,0 +1,244 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import DOMPurify from 'dompurify';
+import { JSDOM } from 'jsdom';
+import * as vscode from 'vscode';
+
+/**
+ * View type that uniquely identifies the Mermaid chat output renderer.
+ */
+const viewType = 'vscode.chatMermaidDiagram';
+
+/**
+ * Mime type used to identify Mermaid diagram data in chat output.
+ */
+const mime = 'text/vnd.mermaid';
+
+const maxFixAttempts = 3;
+
+export function activate(context: vscode.ExtensionContext) {
+
+	// Register tools
+	context.subscriptions.push(
+		vscode.lm.registerTool<{ markup: string }>('renderMermaidDiagram', {
+			invoke: async (options, token) => {
+				let sourceCode = options.input.markup;
+				sourceCode = await runMermaidMarkupFixLoop(sourceCode, token);
+				return writeMermaidToolOutput(sourceCode);
+			},
+		})
+	);
+
+	// Register the chat output renderer for Mermaid diagrams.
+	// This will be invoked with the data generated by the tools.
+	// It can also be invoked when rendering old Mermaid diagrams in the chat history.
+	context.subscriptions.push(
+		vscode.chat.registerChatOutputRenderer(viewType, {
+			async renderChatOutput({ value }, webview, _ctx, _token) {
+				const mermaidSource = new TextDecoder().decode(value);
+
+				// Set the options for the webview
+				const mediaRoot = vscode.Uri.joinPath(context.extensionUri, 'chat-webview-out');
+				webview.options = {
+					enableScripts: true,
+					localResourceRoots: [mediaRoot],
+				};
+
+				// Set the HTML content for the webview
+				const nonce = getNonce();
+				const mermaidScript = vscode.Uri.joinPath(mediaRoot, 'index.js');
+
+				webview.html = `
+					<!DOCTYPE html>
+					<html lang="en">
+
+					<head>
+						<meta charset="UTF-8">
+						<meta name="viewport" content="width=device-width, initial-scale=1.0">
+						<title>Mermaid Diagram</title>
+						<meta http-equiv="Content-Security-Policy" content="default-src 'none'; script-src ${webview.cspSource} 'nonce-${nonce}'; style-src 'self' 'unsafe-inline';" />
+					</head>
+
+					<body>
+						<pre class="mermaid">
+							${escapeHtmlText(mermaidSource)}
+						</pre>
+
+						<script type="module" nonce="${nonce}" src="${webview.asWebviewUri(mermaidScript)}"></script>
+					</body>
+					</html>`;
+			},
+		}));
+}
+
+/**
+ * Lazily load mermaid
+ */
+const getMermaidInstance = (() => {
+	const createMermaidInstance = async () => {
+		// Patch the global window object for mermaid
+
+		const { window } = new JSDOM('');
+		(global as any).window = window;
+		(global as any).DOMPurify = DOMPurify(window);
+		return import('mermaid');
+	};
+
+	let cached: Promise<typeof import('mermaid')> | undefined;
+	return async (): Promise<typeof import('mermaid').default> => {
+		cached ??= createMermaidInstance();
+		return (await cached).default;
+	};
+})();
+
+/**
+ * Tries to fix mermaid syntax errors in a set number of attempts.
+ *
+ * @returns The best effort to fix the Mermaid markup.
+ */
+async function runMermaidMarkupFixLoop(sourceCode: string, token: vscode.CancellationToken): Promise<string> {
+	let attempt = 0;
+	while (attempt < maxFixAttempts) {
+		const result = await validateMermaidMarkup(sourceCode);
+		if (token.isCancellationRequested) {
+			throw new Error('Operation cancelled');
+		}
+
+		if (result.type === 'success') {
+			return sourceCode;
+		}
+
+		attempt++;
+
+		sourceCode = await tryFixingUpMermaidMarkup(sourceCode, result.message, token);
+		if (token.isCancellationRequested) {
+			throw new Error('Operation cancelled');
+		}
+	}
+
+	// Return whatever we have after max attempts
+	return sourceCode;
+}
+
+/**
+ * Validates the syntax of the provided Mermaid markup.
+ */
+async function validateMermaidMarkup(sourceCode: string): Promise<{ type: 'success' } | { type: 'error'; message: string }> {
+	try {
+		const mermaid = await getMermaidInstance();
+		await mermaid.parse(sourceCode);
+		return { type: 'success' };
+	} catch (error) {
+		if (!(error instanceof Error)) {
+			throw error;
+		}
+
+		return { type: 'error', message: error.message };
+	}
+}
+
+/**
+ * Uses a language model to try to fix Mermaid markup based on an error message.
+ */
+async function tryFixingUpMermaidMarkup(sourceCode: string, errorMessage: string, token: vscode.CancellationToken): Promise<string> {
+	const model = await getPreferredLm();
+	if (!model) {
+		console.warn('No suitable model found for fixing Mermaid markup');
+		return sourceCode;
+	}
+
+	if (token.isCancellationRequested) {
+		throw new Error('Operation cancelled');
+	}
+
+	const completion = await model.sendRequest([
+		vscode.LanguageModelChatMessage.Assistant(joinLines(
+			`The user will provide you with the source code for the Mermaid diagram and an error message.`,
+			`Your task is to fix the Mermaid source code based on the error message.`,
+			`Please return the fixed Mermaid source code inside a \`mermaid\` fenced code block. Do not add any comments or explanation.`,
+			`Make sure to return the entire source code.`
+		)),
+		vscode.LanguageModelChatMessage.User(joinLines(
+			`Here is my Mermaid source code:`,
+			``,
+			`\`\`\`mermaid`,
+			`${sourceCode}`,
+			`\`\`\``,
+			``,
+			`And here is the mermaid error message:`,
+			``,
+			errorMessage,
+		)),
+	], {}, token);
+
+	return await parseMermaidMarkupFromChatResponse(completion, token) ?? sourceCode;
+}
+
+async function parseMermaidMarkupFromChatResponse(chatResponse: vscode.LanguageModelChatResponse, token: vscode.CancellationToken): Promise<string | undefined> {
+	const parts: string[] = [];
+	for await (const line of chatResponse.text) {
+		if (token.isCancellationRequested) {
+			throw new Error('Operation cancelled');
+		}
+
+		parts.push(line);
+	}
+
+	const response = parts.join('');
+	const lines = response.split('\n');
+	if (!lines.at(0)?.startsWith('```') || !lines.at(-1)?.endsWith('```')) {
+		console.warn('Invalid response format from model, expected fenced code block');
+		return undefined;
+	}
+
+	return lines.slice(1, -1).join('\n').trim();
+}
+
+
+async function getPreferredLm(): Promise<vscode.LanguageModelChat | undefined> {
+	return (await vscode.lm.selectChatModels({ family: 'gpt-4o-mini' })).at(0)
+		?? (await vscode.lm.selectChatModels({ family: 'gpt-4o' })).at(0)
+		?? (await vscode.lm.selectChatModels({})).at(0);
+}
+
+function writeMermaidToolOutput(sourceCode: string): vscode.LanguageModelToolResult {
+	// Expose the source code as a tool result for the LM
+	const result = new vscode.LanguageModelToolResult([
+		new vscode.LanguageModelTextPart(sourceCode)
+	]);
+
+	// And store custom data in the tool result details to indicate that a custom renderer should be used for it.
+	// In this case we just store the source code as binary data.
+
+	// Add cast to use proposed API
+	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {
+		mime,
+		value: new TextEncoder().encode(sourceCode),
+	};
+
+	return result;
+}
+
+function joinLines(...lines: string[]): string {
+	return lines.join('\n');
+}
+
+function escapeHtmlText(str: string): string {
+	return str
+		.replace(/&/g, '&amp;')
+		.replace(/</g, '&lt;')
+		.replace(/>/g, '&gt;')
+		.replace(/"/g, '&quot;')
+		.replace(/'/g, '&#39;');
+}
+
+function getNonce() {
+	let text = '';
+	const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+	for (let i = 0; i < 64; i++) {
+		text += possible.charAt(Math.floor(Math.random() * possible.length));
+	}
+	return text;
+}

extensions/mermaid-chat-features/tsconfig.json

@@ -0,0 +1,15 @@
+{
+	"extends": "../tsconfig.base.json",
+	"compilerOptions": {
+		"outDir": "./out",
+		"types": []
+	},
+	"include": [
+		"src/**/*",
+		"../../src/vscode-dts/vscode.d.ts",
+		"../../src/vscode-dts/vscode.proposed.chatOutputRenderer.d.ts",
+		"../../src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts",
+		"../../src/vscode-dts/vscode.proposed.languageModelThinkingPart.d.ts",
+		"../../src/vscode-dts/vscode.proposed.chatParticipantPrivate.d.ts"
+	]
+}

extensions/simple-browser/package.json

@@ -57,10 +57,10 @@
     ]
   },
   "scripts": {
-    "compile": "gulp compile-extension:markdown-language-features && npm run build-preview",
-    "watch": "npm run build-preview && gulp watch-extension:markdown-language-features",
+    "compile": "gulp compile-extension:simple-browser && npm run build-preview",
+    "watch": "npm run build-preview && gulp watch-extension:simple-browser",
     "vscode:prepublish": "npm run build-ext && npm run build-preview",
-    "build-ext": "node ../../node_modules/gulp/bin/gulp.js --gulpfile ../../build/gulpfile.extensions.js compile-extension:markdown-language-features ./tsconfig.json",
+    "build-ext": "node ../../node_modules/gulp/bin/gulp.js --gulpfile ../../build/gulpfile.extensions.js compile-extension:simple-browser ./tsconfig.json",
     "build-preview": "node ./esbuild-preview.mjs",
     "compile-web": "npx webpack-cli --config extension-browser.webpack.config --mode none",
     "watch-web": "npx webpack-cli --config extension-browser.webpack.config --mode none --watch --info-verbosity verbose"