Restore variant analysis view on restart of VSCode

This implements persistence for the variant analysis webview, allowing
the webview panel to be restored when VSCode is restarted. It's probably
easier to add this now than to try to add it later.

The basic idea is that there are no real differences when opening the
webview for the first time. However, when VSCode is restarted it will
use the `VariantAnalysisViewSerializer` to restore the webview panel.
In our case this means recreating the `VariantAnalysisView`.

To fully test this, I've also added a mock variant analysis ID as the
state of the webview. This value is now randomly generated when calling
the `codeQL.mockVariantAnalysisView` command. This allows us to test
opening multiple webviews and that the webviews are restored with the
correct state.

See: https://code.visualstudio.com/api/extension-guides/webview#persistence

Author: koesie10

extensions/ql-vscode/.storybook/preview.ts

@@ -33,5 +33,6 @@ export const parameters = {
 };
 
 (window as any).acquireVsCodeApi = () => ({
-  postMessage: action('post-vscode-message')
+  postMessage: action('post-vscode-message'),
+  setState: action('set-vscode-state'),
 });

extensions/ql-vscode/package.json

@@ -63,6 +63,7 @@
     "onCommand:codeQL.quickQuery",
     "onCommand:codeQL.restartQueryServer",
     "onWebviewPanel:resultsView",
+    "onWebviewPanel:codeQL.variantAnalysis",
     "onFileSystem:codeql-zip-archive"
   ],
   "main": "./out/extension",

extensions/ql-vscode/src/abstract-webview.ts

@@ -33,6 +33,11 @@ export abstract class AbstractWebview<ToMessage extends WebviewMessage, FromMess
     super();
   }
 
+  public async restoreView(panel: WebviewPanel): Promise<void> {
+    this.panel = panel;
+    this.setupPanel(panel);
+  }
+
   protected get isShowingPanel() {
     return !!this.panel;
   }
@@ -59,37 +64,43 @@ export abstract class AbstractWebview<ToMessage extends WebviewMessage, FromMess
           ],
         }
       );
-      this.push(
-        this.panel.onDidDispose(
-          () => {
-            this.panel = undefined;
-            this.panelLoaded = false;
-            this.onPanelDispose();
-          },
-          null,
-          ctx.subscriptions
-        )
-      );
-
-      this.panel.webview.html = getHtmlForWebview(
-        ctx,
-        this.panel.webview,
-        config.view,
-        {
-          allowInlineStyles: true,
-        }
-      );
-      this.push(
-        this.panel.webview.onDidReceiveMessage(
-          async (e) => this.onMessage(e),
-          undefined,
-          ctx.subscriptions
-        )
-      );
+      this.setupPanel(this.panel);
     }
     return this.panel;
   }
 
+  protected setupPanel(panel: WebviewPanel): void {
+    const config = this.getPanelConfig();
+
+    this.push(
+      panel.onDidDispose(
+        () => {
+          this.panel = undefined;
+          this.panelLoaded = false;
+          this.onPanelDispose();
+        },
+        null,
+        this.ctx.subscriptions
+      )
+    );
+
+    panel.webview.html = getHtmlForWebview(
+      this.ctx,
+      panel.webview,
+      config.view,
+      {
+        allowInlineStyles: true,
+      }
+    );
+    this.push(
+      panel.webview.onDidReceiveMessage(
+        async (e) => this.onMessage(e),
+        undefined,
+        this.ctx.subscriptions
+      )
+    );
+  }
+
   protected abstract getPanelConfig(): WebviewPanelConfig;
 
   protected abstract onPanelDispose(): void;

extensions/ql-vscode/src/extension.ts

@@ -105,6 +105,7 @@ import { createInitialQueryInfo } from './run-queries-shared';
 import { LegacyQueryRunner } from './legacy-query-server/legacyRunner';
 import { QueryRunner } from './queryRunner';
 import { VariantAnalysisView } from './remote-queries/variant-analysis-view';
