From fbebcbcaea387c7e4115ef8a69482cf7117f6a37 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Wouts?= <f@zenc.io>
Date: Tue, 28 Mar 2023 23:43:19 +1100
Subject: [PATCH] feat(vscode): add a command to reset Preview.js and clear
 cache

This addresses the VS Code implementation for #1488.
---
 daemon/src/client.ts                          | 13 +++
 integrations/vscode/package.json              |  5 +
 integrations/vscode/src/component-detector.ts | 26 +++--
 integrations/vscode/src/index.ts              | 95 +++++++++++--------
 integrations/vscode/src/preview-panel.ts      | 33 +++----
 integrations/vscode/src/preview-server.ts     | 35 +++----
 integrations/vscode/src/start-daemon.ts       | 36 ++++---
 integrations/vscode/src/state.ts              | 69 ++++++++++++++
 integrations/vscode/src/workspaces.ts         | 28 ++++--
 9 files changed, 224 insertions(+), 116 deletions(-)
 create mode 100644 integrations/vscode/src/state.ts

diff --git a/daemon/src/client.ts b/daemon/src/client.ts
index 33c5af87b8e..fc0eb0650f7 100644
--- a/daemon/src/client.ts
+++ b/daemon/src/client.ts
@@ -1,4 +1,5 @@
 import { exclusivePromiseRunner } from "exclusive-promises";
