plumb VS Code getSession through to cell environment

githubnext · Jun 28, 2024 · f2f22d4 · f2f22d4
1 parent 22cec49
commit f2f22d4
Show file tree

Hide file tree

Showing 9 changed files with 272 additions and 129 deletions.
diff --git a/packages/server/package.json b/packages/server/package.json
@@ -26,6 +26,7 @@
     "@babel/parser": "^7.24.4",
     "@babel/traverse": "^7.24.1",
     "@babel/types": "^7.24.0",
+    "@types/vscode": "^1.90.0",
     "birpc": "^0.2.17",
     "cac": "^6.7.14",
     "cors": "^2.8.5",

diff --git a/packages/server/src/client.ts b/packages/server/src/client.ts
@@ -0,0 +1,125 @@
+import { Cells } from "./cells";
+import { executeCell } from "./executeCell";
+import { Runtime } from "./runtime";
+import { Rpc } from "./rpc";
+
+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}"`);
+  }
+}
+
+export class Client {
+  private rpc: Rpc;
+  private cells: Cells;
+  private runtime: Runtime;
+
+  constructor(rpc: Rpc, cells: Cells, runtime: Runtime) {
+    this.rpc = rpc;
+    this.cells = cells;
+    this.runtime = runtime;
+  }
+
+  async executeCells(
+    cells: {
+      path: string;
+      cellId: string;
+      language: string;
+      code?: string;
+    }[],
+    force: boolean,
+    executeDirtyCells: boolean
+  ) {
+    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
+
+    for (const { path, cellId, language, code } of cells) {
+      const ext = extOfLanguage(language);
+      const id = `${path}-cellId=${cellId}.${ext}`;
+      if (code) {
+        this.cells.set(id, { cellId, code, language });
+      }
+
+      this.runtime.invalidateServerModule(id);
+      this.runtime.handleHMRUpdate(id);
+      this.runtime.invalidateRuntimeModule(id, dirtyCells);
+    }
+
+    dirtyCells = dirtyCells.filter(
+      (dirtyCell) =>
+        !cells.some(
+          (cell) =>
+            cell.path === dirtyCell.path && cell.cellId === dirtyCell.cellId
+        )
+    );
+
+    const cellsToExecute = [
+      ...cells.map(({ path, cellId, language }) => ({
+        path,
+        cellId,
+        ext: extOfLanguage(language),
+        force,
+      })),
+      ...(executeDirtyCells
+        ? dirtyCells.map((cell) => ({ ...cell, force: false }))
+        : []),
+    ];
+
+    const executed = await Promise.all(
+      cellsToExecute.map(({ path, cellId, ext, force }) =>
+        executeCell(
+          this.rpc,
+          this.cells,
+          this.runtime,
+          `${path}-cellId=${cellId}.${ext}`,
+          path,
+          cellId,
+          force
+        )
+      )
+    );
+
+    const cellsToMarkDirty = cellsToExecute.filter(
+      ({ force }, i) => !force && !executed[i]
+    );
+    this.rpc.markCellsDirty(cellsToMarkDirty);
+  }
+
+  removeCells(
+    cells: {
+      path: string;
+      cellId: string;
+      language: string;
+    }[]
+  ) {
+    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
+
+    for (const { path, cellId, language } of cells) {
+      const ext = extOfLanguage(language);
+      const id = `${path}-cellId=${cellId}.${ext}`;
+      this.cells.delete(id);
+
+      this.runtime.invalidateServerModule(id);
+      // TODO(jaked) HMR remove?
+      this.runtime.invalidateRuntimeModule(id, dirtyCells);
+    }
+
+    // don't mark cells dirty if they were just removed
+    dirtyCells = dirtyCells.filter(
+      (dirtyCell) =>
+        !cells.some(
+          (cell) =>
+            cell.path === dirtyCell.path && cell.cellId === dirtyCell.cellId
+        )
+    );
+    this.rpc.markCellsDirty(dirtyCells);
+  }
+}
diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts
@@ -1,3 +1,4 @@
+import * as vscode from "vscode";
 export * from "./rpc-types";
 
 export function text(data: string) {
@@ -41,3 +42,33 @@ export function json(obj: object) {
 export function jsonView(obj: object) {
   return { data: JSON.stringify(obj), mime: "application/x-json-view" };
 }
+
+export function getSession(
+  providerId: string,
+  scopes: readonly string[],
+  options: vscode.AuthenticationGetSessionOptions & {
+    /** */ createIfNone: true;
+  }
+): Thenable<vscode.AuthenticationSession>;
+export function getSession(
+  providerId: string,
+  scopes: readonly string[],
+  options: vscode.AuthenticationGetSessionOptions & {
+    /** literal-type defines return type */ forceNewSession:
+      | true
+      | vscode.AuthenticationForceNewSessionOptions;
+  }
+): Thenable<vscode.AuthenticationSession>;
+export function getSession(
+  providerId: string,
+  scopes: readonly string[],
+  options?: vscode.AuthenticationGetSessionOptions
+): Thenable<vscode.AuthenticationSession | undefined>;
+export function getSession(
+  providerId: string,
+  scopes: readonly string[],
+  options?: vscode.AuthenticationGetSessionOptions
+) {
+  // @ts-ignore
+  return global["__vitale_rpc__"].getSession(providerId, scopes, options);
+}
diff --git a/packages/server/src/rpc-types.ts b/packages/server/src/rpc-types.ts
@@ -1,3 +1,5 @@
+import * as vscode from "vscode";
+
 export interface CellOutputItem {
   data: number[]; // Uint8Array
   mime: string;
@@ -49,4 +51,7 @@ export type ClientFunctions = {
     cellId: string,
     cellOutput?: CellOutput
   ) => void;
+
+  // VS Code API
+  getSession: typeof vscode.authentication.getSession;
 };
diff --git a/packages/server/src/rpc.ts b/packages/server/src/rpc.ts
@@ -1,38 +1,15 @@
+import * as vscode from "vscode";
 import { createBirpc, type BirpcReturn } from "birpc";
 import JSON5 from "json5";
 import type WebSocket from "ws";
-import { Cells } from "./cells";
-import { executeCell } from "./executeCell";
 import type { CellOutput, ClientFunctions, ServerFunctions } from "./rpc-types";
-import { Runtime } from "./runtime";
-
-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}"`);
-  }
-}
+import { Client } from "./client";
 
 export class Rpc {
   private clients: Map<
     WebSocket,
     BirpcReturn<ClientFunctions, ServerFunctions>
   > = new Map();
-  private cells: Cells;
-  private runtime: Runtime;
-
-  constructor(cells: Cells, runtime: Runtime) {
-    this.cells = cells;
-    this.runtime = runtime;
-  }
 
   startCellExecution(path: string, cellId: string, force: boolean) {
     return Promise.all(
@@ -70,6 +47,21 @@ export class Rpc {
     );
   }
 
+  async getSession(
+    providerId: string,
+    scopes: readonly string[],
+    options: vscode.AuthenticationGetSessionOptions
+  ) {
+    // TODO(jaked)
+    // it doesn't make sense to call this for all clients
+    const sessions = await Promise.all(
+      Array.from(this.clients.values()).map((client) =>
+        client.getSession(providerId, scopes, options)
+      )
+    );
+    return sessions[0];
+  }
+
   markCellsDirty(cells: { path: string; cellId: string }[]) {
     if (cells.length === 0) {
       return;
@@ -79,8 +71,7 @@ export class Rpc {
     }
   }
 
-  setupClient(ws: WebSocket) {
-    const self = this;
+  setupClient(ws: WebSocket, client: Client) {
     const rpc = createBirpc<ClientFunctions, ServerFunctions>(
       {
         ping: async () => {
@@ -89,14 +80,14 @@ export class Rpc {
         },
         async executeCells(cells, force, executeDirtyCells) {
           try {
-            return self.executeCellsRPC(cells, force, executeDirtyCells);
+            return client.executeCells(cells, force, executeDirtyCells);
           } catch (e) {
             console.error(e);
           }
         },
         async removeCells(cells) {
           try {
-            return self.removeCellsRPC(cells);
+            return client.removeCells(cells);
           } catch (e) {
             console.error(e);
           }
@@ -115,98 +106,4 @@ export class Rpc {
       this.clients.delete(ws);
     });
   }
-
-  private async executeCellsRPC(
-    cells: {
-      path: string;
-      cellId: string;
-      language: string;
-      code?: string;
-    }[],
-    force: boolean,
-    executeDirtyCells: boolean
-  ) {
-    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
-
-    for (const { path, cellId, language, code } of cells) {
-      const ext = extOfLanguage(language);
-      const id = `${path}-cellId=${cellId}.${ext}`;
-      if (code) {
-        this.cells.set(id, { cellId, code, language });
-      }
-
-      this.runtime.invalidateServerModule(id);
-      this.runtime.handleHMRUpdate(id);
-      this.runtime.invalidateRuntimeModule(id, dirtyCells);
-    }
-
-    dirtyCells = dirtyCells.filter(
-      (dirtyCell) =>
-        !cells.some(
-          (cell) =>
-            cell.path === dirtyCell.path && cell.cellId === dirtyCell.cellId
-        )
-    );
-
-    const cellsToExecute = [
-      ...cells.map(({ path, cellId, language }) => ({
-        path,
-        cellId,
-        ext: extOfLanguage(language),
-        force,
-      })),
-      ...(executeDirtyCells
-        ? dirtyCells.map((cell) => ({ ...cell, force: false }))
-        : []),
-    ];
-
-    const executed = await Promise.all(
-      cellsToExecute.map(({ path, cellId, ext, force }) =>
-        executeCell(
-          this,
-          this.cells,
-          this.runtime,
-          `${path}-cellId=${cellId}.${ext}`,
-          path,
-          cellId,
-          force
-        )
-      )
-    );
-
-    const cellsToMarkDirty = cellsToExecute.filter(
-      ({ force }, i) => !force && !executed[i]
-    );
-    this.markCellsDirty(cellsToMarkDirty);
-  }
-
-  private removeCellsRPC(
-    cells: {
-      path: string;
-      cellId: string;
-      language: string;
-    }[]
-  ) {
-    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
-
-    for (const { path, cellId, language } of cells) {
-      const ext = extOfLanguage(language);
-      const id = `${path}-cellId=${cellId}.${ext}`;
-      this.cells.delete(id);
-
-      this.runtime.invalidateServerModule(id);
-      // TODO(jaked) HMR remove?
-      this.runtime.invalidateRuntimeModule(id, dirtyCells);
-    }
-
-    // don't mark cells dirty if they were just removed
-    dirtyCells = dirtyCells.filter(
-      (dirtyCell) =>
-        !cells.some(
-          (cell) =>
-            cell.path === dirtyCell.path && cell.cellId === dirtyCell.cellId
-        )
-    );
-    this.markCellsDirty(dirtyCells);
-  }
 }