Refactor TensorBoard prompt and import tracking and add tests

.eslintignore

@@ -946,3 +946,6 @@ src/test/startPage/mountedWebViewFactory.ts
 src/test/startPage/reactHelpers.ts
 src/test/startPage/webBrowserPanel.ts
 src/test/startPage/webBrowserPanelProvider.ts
+
+# This directory contains vendored code
+src/client/tensorBoard/terminal/*

.prettierignore

@@ -0,0 +1,2 @@
+# Ignore vendored files from https://github.com/xtermjs/xterm.js
+src/client/tensorBoard/terminal/**

package.json

@@ -2035,7 +2035,8 @@
         "vscode-tas-client": "^0.1.15",
         "winreg": "^1.2.4",
         "winston": "^3.2.1",
-        "xml2js": "^0.4.19"
+        "xml2js": "^0.4.19",
+        "xterm": "^4.9.0"
     },
     "devDependencies": {
         "@enonic/fnv-plus": "^1.3.0",

src/client/tensorBoard/serviceRegistry.ts

@@ -9,7 +9,7 @@ import { TensorBoardFileWatcher } from './tensorBoardFileWatcher';
 import { TensorBoardImportTracker } from './tensorBoardImportTracker';
 import { TensorBoardPrompt } from './tensorBoardPrompt';
 import { TensorBoardSessionProvider } from './tensorBoardSessionProvider';
-import { ITensorBoardImportTracker } from './types';
+import { TensorBoardTerminalListener } from './tensorBoardTerminalListener';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingleton<IExtensionSingleActivationService>(
@@ -21,8 +21,12 @@ export function registerTypes(serviceManager: IServiceManager): void {
         TensorBoardFileWatcher,
     );
     serviceManager.addSingleton<TensorBoardPrompt>(TensorBoardPrompt, TensorBoardPrompt);
-    serviceManager.addSingleton<ITensorBoardImportTracker>(ITensorBoardImportTracker, TensorBoardImportTracker);
-    serviceManager.addBinding(ITensorBoardImportTracker, IExtensionSingleActivationService);
+    serviceManager.addSingleton<IExtensionSingleActivationService>(
+        IExtensionSingleActivationService,
+        TensorBoardImportTracker,
+    );
+    serviceManager.addSingleton<TensorBoardTerminalListener>(TensorBoardTerminalListener, TensorBoardTerminalListener);
+    serviceManager.addBinding(TensorBoardTerminalListener, IExtensionSingleActivationService);
     serviceManager.addSingleton<IExtensionSingleActivationService>(
         IExtensionSingleActivationService,
         TensorBoardCodeLensProvider,

src/client/tensorBoard/tensorBoardImportTracker.ts

@@ -3,36 +3,24 @@
 
 import { inject, injectable } from 'inversify';
 import * as path from 'path';
-import { Event, EventEmitter, TextEditor } from 'vscode';
+import { TextEditor } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IDocumentManager } from '../common/application/types';
 import { isTestExecution } from '../common/constants';
 import { IDisposableRegistry } from '../common/types';
 import { getDocumentLines } from '../telemetry/importTracker';
 import { containsTensorBoardImport } from './helpers';
-import { ITensorBoardImportTracker } from './types';
+import { TensorBoardPrompt } from './tensorBoardPrompt';
 
 const testExecution = isTestExecution();
 @injectable()
-export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExtensionSingleActivationService {
-    private pendingChecks = new Map<string, NodeJS.Timer | number>();
-
-    private _onDidImportTensorBoard = new EventEmitter<void>();
-
+export class TensorBoardImportTracker implements IExtensionSingleActivationService {
     constructor(
         @inject(IDocumentManager) private documentManager: IDocumentManager,
         @inject(IDisposableRegistry) private disposables: IDisposableRegistry,
+        @inject(TensorBoardPrompt) private prompt: TensorBoardPrompt,
     ) {}
 
-    // Fires when the active text editor contains a tensorboard import.
-    public get onDidImportTensorBoard(): Event<void> {
-        return this._onDidImportTensorBoard.event;
-    }
-
-    public dispose(): void {
-        this.pendingChecks.clear();
-    }
-
     public async activate(): Promise<void> {
         if (testExecution) {
             await this.activateInternal();
@@ -63,7 +51,7 @@ export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExt
         ) {
             const lines = getDocumentLines(document);
             if (containsTensorBoardImport(lines)) {
-                this._onDidImportTensorBoard.fire();
+                this.prompt.showNativeTensorBoardPrompt().ignoreErrors();
             }
         }
     }

src/client/tensorBoard/tensorBoardPrompt.ts

@@ -5,9 +5,8 @@ import { inject, injectable } from 'inversify';
 import { IApplicationShell, ICommandManager } from '../common/application/types';
 import { Commands } from '../common/constants';
 import { NativeTensorBoard } from '../common/experiments/groups';
-import { IDisposableRegistry, IExperimentService, IPersistentState, IPersistentStateFactory } from '../common/types';
+import { IExperimentService, IPersistentState, IPersistentStateFactory } from '../common/types';
 import { Common, TensorBoard } from '../common/utils/localize';
-import { ITensorBoardImportTracker } from './types';
 
 enum TensorBoardPromptStateKeys {
     ShowNativeTensorBoardPrompt = 'showNativeTensorBoardPrompt',
@@ -17,7 +16,7 @@ enum TensorBoardPromptStateKeys {
 export class TensorBoardPrompt {
     private state: IPersistentState<boolean>;
 
-    private enabled: Promise<boolean>;
+    private enabled: boolean;
 
     private inExperiment: Promise<boolean>;
 
@@ -28,8 +27,6 @@ export class TensorBoardPrompt {
     constructor(
         @inject(IApplicationShell) private applicationShell: IApplicationShell,
         @inject(ICommandManager) private commandManager: ICommandManager,
-        @inject(ITensorBoardImportTracker) private importTracker: ITensorBoardImportTracker,
-        @inject(IDisposableRegistry) private disposableRegistry: IDisposableRegistry,
         @inject(IPersistentStateFactory) private persistentStateFactory: IPersistentStateFactory,
         @inject(IExperimentService) private experimentService: IExperimentService,
     ) {
@@ -39,13 +36,12 @@ export class TensorBoardPrompt {
         );
         this.enabled = this.isPromptEnabled();
         this.inExperiment = this.isInExperiment();
-        this.importTracker.onDidImportTensorBoard(this.showNativeTensorBoardPrompt, this, this.disposableRegistry);
     }
 
     public async showNativeTensorBoardPrompt(): Promise<void> {
         if (
             (await this.inExperiment) &&
-            (await this.enabled) &&
+            this.enabled &&
             this.enabledInCurrentSession &&
             !this.waitingForUserSelection
         ) {
@@ -73,7 +69,7 @@ export class TensorBoardPrompt {
         }
     }
 
-    private async isPromptEnabled(): Promise<boolean> {
+    private isPromptEnabled(): boolean {
         return this.state.value;
     }
 

src/client/tensorBoard/tensorBoardSession.ts

@@ -19,7 +19,7 @@ import { createPromiseFromCancellation } from '../common/cancellation';
 import { traceError, traceInfo } from '../common/logger';
 import { tensorboardLauncher } from '../common/process/internal/scripts';
 import { IProcessServiceFactory, ObservableExecutionResult } from '../common/process/types';
-import { IInstaller, InstallerResponse, Product } from '../common/types';
+import { IDisposableRegistry, IInstaller, InstallerResponse, Product } from '../common/types';
 import { createDeferred, sleep } from '../common/utils/async';
 import { TensorBoard } from '../common/utils/localize';
 import { IInterpreterService } from '../interpreter/contracts';
@@ -47,6 +47,7 @@ export class TensorBoardSession {
         private readonly workspaceService: IWorkspaceService,
         private readonly processServiceFactory: IProcessServiceFactory,
         private readonly commandManager: ICommandManager,
+        private readonly disposables: IDisposableRegistry,
     ) {}
 
     public async initialize(): Promise<void> {
@@ -248,11 +249,15 @@ export class TensorBoardSession {
             this.process?.kill();
             this.process = undefined;
         });
-        webviewPanel.onDidChangeViewState(() => {
-            if (webviewPanel.visible) {
-                this.update();
-            }
-        }, null);
+        webviewPanel.onDidChangeViewState(
+            () => {
+                if (webviewPanel.visible) {
+                    this.update();
+                }
+            },
+            undefined,
+            this.disposables,
+        );
         return webviewPanel;
     }
 
@@ -262,7 +267,7 @@ export class TensorBoardSession {
             <html lang="en">
                 <head>
                     <meta charset="UTF-8">
-                    <meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src ${this.url};">
+                    <meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src http: https:;">
                     <meta name="viewport" content="width=device-width, initial-scale=1.0">
                     <title>TensorBoard</title>
                 </head>
@@ -275,6 +280,7 @@ export class TensorBoardSession {
                                 f.style.width = window.innerWidth / 0.7 + "px";
                             }
                         }
+                        resizeFrame();
                         window.addEventListener('resize', resizeFrame);
                     </script>
                     <iframe

src/client/tensorBoard/tensorBoardSessionProvider.ts

@@ -50,6 +50,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
                 this.workspaceService,
                 this.processServiceFactory,
                 this.commandManager,
+                this.disposables,
             );
             await newSession.initialize();
         } catch (e) {

src/client/tensorBoard/tensorBoardTerminalListener.ts

@@ -0,0 +1,100 @@
+/* eslint-disable import/first */
+import { inject, injectable } from 'inversify';
+import * as vscode from 'vscode';
+import { IExtensionSingleActivationService } from '../activation/types';
+import { IDisposableRegistry, IExperimentService } from '../common/types';
+import { TensorBoardPrompt } from './tensorBoardPrompt';
+import { NativeTensorBoard } from '../common/experiments/groups';
+import { isTestExecution } from '../common/constants';
+// This is a necessary hack for the xterm npm package to load, because
+// it currently expects to be loaded in the browser context and not in Node.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+(global as any).window = undefined;
+// Import from xterm only after setting the global window variable to
+// prevent a ReferenceError being thrown.
+// eslint-disable-next-line import/order
+import { Terminal } from 'xterm';
+
+@injectable()
+export class TensorBoardTerminalListener implements IExtensionSingleActivationService {
+    private disposables: vscode.Disposable[] = [];
+
+    private terminal: Terminal | undefined;
+
+    constructor(
+        @inject(IDisposableRegistry) private disposableRegistry: IDisposableRegistry,
+        @inject(TensorBoardPrompt) private prompt: TensorBoardPrompt,
+        @inject(IExperimentService) private experimentService: IExperimentService,
+    ) {}
+
+    public async activate(): Promise<void> {
+        this.activateInternal().ignoreErrors();
+    }
+
+    private async activateInternal() {
+        if (isTestExecution() || (await this.experimentService.inExperiment(NativeTensorBoard.experiment))) {
+            const terminal = new Terminal({ allowProposedApi: true });
+
+            this.disposables.push(terminal);
+            this.disposables.push(
+                vscode.window.onDidWriteTerminalData((e) => this.handleTerminalData(e).ignoreErrors(), this),
+            );
+            // Only track and parse the active terminal's data since we only care about user input
+            this.disposables.push(vscode.window.onDidChangeActiveTerminal(() => terminal.reset(), this));
+            this.disposables.push(terminal.onCursorMove(() => this.findTensorBoard(terminal.buffer.active.cursorY)));
+            // Only bother tracking one line at a time
+            this.disposables.push(
+                terminal.onLineFeed(() => {
+                    this.findTensorBoard(terminal.buffer.active.cursorY - 1);
+                    terminal.reset();
+                }),
+            );
+            // Add all our disposables to the extension's disposable registry.
+            // We will dispose ourselves as soon as we see a matching terminal
+            // command or when the extension as a whole is disposed, whichever
+            // happens first
+            this.disposableRegistry.push(...this.disposables);
+
+            this.terminal = terminal;
+        }
+    }
+
+    private findTensorBoard(row: number) {
+        const line = this.terminal?.buffer.active.getLine(row);
+        const bufferContents = line?.translateToString(false);
+        if (bufferContents && bufferContents.includes('tensorboard')) {
+            this.complete();
+        }
+    }
+
+    private complete() {
+        this.prompt.showNativeTensorBoardPrompt().ignoreErrors();
+        // Unsubscribe from terminal data events ASAP
+        this.disposables.forEach((d) => {
+            d.dispose();
+        });
+    }
+
+    // This function is called whenever any data is written to a VS Code integrated
+    // terminal. It shows our tensorboard prompt when the user attempts to launch
+    // tensorboard from the active terminal.
+    // onDidWriteTerminalData emits raw data being written to the terminal output.
+    // TerminalDataWriteEvent.data can be a individual single character as user is typing
+    // something into terminal. It can also be a series of characters if the user pastes
+    // a command into the terminal or uses terminal history to fetch past commands.
+    // It can also fire with multiple characters from terminal prompt characters or terminal output.
+    private async handleTerminalData(e: vscode.TerminalDataWriteEvent) {
+        if (!vscode.window.activeTerminal || vscode.window.activeTerminal !== e.terminal) {
+            return;
+        }
+        // In the case of terminal scrollback or a pasted command, we might be lucky enough
+        // to get the whole command in e.data without having to feed it through the parser
+        if (e.data.includes('tensorboard')) {
+            this.complete();
+            return;
+        }
+        // If the user is entering one character at a time, we'll need to buffer individual characters
+        // and handle escape sequences which manipulate the position of the cursor in the buffer
+        this.terminal?.write(e.data);
+    }
+}

src/test/tensorBoard/helpers.ts

@@ -0,0 +1,24 @@
+import * as TypeMoq from 'typemoq';
+import { IApplicationShell, ICommandManager } from '../../client/common/application/types';
+import { IExperimentService, IPersistentStateFactory } from '../../client/common/types';
+import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt';
+import { MockState } from '../interpreters/mocks';
+
+export function createTensorBoardPromptWithMocks(): TensorBoardPrompt {
+    const appShell = TypeMoq.Mock.ofType<IApplicationShell>();
+    const commandManager = TypeMoq.Mock.ofType<ICommandManager>();
+    const persistentStateFactory = TypeMoq.Mock.ofType<IPersistentStateFactory>();
+    const expService = TypeMoq.Mock.ofType<IExperimentService>();
+    const persistentState = new MockState(true);
+    persistentStateFactory
+        .setup((factory) => {
+            factory.createWorkspacePersistentState(TypeMoq.It.isAny(), TypeMoq.It.isAny());
+        })
+        .returns(() => persistentState);
+    return new TensorBoardPrompt(
+        appShell.object,
+        commandManager.object,
+        persistentStateFactory.object,
+        expService.object,
+    );
+}