diff --git a/src/client/datascience/jupyter/jupyterCommandFinder.ts b/src/client/datascience/jupyter/jupyterCommandFinder.ts index 0a3b10da7412..5e1f6d77e7c0 100644 --- a/src/client/datascience/jupyter/jupyterCommandFinder.ts +++ b/src/client/datascience/jupyter/jupyterCommandFinder.ts @@ -51,7 +51,7 @@ type ProgressNotification = Progress<{ message?: string | undefined; increment?: export class JupyterCommandFinder { private readonly processServicePromise: Promise; private jupyterPath?: string; - private readonly commands = new Map(); + private readonly commands = new Map>(); constructor( private readonly interpreterService: IInterpreterService, private readonly executionFactory: IPythonExecutionFactory, @@ -100,13 +100,29 @@ export class JupyterCommandFinder { // Only log telemetry if not already found (meaning the first time) const timer = new StopWatch(); - try { - const result = await this.findBestCommandImpl(command, cancelToken); - this.commands.set(command, result); - return result; - } finally { - sendTelemetryEvent(Telemetry.FindJupyterCommand, timer.elapsedTime, { command }); + const promise = this.findBestCommandImpl(command, cancelToken) + .finally(() => sendTelemetryEvent(Telemetry.FindJupyterCommand, timer.elapsedTime, { command })); + + if (cancelToken) { + let promiseCompleted = false; + promise.finally(() => promiseCompleted = true).ignoreErrors(); + + // If the promise is not pending, then remove the item from cache. + // As the promise would not complete correctly, as its been cancelled. + if (cancelToken.isCancellationRequested && !promiseCompleted) { + this.commands.delete(command); + } + cancelToken.onCancellationRequested(() => { + // If the promise is not pending, then remove the item from cache. + // As the promise would not complete correctly, as its been cancelled. + if (!promiseCompleted) { + this.commands.delete(command); + } + }); } + + this.commands.set(command, promise); + return promise; } /**