CoCalc -- chatOutputRenderer.ts

GitHub Repository: microsoft/vscode
Path: blob/main/extensions/mermaid-chat-features/src/chatOutputRenderer.ts
⁵²²³ views
1
/*---------------------------------------------------------------------------------------------
2
 *  Copyright (c) Microsoft Corporation. All rights reserved.
3
 *  Licensed under the MIT License. See License.txt in the project root for license information.
4
 *--------------------------------------------------------------------------------------------*/
5
import * as vscode from 'vscode';
6
import { MermaidEditorManager } from './editorManager';
7
import { MermaidWebviewManager } from './webviewManager';
8
import { escapeHtmlText } from './util/html';
9
import { generateUuid } from './util/uuid';
10
import { disposeAll } from './util/dispose';
11

12
/**
13
 * Mime type used to identify Mermaid diagram data in chat output.
14
 */
15
const mime = 'text/vnd.mermaid';
16

17
/**
18
 * View type that uniquely identifies the Mermaid chat output renderer.
19
 */
20
const viewType = 'vscode.chat-mermaid-features.chatOutputItem';
21

22
class MermaidChatOutputRenderer implements vscode.ChatOutputRenderer {
23

24
	constructor(
25
		private readonly _extensionUri: vscode.Uri,
26
		private readonly _webviewManager: MermaidWebviewManager
27
	) { }
28

29
	async renderChatOutput({ value }: vscode.ChatOutputDataItem, chatOutputWebview: vscode.ChatOutputWebview, _ctx: unknown, _token: vscode.CancellationToken): Promise<void> {
30
		const webview = chatOutputWebview.webview;
31
		const decoded = decodeMermaidData(value);
32
		const mermaidSource = decoded.source;
33
		const title = decoded.title;
34

35
		// Generate unique ID for this webview
36
		const webviewId = generateUuid();
37

38
		const disposables: vscode.Disposable[] = [];
39

40
		// Register and set as active
41
		disposables.push(this._webviewManager.registerWebview(webviewId, webview, mermaidSource, title, 'chat'));
42

43
		// Listen for messages from the webview
44
		disposables.push(webview.onDidReceiveMessage(message => {
45
			if (message.type === 'openInEditor') {
46
				vscode.commands.executeCommand('_mermaid-chat.openInEditor', { mermaidWebviewId: webviewId });
47
			}
48
		}));
49

50
		// Dispose resources when webview is disposed
51
		chatOutputWebview.onDidDispose(() => {
52
			disposeAll(disposables);
53
		});
54

55
		// Set the options for the webview
56
		const mediaRoot = vscode.Uri.joinPath(this._extensionUri, 'chat-webview-out');
57
		webview.options = {
58
			enableScripts: true,
59
			localResourceRoots: [mediaRoot],
60
		};
61

62
		// Set the HTML content for the webview
63
		const nonce = generateUuid();
64
		const mermaidScript = vscode.Uri.joinPath(mediaRoot, 'index.js');
65
		const codiconsUri = webview.asWebviewUri(vscode.Uri.joinPath(mediaRoot, 'codicon.css'));
66

67
		webview.html = `
68
			<!DOCTYPE html>
69
			<html lang="en">
70

71
			<head>
72
				<meta charset="UTF-8">
73
				<meta name="viewport" content="width=device-width, initial-scale=1.0">
74
				<title>Mermaid Diagram</title>
75
				<meta http-equiv="Content-Security-Policy" content="default-src 'none'; script-src 'nonce-${nonce}'; style-src ${webview.cspSource} 'unsafe-inline'; font-src data:;" />
76
				<link rel="stylesheet" type="text/css" href="${codiconsUri}">
77

78
				<style>
79
					body {
80
						padding: 0;
81
					}
82
					.mermaid {
83
						visibility: hidden;
84
					}
85
					.mermaid.rendered {
86
						visibility: visible;
87
					}
88
					.open-in-editor-btn {
89
						position: absolute;
90
						top: 8px;
91
						right: 8px;
92
						display: flex;
93
						align-items: center;
94
						justify-content: center;
95
						width: 26px;
96
						height: 26px;
97
						background: var(--vscode-editorWidget-background);
98
						color: var(--vscode-icon-foreground);
99
						border: 1px solid var(--vscode-editorWidget-border);
100
						border-radius: 6px;
101
						cursor: pointer;
102
						z-index: 100;
103
						opacity: 0;
104
						transition: opacity 0.2s;
105
					}
106
					body:hover .open-in-editor-btn {
107
						opacity: 1;
108
					}
109
					.open-in-editor-btn:hover {
110
						opacity: 1;
111
						background: var(--vscode-toolbar-hoverBackground);
112
					}
113
				</style>
114
			</head>
115

116
			<body data-vscode-context='${JSON.stringify({ preventDefaultContextMenuItems: true, mermaidWebviewId: webviewId })}' data-vscode-mermaid-webview-id="${webviewId}">
117
				<button class="open-in-editor-btn" title="${vscode.l10n.t('Open in Editor')}"><i class="codicon codicon-open-preview"></i></button>
118
				<pre class="mermaid">
119
					${escapeHtmlText(mermaidSource)}
120
				</pre>
121

122
				<script type="module" nonce="${nonce}" src="${webview.asWebviewUri(mermaidScript)}"></script>
123
			</body>
124
			</html>`;
125
	}
126
}
127

