Wire up a code lens command with the webview (#457)

* Wire up a code lens command with the webview

* fixed import

new-packages/language-server/src/index.ts

@@ -9,7 +9,7 @@ import {
   Provide,
   create as createTypeScriptService,
 } from 'volar-service-typescript';
-import { helloRequest } from './protocol';
+import { getMachineAtIndex } from './protocol';
 
 const projectCache = new WeakMap<Program, XStateProject>();
 
@@ -31,40 +31,58 @@ connection.onInitialize((params) => {
       'ts',
       'tsx',
     ],
-    getServicePlugins: () => [createTypeScriptService(getTsLib())],
+    getServicePlugins: () => {
+      const service = createTypeScriptService(getTsLib());
+      return [
+        service,
+        {
+          create: () => {
+            return {
+              provideCodeLenses: async (textDocument) => {
+                const xstateProject = await getXStateProject(textDocument.uri);
+                if (!xstateProject) {
+                  return [];
+                }
+
+                // TODO: a range is returned here regardless of the extraction status (extraction could error)
+                // DX has to account for this somehow or results with errors have to be ignored (this would be slower but it might be a good tradeoff)
+                return xstateProject
+                  .findMachines(server.env.uriToFileName(textDocument.uri))
+                  .map((range, index) => ({
+                    command: {
+                      title: 'Open Visual Editor',
+                      command: 'stately-xstate/edit-machine',
+                      arguments: [textDocument.uri, index],
+                    },
+                    range,
+                  }));
+              },
+            };
+          },
+        },
+      ];
+    },
     getLanguagePlugins: () => [],
   });
 });
 
