Ensure kernel points to live kernel connection when starting remote kernels

.github/workflows/build-test.yml

@@ -283,8 +283,8 @@ jobs:
       matrix:
         jupyter: [raw, local]
         python: [nonConda, conda]
-        pythonVersion: ['3.9', '3.10']
         webOrDesktop: ['']
+        pythonVersion: ['3.9', '3.10']
         # Whether we're using stable (<empty>) or pre-release versions of Python packages.
         # When installing pre-release versions, we're only focused on jupyter & related packages.
         # Not pre-release versions of pandas, numpy or other such packages that are not core to Jupyter.

news/2 Fixes/10572.md

@@ -0,0 +1 @@
+Ensure the kernel picker points to the live Kernel Connection when starting a new remote kernel.

src/interactive-window/debugger/jupyter/kernelDebugAdapter.ts

@@ -15,6 +15,7 @@ import { getInteractiveCellMetadata } from '../../../interactive-window/helpers'
 import { KernelDebugAdapterBase } from '../../../notebooks/debugger/kernelDebugAdapterBase';
 import { InteractiveCellMetadata } from '../../editor-integration/types';
 import { IDebugService } from '../../../platform/common/application/types';
+import { noop } from '../../../platform/common/utils/misc';
 /**
  * KernelDebugAdapter listens to debug messages in order to translate file requests into real files
  * (Interactive Window generally executes against a real file)
@@ -75,7 +76,11 @@ export class KernelDebugAdapter extends KernelDebugAdapterBase {
         if (message.type === 'response' && this.debugLocationTracker?.onDidSendMessage) {
             this.debugLocationTracker.onDidSendMessage(message);
         }
-        return super.handleMessage(message);
+        const promise = super.handleMessage(message);
+        // The VS Code debugger class does not support an async `handleMessage`, its supposed to be sync.
+        // As a result, any errors here will not be handled and they'll be treated as unhandled errors.
+        promise.catch(noop);
+        return promise;
     }
 
     // Dump content of given cell into a tmp file and return path to file.

src/interactive-window/interactiveWindow.ts

@@ -32,6 +32,7 @@ import { IConfigurationService, InteractiveWindowMode, IsWebExtension, Resource
 import { noop } from '../platform/common/utils/misc';
 import {
     IKernel,
+    IKernelProvider,
     isLocalConnection,
     KernelAction,
     KernelConnectionMetadata,
@@ -41,7 +42,6 @@ import { chainable } from '../platform/common/utils/decorators';
 import { InteractiveCellResultError } from '../platform/errors/interactiveCellResultError';
 import { DataScience } from '../platform/common/utils/localize';
 import { createDeferred, Deferred } from '../platform/common/utils/async';
-import { IServiceContainer } from '../platform/ioc/types';
 import { SysInfoReason } from '../messageTypes';
 import { createOutputWithErrorMessageForDisplay } from '../platform/errors/errorUtils';
 import { INotebookExporter } from '../kernels/jupyter/types';
@@ -54,15 +54,13 @@ import {
     IInteractiveWindowDebuggingManager,
     InteractiveTab
 } from './types';
-import { generateInteractiveCode, isInteractiveInputTab } from './helpers';
+import { generateInteractiveCode, getInteractiveCellMetadata, isInteractiveInputTab } from './helpers';
 import {
     IControllerRegistration,
     IControllerSelection,
     IVSCodeNotebookController
 } from '../notebooks/controllers/types';
 import { DisplayOptions } from '../kernels/displayOptions';
-import { getInteractiveCellMetadata } from './helpers';
-import { KernelConnector } from '../notebooks/controllers/kernelConnector';
 import { getFilePath } from '../platform/common/platform/fs-paths';
 import {
     ICodeGeneratorFactory,
@@ -75,6 +73,7 @@ import { updateNotebookMetadata } from '../kernels/execution/helpers';
 import { chainWithPendingUpdates } from '../kernels/execution/notebookUpdater';
 import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../kernels/telemetry/helper';
 import { generateMarkdownFromCodeLines, parseForComments } from '../platform/common/utils';
+import { IServiceContainer } from '../platform/ioc/types';
 
 /**
  * ViewModel for an interactive window from the Jupyter extension's point of view.
@@ -110,7 +109,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
     private kernelDisposables: Disposable[] = [];
     private _insertSysInfoPromise: Promise<NotebookCell> | undefined;
     private currentKernelInfo: {
-        kernel?: Deferred<IKernel>;
+        kernelStarted?: Deferred<void>;
+        kernel?: IKernel;
         controller?: NotebookController;
         metadata?: KernelConnectionMetadata;
     } = {};
@@ -140,6 +140,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
     private readonly debuggingManager: IInteractiveWindowDebuggingManager;
     private readonly isWebExtension: boolean;
     private readonly commandManager: ICommandManager;
+    private readonly kernelProvider: IKernelProvider;
     private readonly controllerRegistration: IControllerRegistration;
     constructor(
         private readonly serviceContainer: IServiceContainer,
@@ -162,6 +163,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
         this.errorHandler = this.serviceContainer.get<IDataScienceErrorHandler>(IDataScienceErrorHandler);
         this.codeGeneratorFactory = this.serviceContainer.get<ICodeGeneratorFactory>(ICodeGeneratorFactory);
         this.storageFactory = this.serviceContainer.get<IGeneratedCodeStorageFactory>(IGeneratedCodeStorageFactory);
+        this.kernelProvider = this.serviceContainer.get<IKernelProvider>(IKernelProvider);
         this.debuggingManager = this.serviceContainer.get<IInteractiveWindowDebuggingManager>(
             IInteractiveWindowDebuggingManager
         );
@@ -248,42 +250,54 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
             // This cannot happen, but we need to make typescript happy.
             throw new Error('Controller not selected');
         }
-        if (this.currentKernelInfo.kernel) {
-            return this.currentKernelInfo.kernel.promise;
+        if (this.currentKernelInfo.kernelStarted) {
+            await this.currentKernelInfo.kernelStarted.promise;
+            return this.currentKernelInfo.kernel!;
+        }
+        const vscController = this.controllerRegistration.registered.find(
+            (item) => item.controller.id === controller.id
+        );
+        if (!vscController) {
+            // This cannot happen, but we need to make typescript happy.
+            throw new Error('VSCController not available');
         }
-        const kernelPromise = createDeferred<IKernel>();
-        kernelPromise.promise.catch(noop);
-        this.currentKernelInfo = { controller, metadata, kernel: kernelPromise };
+        const kernelStarted = createDeferred<void>();
+        kernelStarted.promise.catch(noop);
+        this.currentKernelInfo = { controller, metadata, kernelStarted };
 
         const sysInfoCell = this.insertSysInfoMessage(metadata, SysInfoReason.Start);
         try {
             // Try creating a kernel
             await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(this.owner, metadata);
 
-            const onStartKernel = (action: KernelAction, k: IKernel) => {
+            const onKernelStarted = async (action: KernelAction, k: IKernel) => {
                 if (action !== 'start' && action !== 'restart') {
                     return;
                 }
                 // Id may be different if the user switched controllers
-                this.currentKernelInfo.controller = k.controller;
-                this.currentKernelInfo.metadata = k.kernelConnectionMetadata;
+                this.currentKernelInfo.kernel = k;
                 !!this.pendingCellAdd && this.setPendingCellAdd(this.pendingCellAdd);
                 this.updateSysInfoMessage(
                     this.getSysInfoMessage(k.kernelConnectionMetadata, SysInfoReason.Start),
                     false,
                     sysInfoCell
                 );
             };
+            const onKernelStartCompleted = async (action: KernelAction, _: unknown, k: IKernel) => {
+                if (action !== 'start' && action !== 'restart') {
+                    return;
+                }
+                // Id may be different if the user switched controllers
+                this.currentKernelInfo.controller = k.controller;
+                this.currentKernelInfo.metadata = k.kernelConnectionMetadata;
+            };
             // When connecting, we need to update the sys info message
             this.updateSysInfoMessage(this.getSysInfoMessage(metadata, SysInfoReason.Start), false, sysInfoCell);
-            const kernel = await KernelConnector.connectToNotebookKernel(
-                metadata,
-                this.serviceContainer,
-                { resource: this.owner, notebook: this.notebookDocument, controller },
+            const kernel = await vscController.connectToKernel(
+                { resource: this.owner, notebook: this.notebookDocument },
                 new DisplayOptions(false),
-                this.internalDisposables,
-                'jupyterExtension',
-                onStartKernel
+                onKernelStarted,
+                onKernelStartCompleted
             );
             this.currentKernelInfo.controller = kernel.controller;
             this.currentKernelInfo.metadata = kernel.kernelConnectionMetadata;
@@ -323,10 +337,11 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
             this.fileInKernel = undefined;
             await this.runInitialization(kernel, this.owner);
             this.finishSysInfoMessage(kernel, sysInfoCell, SysInfoReason.Start);
-            kernelPromise.resolve(kernel);
+            kernelStarted.resolve();
             return kernel;
         } catch (ex) {
-            kernelPromise.reject(ex);
+            kernelStarted.reject(ex);
+            this.currentKernelInfo.kernelStarted = undefined;
             this.currentKernelInfo.kernel = undefined;
             this.disconnectKernel();
             if (this.owner) {
@@ -442,6 +457,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
         message = message.split('\n').join('  \n');
         this.updateSysInfoMessage(message, true, cellPromise);
     }
+
     private listenForControllerSelection() {
         // Ensure we hear about any controller changes so we can update our cached promises
         this.notebookControllerSelection.onControllerSelected(
@@ -451,7 +467,13 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
                 }
 
                 // Clear cached kernel when the selected controller for this document changes
-                if (e.controller.id !== this.currentKernelInfo.controller?.id) {
+                const kernel = this.kernelProvider.get(this.notebookDocument.uri);
+                const isControllerAttachedToTheSameKernel =
+                    kernel && this.currentKernelInfo.kernel && kernel === this.currentKernelInfo.kernel;
+                if (
+                    e.controller.connection.id !== this.currentKernelInfo.kernel?.kernelConnectionMetadata.id &&
+                    !isControllerAttachedToTheSameKernel
+                ) {
                     this.disconnectKernel();
                     this.startKernel(e.controller.controller, e.controller.connection).ignoreErrors();
                 }
@@ -620,6 +642,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
     private disconnectKernel() {
         this.kernelDisposables.forEach((d) => d.dispose());
         this.kernelDisposables = [];
+        this.currentKernelInfo.kernelStarted = undefined;
         this.currentKernelInfo.kernel = undefined;
     }
 

src/kernels/common/delayedFutureExecute.ts

@@ -7,6 +7,7 @@ import type { Kernel, KernelMessage } from '@jupyterlab/services';
 import { traceInfoIfCI } from '../../platform/logging';
 import { createDeferred } from '../../platform/common/utils/async';
 import { CancellationError } from 'vscode';
+import { noop } from '../../platform/common/utils/misc';
 
 // Wraps a future so that a requestExecute on a session will wait for the previous future to finish before actually executing
 export class DelayedFutureExecute
@@ -40,6 +41,8 @@ export class DelayedFutureExecute
         private disposeOnDone?: boolean,
         private metadata?: JSONObject
     ) {
+        // Ensure we don't have any unhandled promises.
+        this.doneDeferred.promise.catch(noop);
         // Setup our request based on the previous link finishing
         previousLink.done.then(() => this.requestExecute()).catch((e) => this.doneDeferred.reject(e));
 

src/kernels/execution/cellExecution.ts

@@ -82,6 +82,7 @@ export class CellExecution implements IDisposable {
     private _completed?: boolean;
     private startTime?: number;
     private endTime?: number;
+    private disposed?: boolean;
     private execution?: NotebookCellExecution;
     private cancelHandled = false;
     private request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg> | undefined;
@@ -217,6 +218,10 @@ export class CellExecution implements IDisposable {
      * Or when execution has been cancelled.
      */
     public dispose() {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
         traceCellMessage(this.cell, 'Execution disposed');
         disposeAllDisposables(this.disposables);
         this.cellExecutionHandler?.dispose();

src/kernels/execution/cellExecutionMessageHandler.ts

@@ -114,6 +114,7 @@ export class CellExecutionMessageHandler implements IDisposable {
     private temporaryExecution?: NotebookCellExecution;
     private previousResultsToRestore?: NotebookCellExecutionSummary;
     private cellHasErrorsInOutput?: boolean;
+    private disposed?: boolean;
 
     public get hasErrorOutput() {
         return this.cellHasErrorsInOutput === true;
@@ -235,7 +236,11 @@ export class CellExecutionMessageHandler implements IDisposable {
      * Or when execution has been cancelled.
      */
     public dispose() {
-        traceCellMessage(this.cell, 'Execution disposed');
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
+        traceCellMessage(this.cell, 'Execution Message Handler disposed');
         disposeAllDisposables(this.disposables);
         this.prompts.forEach((item) => item.dispose());
         this.prompts.clear();

src/kernels/execution/kernelExecution.ts

@@ -36,6 +36,7 @@ import { noop } from '../../platform/common/utils/misc';
  */
 export class BaseKernelExecution<TKernel extends IBaseKernel = IBaseKernel> implements IDisposable {
     protected readonly disposables: IDisposable[] = [];
+    private disposed?: boolean;
     private _interruptPromise?: Promise<InterruptResult>;
     private _restartPromise?: Promise<void>;
     protected get restarting() {
@@ -100,6 +101,10 @@ export class BaseKernelExecution<TKernel extends IBaseKernel = IBaseKernel> impl
         await this._restartPromise;
     }
     public dispose() {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
         traceInfoIfCI(`Dispose KernelExecution`);
         this.disposables.forEach((d) => d.dispose());
     }

src/kernels/kernelConnectionMetadataWrapper.ts

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { KernelConnectionMetadata } from './types';
+
+export class KernelConnectionMetadataProxy {
+    public get kernelModel() {
+        return 'kernelModel' in this.metadata ? this.metadata['kernelModel'] : undefined;
+    }
+    public get interpreter() {
+        return this.metadata.interpreter;
+    }
+    public get kernelSpec() {
+        return 'kernelSpec' in this.metadata ? this.metadata['kernelSpec'] : undefined;
+    }
+    public get kind() {
+        return this.metadata.kind;
+    }
+    public get baseUrl() {
+        return 'baseUrl' in this.metadata ? this.metadata['baseUrl'] : undefined;
+    }
+    public get serverId() {
+        return 'serverId' in this.metadata ? this.metadata['serverId'] : undefined;
+    }
+    public get id() {
+        return this.metadata.id;
+    }
+    private constructor(public metadata: KernelConnectionMetadata) {
+        this.metadata = metadata;
+    }
+    public update(metadata: KernelConnectionMetadata) {
+        this.metadata = metadata;
+    }
+    public toString() {
+        return JSON.stringify(this.toJSON());
+    }
+    public toJSON() {
+        return {
+            kernelModel: this.kernelModel,
+            kernelSpec: this.kernelSpec,
+            interpreter: this.interpreter,
+            kind: this.kind,
+            baseUrl: this.baseUrl,
+            serverId: this.serverId,
+            id: this.id
+        };
+    }
+    static isWrapped(
+        metadata: KernelConnectionMetadata | KernelConnectionMetadataProxy
+    ): metadata is KernelConnectionMetadataProxy {
+        if (metadata instanceof KernelConnectionMetadataProxy) {
+            return true;
+        }
+        return false;
+    }
+    static wrap(metadata: KernelConnectionMetadata): KernelConnectionMetadata {
+        if (metadata instanceof KernelConnectionMetadataProxy) {
+            return metadata;
+        }
+        return new KernelConnectionMetadataProxy(metadata) as unknown as KernelConnectionMetadata;
+    }
+}