Merge pull request #1202 from mjbvz/rising-toucan

Add chat output renderer proposed API sample

.scripts/samples.ts

@@ -294,6 +294,7 @@ export const samples: Sample[] = [
 	{ description: 'authenticationprovider-sample', excludeFromReadme: true, path: 'authenticationprovider-sample', guide: null, apis: [], contributions: [] },
 	{ description: 'configuration-sample', excludeFromReadme: true, path: 'configuration-sample', guide: null, apis: [], contributions: [] },
 	{ description: 'chat-model-provider-sample', excludeFromReadme: true, path: 'chat-model-provider-sample', guide: null, apis: [], contributions: [] },
+	{ description: 'chat-output-renderer-sample', excludeFromReadme: true, path: 'chat-output-renderer-sample', guide: null, apis: [], contributions: [] },
 	{ description: 'contentprovider-sample', excludeFromReadme: true, path: 'contentprovider-sample', guide: null, apis: [], contributions: [] },
 	{ description: 'diagnostic-related-information-sample', excludeFromReadme: true, path: 'diagnostic-related-information-sample', guide: null, apis: [], contributions: [] },
 	{ description: 'document-paste', excludeFromReadme: true, path: 'document-paste', guide: null, apis: [], contributions: [] },

chat-output-renderer-sample/.gitignore

@@ -0,0 +1,5 @@
+out
+node_modules
+.vscode-test/
+*.vsix
+vscode*.d.ts

chat-output-renderer-sample/.vscode/launch.json

@@ -0,0 +1,21 @@
+// A launch configuration that compiles the extension and then opens it inside a new window
+// Use IntelliSense to learn about possible attributes.
+// Hover to view descriptions of existing attributes.
+// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+{
+	"version": "0.2.0",
+	"configurations": [{
+			"name": "Run Extension",
+			"type": "extensionHost",
+			"request": "launch",
+			"runtimeExecutable": "${execPath}",
+			"args": [
+				"--extensionDevelopmentPath=${workspaceFolder}"
+			],
+			"outFiles": [
+				"${workspaceFolder}/out/**/*.js"
+			],
+			"preLaunchTask": "npm: watch"
+		}
+	]
+}

chat-output-renderer-sample/.vscode/tasks.json

@@ -0,0 +1,20 @@
+// See https://go.microsoft.com/fwlink/?LinkId=733558
+// for the documentation about the tasks.json format
+{
+	"version": "2.0.0",
+	"tasks": [
+		{
+			"type": "npm",
+			"script": "watch",
+			"problemMatcher": "$tsc-watch",
+			"isBackground": true,
+			"presentation": {
+				"reveal": "never"
+			},
+			"group": {
+				"kind": "build",
+				"isDefault": true
+			}
+		}
+	]
+}

chat-output-renderer-sample/README.md

@@ -0,0 +1,12 @@
+# Chat Output Renderer sample
+
+This VS Code extension sample demonstrates usage of the proposed chat output renderer API. This API allows extensions to
+contribute custom rendered widgets into VS Code's chat interface. Language models can invoke tools to create these widgets. The widgets are rendered using VS Code's webview API.
+
+## Running the Sample
+
+- Make sure you are using VS Code 1.103 or newer.
+- Run `npm install` in terminal to install dependencies
+- Run the `Run Extension` target in the Debug View. This will:
+	- Start a task `npm: watch` to compile the code
+	- Run the extension in a new VS Code window

chat-output-renderer-sample/eslint.config.mjs

@@ -0,0 +1,44 @@
+/**
+ * ESLint configuration for the project.
+ * 
+ * See https://eslint.style and https://typescript-eslint.io for additional linting options.
+ */
+// @ts-check
+import js from '@eslint/js';
+import tseslint from 'typescript-eslint';
+import stylistic from '@stylistic/eslint-plugin';
+
+export default tseslint.config(
+	{
+		ignores: [
+			'.vscode-test',
+			'out',
+		]
+	},
+	js.configs.recommended,
+	...tseslint.configs.recommended,
+	...tseslint.configs.stylistic,
+	{
+		plugins: {
+			'@stylistic': stylistic
+		},
+		rules: {
+			'curly': 'warn',
+			'@stylistic/semi': ['warn', 'always'],
+			'@typescript-eslint/no-empty-function': 'off',
+			'@typescript-eslint/naming-convention': [
+				'warn',
+				{
+					'selector': 'import',
+					'format': ['camelCase', 'PascalCase']
+				}
+			],
+			'@typescript-eslint/no-unused-vars': [
+				'error',
+				{
+					'argsIgnorePattern': '^_'
+				}
+			]
+		}
+	}
+);

chat-output-renderer-sample/package.json