128

129
export function registerChatSupport(
130
	context: vscode.ExtensionContext,
131
	webviewManager: MermaidWebviewManager,
132
	editorManager: MermaidEditorManager
133
): vscode.Disposable {
134
	const disposables: vscode.Disposable[] = [];
135

136
	disposables.push(
137
		vscode.commands.registerCommand('_mermaid-chat.openInEditor', (ctx?: { mermaidWebviewId?: string }) => {
138
			const webviewInfo = ctx?.mermaidWebviewId ? webviewManager.getWebview(ctx.mermaidWebviewId) : webviewManager.activeWebview;
139
			if (webviewInfo) {
140
				editorManager.openPreview(webviewInfo.mermaidSource, webviewInfo.title);
141
			}
142
		})
143
	);
144

145
	// Register lm tools
146
	disposables.push(
147
		vscode.lm.registerTool<{ markup: string; title?: string }>('renderMermaidDiagram', {
148
			invoke: async (options, _token) => {
149
				const sourceCode = options.input.markup;
150
				const title = options.input.title;
151
				return writeMermaidToolOutput(sourceCode, title);
152
			},
153
		})
154
	);
155

156
	// Register the chat output renderer for Mermaid diagrams.
157
	// This will be invoked with the data generated by the tools.
158
	// It can also be invoked when rendering old Mermaid diagrams in the chat history.
159
	const renderer = new MermaidChatOutputRenderer(context.extensionUri, webviewManager);
160
	disposables.push(vscode.chat.registerChatOutputRenderer(viewType, renderer));
161

162
	return vscode.Disposable.from(...disposables);
163
}
164

165
function writeMermaidToolOutput(sourceCode: string, title: string | undefined): vscode.LanguageModelToolResult {
166
	// Expose the source code as a markdown mermaid code block
167
	const fence = getFenceForContent(sourceCode);
168
	const result = new vscode.LanguageModelToolResult([
169
		new vscode.LanguageModelTextPart(`${fence}mermaid\n${sourceCode}\n${fence}`)
170
	]);
171

172
	// And store custom data in the tool result details to indicate that a custom renderer should be used for it.
173
	// Encode source and optional title as JSON.
174
	const data = JSON.stringify({ source: sourceCode, title });
175
	// Add cast to use proposed API
176
	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {
177
		mime,
178
		value: new TextEncoder().encode(data),
179
	};
180

181
	return result;
182
}
183

184
function getFenceForContent(content: string): string {
185
	const backtickMatch = content.matchAll(/`+/g);
186
	if (!backtickMatch) {
187
		return '```';
188
	}
189

190
	const maxBackticks = Math.max(...Array.from(backtickMatch, s => s[0].length));
191
	return '`'.repeat(Math.max(3, maxBackticks + 1));
192
}
193

194
interface MermaidData {
195
	readonly title: string | undefined;
196
	readonly source: string;
197
}
198

199
function decodeMermaidData(value: Uint8Array): MermaidData {
200
	const text = new TextDecoder().decode(value);
201

202
	// Try to parse as JSON (new format with title), fall back to plain text (legacy format)
203
	try {
204
		const parsed = JSON.parse(text);
205
		if (typeof parsed === 'object' && typeof parsed.source === 'string') {
206
			return { title: parsed.title, source: parsed.source };
207
		}
208
	} catch {
209
		// Not JSON, treat as legacy plain text format
210
	}
211

212
	return { title: undefined, source: text };
213
}
214

215
Product

Resources

Company