-connection.onInitialized(() => {
-  server.initialized();
-});
-
-connection.onRequest(helloRequest, async ({ textDocument, name }) => {
-  if (!textDocument) {
-    connection.console.log(`Hello, ${name}!`);
-    return true;
-  }
-  const tsProgram = (
-    await getTypeScriptLanguageService(textDocument.uri)
-  ).getProgram();
+connection.onRequest(getMachineAtIndex, async ({ uri, machineIndex }) => {
+  const xstateProject = await getXStateProject(uri);
 
-  if (!tsProgram) {
-    return false;
+  if (!xstateProject) {
+    return;
   }
 
-  const xstateProject = getOrCreateXStateProject(
-    await getTypeScriptModule(textDocument.uri),
-    tsProgram,
-  );
-
-  const machines = xstateProject.extractMachines(
-    server.env.uriToFileName(textDocument.uri),
-  );
+  // TODO: it would be faster to extract a single machine instead of all of them
+  const [digraph] = xstateProject.extractMachines(
+    server.env.uriToFileName(uri),
+  )[machineIndex];
 
-  connection.console.log(JSON.stringify(machines));
+  return digraph;
+});
 
-  return true;
+connection.onInitialized(() => {
+  server.initialized();
 });
 
 connection.onShutdown(() => {
@@ -93,15 +111,21 @@ async function getTypeScriptLanguageService(uri: string) {
     );
 }
 
-function getOrCreateXStateProject(
-  ts: typeof import('typescript'),
-  tsProgram: Program,
-) {
+async function getXStateProject(uri: string) {
+  const tsProgram = (await getTypeScriptLanguageService(uri)).getProgram();
+
+  if (!tsProgram) {
+    return;
+  }
+
   const existing = projectCache.get(tsProgram);
   if (existing) {
     return existing;
   }
-  const xstateProject = createProject(ts, tsProgram);
+  const xstateProject = createProject(
+    await getTypeScriptModule(uri),
+    tsProgram,
+  );
   projectCache.set(tsProgram, xstateProject);
   return xstateProject;
 }

new-packages/language-server/src/protocol.ts

@@ -1,10 +1,11 @@
+import { ExtractorDigraphDef } from '@xstate/ts-project';
 import * as vscode from 'vscode-languageserver-protocol';
 
-export const helloRequest = new vscode.RequestType<
+export const getMachineAtIndex = new vscode.RequestType<
   {
-    textDocument?: vscode.TextDocumentIdentifier | undefined;
-    name: string;
+    uri: string;
+    machineIndex: number;
   },
-  boolean,
+  ExtractorDigraphDef | undefined,
   never
->('xstate/hello');
+>('stately-xstate/get-machine-at-index');

new-packages/ts-project/src/index.ts

@@ -1,12 +1,13 @@
 import type { CallExpression, Program, SourceFile } from 'typescript';
 import { extractState } from './state';
-import {
+import type {
   ExtractionContext,
   ExtractionError,
   ExtractorDigraphDef,
   TreeNode,
   XStateVersion,
 } from './types';
+export { ExtractorDigraphDef };
 
 function findCreateMachineCalls(
   ts: typeof import('typescript'),
@@ -138,12 +139,34 @@ export interface TSProjectOptions {
   xstateVersion?: XStateVersion | undefined;
 }
 
+interface Position {
+  line: number;
+  character: number;
+}
+
+interface Range {
+  start: Position;
+  end: Position;
+}
+
 export function createProject(
   ts: typeof import('typescript'),
   tsProgram: Program,
   { xstateVersion = '5' }: TSProjectOptions = {},
 ) {
   return {
+    findMachines: (fileName: string): Range[] => {
+      const sourceFile = tsProgram.getSourceFile(fileName);
+      if (!sourceFile) {
+        return [];
+      }
+      return findCreateMachineCalls(ts, sourceFile).map((call) => {
+        return {
+          start: sourceFile.getLineAndCharacterOfPosition(call.getStart()),
+          end: sourceFile.getLineAndCharacterOfPosition(call.getEnd()),
+        };
+      });
+    },
     extractMachines(fileName: string) {
       const sourceFile = tsProgram.getSourceFile(fileName);
       if (!sourceFile) {

new-packages/vscode-xstate/package.json

@@ -38,13 +38,6 @@
           }
         }
       }
-    ],
-    "commands": [
-      {
-        "command": "xstate.hello",
-        "title": "Hello from XState!",
-        "category": "XState"
-      }
     ]
   }
 }

new-packages/vscode-xstate/src/languageClient.ts

@@ -1,8 +1,16 @@
 import { RequestType, getTsdk } from '@volar/vscode';
 import { LanguageClient, TransportKind } from '@volar/vscode/node.js';
-import { helloRequest } from '@xstate/language-server/protocol';
+import { getMachineAtIndex } from '@xstate/language-server/protocol';
+import type { ExtractorDigraphDef } from '@xstate/ts-project';
 import * as vscode from 'vscode';
-import { ActorRef, Snapshot, fromCallback, fromPromise, setup } from 'xstate';
+import {
+  ActorRef,
+  Snapshot,
+  assertEvent,
+  fromCallback,
+  fromPromise,
+  setup,
+} from 'xstate';
 import { assertExtendedEvent } from './utils';
 import { webviewLogic } from './webview';
 
@@ -14,7 +22,7 @@ type LanguageClientInput = {
 export type LanguageClientEvent =
   | { type: 'SERVER_ENABLED_CHANGE' }
   | { type: 'WEBVIEW_CLOSED' }
-  | { type: 'OPEN_MACHINE' };
+  | { type: 'OPEN_MACHINE'; digraph: ExtractorDigraphDef };
 
 export const languageClientMachine = setup({
   types: {} as {
@@ -51,22 +59,34 @@ export const languageClientMachine = setup({
     ),
     registerCommands: fromCallback(
       ({
-        input: { languageClient },
+        input: { languageClient, parent },
       }: {
-        input: { languageClient: LanguageClient };
+        input: {
+          languageClient: LanguageClient;
+          parent: ActorRef<Snapshot<unknown>, LanguageClientEvent>;
+        };
       }) => {
         const disposable = vscode.commands.registerCommand(
-          'xstate.hello',
-          () => {
-            const activeTextEditorUri =
-              vscode.window.activeTextEditor?.document.uri.toString();
-            sendRequest(languageClient, helloRequest, {
-              name: 'Andarist',
-              textDocument: activeTextEditorUri
-                ? {
-                    uri: activeTextEditorUri,
-                  }
-                : undefined,
+          'stately-xstate/edit-machine',
+          (uri: string, machineIndex: number) => {
+            // TODO: actually use this token to cancel redundant requests
+            const tokenSource = new vscode.CancellationTokenSource();
+            sendRequest(
+              languageClient,
+              getMachineAtIndex,
+              {
+                uri,
+                machineIndex,
+              },
+              tokenSource.token,
+            ).then((digraph) => {
+              if (!digraph) {
+                return;
+              }
+              parent.send({
+                type: 'OPEN_MACHINE',
+                digraph,
+              });
             });
           },
         );
@@ -155,10 +175,10 @@ export const languageClientMachine = setup({
         },
         active: {
           invoke: {
-            // TODO: what about commands executed before we get here?
             src: 'registerCommands',
-            input: ({ context }) => ({
+            input: ({ context, self }) => ({
               languageClient: context.languageClient,
+              parent: self,
             }),
           },
           initial: 'webviewClosed',
@@ -172,10 +192,14 @@ export const languageClientMachine = setup({
               invoke: {
                 src: 'webviewLogic',
                 id: 'webview',
-                input: ({ context, self }) => ({
-                  extensionContext: context.extensionContext,
-                  parent: self,
-                }),
+                input: ({ context, event, self }) => {
+                  assertEvent(event, 'OPEN_MACHINE');
+                  return {
+                    extensionContext: context.extensionContext,
+                    parent: self,
+                    digraph: event.digraph,
+                  };
+                },
               },
               on: {
                 WEBVIEW_CLOSED: 'webviewClosed',
@@ -202,9 +226,10 @@ export const languageClientMachine = setup({
  * This mitigates the problem by defining 1 focused signature.
  */
 function sendRequest<P, R, E>(
-  client: LanguageClient | undefined,
+  client: LanguageClient,
   request: RequestType<P, R, E>,
   params: P,
-): Promise<R> | undefined {
-  return client?.sendRequest(request, params);
+  token: vscode.CancellationToken,
+): Promise<R> {
+  return client.sendRequest(request, params, token);
 }

new-packages/vscode-xstate/src/webview.ts

@@ -1,3 +1,4 @@
+import type { ExtractorDigraphDef } from '@xstate/ts-project';
 import * as vscode from 'vscode';
 import { ActorRef, EventObject, Snapshot, fromCallback } from 'xstate';
 import { LanguageClientEvent } from './languageClient';
@@ -25,6 +26,11 @@ function createWebviewPanel() {
 async function getWebviewHtml(
   extensionContext: vscode.ExtensionContext,
   webviewPanel: vscode.WebviewPanel,
+  {
+    digraph,
+  }: {
+    digraph: ExtractorDigraphDef;
+  },
 ) {
   const bundledEditorRootUri = vscode.Uri.joinPath(
     vscode.Uri.file(extensionContext.extensionPath),
@@ -43,15 +49,18 @@ async function getWebviewHtml(
   const theme =
     vscode.workspace.getConfiguration('xstate').get('theme') ?? 'dark';
 
-  const initialDataScript = `<script>window.__params = ${JSON.stringify({
-    themeKind:
-      theme === 'auto'
-        ? vscode.window.activeColorTheme.kind === vscode.ColorThemeKind.Dark
-          ? 'dark'
-          : 'light'
-        : theme,
-    distinctId: `vscode:${vscode.env.machineId}`,
-  })}</script>`;
+  const initialDataScript = `<script>window.__vscodeInitParams = ${JSON.stringify(
+    {
+      themeKind:
+        theme === 'auto'
+          ? vscode.window.activeColorTheme.kind === vscode.ColorThemeKind.Dark
+            ? 'dark'
+            : 'light'
+          : theme,
+      distinctId: `vscode:${vscode.env.machineId}`,
+      digraph,
+    },
+  )}</script>`;
 
   return htmlContent.replace('<head>', `<head>${baseTag}${initialDataScript}`);
 }
@@ -61,8 +70,9 @@ export const webviewLogic = fromCallback<
   {
     extensionContext: vscode.ExtensionContext;
     parent: ActorRef<Snapshot<unknown>, LanguageClientEvent>;
+    digraph: ExtractorDigraphDef;
   }
->(({ input: { extensionContext, parent }, receive }) => {
+>(({ input: { extensionContext, parent, digraph }, receive }) => {
   let canceled = false;
   const webviewPanel = createWebviewPanel();
 
@@ -80,7 +90,9 @@ export const webviewLogic = fromCallback<
   );
 
   (async () => {
-    const html = await getWebviewHtml(extensionContext, webviewPanel);
+    const html = await getWebviewHtml(extensionContext, webviewPanel, {
+      digraph,
+    });
     if (canceled) {
       return;
     }