Restore variant analysis view on restart of VSCode

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

@@ -104,6 +104,8 @@ import { LogScannerService } from './log-insights/log-scanner-service';
 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';
 import { VariantAnalysis } from './remote-queries/shared/variant-analysis';
 import {
   VariantAnalysis as VariantAnalysisApiResponse,
@@ -176,6 +178,7 @@ export interface CodeQLExtensionInterface {
   readonly distributionManager: DistributionManager;
   readonly databaseManager: DatabaseManager;
   readonly databaseUI: DatabaseUI;
+  readonly variantAnalysisManager: VariantAnalysisManager;
   readonly dispose: () => void;
 }
 
@@ -386,14 +389,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(codeQlExtension.variantAnalysisManager);
+
+  return codeQlExtension;
 }
 
 async function activateWithInstalledDistribution(
@@ -934,7 +944,10 @@ async function activateWithInstalledDistribution(
 
   ctx.subscriptions.push(
     commandRunner('codeQL.mockVariantAnalysisView', async () => {
-      await variantAnalysisManager.showView(1);
+      // Generate a random variant analysis ID for testing
+      const variantAnalysisId: number = Math.floor(Math.random() * 1000000);
+
+      await variantAnalysisManager.showView(variantAnalysisId);
     })
   );
 
@@ -1159,6 +1172,7 @@ async function activateWithInstalledDistribution(
     distributionManager,
     databaseManager: dbm,
     databaseUI,
+    variantAnalysisManager,
     dispose: () => {
       ctx.subscriptions.forEach(d => d.dispose());
     }

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

@@ -440,6 +440,15 @@ export interface SetVariantAnalysisMessage {
   variantAnalysis: VariantAnalysis;
 }
 
+export type StopVariantAnalysisMessage = {
+  t: 'stopVariantAnalysis';
+  variantAnalysisId: number;
+}
+
+export type VariantAnalysisState = {
+  variantAnalysisId: number;
+}
+
 export interface SetRepoResultsMessage {
   t: 'setRepoResults';
   repoResults: VariantAnalysisScannedRepositoryResult[];
@@ -456,4 +465,5 @@ export type ToVariantAnalysisMessage =
   | SetRepoStatesMessage;
 
 export type FromVariantAnalysisMessage =
-  | ViewLoadedMsg;
+  | ViewLoadedMsg
+  | StopVariantAnalysisMessage;

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

@@ -0,0 +1,48 @@
+import { ExtensionContext, WebviewPanel, WebviewPanelSerializer } from 'vscode';
+import { VariantAnalysisView } from './variant-analysis-view';
+import { VariantAnalysisState } from '../pure/interface-types';
+import { VariantAnalysisViewManager } from './variant-analysis-view-manager';
+
+export class VariantAnalysisViewSerializer implements WebviewPanelSerializer {
+  private resolvePromises: ((value: VariantAnalysisViewManager<VariantAnalysisView>) => void)[] = [];
+
+  private manager?: VariantAnalysisViewManager<VariantAnalysisView>;
+
+  public constructor(
+    private readonly ctx: ExtensionContext,
+  ) { }
+
+  onExtensionLoaded(manager: VariantAnalysisViewManager<VariantAnalysisView>): void {
+    this.manager = manager;
+
+    this.resolvePromises.forEach((resolve) => resolve(manager));
+    this.resolvePromises = [];
+  }
+
+  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;
+
+    const manager = await this.waitForExtensionFullyLoaded();
+
+    const view = new VariantAnalysisView(this.ctx, variantAnalysisState.variantAnalysisId, manager);
+    await view.restoreView(webviewPanel);
+  }
+
+  private waitForExtensionFullyLoaded(): Promise<VariantAnalysisViewManager<VariantAnalysisView>> {
+    if (this.manager) {
+      return Promise.resolve(this.manager);
+    }
+
+    return new Promise<VariantAnalysisViewManager<VariantAnalysisView>>((resolve) => {
+      this.resolvePromises.push(resolve);
+    });
+  }
+}