@@ -0,0 +1,89 @@
+{
+	"name": "chat-output-renderer-sample",
+	"displayName": "chat-output-renderer-sample",
+	"description": "Sample chat output renderer extension that renders Mermaid diagrams.",
+	"version": "0.0.1",
+	"publisher": "vscode-samples",
+	"repository": "https://github.com/Microsoft/vscode-extension-samples/helloworld-sample",
+	"engines": {
+		"vscode": "^1.103.0"
+	},
+	"categories": [
+		"Other"
+	],
+	"enabledApiProposals": [
+		"chatOutputRenderer"
+	],
+	"activationEvents": [],
+	"main": "./out/extension.js",
+	"contributes": {
+		"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."
+						}
+					}
+				}
+			},
+			{
+				"name": "createMermaidDiagram",
+				"displayName": "Mermaid Creator",
+				"toolReferenceName": "createMermaidDiagram",
+				"canBeReferencedInPrompt": true,
+				"modelDescription": "Creates a Mermaid diagram from a description and renders for the user.",
+				"userDescription": "Creates and renders Mermaid.js diagrams.",
+				"inputSchema": {
+					"type": "object",
+					"properties": {
+						"description": {
+							"type": "string",
+							"description": "A plain text description of the Mermaid diagram to create. The description should be detailed enough for the model to generate a valid Mermaid diagram markup."
+						}
+					}
+				}
+			}
+		],
+		"chatOutputRenderer": [
+			{
+				"mimeTypes": [
+					"application/vnd.chat-output-renderer.mermaid"
+				]
+			}
+		]
+	},
+	"scripts": {
+		"vscode:prepublish": "npm run compile",
+		"compile": "tsc -p ./",
+		"lint": "eslint",
+		"watch": "tsc -watch -p ./",
+		"download-api": "dts dev",
+		"postdownload-api": "dts main",
+		"postinstall": "npm run download-api"
+	},
+	"devDependencies": {
+		"@eslint/js": "^9.13.0",
+		"@stylistic/eslint-plugin": "^2.9.0",
+		"@types/jsdom": "^21.1.7",
+		"@types/node": "^20",
+		"@types/vscode": "^1.73.0",
+		"@vscode/dts": "^0.4.0",
+		"eslint": "^9.13.0",
+		"typescript": "^5.8.2",
+		"typescript-eslint": "^8.16.0"
+	},
+	"dependencies": {
+		"dompurify": "^3.2.6",
+		"jsdom": "^26.1.0",
+		"mermaid": "^11.9.0"
+	}
+}

chat-output-renderer-sample/src/extension.ts

@@ -0,0 +1,295 @@
+import * as vscode from 'vscode';
+import { JSDOM } from 'jsdom';
+import * as DOMPurify from 'dompurify';
+
+/**
+ * Mime type used to identify Mermaid diagram data in chat output.
+ */
+const mime = 'application/vnd.chat-output-renderer.mermaid';
+
+const maxFixAttempts = 3;
+
+export function activate(context: vscode.ExtensionContext) {
+
+	// Register our tools
+
+	// The first tool takes mermaid markup as input
+	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);
+			},
+		})
+	);
+
+	// The second tool generates mermaid markup based on a description
+	context.subscriptions.push(
+		vscode.lm.registerTool<{ description: string }>('createMermaidDiagram', {
+			invoke: async (options, token) => {
+				const description = options.input.description;
+
+				let sourceCode = await generateMermaidDiagram(description, token);
+				if (!sourceCode) {
+					throw new Error('Failed to generate Mermaid diagram from description');
+				}
+
+				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(mime, {
+			async renderChatOutput(data, webview, _ctx, _token) {
+				const mermaidSource = new TextDecoder().decode(data);
+
+				// Set the options for the webview
+				const mermaidDist = vscode.Uri.joinPath(context.extensionUri, 'node_modules', 'mermaid', 'dist');
+				webview.options = {
+					enableScripts: true,
+					localResourceRoots: [mermaidDist],
+				};
+
+				// Set the HTML content for the webview
+				const nonce = getNonce();
+				const mermaidEsmUri = vscode.Uri.joinPath(mermaidDist, 'mermaid.esm.mjs');
+
+				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}">
+							import mermaid from '${escapeForScriptBlock(webview.asWebviewUri(mermaidEsmUri).toString())}';
+							mermaid.initialize({ startOnLoad: true });
+						</script>
+					</body>
+					</html>`;
+			},
+		}));
+}
+
+/**
+ * Lazily load mermaid
+ */
+const getMermaidInstance = (() => {
+	const createMermaidInstance = async () => {
+		// Patch the global window object for mermaid
+
+		const { window } = new JSDOM("");
+		// eslint-disable-next-line @typescript-eslint/no-explicit-any
+		(global as any).window = window;
+		// eslint-disable-next-line @typescript-eslint/no-explicit-any
+		(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();
+}
+
+/**
+ * Uses a language model to generate Mermaid markup based on a description of the diagram.
+ */
+async function generateMermaidDiagram(description: string, token: vscode.CancellationToken): Promise<string | undefined> {
+	const model = await getPreferredLm();
+	if (!model) {
+		throw new Error('No suitable model found for generating Mermaid diagram');
+	}
+
+	const completion = await model.sendRequest([
+		vscode.LanguageModelChatMessage.Assistant(joinLines(
+			`The user will provide you with a description for a Mermaid diagram.`,
+			`Your task is to generate the Mermaid source code based on the description.`,
+			`Please return the 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(
+			`Please create a Mermaid diagram based on this description:`,
+			`${description}`,
+		)),
+	], {}, token);
+
+	return parseMermaidMarkupFromChatResponse(completion, token);
+
+}
+
+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 escapeForScriptBlock(str: string): string {
+	return str
+		.replace(/\\/g, '\\\\')
+		.replace(/'/g, "\\'")
+		.replace(/"/g, '\\"')
+		.replace(/\r/g, '\\r')
+		.replace(/\n/g, '\\n')
+		.replace(/<\/script>/gi, '<\\/script>');
+}
+
+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;
+}

chat-output-renderer-sample/tsconfig.json

@@ -0,0 +1,18 @@
+{
+	"compilerOptions": {
+		"module": "commonjs",
+		"target": "es2020",
+		"lib": [
+			"es2020"
+		],
+		"outDir": "out",
+		"sourceMap": true,
+		"strict": true,
+		"rootDir": "src",
+		"skipLibCheck": true
+	},
+	"exclude": [
+		"node_modules",
+		".vscode-test"
+	]
+}