Fixes to flaky Installer Prompt tests (#9358)

src/client/common/utils/async.ts

@@ -138,9 +138,9 @@ export function createDeferred<T>(scope: any = null): Deferred<T> {
     return new DeferredImpl<T>(scope);
 }
 
-export function createDeferredFrom<T>(...promises: Promise<T>[]): Deferred<T> {
+export function createDeferredFrom<T>(promise: Promise<T>): Deferred<T> {
     const deferred = createDeferred<T>();
-    Promise.all<T>(promises)
+    promise
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         .then(deferred.resolve.bind(deferred) as any)
         // eslint-disable-next-line @typescript-eslint/no-explicit-any

src/client/telemetry/index.ts

@@ -18,7 +18,7 @@ import {
 } from '../datascience/constants';
 import { ResourceSpecificTelemetryProperties } from '../datascience/telemetry/types';
 import { ExportFormat } from '../datascience/export/types';
-import { InterruptResult } from '../datascience/types';
+import { InterruptResult, KernelInterpreterDependencyResponse } from '../datascience/types';
 import { CheckboxState, EventName, PlatformErrors, SliceOperationSource } from './constants';
 import { noop } from '../common/utils/misc';
 import { isPromise } from 'rxjs/internal-compatibility';
@@ -846,6 +846,7 @@ export interface IEventNamePropertyMapping {
          */
         isModulePresent?: 'true' | undefined;
         action:
+            | 'cancelled' // User cancelled the installation or closed the notebook or the like.
             | 'displayed' // Install prompt may have been displayed.
             | 'prompted' // Install prompt was displayed.
             | 'installed' // Installation disabled (this is what python extension returns).
@@ -1309,7 +1310,7 @@ export interface IEventNamePropertyMapping {
     [Telemetry.RawKernelSessionStartTimeout]: never | undefined;
     [Telemetry.RawKernelSessionStartUserCancel]: never | undefined;
     [Telemetry.RawKernelSessionStartNoIpykernel]: {
-        reason: number;
+        reason: KernelInterpreterDependencyResponse;
     } & TelemetryErrorProperties;
     /**
      * This event is sent when the underlying kernelProcess for a

src/extension/errors/errorHandler.ts

@@ -39,6 +39,7 @@ import { JupyterKernelDependencyError } from './jupyterKernelDependencyError';
 import { WrappedError, BaseKernelError, WrappedKernelError, BaseError } from './types';
 import { noop } from '../../client/common/utils/misc';
 import { EnvironmentType } from '../../client/pythonEnvironments/info';
+import { KernelDeadError } from './kernelDeadError';
 
 @injectable()
 export class DataScienceErrorHandler implements IDataScienceErrorHandler {
@@ -97,7 +98,11 @@ export class DataScienceErrorHandler implements IDataScienceErrorHandler {
         context: 'start' | 'restart' | 'interrupt' | 'execution'
     ) {
         error = WrappedError.unwrap(error);
-        if (error instanceof JupyterKernelDependencyError) {
+        if (error instanceof KernelDeadError) {
+            // When we get this we've already asked the user to restart the kernel,
+            // No need to display errors in each cell.
+            return '';
+        } else if (error instanceof JupyterKernelDependencyError) {
             return getIPyKernelMissingErrorMessageForCell(error.kernelConnectionMetadata) || error.message;
         } else if (error instanceof JupyterInstallError) {
             return getJupyterMissingErrorMessageForCell(error) || error.message;

src/extension/errors/kernelDeadError.ts

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { DataScience } from '../../client/common/utils/localize';
+import { getDisplayNameOrNameOfKernelConnection } from '../../kernels/helpers';
+import { KernelConnectionMetadata } from '../../kernels/types';
+import { WrappedKernelError } from './types';
+
+export class KernelDeadError extends WrappedKernelError {
+    constructor(kernelConnectionMetadata: KernelConnectionMetadata) {
+        super(
+            DataScience.kernelDiedWithoutError().format(
+                getDisplayNameOrNameOfKernelConnection(kernelConnectionMetadata)
+            ),
+            undefined,
+            kernelConnectionMetadata
+        );
+    }
+}

src/kernels/helpers.ts

@@ -31,7 +31,7 @@ import { traceError, traceInfo, traceInfoIfCI, traceVerbose, traceWarning } from
 import { getDisplayPath } from '../client/common/platform/fs-paths';
 import { IPythonExecutionFactory } from '../client/common/process/types';
 import { IPathUtils, IConfigurationService, Resource, IMemento, GLOBAL_MEMENTO } from '../client/common/types';
-import { createDeferred } from '../client/common/utils/async';
+import { createDeferred, createDeferredFromPromise, Deferred } from '../client/common/utils/async';
 import { DataScience } from '../client/common/utils/localize';
 import { getResourceType } from '../client/datascience/common';
 import { Settings } from '../client/datascience/constants';
@@ -44,7 +44,8 @@ import {
     IDataScienceErrorHandler,
     IRawNotebookProvider,
     KernelInterpreterDependencyResponse,
-    IJupyterKernelSpec
+    IJupyterKernelSpec,
+    IDisplayOptions
 } from '../client/datascience/types';
 import { IServiceContainer } from '../client/ioc/types';
 import {
@@ -71,6 +72,8 @@ import { isPythonNotebook } from '../notebooks/helpers';
 import { INotebookControllerManager } from '../notebooks/types';
 import { PreferredRemoteKernelIdProvider } from './raw/finder/preferredRemoteKernelIdProvider';
 import { findNotebookEditor, selectKernel } from '../notebooks/controllers/kernelSelector';
+import { DisplayOptions } from '../client/datascience/displayOptions';
+import { KernelDeadError } from '../extension/errors/kernelDeadError';
 
 // Helper functions for dealing with kernels and kernelspecs
 
@@ -1975,10 +1978,10 @@ export async function handleKernelError(
     return resultController;
 }
 
-function convertContextToFunction(context: 'start' | 'interrupt' | 'restart') {
+function convertContextToFunction(context: 'start' | 'interrupt' | 'restart', options?: IDisplayOptions) {
     switch (context) {
         case 'start':
-            return (k: IKernel) => k.start();
+            return (k: IKernel) => k.start(options);
 
         case 'interrupt':
             return (k: IKernel) => k.interrupt();
@@ -1988,17 +1991,122 @@ function convertContextToFunction(context: 'start' | 'interrupt' | 'restart') {
     }
 }
 
+const connections = new WeakMap<
+    NotebookDocument,
+    {
+        kernel: Deferred<{
+            kernel: IKernel;
+            controller: VSCodeNotebookController;
+            deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
+        }>;
+        options: IDisplayOptions;
+    }
+>();
+
+async function verifyKernelState(
+    serviceContainer: IServiceContainer,
+    resource: Resource,
+    notebook: NotebookDocument,
+    options: IDisplayOptions = new DisplayOptions(false),
+    promise: Promise<{
+        kernel: IKernel;
+        controller: VSCodeNotebookController;
+        deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
+    }>
+): Promise<IKernel> {
+    const { kernel, controller, deadKernelAction } = await promise;
+    // Before returning, but without disposing the kernel, double check it's still valid
+    // If a restart didn't happen, then we can't connect. Throw an error.
+    // Do this outside of the loop so that subsequent calls will still ask because the kernel isn't disposed
+    if (kernel.status === 'dead' || (kernel.status === 'terminating' && !kernel.disposed && !kernel.disposing)) {
+        // If the kernel is dead, then remove the cached promise, & try to get the kernel again.
+        // At that point, it will get restarted.
+        if (connections.get(notebook)?.kernel.promise === promise) {
+            connections.delete(notebook);
+        }
+        if (deadKernelAction === 'deadKernelWasNoRestarted') {
+            throw new KernelDeadError(kernel.kernelConnectionMetadata);
+        } else if (deadKernelAction === 'deadKernelWasRestarted') {
+            return kernel;
+        }
+        // Kernel is dead and we didn't prompt the user to restart it, hence re-run the code that will prompt the user for a restart.
+        return wrapKernelMethod(controller, 'start', serviceContainer, resource, notebook, options);
+    }
+    return kernel;
+}
+
 export async function wrapKernelMethod(
     initialController: VSCodeNotebookController,
     initialContext: 'start' | 'interrupt' | 'restart',
     serviceContainer: IServiceContainer,
     resource: Resource,
-    notebook: NotebookDocument
-) {
+    notebook: NotebookDocument,
+    options: IDisplayOptions = new DisplayOptions(false)
+): Promise<IKernel> {
+    let currentPromise = connections.get(notebook);
+    if (!options.disableUI && currentPromise?.options.disableUI) {
+        currentPromise.options.disableUI = false;
+    }
+    // If the current kernel has been disposed or in the middle of being disposed, then create another one.
+    // But do that only if we require a UI, else we can just use the current one.
+    if (
+        !options.disableUI &&
+        currentPromise?.kernel.resolved &&
+        (currentPromise?.kernel.value?.kernel?.disposed || currentPromise?.kernel.value?.kernel?.disposing)
+    ) {
+        connections.delete(notebook);
+        currentPromise = undefined;
+    }
+
+    // Wrap the kernel method again to interrupt/restart this kernel.
+    if (currentPromise && initialContext !== 'restart' && initialContext !== 'interrupt') {
+        return verifyKernelState(serviceContainer, resource, notebook, options, currentPromise.kernel.promise);
+    }
+
+    const promise = wrapKernelMethodImpl(
+        initialController,
+        initialContext,
+        serviceContainer,
+        resource,
+        notebook,
+        options
+    );
+    const deferred = createDeferredFromPromise(promise);
+    // If the kernel gets disposed or we fail to create the kernel, then ensure we remove the cached result.
+    promise
+        .then((result) => {
+            result.kernel.onDisposed(() => {
+                if (connections.get(notebook)?.kernel === deferred) {
+                    connections.delete(notebook);
+                }
+            });
+        })
+        .catch(() => {
+            if (connections.get(notebook)?.kernel === deferred) {
+                connections.delete(notebook);
+            }
+        });
+
+    connections.set(notebook, { kernel: deferred, options });
+    return verifyKernelState(serviceContainer, resource, notebook, options, deferred.promise);
+}
+
+export async function wrapKernelMethodImpl(
+    initialController: VSCodeNotebookController,
+    initialContext: 'start' | 'interrupt' | 'restart',
+    serviceContainer: IServiceContainer,
+    resource: Resource,
+    notebook: NotebookDocument,
+    options: IDisplayOptions = new DisplayOptions(false)
+): Promise<{
+    kernel: IKernel;
+    controller: VSCodeNotebookController;
+    deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
+}> {
     const kernelProvider = serviceContainer.get<IKernelProvider>(IKernelProvider);
     let kernel: IKernel | undefined;
     let controller: VSCodeNotebookController = initialController;
-    let currentMethod = convertContextToFunction(initialContext);
+    let currentMethod = convertContextToFunction(initialContext, options);
     let context = initialContext;
     while (kernel === undefined) {
         // Try to create the kernel (possibly again)
@@ -2008,46 +2116,52 @@ export async function wrapKernelMethod(
             resourceUri: resource
         });
 
+        const isKernelDead = (k: IKernel) =>
+            k.status === 'dead' || (k.status === 'terminating' && !k.disposed && !k.disposing);
+
         try {
-            await currentMethod(kernel);
+            // If the kernel is dead, ask the user if they want to restart.
+            // We need to perform this check first, as its possible we'd call this method for dead kernels.
+            // & if the kernel is dead, prompt to restart.
+            if (initialContext !== 'restart' && isKernelDead(kernel) && !options.disableUI) {
+                const restarted = await notifyAndRestartDeadKernel(kernel, serviceContainer);
+                return {
+                    kernel,
+                    controller,
+                    deadKernelAction: restarted ? 'deadKernelWasRestarted' : 'deadKernelWasNoRestarted'
+                };
+            } else {
+                await currentMethod(kernel);
 
-            // If the kernel is dead, ask the user if they want to restart
-            if (
-                kernel.status === 'dead' ||
-                (kernel.status === 'terminating' && !kernel.disposed && !kernel.disposing)
-            ) {
-                await notifyAndRestartDeadKernel(kernel, serviceContainer);
+                // If the kernel is dead, ask the user if they want to restart
+                if (isKernelDead(kernel) && !options.disableUI) {
+                    await notifyAndRestartDeadKernel(kernel, serviceContainer);
+                }
             }
         } catch (error) {
+            if (options.disableUI) {
+                throw error;
+            }
             controller = await handleKernelError(serviceContainer, error, context, resource, kernel, controller);
 
             // When we wrap around, update the current method to start. This
             // means if we're handling a restart or an interrupt that fails, we move onto trying to start the kernel.
-            currentMethod = (k) => k.start();
+            currentMethod = (k) => k.start(options);
             context = 'start';
 
             // Since an error occurred, we have to try again (controller may have switched so we have to pick a new kernel)
             kernel = undefined;
         }
     }
-    // Before returning, but without disposing the kernel, double check it's still valid
-    // If a restart didn't happen, then we can't connect. Throw an error.
-    // Do this outside of the loop so that subsequent calls will still ask because the kernel isn't disposed
-    if (kernel.status === 'dead' || (kernel.status === 'terminating' && !kernel.disposed && !kernel.disposing)) {
-        throw new Error(
-            DataScience.kernelDiedWithoutError().format(
-                getDisplayNameOrNameOfKernelConnection(kernel.kernelConnectionMetadata)
-            )
-        );
-    }
-    return kernel;
+    return { kernel, controller };
 }
 
 export async function connectToKernel(
     initialController: VSCodeNotebookController,
     serviceContainer: IServiceContainer,
     resource: Resource,
-    notebook: NotebookDocument
+    notebook: NotebookDocument,
+    options: IDisplayOptions = new DisplayOptions(false)
 ): Promise<IKernel> {
-    return wrapKernelMethod(initialController, 'start', serviceContainer, resource, notebook);
+    return wrapKernelMethod(initialController, 'start', serviceContainer, resource, notebook, options);
 }

src/kernels/installer/moduleInstaller.ts

@@ -55,25 +55,31 @@ export abstract class ModuleInstaller implements IModuleInstaller {
         const environmentService = this.serviceContainer.get<IEnvironmentVariablesService>(
             IEnvironmentVariablesService
         );
+        if (cancelTokenSource.token.isCancellationRequested) {
+            return;
+        }
         const install = async (
             progress: Progress<{
                 message?: string | undefined;
                 increment?: number | undefined;
             }>,
             token: CancellationToken
         ) => {
+            const deferred = createDeferred();
             // When the progress is canceled notify caller
             token.onCancellationRequested(() => {
                 cancelTokenSource.cancel();
+                deferred.resolve();
             });
-            // Args can be for a specific exe or for the interpreter. Both need to
-            // use an activated environment though
-            const deferred = createDeferred();
+
             let observable: ObservableExecutionResult<string> | undefined;
 
             // Some installers only work with shellexec
             if (args.useShellExec) {
                 const proc = await procFactory.create(undefined);
+                if (cancelTokenSource.token.isCancellationRequested) {
+                    return;
+                }
                 try {
                     const results = await proc.shellExec(args.args.join(' '), { cwd: args.cwd });
                     traceInfo(results.stdout);
@@ -82,15 +88,26 @@ export abstract class ModuleInstaller implements IModuleInstaller {
                     deferred.reject(ex);
                 }
             } else if (args.exe) {
+                // Args can be for a specific exe or for the interpreter. Both need to
+                // use an activated environment though
                 // For the exe, just figure out the environment variables.
                 const envVars = await activationHelper.getActivatedEnvironmentVariables(undefined, interpreter, false);
+                if (cancelTokenSource.token.isCancellationRequested) {
+                    return;
+                }
                 const env = { ...process.env };
                 environmentService.mergeVariables(envVars || {}, env);
                 environmentService.mergePaths(envVars || {}, env);
                 const proc = await procFactory.create(undefined);
+                if (cancelTokenSource.token.isCancellationRequested) {
+                    return;
+                }
                 observable = proc.execObservable(args.exe, args.args, { encoding: 'utf-8', token, env, cwd: args.cwd });
             } else {
                 const proc = await pythonFactory.createActivatedEnvironment({ interpreter });
+                if (cancelTokenSource.token.isCancellationRequested) {
+                    return;
+                }
                 observable = proc.execObservable(args.args, {
                     encoding: 'utf-8',
                     token,