From 2076b05d9082e9413213c0d01574e040fbe44f71 Mon Sep 17 00:00:00 2001
From: Noam Zaks <imnoamzaks@gmail.com>
Date: Sun, 23 Jun 2024 17:35:04 +0300
Subject: [PATCH] Persist webview preview through vscode restarts (#300)

* feat: persist webview preview through vscode restarts

* make sure code works without VSCODE_STATE

* Apply suggestions from code review

* Comment backtick replacement

* Use base64 to solve security stuff

* fix: minor fix about non-ascii characters

---------

Co-authored-by: mgt <mgt@oi-wiki.org>
---
 addons/frontend/index.html     | 18 ++++++++++++++
 addons/frontend/src/main.js    |  4 ++++
 addons/vscode/package.json     |  4 +++-
 addons/vscode/src/extension.ts | 43 +++++++++++++++++++++++++++++++---
 4 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/addons/frontend/index.html b/addons/frontend/index.html
index d1bac96..b3302eb 100644
--- a/addons/frontend/index.html
+++ b/addons/frontend/index.html
@@ -15,6 +15,24 @@
     <style id="preview-theme-styles"></style>
     <!-- before all of code to avoid rerender by style replacement -->
     <script>
+      if (typeof acquireVsCodeApi !== "undefined") {
+        function base64ToString(base64) {
+          const binString = atob(base64);
+          const arr = Uint8Array.from(binString, (m) => m.codePointAt(0));
+          return new TextDecoder().decode(arr);
+        }
+        /// The string here is a placeholder 
+        /// It will be replaced by the actual preview mode. 
+        let state = `preview-arg:state:`;
+        state = state.replace("preview-arg:state:", "");
+        console.log("state", state);
+        if (state !== "") {
+          /// Set it later when acquiring the VSCode API
+          window.vscode_state = JSON.parse(base64ToString(state));
+          console.log("vscode_state", window.vscode_state);
+        }
+      }
+
       /// https://stackoverflow.com/questions/13586999/color-difference-similarity-between-two-values-with-js
       function deltaE(rgbA, rgbB) {
         let labA = rgb2lab(rgbA);
diff --git a/addons/frontend/src/main.js b/addons/frontend/src/main.js
index e950eaa..5643619 100644
--- a/addons/frontend/src/main.js
+++ b/addons/frontend/src/main.js
@@ -92,6 +92,10 @@ function setupVscodeChannel(nextWs) {
     if (vscodeAPI?.postMessage) {
         vscodeAPI.postMessage({ type: 'started' });
     }
+    if (vscodeAPI?.setState && window.vscode_state) {
+        vscodeAPI.setState(window.vscode_state);
+    }
+
 
     // Handle messages sent from the extension to the webview
     window.addEventListener('message', event => {
diff --git a/addons/vscode/package.json b/addons/vscode/package.json
index d095be0..4644c7a 100644
--- a/addons/vscode/package.json
+++ b/addons/vscode/package.json
@@ -16,7 +16,9 @@
   "categories": [
     "Other"
   ],
-  "activationEvents": [],
+  "activationEvents": [
+    "onWebviewPanel:typst-preview"
+  ],
   "main": "./out/src/extension.js",
   "contributes": {
     "icons": {
diff --git a/addons/vscode/src/extension.ts b/addons/vscode/src/extension.ts
index fe79d4f..55d3a34 100644
--- a/addons/vscode/src/extension.ts
+++ b/addons/vscode/src/extension.ts
@@ -339,6 +339,7 @@ interface LaunchTask {
 	activeEditor: vscode.TextEditor,
 	bindDocument: vscode.TextDocument,
 	mode: 'doc' | 'slide',
+	webviewPanel?: vscode.WebviewPanel
 }
 
 interface LaunchInBrowserTask extends LaunchTask {
@@ -357,6 +358,7 @@ const launchPreview = async (task: LaunchInBrowserTask | LaunchInWebViewTask) =>
 		outputChannel,
 		activeEditor,
 		bindDocument,
+		webviewPanel
 	} = task;
 	const filePath = bindDocument.uri.fsPath;
 
@@ -462,7 +464,7 @@ const launchPreview = async (task: LaunchInBrowserTask | LaunchInWebViewTask) =>
 	async function launchPreviewInWebView() {
 		const basename = path.basename(activeEditor.document.fileName);
 		// Create and show a new WebView
-		const panel = vscode.window.createWebviewPanel(
+		const panel = webviewPanel !== undefined ? webviewPanel : vscode.window.createWebviewPanel(
 			'typst-preview', // 标识符
 			`${basename} (Preview)`, // 面板标题
 			getTargetViewColumn(activeEditor.viewColumn),
@@ -490,10 +492,13 @@ const launchPreview = async (task: LaunchInBrowserTask | LaunchInWebViewTask) =>
 				.toString()}/typst-webview-assets`
 		);
 		const previewMode = task.mode === 'doc' ? "Doc" : "Slide";
+		const previewState = { mode: task.mode, fsPath: bindDocument.uri.fsPath };
+		const previewStateEncoded = Buffer.from(JSON.stringify(previewState), 'utf-8').toString('base64');
 		html = html.replace(
 			"preview-arg:previewMode:Doc",
 			`preview-arg:previewMode:${previewMode}`
-		);
+		).replace("preview-arg:state:", `preview-arg:state:${previewStateEncoded}`);
+
 		panel.webview.html = html.replace("ws://127.0.0.1:23625", `ws://127.0.0.1:${dataPlanePort}`);
 		// 虽然配置的是 http，但是如果是桌面客户端，任何 tcp 连接都支持，这也就包括了 ws
 		// https://code.visualstudio.com/api/advanced-topics/remote-extensions#forwarding-localhost
@@ -761,6 +766,37 @@ export class OutlineItem extends vscode.TreeItem {
 	contextValue = 'outline-item';
 }
 
+class TypstPreviewSerializer implements vscode.WebviewPanelSerializer {
+	context: vscode.ExtensionContext;
+	outputChannel: vscode.OutputChannel;
+
+	constructor(context: vscode.ExtensionContext, outputChannel: vscode.OutputChannel) {
+		this.context = context;
+		this.outputChannel = outputChannel;
+	}
+
+	async deserializeWebviewPanel(webviewPanel: vscode.WebviewPanel, state: any) {
+		const activeEditor = vscode.window.visibleTextEditors.find(editor => editor.document.uri.fsPath === state.fsPath);
+
+		if (!activeEditor) {
+			return;
+		}
+
+		const bindDocument = activeEditor.document;
+		const mode = state.mode;
+
+		launchPreview({
+			kind: "webview",
+			context: this.context,
+			outputChannel: this.outputChannel,
+			activeEditor,
+			bindDocument,
+			mode,
+			webviewPanel
+		});
+	}
+}
+
 let statusBarItem: vscode.StatusBarItem;
 // This method is called when your extension is activated
 // Your extension is activated the very first time the command is executed
@@ -854,8 +890,9 @@ export function activate(context: vscode.ExtensionContext) {
 	let showLogDisposable = vscode.commands.registerCommand('typst-preview.showLog', async () => {
 		outputChannel.show();
 	});
+	let serializerDisposable = vscode.window.registerWebviewPanelSerializer('typst-preview', new TypstPreviewSerializer(context, outputChannel));
 
-	context.subscriptions.push(webviewDisposable, browserDisposable, webviewSlideDisposable, browserSlideDisposable, syncDisposable, showLogDisposable, statusBarItem, revealDocumentDisposable, awaitTreeDisposable);
+	context.subscriptions.push(webviewDisposable, browserDisposable, webviewSlideDisposable, browserSlideDisposable, syncDisposable, showLogDisposable, statusBarItem, revealDocumentDisposable, awaitTreeDisposable, serializerDisposable);
 	process.on('SIGINT', () => {
 		for (const serverProcess of serverProcesses) {
 			serverProcess.kill();