Swap controllers once kernel starts

Author: DonJayamanne

.github/workflows/build-test.yml

@@ -291,8 +291,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.

TELEMETRY.md

@@ -4311,7 +4311,7 @@ No properties for event
             const telemetryEvent = isLocalConnection(this.kernelConnection)
                 ? Telemetry.SelectLocalJupyterKernel
                 : Telemetry.SelectRemoteJupyterKernel;
-            sendKernelTelemetryEvent(document.uri, telemetryEvent);
+            sendKernelTelemetryEvent(notebook.uri, telemetryEvent);
             this.notebookApi.notebookEditors
 ```
 
@@ -4338,9 +4338,9 @@ No properties for event
             const telemetryEvent = isLocalConnection(this.kernelConnection)
                 ? Telemetry.SelectLocalJupyterKernel
                 : Telemetry.SelectRemoteJupyterKernel;
-            sendKernelTelemetryEvent(document.uri, telemetryEvent);
+            sendKernelTelemetryEvent(notebook.uri, telemetryEvent);
             this.notebookApi.notebookEditors
-                .filter((editor) => editor.notebook === document)
+                .filter((editor) => editor.notebook === notebook)
 ```
 
 </details>
@@ -8959,7 +8959,7 @@ No properties for event
             default:
             // We don't know as its the default kernel on Jupyter server.
         }
-        sendKernelTelemetryEvent(document.uri, Telemetry.SwitchKernel);
+        sendKernelTelemetryEvent(notebook.uri, Telemetry.SwitchKernel);
         // If we have an existing kernel, then we know for a fact the user is changing the kernel.
         // Else VSC is just setting a kernel for a notebook after it has opened.
         if (existingKernel) {

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/interactiveWindow.ts

@@ -33,6 +33,7 @@ import { IConfigurationService, InteractiveWindowMode, IsWebExtension, Resource
 import { noop } from '../platform/common/utils/misc';
 import {
     IKernel,
+    IKernelProvider,
     isLocalConnection,
     KernelAction,
     KernelConnectionMetadata,
@@ -42,7 +43,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';
@@ -55,15 +55,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,
@@ -76,6 +74,8 @@ 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';
+import { KernelConnector } from '../notebooks/controllers/kernelConnector';
 
 /**
  * ViewModel for an interactive window from the Jupyter extension's point of view.
@@ -111,7 +111,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;
     } = {};
@@ -141,6 +142,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,
@@ -163,6 +165,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
         );
@@ -249,32 +252,40 @@ 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 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
             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(
@@ -284,7 +295,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
                 new DisplayOptions(false),
                 this.internalDisposables,
                 'jupyterExtension',
-                onStartKernel
+                onKernelStarted,
+                onKernelStartCompleted
             );
             this.currentKernelInfo.controller = kernel.controller;
             this.currentKernelInfo.metadata = kernel.kernelConnectionMetadata;
@@ -324,10 +336,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) {
@@ -443,6 +456,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(
@@ -452,7 +466,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();
                 }
@@ -621,6 +641,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/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;
+    }
+}

src/kernels/kernelProvider.base.ts

@@ -3,7 +3,7 @@
 
 'use strict';
 import type { KernelMessage } from '@jupyterlab/services';
-import { Event, EventEmitter, NotebookDocument, Uri } from 'vscode';
+import { Event, EventEmitter, NotebookController, NotebookDocument, Uri } from 'vscode';
 import { IVSCodeNotebook } from '../platform/common/application/types';
 import { traceInfoIfCI, traceVerbose, traceWarning } from '../platform/logging';
 import { getDisplayPath } from '../platform/common/platform/fs-paths';
@@ -15,8 +15,11 @@ import {
     IKernel,
     KernelOptions,
     IThirdPartyKernelProvider,
-    ThirdPartyKernelOptions
+    ThirdPartyKernelOptions,
+    KernelConnectionMetadata
 } from './types';
+import { NotebookControllerWrapper } from './notebookControllerWrapper';
+import { KernelConnectionMetadataProxy } from './kernelConnectionMetadataWrapper';
 
 /**
  * Provides kernels to the system. Generally backed by a URI or a notebook object.
@@ -137,6 +140,26 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider {
         }
         this.kernelsByNotebook.delete(notebook);
     }
+    updateKernel(kernel: IKernel, metadata: KernelConnectionMetadata, controller: NotebookController): void {
+        const controllerWrapper = kernel.controller;
+        if (NotebookControllerWrapper.isWrapped(controllerWrapper)) {
+            controllerWrapper.update(controller);
+        } else {
+            throw new Error('IKernel instantiated without passing a NotebookControllerWrapper');
+        }
+        const metadataWrapper = kernel.kernelConnectionMetadata;
+        if (KernelConnectionMetadataProxy.isWrapped(metadataWrapper)) {
+            metadataWrapper.update(metadata);
+        } else {
+            throw new Error('IKernel instantiated without passing a KernelConnectionMetadataWrapper');
+        }
+
+        const info = this.kernelsByNotebook.get(kernel.notebook);
+        if (info) {
+            info.options.controller = controller;
+            info.options.metadata = metadata;
+        }
+    }
 }
 
 export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelProvider {

src/kernels/kernelProvider.node.ts

@@ -27,6 +27,8 @@ import {
     KernelOptions,
     ThirdPartyKernelOptions
 } from './types';
+import { KernelConnectionMetadataProxy } from './kernelConnectionMetadataWrapper';
+import { NotebookControllerWrapper } from './notebookControllerWrapper';
 
 /**
  * Node version of a kernel provider. Needed in order to create the node version of a kernel.
@@ -65,12 +67,12 @@ export class KernelProvider extends BaseCoreKernelProvider {
             uri,
             resourceUri,
             notebook,
-            options.metadata,
+            KernelConnectionMetadataProxy.wrap(options.metadata),
             this.notebookProvider,
             waitForIdleTimeout,
             interruptTimeout,
             this.appShell,
-            options.controller,
+            NotebookControllerWrapper.wrap(options.controller),
             this.configService,
             this.outputTracker,
             this.statusProvider,

src/kernels/kernelProvider.web.ts

@@ -25,6 +25,8 @@ import {
     KernelOptions,
     ThirdPartyKernelOptions
 } from './types';
+import { NotebookControllerWrapper } from './notebookControllerWrapper';
+import { KernelConnectionMetadataProxy } from './kernelConnectionMetadataWrapper';
 
 /**
  * Web version of a kernel provider. Needed in order to create the web version of a kernel.
@@ -62,12 +64,12 @@ export class KernelProvider extends BaseCoreKernelProvider {
             uri,
             resourceUri,
             notebook,
-            options.metadata,
+            KernelConnectionMetadataProxy.wrap(options.metadata),
             this.notebookProvider,
             waitForIdleTimeout,
             interruptTimeout,
             this.appShell,
-            options.controller,
+            NotebookControllerWrapper.wrap(options.controller),
             this.configService,
             this.outputTracker,
             this.statusProvider,