+import { VariantAnalysisViewSerializer } from './remote-queries/variant-analysis-view-serializer';
 
 /**
  * extension.ts
@@ -381,14 +382,21 @@ export async function activate(ctx: ExtensionContext): Promise<CodeQLExtensionIn
     allowAutoUpdating: true
   })));
 
-  return await installOrUpdateThenTryActivate({
+  const variantAnalysisViewSerializer = new VariantAnalysisViewSerializer(ctx);
+  Window.registerWebviewPanelSerializer(VariantAnalysisView.viewType, variantAnalysisViewSerializer);
+
+  const codeQlExtension = await installOrUpdateThenTryActivate({
     isUserInitiated: !!ctx.globalState.get(shouldUpdateOnNextActivationKey),
     shouldDisplayMessageWhenNoUpdates: false,
 
     // only auto update on startup if the user has previously requested an update
     // otherwise, ask user to accept the update
     allowAutoUpdating: !!ctx.globalState.get(shouldUpdateOnNextActivationKey)
   });
+
+  variantAnalysisViewSerializer.onExtensionLoaded();
+
+  return codeQlExtension;
 }
 
 async function activateWithInstalledDistribution(
@@ -909,8 +917,11 @@ async function activateWithInstalledDistribution(
 
   ctx.subscriptions.push(
     commandRunner('codeQL.mockVariantAnalysisView', async () => {
-      const variantAnalysisView = new VariantAnalysisView(ctx);
-      variantAnalysisView.openView();
+      // Generate a random variant analysis ID for testing
+      const variantAnalysisId: number = Math.floor(Math.random() * 1000000);
+
+      const variantAnalysisView = new VariantAnalysisView(ctx, variantAnalysisId);
+      void variantAnalysisView.openView();
     })
   );
 

extensions/ql-vscode/src/pure/interface-types.ts

@@ -2,6 +2,7 @@ import * as sarif from 'sarif';
 import { AnalysisResults } from '../remote-queries/shared/analysis-result';
 import { AnalysisSummary, RemoteQueryResult } from '../remote-queries/shared/remote-query-result';
 import { RawResultSet, ResultRow, ResultSetSchema, Column, ResolvableLocationValue } from './bqrs-cli-types';
+import { VariantAnalysis } from '../remote-queries/shared/variant-analysis';
 
 /**
  * This module contains types and code that are shared between
@@ -429,3 +430,24 @@ export interface CopyRepoListMessage {
   t: 'copyRepoList';
   queryId: string;
 }
+
+export interface SetVariantAnalysisMessage {
+  t: 'setVariantAnalysis';
+  variantAnalysis: VariantAnalysis;
+}
+
+export type ToVariantAnalysisMessage =
+  | SetVariantAnalysisMessage;
+
+export type StopVariantAnalysisMessage = {
+  t: 'stopVariantAnalysis';
+  variantAnalysisId: number;
+}
+
+export type FromVariantAnalysisMessage =
+  | ViewLoadedMsg
+  | StopVariantAnalysisMessage;
+
+export type VariantAnalysisState = {
+  variantAnalysisId: number;
+}

extensions/ql-vscode/src/remote-queries/variant-analysis-view-serializer.ts

@@ -0,0 +1,44 @@
+import { ExtensionContext, WebviewPanel, WebviewPanelSerializer } from 'vscode';
+import { VariantAnalysisView } from './variant-analysis-view';
+import { VariantAnalysisState } from '../pure/interface-types';
+
+export class VariantAnalysisViewSerializer implements WebviewPanelSerializer {
+  private extensionLoaded = false;
+  private readonly resolvePromises: (() => void)[] = [];
+
+  public constructor(
+    private readonly ctx: ExtensionContext
+  ) { }
+
+  onExtensionLoaded(): void {
+    this.extensionLoaded = true;
+    this.resolvePromises.forEach((resolve) => resolve());
+  }
+
+  async deserializeWebviewPanel(webviewPanel: WebviewPanel, state: unknown): Promise<void> {
+    if (!state || typeof state !== 'object') {
+      return;
+    }
+
+    if (!('variantAnalysisId' in state)) {
+      return;
+    }
+
+    const variantAnalysisState: VariantAnalysisState = state as VariantAnalysisState;
+
+    await this.waitForExtensionFullyLoaded();
+
+    const view = new VariantAnalysisView(this.ctx, variantAnalysisState.variantAnalysisId);
+    await view.restoreView(webviewPanel);
+  }
+
+  private waitForExtensionFullyLoaded(): Promise<void> {
+    if (this.extensionLoaded) {
+      return Promise.resolve();
+    }
+
+    return new Promise((resolve) => {
+      this.resolvePromises.push(resolve);
+    });
+  }
+}