diff --git a/extensions/mermaid-chat-features/package.json b/extensions/mermaid-chat-features/package.json index 690bb9ce086..16b6a03ce48 100644 --- a/extensions/mermaid-chat-features/package.json +++ b/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." } } } diff --git a/extensions/mermaid-chat-features/src/chatOutputRenderer.ts b/extensions/mermaid-chat-features/src/chatOutputRenderer.ts index 4bf4d9de7e7..c6aacbd748e 100644 --- a/extensions/mermaid-chat-features/src/chatOutputRenderer.ts +++ b/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 { 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 }; +} diff --git a/extensions/mermaid-chat-features/src/editorManager.ts b/extensions/mermaid-chat-features/src/editorManager.ts index 8482c7fc8b5..5ee5cec29da 100644 --- a/extensions/mermaid-chat-features/src/editorManager.ts +++ b/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) { diff --git a/extensions/mermaid-chat-features/src/webviewManager.ts b/extensions/mermaid-chat-features/src/webviewManager.ts index 40c56993a88..253ae92b959 100644 --- a/extensions/mermaid-chat-features/src/webviewManager.ts +++ b/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);