microsoft · Don Jayamanne (DonJayamanne) · Dec 16, 2019 · Dec 16, 2019 · Dec 16, 2019 · Dec 16, 2019
diff --git a/src/client/datascience/jupyter/jupyterExecution.ts b/src/client/datascience/jupyter/jupyterExecution.ts
@@ -18,6 +18,7 @@ import { JupyterCommands, Telemetry } from '../constants';
 import {
     IConnection,
     IJupyterExecution,
+    IJupyterSessionManagerFactory,
     INotebookServer,
     INotebookServerLaunchInfo,
     INotebookServerOptions
@@ -136,7 +137,9 @@ export class JupyterExecutionBase implements IJupyterExecution {
 
                     // In a remote situation, figure out a kernel spec too.
                     if (!kernelSpecInterpreter.kernelSpec && connection) {
-                        kernelSpecInterpreter = await this.kernelSelector.getKernelForRemoteConnection(connection, options?.metadata, cancelToken);
+                        const sessionManagerFactory = this.serviceContainer.get<IJupyterSessionManagerFactory>(IJupyterSessionManagerFactory);
+                        const sessionManager = await sessionManagerFactory.create(connection);
+                        kernelSpecInterpreter = await this.kernelSelector.getKernelForRemoteConnection(sessionManager, options?.metadata, cancelToken);
                     }
 
                     // Populate the launch info that we are starting our server with

diff --git a/src/client/datascience/jupyter/jupyterSession.ts b/src/client/datascience/jupyter/jupyterSession.ts
@@ -17,6 +17,7 @@ import { noop } from '../../common/utils/misc';
 import { IConnection, IJupyterKernel, IJupyterKernelSpec, IJupyterSession } from '../types';
 import { JupyterWaitForIdleError } from './jupyterWaitForIdleError';
 import { JupyterKernelPromiseFailedError } from './kernels/jupyterKernelPromiseFailedError';
+import { KernelSelector } from './kernels/kernelSelector';
 
 export class JupyterSession implements IJupyterSession {
     private session: Session.ISession | undefined;
@@ -31,7 +32,8 @@ export class JupyterSession implements IJupyterSession {
         private serverSettings: ServerConnection.ISettings,
         private kernelSpec: IJupyterKernelSpec | IJupyterKernel & Partial<IJupyterKernelSpec> | undefined,
         private sessionManager: SessionManager,
-        private contentsManager: ContentsManager
+        private contentsManager: ContentsManager,
+        private readonly kernelSelector: KernelSelector
     ) {
         this.statusHandler = this.onStatusChanged.bind(this);
     }
@@ -101,6 +103,7 @@ export class JupyterSession implements IJupyterSession {
             if (!this.session) {
                 throw new Error(localize.DataScience.sessionDisposed());
             }
+            this.kernelSelector.removeKernelFromIgnoreList(this.session.kernel);
             traceInfo(`Got new session ${this.session.kernel.id}`);
 
             // Rewire our status changed event.
@@ -262,6 +265,7 @@ export class JupyterSession implements IJupyterSession {
             try {
                 result = await this.createSession(serverSettings, contentsManager, cancelToken);
                 await this.waitForIdleOnSession(result, 30000);
+                this.kernelSelector.addKernelToIgnoreList(result.kernel);
                 return result;
             } catch (exc) {
                 traceInfo(`Error waiting for restart session: ${exc}`);

diff --git a/src/client/datascience/jupyter/jupyterSessionManager.ts b/src/client/datascience/jupyter/jupyterSessionManager.ts
@@ -20,6 +20,7 @@ import {
 import { JupyterSession } from './jupyterSession';
 import { createJupyterWebSocket } from './jupyterWebSocket';
 import { JupyterKernelSpec } from './kernels/jupyterKernelSpec';
+import { KernelSelector } from './kernels/kernelSelector';
 
 export class JupyterSessionManager implements IJupyterSessionManager {
 
@@ -31,7 +32,8 @@ export class JupyterSessionManager implements IJupyterSessionManager {
     constructor(
         private jupyterPasswordConnect: IJupyterPasswordConnect,
         private config: IConfigurationService,
-        private failOnPassword: boolean | undefined
+        private failOnPassword: boolean | undefined,
+        private kernelSelector: KernelSelector
     ) {
     }
 
@@ -85,7 +87,7 @@ export class JupyterSessionManager implements IJupyterSessionManager {
             throw new Error(localize.DataScience.sessionDisposed());
         }
         // Create a new session and attempt to connect to it
-        const session = new JupyterSession(this.connInfo, this.serverSettings, kernelSpec, this.sessionManager, this.contentsManager);
+        const session = new JupyterSession(this.connInfo, this.serverSettings, kernelSpec, this.sessionManager, this.contentsManager, this.kernelSelector);
         try {
             await session.connect(cancelToken);
         } finally {

diff --git a/src/client/datascience/jupyter/jupyterSessionManagerFactory.ts b/src/client/datascience/jupyter/jupyterSessionManagerFactory.ts
@@ -6,13 +6,15 @@ import { inject, injectable } from 'inversify';
 import { IConfigurationService } from '../../common/types';
 import { IConnection, IJupyterPasswordConnect, IJupyterSessionManager, IJupyterSessionManagerFactory } from '../types';
 import { JupyterSessionManager } from './jupyterSessionManager';
+import { KernelSelector } from './kernels/kernelSelector';
 
 @injectable()
 export class JupyterSessionManagerFactory implements IJupyterSessionManagerFactory {
 
     constructor(
         @inject(IJupyterPasswordConnect) private jupyterPasswordConnect: IJupyterPasswordConnect,
-        @inject(IConfigurationService) private config: IConfigurationService
+        @inject(IConfigurationService) private config: IConfigurationService,
+        @inject(KernelSelector) private kernelSelector: KernelSelector
     ) {
     }
 
@@ -22,7 +24,7 @@ export class JupyterSessionManagerFactory implements IJupyterSessionManagerFacto
      * @param failOnPassword - whether or not to fail the creation if a password is required.
      */
     public async create(connInfo: IConnection, failOnPassword?: boolean): Promise<IJupyterSessionManager> {
-        const result = new JupyterSessionManager(this.jupyterPasswordConnect, this.config, failOnPassword);
+        const result = new JupyterSessionManager(this.jupyterPasswordConnect, this.config, failOnPassword, this.kernelSelector);
         await result.initialize(connInfo);
         return result;
     }

diff --git a/src/client/datascience/jupyter/kernels/kernelSelector.ts b/src/client/datascience/jupyter/kernels/kernelSelector.ts
@@ -3,6 +3,7 @@
 import '../../../common/extensions';
 
 import { nbformat } from '@jupyterlab/coreutils';
+import { Kernel } from '@jupyterlab/services';
 import { inject, injectable } from 'inversify';
 import { CancellationToken } from 'vscode-jsonrpc';
 
@@ -14,7 +15,7 @@ import { noop } from '../../../common/utils/misc';
 import { IInterpreterService, PythonInterpreter } from '../../../interpreter/contracts';
 import { sendTelemetryEvent } from '../../../telemetry';
 import { Telemetry } from '../../constants';
-import { IConnection, IJupyterKernel, IJupyterKernelSpec, IJupyterSessionManager, IJupyterSessionManagerFactory } from '../../types';
+import { IJupyterKernel, IJupyterKernelSpec, IJupyterSessionManager } from '../../types';
 import { KernelSelectionProvider } from './kernelSelections';
 import { KernelService } from './kernelService';
 import { IKernelSpecQuickPickItem } from './types';
@@ -40,14 +41,43 @@ export type KernelSpecInterpreter = {
 
 @injectable()
 export class KernelSelector {
+    /**
+     * List of ids of kernels that should be hiddenn from the kernel picker.
+     *
+     * @private
+     * @type {new Set<string>}
+     * @memberof KernelSelector
+     */
+    private readonly kernelIdsToHide = new Set<string>();
     constructor(
         @inject(KernelSelectionProvider) private readonly selectionProvider: KernelSelectionProvider,
         @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
         @inject(KernelService) private readonly kernelService: KernelService,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
-        @inject(IInstaller) private readonly installer: IInstaller,
-        @inject(IJupyterSessionManagerFactory) private readonly jupyterSessionManagerFactory: IJupyterSessionManagerFactory
-    ) { }
+        @inject(IInstaller) private readonly installer: IInstaller
+    ) {}
+
+    /**
+     * Ensure kernels such as those associated with the restart session are not displayed in the kernel picker.
+     *
+     * @param {Kernel.IKernelConnection} kernel
+     * @memberof KernelSelector
+     */
+    public addKernelToIgnoreList(kernel: Kernel.IKernelConnection): void {
+        this.kernelIdsToHide.add(kernel.id);
+        this.kernelIdsToHide.add(kernel.clientId);
+    }
+    /**
+     * Opposite of the add counterpart.
+     *
+     * @param {Kernel.IKernelConnection} kernel
+     * @memberof KernelSelector
+     */
+    public removeKernelFromIgnoreList(kernel: Kernel.IKernelConnection): void {
+        this.kernelIdsToHide.delete(kernel.id);
+        this.kernelIdsToHide.delete(kernel.clientId);
+    }
+
     /**
      * Selects a kernel from a remote session.
      *
@@ -57,7 +87,8 @@ export class KernelSelector {
      * @memberof KernelSelector
      */
     public async selectRemoteKernel(session: IJupyterSessionManager, cancelToken?: CancellationToken, currentKernel?: IJupyterKernelSpec | IJupyterKernel & Partial<IJupyterKernelSpec>): Promise<KernelSpecInterpreter> {
-        const suggestions = await this.selectionProvider.getKernelSelectionsForRemoteSession(session, cancelToken);
+        let suggestions = await this.selectionProvider.getKernelSelectionsForRemoteSession(session, cancelToken);
+        suggestions = suggestions.filter(item => !this.kernelIdsToHide.has(item.selection.kernelModel?.id || ''));
         return this.selectKernel(suggestions, session, cancelToken, currentKernel);
     }
     /**
@@ -69,7 +100,8 @@ export class KernelSelector {
      * @memberof KernelSelector
      */
     public async selectLocalKernel(session?: IJupyterSessionManager, cancelToken?: CancellationToken, currentKernel?: IJupyterKernelSpec | IJupyterKernel & Partial<IJupyterKernelSpec>): Promise<KernelSpecInterpreter> {
-        const suggestions = await this.selectionProvider.getKernelSelectionsForLocalSession(session, cancelToken);
+        let suggestions = await this.selectionProvider.getKernelSelectionsForLocalSession(session, cancelToken);
+        suggestions = suggestions.filter(item => !this.kernelIdsToHide.has(item.selection.kernelModel?.id || ''));
         return this.selectKernel(suggestions, session, cancelToken, currentKernel);
     }
     /**
@@ -133,11 +165,10 @@ export class KernelSelector {
      */
     // tslint:disable-next-line: cyclomatic-complexity
     public async getKernelForRemoteConnection(
-        connInfo: IConnection,
+        sessionManager?: IJupyterSessionManager,
         notebookMetadata?: nbformat.INotebookMetadata,
         cancelToken?: CancellationToken
     ): Promise<KernelSpecInterpreter> {
-        const sessionManager = await this.jupyterSessionManagerFactory.create(connInfo);
         const [interpreter, specs] = await Promise.all([this.interpreterService.getActiveInterpreter(undefined), this.kernelService.getKernelSpecs(sessionManager, cancelToken)]);
         let bestMatch: IJupyterKernelSpec | undefined;
         let bestScore = 0;