Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
microsoft
GitHub Repository: microsoft/vscode
 Path: blob/main/extensions/mermaid-chat-features/src/extension.ts
3520 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


7
/**

8
 * View type that uniquely identifies the Mermaid chat output renderer.

9
 */

10
const viewType = 'vscode.chatMermaidDiagram';

11


12
/**

13
 * Mime type used to identify Mermaid diagram data in chat output.

14
 */

15
const mime = 'text/vnd.mermaid';

16


17
export function activate(context: vscode.ExtensionContext) {

18


19
	// Register tools

20
	context.subscriptions.push(

21
		vscode.lm.registerTool<{ markup: string }>('renderMermaidDiagram', {

22
			invoke: async (options, _token) => {

23
				const sourceCode = options.input.markup;

24
				return writeMermaidToolOutput(sourceCode);

25
			},

26
		})

27
	);

28


29
	// Register the chat output renderer for Mermaid diagrams.

30
	// This will be invoked with the data generated by the tools.

31
	// It can also be invoked when rendering old Mermaid diagrams in the chat history.

32
	context.subscriptions.push(

33
		vscode.chat.registerChatOutputRenderer(viewType, {

34
			async renderChatOutput({ value }, webview, _ctx, _token) {

35
				const mermaidSource = new TextDecoder().decode(value);

36


37
				// Set the options for the webview

38
				const mediaRoot = vscode.Uri.joinPath(context.extensionUri, 'chat-webview-out');

39
				webview.options = {

40
					enableScripts: true,

41
					localResourceRoots: [mediaRoot],

42
				};

43


44
				// Set the HTML content for the webview

45
				const nonce = getNonce();

46
				const mermaidScript = vscode.Uri.joinPath(mediaRoot, 'index.js');

47


48
				webview.html = `

49
					<!DOCTYPE html>

50
					<html lang="en">

51


52
					<head>

53
						<meta charset="UTF-8">

54
						<meta name="viewport" content="width=device-width, initial-scale=1.0">

55
						<title>Mermaid Diagram</title>

56
						<meta http-equiv="Content-Security-Policy" content="default-src 'none'; script-src ${webview.cspSource} 'nonce-${nonce}'; style-src 'self' 'unsafe-inline';" />

57
					</head>

58


59
					<body>

60
						<pre class="mermaid">

61
							${escapeHtmlText(mermaidSource)}

62
						</pre>

63


64
						<script type="module" nonce="${nonce}" src="${webview.asWebviewUri(mermaidScript)}"></script>

65
					</body>

66
					</html>`;

67
			},

68
		}));

69
}

70


71


72
function writeMermaidToolOutput(sourceCode: string): vscode.LanguageModelToolResult {

73
	// Expose the source code as a tool result for the LM

74
	const result = new vscode.LanguageModelToolResult([

75
		new vscode.LanguageModelTextPart(sourceCode)

76
	]);

77


78
	// And store custom data in the tool result details to indicate that a custom renderer should be used for it.

79
	// In this case we just store the source code as binary data.

80


81
	// Add cast to use proposed API

82
	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {

83
		mime,

84
		value: new TextEncoder().encode(sourceCode),

85
	};

86


87
	return result;

88
}

89


90
function escapeHtmlText(str: string): string {

91
	return str

92
		.replace(/&/g, '&amp;')

93
		.replace(/</g, '&lt;')

94
		.replace(/>/g, '&gt;')

95
		.replace(/"/g, '&quot;')

96
		.replace(/'/g, '&#39;');

97
}

98


99
function getNonce() {

100
	let text = '';

101
	const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';

102
	for (let i = 0; i < 64; i++) {

103
		text += possible.charAt(Math.floor(Math.random() * possible.length));

104
	}

105
	return text;

106
}

107


108