fix: Disposal Logic + Error toast

package.json

@@ -46,7 +46,7 @@
     "test:lint": "eslint . --ext ts",
     "test:unit": "vitest --reporter=default --reporter=junit --outputFile=./test-reports/vitest.junit.xml",
     "test": "npm run test:unit",
-    "vscode:prepublish": "compile:prod",
+    "vscode:prepublish": "npm run compile:prod",
     "watch:esbuild": "node scripts/esbuild.js --watch",
     "watch:tsc": "tsc --build ./tsconfig.json --watch",
     "watch": "npm-run-all -p watch:*"

src/controllers/ConnectionController.ts

@@ -13,14 +13,14 @@ import {
   createConnectionQuickPick,
   createConnectionQuickPickOptions,
   createConnectStatusBarItem,
-  getConnectionsForConsoleType,
   getConsoleType,
   getEditorForUri,
   isSupportedLanguageId,
   Logger,
   updateConnectionStatusBarItem,
 } from '../util';
 import { UnsupportedConsoleTypeError } from '../common';
+import { getConnectionsForConsoleType } from '../services';
 
 const logger = new Logger('ConnectionController');
 
@@ -153,6 +153,7 @@ export class ConnectionController implements Disposable {
       // disconnect from it.
       if (err instanceof UnsupportedConsoleTypeError && newConnectionUrl) {
         this._serverManager.disconnectFromServer(newConnectionUrl);
+        this._toaster.error(err.message);
       }
 
       throw err;

src/controllers/ExtensionController.ts

@@ -117,11 +117,14 @@ export class ExtensionController implements Disposable {
   readonly _config: IConfigService;
 
   private _connectionController: ConnectionController | null = null;
-  private _coreClientCache: URLMap<CoreAuthenticatedClient> | null = null;
+  private _coreClientCache: URLMap<
+    CoreAuthenticatedClient & Disposable
+  > | null = null;
   private _coreClientFactory: ICoreClientFactory | null = null;
   private _coreJsApiCache: IAsyncCacheService<URL, typeof DhcType> | null =
     null;
-  private _dheClientCache: URLMap<DheAuthenticatedClient> | null = null;
+  private _dheClientCache: URLMap<DheAuthenticatedClient & Disposable> | null =
+    null;
   private _dheClientFactory: IDheClientFactory | null = null;
   private _dheServiceCache: IAsyncCacheService<URL, IDheService> | null = null;
   private _panelController: PanelController | null = null;
@@ -161,6 +164,7 @@ export class ExtensionController implements Disposable {
     const codelensProvider = new RunCommandCodeLensProvider();
 
     this._context.subscriptions.push(
+      codelensProvider,
       vscode.languages.registerCodeLensProvider('groovy', codelensProvider),
       vscode.languages.registerCodeLensProvider('python', codelensProvider)
     );
@@ -329,22 +333,43 @@ export class ExtensionController implements Disposable {
 
     this._coreClientFactory = async (
       url: URL
-    ): Promise<CoreUnauthenticatedClient> => {
+    ): Promise<CoreUnauthenticatedClient & Disposable> => {
       assertDefined(this._coreJsApiCache, 'coreJsApiCache');
       const dhc = await this._coreJsApiCache.get(url);
-      return new dhc.CoreClient(url.toString()) as CoreUnauthenticatedClient;
+
+      const client = new dhc.CoreClient(
+        url.toString()
+      ) as CoreUnauthenticatedClient;
+
+      // Attach a dispose method so that client caches can dispose of the client
+      return Object.assign(client, {
+        dispose: async () => {
+          client.disconnect();
+        },
+      });
     };
 
     this._dheClientFactory = async (
       url: URL
-    ): Promise<DheUnauthenticatedClient> => {
+    ): Promise<DheUnauthenticatedClient & Disposable> => {
       assertDefined(this._dheJsApiCache, 'dheJsApiCache');
       const dhe = await this._dheJsApiCache.get(url);
-      return createDheClient(dhe, getWsUrl(url));
+
+      const client = await createDheClient(dhe, getWsUrl(url));
+
+      // Attach a dispose method so that client caches can dispose of the client
+      return Object.assign(client, {
+        dispose: async () => {
+          client.disconnect();
+        },
+      });
     };
 
-    this._coreClientCache = new URLMap();
-    this._dheClientCache = new URLMap();
+    this._coreClientCache = new URLMap<CoreAuthenticatedClient & Disposable>();
+    this._context.subscriptions.push(this._coreClientCache);
+
+    this._dheClientCache = new URLMap<DheAuthenticatedClient & Disposable>();
+    this._context.subscriptions.push(this._dheClientCache);
 
     this._panelService = new PanelService();
     this._context.subscriptions.push(this._panelService);
@@ -382,6 +407,7 @@ export class ExtensionController implements Disposable {
 
     this._serverManager.onDidDisconnect(
       serverUrl => {
+        this._panelService?.clearServerData(serverUrl);
         this._outputChannel?.appendLine(
           `Disconnected from server: '${serverUrl}'.`
         );
@@ -529,7 +555,10 @@ export class ExtensionController implements Disposable {
     this._context.subscriptions.push(
       this._serverTreeView,
       this._serverConnectionTreeView,
-      this._serverConnectionPanelTreeView
+      this._serverConnectionPanelTreeView,
+      this._serverTreeProvider,
+      this._serverConnectionTreeProvider,
+      this._serverConnectionPanelTreeProvider
     );
   };
 

src/controllers/PipServerController.ts

@@ -13,8 +13,9 @@ import {
   PIP_SERVER_SUPPORTED_PLATFORMS,
   PYTHON_ENV_WAIT,
 } from '../common';
-import { pollUntilTrue, waitFor } from '../util/promiseUtils';
+import { waitFor } from '../util';
 import { isDhcServerRunning } from '../dh/dhc';
+import { pollUntilTrue } from '../services';
 
 const logger = new Logger('PipServerController');
 

src/dh/dhc.ts

@@ -20,16 +20,6 @@ export type ConnectionAndSession<TConnection, TSession> = {
   session: TSession;
 };
 
-// TODO: https://github.com/deephaven/deephaven-core/issues/5911 to address the
-// underlying issue of jsapi-types being unaware of `dhinternal`. Once that is
-// addressed, this can be removed.
-declare global {
-  // eslint-disable-next-line no-unused-vars
-  namespace dhinternal.io.deephaven.proto.ticket_pb {
-    export type TypedTicket = unknown;
-  }
-}
-
 /**
  * Download the DH Core jsapi from a running server and return the `dh` object.
  * @param serverUrl URL of the DH Core server to download the api from.

src/providers/RunCommandCodeLensProvider.ts

@@ -1,11 +1,12 @@
 import * as vscode from 'vscode';
 import { ICON_ID } from '../common';
+import type { Disposable } from '../types';
 
 /**
  * Provides inline editor code lenses for running Deephaven code.
  */
 export class RunCommandCodeLensProvider
-  implements vscode.CodeLensProvider<vscode.CodeLens>
+  implements vscode.CodeLensProvider<vscode.CodeLens>, Disposable
 {
   constructor() {
     vscode.workspace.onDidChangeConfiguration(() => {
@@ -22,6 +23,10 @@ export class RunCommandCodeLensProvider
   readonly onDidChangeCodeLenses: vscode.Event<void> =
     this._onDidChangeCodeLenses.event;
 
+  dispose = async (): Promise<void> => {
+    this._onDidChangeCodeLenses.dispose();
+  };
+
   provideCodeLenses(
     document: vscode.TextDocument,
     _token: vscode.CancellationToken

src/providers/ServerConnectionPanelTreeProvider.ts

@@ -11,6 +11,7 @@ import {
   getPanelVariableTreeItem,
   sortByStringProp,
 } from '../util';
+import { getFirstSupportedConsoleType } from '../services';
 
 export class ServerConnectionPanelTreeProvider extends TreeDataProviderBase<ServerConnectionPanelNode> {
   constructor(serverManager: IServerManager, panelService: IPanelService) {
@@ -31,7 +32,10 @@ export class ServerConnectionPanelTreeProvider extends TreeDataProviderBase<Serv
       return getPanelVariableTreeItem(connectionOrVariable);
     }
 
-    return getPanelConnectionTreeItem(connectionOrVariable);
+    return getPanelConnectionTreeItem(
+      connectionOrVariable,
+      getFirstSupportedConsoleType
+    );
   };
 
   getChildren = (

src/providers/TreeDataProviderBase.ts

@@ -1,11 +1,11 @@
 import * as vscode from 'vscode';
-import type { IServerManager } from '../types';
+import type { Disposable, IServerManager } from '../types';
 
 /**
  * Base class for tree view data providers.
  */
 export abstract class TreeDataProviderBase<T>
-  implements vscode.TreeDataProvider<T>
+  implements vscode.TreeDataProvider<T>, Disposable
 {
   constructor(serverManager: IServerManager) {
     this.serverManager = serverManager;
@@ -27,6 +27,10 @@ export abstract class TreeDataProviderBase<T>
 
   abstract getChildren(element?: T | undefined): vscode.ProviderResult<T[]>;
 
+  dispose = async (): Promise<void> => {
+    this._onDidChangeTreeData.dispose();
+  };
+
   refresh(): void {
     this._onDidChangeTreeData.fire();
   }

src/services/DhcService.ts

@@ -178,7 +178,7 @@ export class DhcService implements IDhcService {
       try {
         const { cn, session } = await this.initSessionPromise;
 
-        cn.subscribeToFieldUpdates(changes => {
+        const fieldUpdateSubscription = cn.subscribeToFieldUpdates(changes => {
           this.panelService.updateVariables(
             this.serverUrl,
             changes as VariableChanges
@@ -198,15 +198,15 @@ export class DhcService implements IDhcService {
             panelVariablesToUpdate
           );
         });
+        this.subscriptions.push(fieldUpdateSubscription);
 
         // TODO: Use constant 'disconnect' event name
-        this.subscriptions.push(
-          cn.addEventListener('disconnect', () => {
-            this.clearCaches();
-          })
-        );
+        const disconnectSubscription = cn.addEventListener('disconnect', () => {
+          this.clearCaches();
+        });
+        this.subscriptions.push(disconnectSubscription);
 
-        session.onLogMessage(logItem => {
+        const logMessageSubscription = session.onLogMessage(logItem => {
           // TODO: Should this pull log level from config somewhere?
           if (logItem.logLevel !== 'INFO') {
             const date = new Date(logItem.micros / 1000);
@@ -217,6 +217,7 @@ export class DhcService implements IDhcService {
             );
           }
         });
+        this.subscriptions.push(logMessageSubscription);
       } catch (err) {}
     }
 

src/services/DheService.ts

@@ -256,8 +256,10 @@ export class DheService implements IDheService {
     const querySerials = [...this._querySerialSet];
 
     this._querySerialSet.clear();
-    this._workerInfoMap.clear();
 
-    await this._disposeQueries(querySerials);
+    await Promise.all([
+      this._workerInfoMap.dispose(),
+      this._disposeQueries(querySerials),
+    ]);
   };
 }

src/services/PanelService.ts

@@ -22,8 +22,25 @@ export class PanelService implements IPanelService, Disposable {
   private readonly _cnPanelMap: URLMap<VariablePanelMap>;
   private readonly _cnVariableMap: URLMap<VariableMap>;
 
+  /**
+   * Clear panel data for the given connection url.
+   * @param url The connection url.
+   */
+  clearServerData = (url: URL): void => {
+    this._cnPanelMap.delete(url);
+    this._cnVariableMap.delete(url);
+  };
+
+  /**
+   * Cleanup resources.
+   */
   dispose = async (): Promise<void> => {
-    this._cnPanelMap.clear();
+    this._onDidUpdate.dispose();
+
+    await Promise.all([
+      this._cnPanelMap.dispose(),
+      this._cnVariableMap.dispose(),
+    ]);
   };
 
   /**