feat: Connect to DHE worker

package.json

@@ -160,6 +160,10 @@
         "title": "Deephaven: Refresh Server Connection Tree",
         "icon": "$(refresh)"
       },
+      {
+        "command": "vscode-deephaven.requestDheUserCredentials",
+        "title": "Deephaven: Request DHE User Credentials"
+      },
       {
         "command": "vscode-deephaven.startServer",
         "title": "Deephaven: Start Server",
@@ -624,6 +628,10 @@
         {
           "command": "vscode-deephaven.refreshVariablePanels",
           "when": "false"
+        },
+        {
+          "command": "vscode-deephaven.requestDheUserCredentials",
+          "when": "false"
         }
       ],
       "editor/context": [
@@ -665,7 +673,7 @@
       "view/item/context": [
         {
           "command": "vscode-deephaven.connectToServer",
-          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerDisconnected || viewItem == isServerRunningDisconnected)",
+          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerDisconnected || viewItem == isServerRunningDisconnected || viewItem == isDHEServerRunning)",
           "group": "inline"
         },
         {
@@ -685,7 +693,7 @@
         },
         {
           "command": "vscode-deephaven.openInBrowser",
-          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerConnected || viewItem == isServerRunningConnected || viewItem == isServerRunningDisconnected || viewItem == isManagedServerDisconnected)",
+          "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerConnected || viewItem == isServerRunningConnected || viewItem == isServerRunningDisconnected || viewItem == isManagedServerDisconnected || viewItem == isDHEServerRunning)",
           "group": "inline"
         },
         {
@@ -730,6 +738,7 @@
     }
   },
   "devDependencies": {
+    "@deephaven-enterprise/jsapi-types": "^1.20240723.95-beta",
     "@deephaven/jsapi-types": "^1.0.0-dev0.35.3",
     "@types/node": "22.5.4",
     "@types/vscode": "^1.91.0",
@@ -757,7 +766,11 @@
     "wdio-vscode-service": "^6.1.0"
   },
   "dependencies": {
+    "@deephaven-enterprise/query-utils": "^1.20240723.95-beta",
     "@deephaven/require-jsapi": "file:./packages/require-jsapi",
-    "ws": "^8.18.0"
+    "nanoid": "^5.0.7"
+  },
+  "overrides": {
+    "event-target-shim": "^6.0.2"
   }
 }

packages/require-jsapi/package.json

@@ -15,5 +15,11 @@
   "sideEffects": false,
   "scripts": {
     "build": "tsc --build"
+  },
+  "dependencies": {
+    "ws": "^8.18.0"
+  },
+  "devDependencies": {
+    "@deephaven-enterprise/jsapi-types": "^1.20240723.95-beta"
   }
 }

packages/require-jsapi/src/dhe.ts

@@ -1,4 +1,8 @@
-import { hasStatusCode } from './serverUtils';
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import type { EnterpriseDhType as DheType } from '@deephaven-enterprise/jsapi-types';
+import { polyfillDh } from './polyfill';
+import { downloadFromURL, hasStatusCode } from './serverUtils';
 
 /**
  * Check if a given server is running by checking if the `irisapi/irisapi.nocache.js`
@@ -15,3 +19,51 @@ export async function isDheServerRunning(serverUrl: URL): Promise<boolean> {
     return false;
   }
 }
+
+/**
+ * Polyfill browser apis, download jsapi to a local directory, and return the
+ * default export.
+ * @param serverUrl URL of the server to download the jsapi from
+ * @param storageDir Directory to save the downloaded jsapi
+ * @returns Default export of downloaded jsapi
+ */
+export async function initDheApi(
+  serverUrl: URL,
+  storageDir: string
+): Promise<DheType> {
+  polyfillDh();
+  return getDhe(serverUrl, true, storageDir);
+}
+
+declare global {
+  export const iris: DheType;
+}
+
+/**
+ * Returns a CJS module that exports the Deephaven JS API. If the download flag
+ * is true, the API is downloaded from the server and saved to a local storage
+ * directory. If it is false, it will be returned from the local storage directory.
+ * @param serverUrl URL of the server to download the jsapi from
+ * @param download Whether to download the jsapi from the server
+ * @param storageDir Directory to save the downloaded jsapi and load it from
+ * @returns Default export of downloaded jsapi
+ */
+async function getDhe(
+  serverUrl: URL,
+  download: boolean,
+  storageDir: string
+): Promise<DheType> {
+  const dheFilePath = path.join(storageDir, 'irisapi.nocache.js');
+
+  if (download) {
+    const dhe = await downloadFromURL(
+      path.join(serverUrl.toString(), 'irisapi/irisapi.nocache.js')
+    );
+
+    fs.writeFileSync(dheFilePath, dhe);
+  }
+
+  require(dheFilePath);
+
+  return iris;
+}

