Add a title for generated mermaid diagrams and use this in the editor

extensions/mermaid-chat-features/package.json

@@ -102,6 +102,10 @@
             "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."
+            },
+            "title": {
+              "type": "string",
+              "description": "A short title that describes the diagram."
             }
           }
         }

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

@@ -28,7 +28,9 @@ class MermaidChatOutputRenderer implements vscode.ChatOutputRenderer {
 
 	async renderChatOutput({ value }: vscode.ChatOutputDataItem, chatOutputWebview: vscode.ChatOutputWebview, _ctx: unknown, _token: vscode.CancellationToken): Promise<void> {
 		const webview = chatOutputWebview.webview;
-		const mermaidSource = new TextDecoder().decode(value);
+		const decoded = decodeMermaidData(value);
+		const mermaidSource = decoded.source;
+		const title = decoded.title;
 
 		// Generate unique ID for this webview
 		const webviewId = generateUuid();
@@ -36,7 +38,7 @@ class MermaidChatOutputRenderer implements vscode.ChatOutputRenderer {
 		const disposables: vscode.Disposable[] = [];
 
 		// Register and set as active
-		disposables.push(this._webviewManager.registerWebview(webviewId, webview, mermaidSource, 'chat'));
+		disposables.push(this._webviewManager.registerWebview(webviewId, webview, mermaidSource, title, 'chat'));
 
 		// Listen for messages from the webview
 		disposables.push(webview.onDidReceiveMessage(message => {
@@ -135,17 +137,18 @@ export function registerChatSupport(
 		vscode.commands.registerCommand('_mermaid-chat.openInEditor', (ctx?: { mermaidWebviewId?: string }) => {
 			const webviewInfo = ctx?.mermaidWebviewId ? webviewManager.getWebview(ctx.mermaidWebviewId) : webviewManager.activeWebview;
 			if (webviewInfo) {
-				editorManager.openPreview(webviewInfo.mermaidSource);
+				editorManager.openPreview(webviewInfo.mermaidSource, webviewInfo.title);
 			}
 		})
 	);
 
 	// Register lm tools
 	disposables.push(
-		vscode.lm.registerTool<{ markup: string }>('renderMermaidDiagram', {
+		vscode.lm.registerTool<{ markup: string; title?: string }>('renderMermaidDiagram', {
 			invoke: async (options, _token) => {
 				const sourceCode = options.input.markup;
-				return writeMermaidToolOutput(sourceCode);
+				const title = options.input.title;
+				return writeMermaidToolOutput(sourceCode, title);
 			},
 		})
 	);
@@ -159,19 +162,52 @@ export function registerChatSupport(
 	return vscode.Disposable.from(...disposables);
 }
 
-function writeMermaidToolOutput(sourceCode: string): vscode.LanguageModelToolResult {
-	// Expose the source code as a tool result for the LM
+function writeMermaidToolOutput(sourceCode: string, title: string | undefined): vscode.LanguageModelToolResult {
+	// Expose the source code as a markdown mermaid code block
+	const fence = getFenceForContent(sourceCode);
 	const result = new vscode.LanguageModelToolResult([
-		new vscode.LanguageModelTextPart(sourceCode)
+		new vscode.LanguageModelTextPart(`${fence}mermaid\n${sourceCode}\n${fence}`)
 	]);
 
 	// 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.
+	// Encode source and optional title as JSON.
+	const data = JSON.stringify({ source: sourceCode, title });
 	// Add cast to use proposed API
 	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {
 		mime,
-		value: new TextEncoder().encode(sourceCode),
+		value: new TextEncoder().encode(data),
 	};
 
 	return result;
 }
+
+function getFenceForContent(content: string): string {
+	const backtickMatch = content.matchAll(/`+/g);
+	if (!backtickMatch) {
+		return '```';
+	}
+
+	const maxBackticks = Math.max(...Array.from(backtickMatch, s => s[0].length));
+	return '`'.repeat(Math.max(3, maxBackticks + 1));
+}
+
+interface MermaidData {
+	readonly title: string | undefined;
+	readonly source: string;
+}
+
+function decodeMermaidData(value: Uint8Array): MermaidData {
+	const text = new TextDecoder().decode(value);
+
+	// Try to parse as JSON (new format with title), fall back to plain text (legacy format)
+	try {
+		const parsed = JSON.parse(text);
+		if (typeof parsed === 'object' && typeof parsed.source === 'string') {
+			return { title: parsed.title, source: parsed.source };
+		}
+	} catch {
+		// Not JSON, treat as legacy plain text format
+	}
+
+	return { title: undefined, source: text };
+}

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

@@ -36,7 +36,7 @@ export class MermaidEditorManager extends Disposable implements vscode.WebviewPa
 	 *
 	 * If a preview already exists for this diagram, it will be revealed instead of creating a new one.
 	 */
-	public openPreview(mermaidSource: string): void {
+	public openPreview(mermaidSource: string, title?: string): void {
 		const webviewId = getWebviewId(mermaidSource);
 		const existingPreview = this._previews.get(webviewId);
 		if (existingPreview) {
@@ -47,6 +47,7 @@ export class MermaidEditorManager extends Disposable implements vscode.WebviewPa
 		const preview = MermaidPreview.create(
 			webviewId,
 			mermaidSource,
+			title,
 			this._extensionUri,
 			this._webviewManager,
 			vscode.ViewColumn.Active);
@@ -126,13 +127,14 @@ class MermaidPreview extends Disposable {
 	public static create(
 		diagramId: string,
 		mermaidSource: string,
+		title: string | undefined,
 		extensionUri: vscode.Uri,
 		webviewManager: MermaidWebviewManager,
 		viewColumn: vscode.ViewColumn
 	): MermaidPreview {
 		const webviewPanel = vscode.window.createWebviewPanel(
 			mermaidEditorViewType,
-			'', // Filled in later
+			title ?? vscode.l10n.t('Mermaid Diagram'),
 			viewColumn,
 			{
 				retainContextWhenHidden: false,
@@ -161,7 +163,6 @@ class MermaidPreview extends Disposable {
 	) {
 		super();
 
-		this._webviewPanel.title = vscode.l10n.t('Mermaid Diagram'); // TODO: Can we generate a better title from the content?
 		this._webviewPanel.iconPath = new vscode.ThemeIcon('graph');
 
 		this._webviewPanel.webview.options = {
@@ -174,7 +175,7 @@ class MermaidPreview extends Disposable {
 		this._webviewPanel.webview.html = this._getHtml();
 
 		// Register with the webview manager
-		this._register(this._webviewManager.registerWebview(this.diagramId, this._webviewPanel.webview, this._mermaidSource, 'editor'));
+		this._register(this._webviewManager.registerWebview(this.diagramId, this._webviewPanel.webview, this._mermaidSource, undefined, 'editor'));
 
 		this._register(this._webviewPanel.onDidChangeViewState(e => {
 			if (e.webviewPanel.active) {

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

@@ -8,6 +8,7 @@ export interface MermaidWebviewInfo {
 	readonly id: string;
 	readonly webview: vscode.Webview;
 	readonly mermaidSource: string;
+	readonly title: string | undefined;
 	readonly type: 'chat' | 'editor';
 }
 
@@ -27,7 +28,7 @@ export class MermaidWebviewManager {
 		return this._activeWebviewId ? this._webviews.get(this._activeWebviewId) : undefined;
 	}
 
-	public registerWebview(id: string, webview: vscode.Webview, mermaidSource: string, type: 'chat' | 'editor'): vscode.Disposable {
+	public registerWebview(id: string, webview: vscode.Webview, mermaidSource: string, title: string | undefined, type: 'chat' | 'editor'): vscode.Disposable {
 		if (this._webviews.has(id)) {
 			throw new Error(`Webview with id ${id} is already registered.`);
 		}
@@ -36,6 +37,7 @@ export class MermaidWebviewManager {
 			id,
 			webview,
 			mermaidSource,
+			title,
 			type
 		};
 		this._webviews.set(id, info);