+import { existsSync, readFileSync, unlinkSync } from "fs";
 import http from "http";
 import type {
   AnalyzeFileRequest,
@@ -93,6 +94,18 @@ export function createClient(baseUrl: string): Client {
   return client;
 }
 
+export function destroyDaemon(lockFilePath: string) {
+  if (existsSync(lockFilePath)) {
+    const pid = parseInt(readFileSync(lockFilePath, "utf8"));
+    try {
+      process.kill(pid, "SIGKILL");
+    } catch {
+      // The daemon was already dead.
+    }
+    unlinkSync(lockFilePath);
+  }
+}
+
 export interface Client {
   waitForReady(): Promise<void>;
   info(): Promise<InfoResponse>;
diff --git a/integrations/vscode/package.json b/integrations/vscode/package.json
index 5f9d73ec910..8ac957629d5 100644
--- a/integrations/vscode/package.json
+++ b/integrations/vscode/package.json
@@ -65,6 +65,11 @@
         "title": "Open Preview.js",
         "category": "Preview.js",
         "icon": "./logo.png"
+      },
+      {
+        "command": "previewjs.reset",
+        "title": "Reset Preview.js after clearing cache",
+        "category": "Preview.js"
       }
     ],
     "keybindings": [
diff --git a/integrations/vscode/src/component-detector.ts b/integrations/vscode/src/component-detector.ts
index 960e56b3004..bc7ac46a333 100644
--- a/integrations/vscode/src/component-detector.ts
+++ b/integrations/vscode/src/component-detector.ts
@@ -1,33 +1,29 @@
 import type { AnalyzeFileResponse, Client } from "@previewjs/daemon/client";
 import type vscode from "vscode";
-import type { createWorkspaceGetter } from "./workspaces";
+import { WorkspaceGetter } from "./workspaces";
 
 export function createComponentDetector(
-  previewjsInitPromise: Promise<Client | null>,
-  getWorkspaceId: ReturnType<typeof createWorkspaceGetter>
-) {
-  async function getComponents(
+  client: Client,
+  getWorkspaceId: WorkspaceGetter
+): ComponentDetector {
+  return async function (
     document?: vscode.TextDocument
   ): Promise<AnalyzeFileResponse["components"]> {
     if (!document || !document.fileName) {
       return [];
     }
-    const previewjsClient = await previewjsInitPromise;
-    if (!previewjsClient) {
-      return [];
-    }
-    const workspaceId = await getWorkspaceId(previewjsClient, document);
+    const workspaceId = await getWorkspaceId(document);
     if (!workspaceId) {
       return [];
     }
-    const { components } = await previewjsClient.analyzeFile({
+    const { components } = await client.analyzeFile({
       workspaceId,
       absoluteFilePath: document.fileName,
     });
     return components;
-  }
-
-  return {
-    getComponents,
   };
 }
+
+export type ComponentDetector = (
+  document?: vscode.TextDocument
+) => Promise<AnalyzeFileResponse["components"]>;
diff --git a/integrations/vscode/src/index.ts b/integrations/vscode/src/index.ts
index c28777885e5..d9d3d6e0d90 100644
--- a/integrations/vscode/src/index.ts
+++ b/integrations/vscode/src/index.ts
@@ -1,11 +1,12 @@
+import { destroyDaemon } from "@previewjs/daemon/client";
+import fs from "fs";
 import path from "path";
 import vscode from "vscode";
 import { clientId } from "./client-id";
-import { createComponentDetector } from "./component-detector";
 import { closePreviewPanel, updatePreviewPanel } from "./preview-panel";
 import { ensurePreviewServerStarted } from "./preview-server";
-import { ensureDaemonRunning } from "./start-daemon";
-import { createWorkspaceGetter } from "./workspaces";
+import { daemonLockFilePath } from "./start-daemon";
+import { createState } from "./state";
 
 const codeLensLanguages = [
   "javascript",
@@ -38,24 +39,7 @@ let dispose = async () => {
 
 export async function activate() {
   const outputChannel = vscode.window.createOutputChannel("Preview.js");
-  const getWorkspaceId = createWorkspaceGetter(outputChannel);
-  const previewjsInitPromise = ensureDaemonRunning(outputChannel)
-    .catch((e) => {
-      outputChannel.appendLine(e.stack);
-      return null;
-    })
-    .then((p) => {
-      if (!p) {
-        outputChannel.appendLine("Preview.js daemon could not be started.");
-        outputChannel.show();
-        return null;
-      }
-      return p;
-    });
-  const componentDetector = createComponentDetector(
-    previewjsInitPromise,
-    getWorkspaceId
-  );
+  let currentState = createState({ outputChannel });
 
   const config = vscode.workspace.getConfiguration();
   let focusedOutputChannelForError = false;
@@ -80,16 +64,23 @@ export async function activate() {
   }
 
   dispose = async () => {
-    const previewjsClient = await previewjsInitPromise;
-    await previewjsClient?.updateClientStatus({
-      clientId,
-      alive: false,
-    });
+    const state = await currentState;
+    if (state) {
+      await state.client.updateClientStatus({
+        clientId,
+        alive: false,
+      });
+      await closePreviewPanel(state);
+    }
     outputChannel.dispose();
   };
 
   vscode.window.onDidChangeActiveTextEditor(async (e) => {
-    const components = await componentDetector.getComponents(e?.document);
+    const state = await currentState;
+    if (!state) {
+      return;
+    }
+    const components = await state.getComponents(e?.document);
     vscode.commands.executeCommand(
       "setContext",
       "previewjs.componentsDetected",
@@ -100,7 +91,11 @@ export async function activate() {
   if (config.get("previewjs.codelens", true)) {
     vscode.languages.registerCodeLensProvider(codeLensLanguages, {
       provideCodeLenses: catchErrors(async (document: vscode.TextDocument) => {
-        const components = await componentDetector.getComponents(document);
+        const state = await currentState;
+        if (!state) {
+          return;
+        }
+        const components = await state.getComponents(document);
         return components.map((c) => {
           const start = document.positionAt(c.start + 2);
           const lens = new vscode.CodeLens(new vscode.Range(start, start));
@@ -126,27 +121,49 @@ export async function activate() {
       document: vscode.TextDocument,
       saved = false
     ) {
-      const previewjsClient = await previewjsInitPromise;
+      const state = await currentState;
       if (
-        !previewjsClient ||
+        !state ||
         !path.isAbsolute(document.fileName) ||
         !watchedExtensions.has(path.extname(document.fileName))
       ) {
         return;
       }
-      previewjsClient.updatePendingFile({
+      state.client.updatePendingFile({
         absoluteFilePath: document.fileName,
         utf8Content: saved ? null : document.getText(),
       });
     }
   }
 
+  vscode.commands.registerCommand(
+    "previewjs.reset",
+    catchErrors(async () => {
+      outputChannel.appendLine("Resetting Preview.js...");
+      const state = await currentState;
+      if (state) {
+        state.dispose();
+      }
+      currentState = Promise.resolve(null);
+      destroyDaemon(daemonLockFilePath);
+      if (state) {
+        for (const rootDirPath of Object.values(state.workspaces)) {
+          fs.rmSync(path.join(rootDirPath, "node_modules", ".previewjs"), {
+            recursive: true,
+            force: true,
+          });
+        }
+      }
+      currentState = createState({ outputChannel });
+    })
+  );
+
   vscode.commands.registerCommand(
     "previewjs.open",
     catchErrors(
       async (document?: vscode.TextDocument, componentId?: string) => {
-        const previewjsClient = await previewjsInitPromise;
-        if (!previewjsClient) {
+        const state = await currentState;
+        if (!state) {
           vscode.window.showErrorMessage(
             "Preview.js was unable to start successfully. Please check Preview.js output panel and consider filing a bug at https://github.com/fwouts/previewjs/issues."
           );
@@ -165,7 +182,7 @@ export async function activate() {
             return;
           }
         }
-        const workspaceId = await getWorkspaceId(previewjsClient, document);
+        const workspaceId = await state.getWorkspaceId(document);
         if (!workspaceId) {
           vscode.window.showErrorMessage(
             `No compatible workspace detected from ${document.fileName}`
@@ -176,7 +193,7 @@ export async function activate() {
           const offset = editor?.selection.active
             ? document.offsetAt(editor.selection.active)
             : 0;
-          const components = await componentDetector.getComponents(document);
+          const components = await state.getComponents(document);
           const component =
             components.find((c) => offset >= c.start && offset <= c.end) ||
             components[0];
@@ -188,11 +205,8 @@ export async function activate() {
           }
           componentId = component.componentId;
         }
-        const preview = await ensurePreviewServerStarted(
-          previewjsClient,
-          workspaceId
-        );
-        updatePreviewPanel(previewjsClient, preview.url, componentId);
+        const preview = await ensurePreviewServerStarted(state, workspaceId);
+        updatePreviewPanel(state, preview.url, componentId);
       }
     )
   );
@@ -200,7 +214,6 @@ export async function activate() {
 
 export async function deactivate() {
   await dispose();
-  await closePreviewPanel();
   dispose = async () => {
     // Do nothing.
   };
diff --git a/integrations/vscode/src/preview-panel.ts b/integrations/vscode/src/preview-panel.ts
index d7bfb8c9aee..0c511b460eb 100644
--- a/integrations/vscode/src/preview-panel.ts
+++ b/integrations/vscode/src/preview-panel.ts
@@ -1,17 +1,14 @@
-import type { Client } from "@previewjs/daemon/client";
 import vscode from "vscode";
-import type { WebviewPanel } from "vscode";
 import { ensurePreviewServerStopped } from "./preview-server";
-
-let previewPanel: WebviewPanel | null = null;
+import { PreviewJsState } from "./state";
 
 export function updatePreviewPanel(
-  client: Client,
+  state: PreviewJsState,
   previewBaseUrl: string,
   componentId: string
 ) {
-  if (!previewPanel) {
-    previewPanel = vscode.window.createWebviewPanel(
+  if (!state.previewPanel) {
+    state.previewPanel = vscode.window.createWebviewPanel(
       "preview", // Identifies the type of the webview. Used internally
       "Preview", // Title of the panel displayed to the user
       vscode.ViewColumn.Two, // Editor column to show the new webview panel in.
@@ -20,16 +17,16 @@ export function updatePreviewPanel(
         retainContextWhenHidden: true,
       }
     );
-    previewPanel.webview.onDidReceiveMessage((message) => {
+    state.previewPanel.webview.onDidReceiveMessage((message) => {
       if (message.command === "open-browser") {
         vscode.env.openExternal(vscode.Uri.parse(message.url));
       }
     });
-    previewPanel.onDidDispose(() => {
-      previewPanel = null;
-      ensurePreviewServerStopped(client).catch(console.error);
+    state.previewPanel.onDidDispose(() => {
+      state.previewPanel = null;
+      ensurePreviewServerStopped(state).catch(console.error);
     });
-    previewPanel.webview.html = `<!DOCTYPE html>
+    state.previewPanel.webview.html = `<!DOCTYPE html>
   <html>
     <head>
       <style>
@@ -78,18 +75,18 @@ export function updatePreviewPanel(
     </body>
   </html>`;
   } else {
-    previewPanel.webview.postMessage({
+    state.previewPanel.webview.postMessage({
       kind: "navigate",
       previewBaseUrl,
       componentId,
     });
   }
-  previewPanel.reveal(vscode.ViewColumn.Beside, true);
+  state.previewPanel.reveal(vscode.ViewColumn.Beside, true);
 }
 
-export function closePreviewPanel() {
-  if (previewPanel) {
-    previewPanel.dispose();
-    previewPanel = null;
+export function closePreviewPanel(state: PreviewJsState) {
+  if (state.previewPanel) {
+    state.previewPanel.dispose();
+    state.previewPanel = null;
   }
 }
diff --git a/integrations/vscode/src/preview-server.ts b/integrations/vscode/src/preview-server.ts
index 162eef100ad..4ba16e54ce3 100644
--- a/integrations/vscode/src/preview-server.ts
+++ b/integrations/vscode/src/preview-server.ts
@@ -1,45 +1,38 @@
-import type { Client } from "@previewjs/daemon/client";
 import { exclusivePromiseRunner } from "exclusive-promises";
-
-let currentPreview:
-  | {
-      workspaceId: string;
-      url: string;
-    }
-  | undefined = undefined;
+import { PreviewJsState } from "./state";
 
 const locking = exclusivePromiseRunner();
 export async function ensurePreviewServerStarted(
-  client: Client,
+  state: PreviewJsState,
   workspaceId: string
 ) {
   return locking(async () => {
-    if (currentPreview?.workspaceId !== workspaceId) {
-      if (currentPreview) {
-        await client.stopPreview({
-          workspaceId: currentPreview.workspaceId,
+    if (state.currentPreview?.workspaceId !== workspaceId) {
+      if (state.currentPreview) {
+        await state.client.stopPreview({
+          workspaceId: state.currentPreview.workspaceId,
         });
       }
-      const { url } = await client.startPreview({
+      const { url } = await state.client.startPreview({
         workspaceId,
       });
-      currentPreview = {
+      state.currentPreview = {
         workspaceId,
         url,
       };
     }
-    return currentPreview;
+    return state.currentPreview;
   });
 }
 
-export async function ensurePreviewServerStopped(client: Client) {
+export async function ensurePreviewServerStopped(state: PreviewJsState) {
   return locking(async () => {
-    if (!currentPreview) {
+    if (!state.currentPreview) {
       return;
     }
-    await client.stopPreview({
-      workspaceId: currentPreview.workspaceId,
+    await state.client.stopPreview({
+      workspaceId: state.currentPreview.workspaceId,
     });
-    currentPreview = undefined;
+    state.currentPreview = null;
   });
 }
diff --git a/integrations/vscode/src/start-daemon.ts b/integrations/vscode/src/start-daemon.ts
index 14508370286..d9d228763a1 100644
--- a/integrations/vscode/src/start-daemon.ts
+++ b/integrations/vscode/src/start-daemon.ts
@@ -1,8 +1,15 @@
-import { createClient } from "@previewjs/daemon/client";
 import type { Client } from "@previewjs/daemon/client";
-import { execa } from "execa";
+import { createClient } from "@previewjs/daemon/client";
 import type { ExecaChildProcess, ExecaReturnValue, Options } from "execa";
-import { closeSync, openSync, readFileSync, utimesSync, watch } from "fs";
+import { execa } from "execa";
+import {
+  FSWatcher,
+  closeSync,
+  openSync,
+  readFileSync,
+  utimesSync,
+  watch,
+} from "fs";
 import path from "path";
 import stripAnsi from "strip-ansi";
 import type { OutputChannel } from "vscode";
@@ -11,11 +18,15 @@ import { clientId } from "./client-id";
 
 const port = process.env.PREVIEWJS_PORT || "9315";
 const logsPath = path.join(__dirname, "daemon.log");
-const daemonLockFilePath = path.join(__dirname, "process.lock");
+
+export const daemonLockFilePath = path.join(__dirname, "process.lock");
 
 export async function ensureDaemonRunning(
   outputChannel: OutputChannel
-): Promise<Client | null> {
+): Promise<{
+  client: Client;
+  watcher: FSWatcher;
+} | null> {
   const client = createClient(`http://localhost:${port}`);
   if (!(await startDaemon(outputChannel))) {
     return null;
@@ -24,12 +35,15 @@ export async function ensureDaemonRunning(
   // daemon running already (e.g. from another workspace) because of the lock file. This is
   // fine and working by design.
   const ready = streamDaemonLogs(outputChannel);
-  await ready;
+  const watcher = await ready;
   await client.updateClientStatus({
     clientId,
     alive: true,
   });
-  return client;
+  return {
+    client,
+    watcher,
+  };
 }
 
 // Important: we wrap daemonProcess into a Promise so that awaiting startDaemon()
@@ -125,8 +139,8 @@ async function startDaemon(outputChannel: OutputChannel): Promise<{
   return { daemonProcess };
 }
 
-function streamDaemonLogs(outputChannel: OutputChannel) {
-  const ready = new Promise<void>((resolve) => {
+function streamDaemonLogs(outputChannel: OutputChannel): Promise<FSWatcher> {
+  const ready = new Promise<FSWatcher>((resolve) => {
     let lastKnownLogsLength = 0;
     let resolved = false;
     // Ensure file exists before watching.
@@ -190,14 +204,14 @@ function streamDaemonLogs(outputChannel: OutputChannel) {
           );
         }
         if (!resolved && logsContent.includes("[ready]")) {
-          resolve();
+          resolve(watcher);
           resolved = true;
         }
       } catch (e: any) {
         // Fine, ignore. It just means log streaming is broken.
       }
     };
-    watch(
+    const watcher = watch(
       logsPath,
       {
         persistent: false,
diff --git a/integrations/vscode/src/state.ts b/integrations/vscode/src/state.ts
new file mode 100644
index 00000000000..ba9ac9a185e
--- /dev/null
+++ b/integrations/vscode/src/state.ts
@@ -0,0 +1,69 @@
+import { Client } from "@previewjs/daemon/client";
+import vscode, { WebviewPanel } from "vscode";
+import {
+  ComponentDetector,
+  createComponentDetector,
+} from "./component-detector";
+import { ensureDaemonRunning } from "./start-daemon";
+import { Workspaces, createWorkspaceGetter } from "./workspaces";
+
+export async function createState({
+  outputChannel,
+}: {
+  outputChannel: vscode.OutputChannel;
+}): Promise<PreviewJsState | null> {
+  const daemon = await ensureDaemonRunning(outputChannel)
+    .catch((e) => {
+      outputChannel.appendLine(e.stack);
+      return null;
+    })
+    .then((p) => {
+      if (!p) {
+        outputChannel.appendLine("Preview.js daemon could not be started.");
+        outputChannel.show();
+        return null;
+      }
+      return p;
+    });
+
+  if (!daemon) {
+    return null;
+  }
+
+  const workspaces: Workspaces = {};
+  const getWorkspaceId = createWorkspaceGetter(
+    daemon.client,
+    outputChannel,
+    workspaces
+  );
+  const getComponents = createComponentDetector(daemon.client, getWorkspaceId);
+  const state: PreviewJsState = {
+    client: daemon.client,
+    dispose: () => {
+      daemon.watcher.close();
+      if (state.previewPanel) {
+        state.previewPanel.dispose();
+        state.previewPanel = null;
+      }
+    },
+    workspaces,
+    getComponents,
+    getWorkspaceId,
+    previewPanel: null,
+    currentPreview: null,
+  };
+  return state;
+}
+
+export type PreviewJsState = {
+  client: Client;
+  dispose: () => void;
+  workspaces: Workspaces;
+  previewPanel: WebviewPanel | null;
+  currentPreview: {
+    workspaceId: string;
+    url: string;
+  } | null;
+  getComponents: ComponentDetector;
+  getWorkspaceId: (document: vscode.TextDocument) => Promise<string | null>;
+};
diff --git a/integrations/vscode/src/workspaces.ts b/integrations/vscode/src/workspaces.ts
index 8b0127668bd..53eb1b46557 100644
--- a/integrations/vscode/src/workspaces.ts
+++ b/integrations/vscode/src/workspaces.ts
@@ -3,28 +3,36 @@ import path from "path";
 import type vscode from "vscode";
 import type { OutputChannel } from "vscode";
 
-export function createWorkspaceGetter(outputChannel: OutputChannel) {
-  const workspaceIds = new Set<string>();
-
-  return async function (
-    previewjsClient: Client,
-    document: vscode.TextDocument
-  ) {
+export function createWorkspaceGetter(
+  client: Client,
+  outputChannel: OutputChannel,
+  workspaces: Workspaces
+): WorkspaceGetter {
+  return async function (document: vscode.TextDocument) {
     if (!path.isAbsolute(document.fileName)) {
       return null;
     }
-    const workspace = await previewjsClient.getWorkspace({
+    const workspace = await client.getWorkspace({
       absoluteFilePath: document.fileName,
     });
     if (!workspace.workspaceId) {
       return null;
     }
-    if (!workspaceIds.has(workspace.workspaceId)) {
+    if (!workspaces[workspace.workspaceId]) {
       outputChannel.appendLine(
         `✨ Created Preview.js workspace for: ${workspace.rootDirPath}`
       );
-      workspaceIds.add(workspace.workspaceId);
+      workspaces[workspace.workspaceId] = workspace.rootDirPath;
     }
     return workspace.workspaceId;
   };
 }
+
+export type Workspaces = Record<
+  string /* workspaceId */,
+  string /* rootDirPath */
+>;
+
+export type WorkspaceGetter = (
+  document: vscode.TextDocument
+) => Promise<string | null>;