packages/require-jsapi/src/polyfill.ts

@@ -1,24 +1,24 @@
 import ws from 'ws';
 
+class Event {
+  type: string;
+  detail: unknown;
+
+  constructor(type: string, dict: { detail: unknown }) {
+    this.type = type;
+    if (dict) {
+      this.detail = dict.detail;
+    }
+  }
+}
+
 export class CustomEvent extends Event {
   constructor(...args: ConstructorParameters<typeof Event>) {
     super(...args);
   }
 }
 
 export function polyfillDh(): void {
-  class Event {
-    type: string;
-    detail: unknown;
-
-    constructor(type: string, dict: { detail: unknown }) {
-      this.type = type;
-      if (dict) {
-        this.detail = dict.detail;
-      }
-    }
-  }
-
   // Copilot will look for `window.document.currentScript` if it finds `window`.
   // Since we are polyfilling `window` below, we also need to set `document` to
   // avoid a "Cannot read properties of undefined (reading 'currentScript')"

src/common/commands.ts

@@ -10,6 +10,7 @@ export const OPEN_VARIABLE_PANELS_CMD = `${EXTENSION_ID}.openVariablePanels`;
 export const REFRESH_SERVER_TREE_CMD = `${EXTENSION_ID}.refreshServerTree`;
 export const REFRESH_SERVER_CONNECTION_TREE_CMD = `${EXTENSION_ID}.refreshServerConnectionTree`;
 export const REFRESH_VARIABLE_PANELS_CMD = `${EXTENSION_ID}.refreshVariablePanels`;
+export const REQUEST_DHE_USER_CREDENTIALS_CMD = `${EXTENSION_ID}.requestDheUserCredentials`;
 export const RUN_CODE_COMMAND = `${EXTENSION_ID}.runCode`;
 export const RUN_SELECTION_COMMAND = `${EXTENSION_ID}.runSelection`;
 export const SELECT_CONNECTION_COMMAND = `${EXTENSION_ID}.selectConnection`;

src/common/constants.ts

@@ -12,6 +12,15 @@ export const CONFIG_KEY = {
 export const DEFAULT_CONSOLE_TYPE = 'python' as const;
 // export const DHFS_SCHEME = 'dhfs';
 
+// 600 seconds is based on default `auto_delete_timeout` in
+// `ControllerClient.make_temporary_config`
+export const DEFAULT_TEMPORARY_QUERY_AUTO_TIMEOUT_MS = 600000 as const;
+export const DEFAULT_TEMPORARY_QUERY_TIMEOUT_MS = 60000 as const;
+
+export const INTERACTIVE_CONSOLE_QUERY_TYPE = 'InteractiveConsole';
+export const INTERACTIVE_CONSOLE_TEMPORARY_QUEUE_NAME =
+  'InteractiveConsoleTemporaryQueue';
+
 export const PYTHON_ENV_WAIT = 1500 as const;
 
 export const PIP_SERVER_STATUS_CHECK_INTERVAL = 3000;
@@ -88,6 +97,7 @@ export const PIP_SERVER_STATUS_DIRECTORY = 'pip-server-status';
 
 export const SERVER_TREE_ITEM_CONTEXT = {
   canStartServer: 'canStartServer',
+  isDHEServerRunning: 'isDHEServerRunning',
   isManagedServerConnected: 'isManagedServerConnected',
   isManagedServerConnecting: 'isManagedServerConnecting',
   isManagedServerDisconnected: 'isManagedServerDisconnected',
@@ -97,3 +107,13 @@ export const SERVER_TREE_ITEM_CONTEXT = {
 } as const;
 
 export type ServerTreeItemContextValue = keyof typeof SERVER_TREE_ITEM_CONTEXT;
+
+export const DEEPHAVEN_POST_MSG = {
+  loginOptionsRequest: 'io.deephaven.message.LoginOptions.request',
+  sessionDetailsRequest: 'io.deephaven.message.SessionDetails.request',
+} as const;
+
+export const VSCODE_POST_MSG = {
+  loginOptionsResponse: 'vscode-ext.loginOptions',
+  sessionDetailsResponse: 'vscode-ext.sessionDetails',
+} as const;

src/controllers/ConnectionController.ts

@@ -14,6 +14,7 @@ import {
   createConnectionQuickPickOptions,
   createConnectStatusBarItem,
   getConnectionsForConsoleType,
+  getConsoleType,
   getEditorForUri,
   isSupportedLanguageId,
   Logger,
@@ -82,9 +83,9 @@ export class ConnectionController implements Disposable {
 
     this._serverManager.onDidServerStatusChange(
       server => {
-        // Auto connect to servers managed by extension when they start
+        // Auto connect to pip servers managed by extension when they start
         if (server.isManaged && server.isRunning) {
-          this._serverManager.connectToServer(server.url);
+          this._serverManager.connectToServer(server.url, 'python');
         }
       },
       undefined,
@@ -131,7 +132,8 @@ export class ConnectionController implements Disposable {
 
     if ('url' in connectionOrServer) {
       const cn = await this._serverManager.connectToServer(
-        connectionOrServer.url
+        connectionOrServer.url,
+        getConsoleType(editor.document.languageId)
       );
 
       if (cn == null) {
@@ -228,7 +230,12 @@ export class ConnectionController implements Disposable {
   };
 
   /**
-   * Prompt user to select a connection and apply the selection.
+   * Prompt user to select a connection and apply the selection. The options
+   * presented to the user consist of:
+   * 1. Active connections that support the console type of the active editor.
+   * 2. A list of running servers composed of:
+   *   - DHC servers that don't yet have a connection
+   *   - All running DHE servers
    */
   onPromptUserToSelectConnection = async (): Promise<boolean> => {
     assertDefined(vscode.window.activeTextEditor, 'activeTextEditor');
@@ -252,9 +259,19 @@ export class ConnectionController implements Disposable {
       await updateStatusPromise;
     }
 
-    const runningServersWithoutConnections = this._serverManager.getServers({
+    // Only include DHC servers that don't have any connections yet since their
+    // single connection will be included in the `connectionsForConsoleType` list.
+    const runningDHCServersWithoutConnections = this._serverManager.getServers({
       isRunning: true,
       hasConnections: false,
+      type: 'DHC',
+    });
+
+    // Since DHE servers allow creating multiple workers, always include the
+    // server in the list regardless of how many worker connections already exist.
+    const runningDHEServers = this._serverManager.getServers({
+      isRunning: true,
+      type: 'DHE',
     });
 
     const connectionsForConsoleType: ConnectionState[] =
@@ -272,7 +289,7 @@ export class ConnectionController implements Disposable {
     try {
       selectedCnResult = await createConnectionQuickPick(
         createConnectionQuickPickOptions(
-          runningServersWithoutConnections,
+          [...runningDHCServersWithoutConnections, ...runningDHEServers],
           connectionsForConsoleType,
           editor.document.languageId,
           editorActiveConnectionUrl

src/controllers/ControllerBase.ts

@@ -0,0 +1,31 @@
+import * as vscode from 'vscode';
+import type { Disposable } from '../types';
+
+/**
+ * Base controller class.
+ */
+export abstract class ControllerBase implements Disposable {
+  protected readonly disposables: Disposable[] = [];
+
+  /**
+   * Register a command and add it's subscription to the disposables list.
+   */
+  registerCommand = (
+    ...args: Parameters<typeof vscode.commands.registerCommand>
+  ): void => {
+    const cmd = vscode.commands.registerCommand(...args);
+    this.disposables.push(cmd);
+  };
+
+  /**
+   * Dispose all registered disposables.
+   */
+  async dispose(): Promise<void> {
+    const disposing = this.disposables.map(subscription =>
+      subscription.dispose()
+    );
+    this.disposables.length = 0;
+
+    await Promise.all(disposing);
+  }
+}