setting to auto-rerun cells when deps change

packages/server/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@githubnext/vitale",
-  "version": "0.0.5",
+  "version": "0.0.6",
   "description": "",
   "license": "MIT",
   "type": "module",

packages/server/src/server.ts

@@ -46,6 +46,21 @@ function isHTMLElementLike(obj: unknown): obj is PossibleHTML {
 
 const cellIdRegex = /^([^?]+\.vnb)\?cellId=([a-zA-z0-9_-]{21})\.([a-z]+)$/;
 
+function extOfLanguage(language: string): string {
+  switch (language) {
+    case "typescriptreact":
+      return "tsx";
+    case "typescript":
+      return "ts";
+    case "javascriptreact":
+      return "jsx";
+    case "javascript":
+      return "js";
+    default:
+      throw new Error(`unknown language "${language}"`);
+  }
+}
+
 class VitaleDevServer {
   static async construct(options: Options) {
     const cells: Map<string, SourceDescription> = new Map();
@@ -216,14 +231,18 @@ class VitaleDevServer {
     );
   }
 
-  private invalidateModule(id: string) {
+  private invalidateModule(id: string, dirty: boolean = true) {
     const mod = this.viteRuntime.moduleCache.get(id);
     this.viteRuntime.moduleCache.delete(id);
 
-    const match = cellIdRegex.exec(id);
-    if (match) {
-      const [_, path, cellId] = match;
-      this.executeCell(id, path, cellId);
+    if (dirty) {
+      const match = cellIdRegex.exec(id);
+      if (match) {
+        const [_, path, cellId] = match;
+        Array.from(this.clients.values()).forEach((client) =>
+          client.dirtyCell(path, cellId)
+        );
+      }
     }
 
     for (const dep of mod.importers ?? []) {
@@ -237,27 +256,16 @@ class VitaleDevServer {
     language: string,
     code: string
   ) {
-    const ext = (() => {
-      switch (language) {
-        case "typescriptreact":
-          return "tsx";
-        case "typescript":
-          return "ts";
-        case "javascriptreact":
-          return "jsx";
-        case "javascript":
-          return "js";
-        default:
-          throw new Error(`unknown language "${language}"`);
-      }
-    })();
+    const ext = extOfLanguage(language);
     const id = `${path}?cellId=${cellId}.${ext}`;
     this.cells.set(id, rewrite(code, language, cellId));
 
     const mod = this.viteServer.moduleGraph.getModuleById(id);
     if (mod) this.viteServer.moduleGraph.invalidateModule(mod);
 
-    this.invalidateModule(id);
+    this.invalidateModule(id, false);
+
+    this.executeCell(id, path, cellId);
   }
 
   private setupClient(ws: WebSocket) {

packages/server/src/types.ts

@@ -21,6 +21,7 @@ export type ServerFunctions = {
 };
 
 export type ClientFunctions = {
+  dirtyCell: (path: string, cellId: string) => void;
   startCellExecution: (path: string, cellId: string) => void;
   endCellExecution: (
     path: string,

packages/vscode/package.json

@@ -2,7 +2,7 @@
   "name": "vitale-vscode",
   "displayName": "Vitale",
   "description": "",
-  "version": "0.0.5",
+  "version": "0.0.6",
   "license": "MIT",
   "engines": {
     "vscode": "^1.87.0"
@@ -23,6 +23,12 @@
         "command": "vitale.restartKernel",
         "category": "Vitale",
         "icon": "$(refresh)"
+      },
+      {
+        "title": "Run Dirty",
+        "command": "vitale.runDirty",
+        "category": "Vitale",
+        "icon": "$(run-all)"
       }
     ],
     "menus": {
@@ -31,6 +37,11 @@
           "command": "vitale.restartKernel",
           "group": "navigation",
           "when": "notebookType == 'vitale-notebook'"
+        },
+        {
+          "command": "vitale.runDirty",
+          "group": "navigation",
+          "when": "notebookType == 'vitale-notebook'"
         }
       ]
     },
@@ -62,7 +73,18 @@
           "application/x-vitale"
         ]
       }
-    ]
+    ],
+    "configuration": {
+      "title": "Vitale",
+      "properties": {
+        "vitale.rerunCellsWhenDirty": {
+          "type": "boolean",
+          "default": true,
+          "description": "Automatically re-run cells when their dependencies change",
+          "scope": "resource"
+        }
+      }
+    }
   },
   "scripts": {
     "package": "vsce package --no-dependencies",

packages/vscode/src/cellStatusBarItemProvider.ts

@@ -0,0 +1,22 @@
+import * as vscode from "vscode";
+
+export class NotebookCellStatusBarItemProvider
+  implements vscode.NotebookCellStatusBarItemProvider
+{
+  provideCellStatusBarItems(cell: vscode.NotebookCell) {
+    // TODO(jaked)
+    // provideCellStatusBarItems is called on cell edits
+    // but the cell.document.dirty flag is always false
+    if (cell.metadata.dirty || cell.document.isDirty) {
+      const item = new vscode.NotebookCellStatusBarItem(
+        "$(circle-filled)",
+        vscode.NotebookCellStatusBarAlignment.Right
+      );
+      // TODO(jaked)
+      // should disable this while the cell is executing
+      item.command = "notebook.cell.execute";
+      return item;
+    }
+    return;
+  }
+}

packages/vscode/src/controller.ts

@@ -21,6 +21,11 @@ function cellOutputToNotebookCellOutput(cellOutput: CellOutput) {
   );
 }
 
+function getRerunCellsWhenDirty(uri: vscode.Uri) {
+  const config = vscode.workspace.getConfiguration("vitale", uri);
+  return config.get("rerunCellsWhenDirty", true);
+}
+
 type Client = BirpcReturn<ServerFunctions, ClientFunctions>;
 
 type State =
@@ -76,12 +81,12 @@ export class NotebookController {
     this.run("idle");
   }
 
-  private resolve(client: Client) {
+  private resolveClient(client: Client) {
     this._clientWaiters.forEach((waiter) => waiter.resolve(client));
     this._clientWaiters = [];
   }
 
-  private reject(error: Error) {
+  private rejectClient(error: Error) {
     this._clientWaiters.forEach((waiter) => waiter.reject(error));
     this._clientWaiters = [];
   }
@@ -116,14 +121,15 @@ export class NotebookController {
       this.run("idle");
     });
     process.on("error", (error) => {
-      this.reject(error);
+      this.rejectClient(error);
       this.run("start-failed");
     });
   }
 
   private makeClient(ws: WebSocket) {
     return createBirpc<ServerFunctions, ClientFunctions>(
       {
+        dirtyCell: this.dirtyCell.bind(this),
         startCellExecution: this.startCellExecution.bind(this),
         endCellExecution: this.endCellExecution.bind(this),
       },
@@ -173,7 +179,7 @@ export class NotebookController {
 
   run(state: State) {
     if (this._state === "disposed") {
-      this.reject(new Error("disposed"));
+      this.rejectClient(new Error("disposed"));
       return;
     }
     this._state = state;
@@ -186,7 +192,7 @@ export class NotebookController {
         break;
 
       case "start-failed":
-        this.reject(new Error(`Couldn't start Vitale server`));
+        this.rejectClient(new Error(`Couldn't start Vitale server`));
         vscode.window.showErrorMessage(
           `Couldn't start Vitale server; is @githubnext/vitale installed?`
         );
@@ -204,14 +210,14 @@ export class NotebookController {
         break;
 
       case "connect-failed":
-        this.reject(new Error(`Couldn't connect to Vitale server`));
+        this.rejectClient(new Error(`Couldn't connect to Vitale server`));
         vscode.window.showErrorMessage(`Couldn't connect to Vitale server`);
         break;
 
       case "connected":
         this._tries = RECONNECT_TRIES;
         this._client = this.makeClient(this._websocket!);
-        this.resolve(this._client);
+        this.resolveClient(this._client);
         break;
 
       default:
@@ -222,11 +228,21 @@ export class NotebookController {
   }
 
   restartKernel() {
+    // TODO(jaked)
+    // should clear outputs and dirty all cells
     if (this._process && this._process.pid) {
       kill(this._process.pid);
     }
   }
 
+  async runDirty(notebookUri: string) {
+    const notebook = await vscode.workspace.openNotebookDocument(
+      vscode.Uri.parse(notebookUri)
+    );
+    const cells = notebook.getCells().filter((cell) => cell.metadata.dirty);
+    this._executeAll(cells, notebook);
+  }
+
   private makeClientPromise() {
     return new Promise<Client>((resolve, reject) => {
       this._clientWaiters.push({ resolve, reject });
@@ -269,12 +285,33 @@ export class NotebookController {
     this._controller.dispose();
   }
 
+  private setCellDirty(cell: vscode.NotebookCell, dirty: boolean) {
+    const metadata = { ...(cell.metadata ?? {}), dirty };
+    const edit = new vscode.WorkspaceEdit();
+    edit.set(cell.notebook.uri, [
+      vscode.NotebookEdit.updateCellMetadata(cell.index, metadata),
+    ]);
+    vscode.workspace.applyEdit(edit);
+  }
+
+  private async dirtyCell(path: string, id: string) {
+    const uri = vscode.Uri.file(path);
+    const notebook = await vscode.workspace.openNotebookDocument(uri);
+    const cell = notebook.getCells().find((cell) => cell.metadata.id === id);
+    if (cell) {
+      this.setCellDirty(cell, true);
+      if (getRerunCellsWhenDirty(uri)) {
+        this._executeAll([cell], notebook);
+      }
+    }
+  }
+
   private async startCellExecution(path: string, id: string) {
     console.log(`client startCellExecution`, path, id);
-    const document = await vscode.workspace.openNotebookDocument(
+    const notebook = await vscode.workspace.openNotebookDocument(
       vscode.Uri.file(path)
     );
-    const cell = document.getCells().find((cell) => cell.metadata.id === id);
+    const cell = notebook.getCells().find((cell) => cell.metadata.id === id);
     if (cell) {
       const execution = this._controller.createNotebookCellExecution(cell);
       execution.token.onCancellationRequested(() => {
@@ -312,16 +349,17 @@ export class NotebookController {
       execution.replaceOutput(notebookCellOutput);
       execution.end(true, Date.now());
       this._executions.delete(key);
+
+      this.setCellDirty(execution.cell, false);
     }
   }
 
   private _executeAll(
     cells: vscode.NotebookCell[],
-    _notebook: vscode.NotebookDocument,
-    _controller: vscode.NotebookController
+    notebook: vscode.NotebookDocument
   ): void {
     for (const cell of cells) {
-      this._doExecution(_notebook.uri.fsPath, cell);
+      this._doExecution(notebook.uri.fsPath, cell);
     }
   }
 

packages/vscode/src/extension.ts

@@ -1,6 +1,7 @@
 import * as vscode from "vscode";
 import { NotebookSerializer } from "./serializer";
 import { NotebookController } from "./controller";
+import { NotebookCellStatusBarItemProvider } from "./cellStatusBarItemProvider";
 import { handleDidChangeNotebookDocument } from "./handleDidChangeNotebookDocument";
 
 export function activate(context: vscode.ExtensionContext) {
@@ -12,6 +13,12 @@ export function activate(context: vscode.ExtensionContext) {
     vscode.commands.registerCommand("vitale.restartKernel", () => {
       controller.restartKernel();
     }),
+    vscode.commands.registerCommand(
+      "vitale.runDirty",
+      (ctx: { notebookEditor: { notebookUri: string } }) => {
+        controller.runDirty(ctx.notebookEditor.notebookUri);
+      }
+    ),
     vscode.workspace.registerNotebookSerializer(
       "vitale-notebook",
       new NotebookSerializer(),
@@ -20,7 +27,11 @@ export function activate(context: vscode.ExtensionContext) {
     vscode.workspace.onDidChangeNotebookDocument(
       handleDidChangeNotebookDocument
     ),
-    controller
+    controller,
+    vscode.notebooks.registerNotebookCellStatusBarItemProvider(
+      "vitale-notebook",
+      new NotebookCellStatusBarItemProvider()
+    )
   );
 }
 

packages/vscode/src/handleDidChangeNotebookDocument.ts

@@ -13,7 +13,6 @@ function uniqueId(doc: NotebookDocument): string {
 export function handleDidChangeNotebookDocument(
   e: NotebookDocumentChangeEvent
 ) {
-  const edit = new WorkspaceEdit();
   const edits: NotebookEdit[] = [];
   for (const contentChange of e.contentChanges) {
     for (const cell of contentChange.addedCells) {
@@ -23,6 +22,7 @@ export function handleDidChangeNotebookDocument(
     }
   }
   if (edits.length > 0) {
+    const edit = new WorkspaceEdit();
     edit.set(e.notebook.uri, edits);
     workspace.applyEdit(edit);
   }

packages/vscode/src/serializer.ts

@@ -3,6 +3,7 @@ import JSON5 from "json5";
 
 interface NotebookCellMetadata {
   id: string;
+  dirty: boolean;
 }
 
 interface NotebookCell {
@@ -42,7 +43,7 @@ export class NotebookSerializer implements vscode.NotebookSerializer {
         item.value,
         item.language
       );
-      cellData.metadata = item.metadata;
+      cellData.metadata = { ...item.metadata, dirty: true };
       return cellData;
     });