Swap controllers once kernel starts

Author: DonJayamanne

.github/workflows/build-test.yml

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

@@ -4141,7 +4141,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
 ```
 
@@ -4168,9 +4168,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>
@@ -8871,7 +8871,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

@@ -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,11 +54,9 @@ import {
     IInteractiveWindowDebuggingManager,
     InteractiveTab
 } from './types';
-import { generateInteractiveCode, isInteractiveInputTab } from './helpers';
+import { generateInteractiveCode, getInteractiveCellMetadata, isInteractiveInputTab } from './helpers';
 import { 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,
@@ -71,6 +69,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.
@@ -105,7 +105,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;
     } = {};
@@ -135,6 +136,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
     private readonly debuggingManager: IInteractiveWindowDebuggingManager;
     private readonly isWebExtension: boolean;
     private readonly commandManager: ICommandManager;
+    private readonly kernelProvider: IKernelProvider;
     constructor(
         private readonly serviceContainer: IServiceContainer,
         private _owner: Resource,
@@ -156,6 +158,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
         );
@@ -229,31 +232,39 @@ 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.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(
@@ -264,7 +275,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;
@@ -304,10 +316,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) {
@@ -423,26 +436,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
         message = message.split('\n').join('  \n');
         this.updateSysInfoMessage(message, true, cellPromise);
     }
-    private registerControllerChangeListener(controller: IVSCodeNotebookController) {
-        const controllerChangeListener = controller.controller.onDidChangeSelectedNotebooks(
-            (selectedEvent: { notebook: NotebookDocument; selected: boolean }) => {
-                // Controller was deselected for this InteractiveWindow's NotebookDocument
-                if (selectedEvent.selected === false && selectedEvent.notebook === this.notebookEditor.notebook) {
-                    controllerChangeListener.dispose();
-                    this.disconnectKernel();
-                }
-            },
-            this,
-            this.internalDisposables
-        );
-    }
 
     private listenForControllerSelection() {
-        const controller = this.notebookControllerSelection.getSelected(this.notebookEditor.notebook);
-        if (controller !== undefined) {
-            this.registerControllerChangeListener(controller);
-        }
-
         // Ensure we hear about any controller changes so we can update our cached promises
         this.notebookControllerSelection.onControllerSelected(
             (e: { notebook: NotebookDocument; controller: IVSCodeNotebookController }) => {
@@ -451,8 +446,13 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
                 }
 
                 // Clear cached kernel when the selected controller for this document changes
-                this.registerControllerChangeListener(e.controller);
-                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();
                 }
@@ -615,6 +615,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/ipywidgets/commonMessageCoordinator.ts

@@ -4,9 +4,9 @@
 'use strict';
 
 import type { KernelMessage } from '@jupyterlab/services';
-import { Event, EventEmitter, NotebookDocument } from 'vscode';
-import { IApplicationShell, ICommandManager } from '../../platform/common/application/types';
-import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
+import { Event, EventEmitter, NotebookDocument, notebooks } from 'vscode';
+import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../platform/common/application/types';
+import { IPyWidgetRendererId, STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
 import { traceVerbose, traceError, traceInfo, traceInfoIfCI } from '../../platform/logging';
 import {
     IDisposableRegistry,
@@ -65,7 +65,9 @@ export class CommonMessageCoordinator {
     private readonly configService: IConfigurationService;
     private webview: IWebviewCommunication | undefined;
     private modulesForWhichWeHaveDisplayedWidgetErrorMessage = new Set<string>();
-
+    private kernelProvider: IKernelProvider;
+    private queuedMessages: { type: string; payload: unknown }[] = [];
+    private readyMessageReceived?: boolean;
     public constructor(
         private readonly document: NotebookDocument,
         private readonly serviceContainer: IServiceContainer
@@ -75,6 +77,7 @@ export class CommonMessageCoordinator {
         this.appShell = this.serviceContainer.get<IApplicationShell>(IApplicationShell, IApplicationShell);
         this.commandManager = this.serviceContainer.get<ICommandManager>(ICommandManager);
         this.configService = this.serviceContainer.get<IConfigurationService>(IConfigurationService);
+        this.kernelProvider = this.serviceContainer.get<IKernelProvider>(IKernelProvider);
     }
 
     public dispose() {
@@ -103,6 +106,12 @@ export class CommonMessageCoordinator {
                             response: webview.asWebviewUri(e.payload)
                         });
                     } else {
+                        if (!webview.isReady || !this.readyMessageReceived) {
+                            // Web view is not yet ready to receive messages, hence queue these to be sent later.
+                            this.queuedMessages.push({ type: e.message, payload: e.payload });
+                            return;
+                        }
+                        this.sendPendingWebViewMessages();
                         webview.postMessage({ type: e.message, payload: e.payload }).then(noop, noop);
                     }
                 },
@@ -113,10 +122,55 @@ export class CommonMessageCoordinator {
                 (m) => {
                     traceInfoIfCI(`${ConsoleForegroundColors.Green}Widget Coordinator received ${m.type}`);
                     this.onMessage(m.type, m.payload);
-
+                    const kernel = this.kernelProvider.get(this.document.uri);
                     // Special case the WidgetManager loaded message. It means we're ready
                     // to use a kernel. (IPyWidget Dispatcher uses this too)
                     if (m.type === IPyWidgetMessages.IPyWidgets_Ready) {
+                        if (kernel?.kernelConnectionMetadata.kind === 'startUsingRemoteKernelSpec') {
+                            traceInfoIfCI(
+                                'Web view is not ready to receive widget messages (kernel points to remote kernel spec)'
+                            );
+                            const nbEditor = this.serviceContainer
+                                .get<IVSCodeNotebook>(IVSCodeNotebook)
+                                .notebookEditors.find((item) => item.notebook === this.document);
+                            // With remote kernel specs, once the kernel is ready we create a live kernel controller and
+                            // switch to that. At that point the webview also changes, hence
+                            // there's no need to render anything while were in this state.
+                            notebooks
+                                .createRendererMessaging(IPyWidgetRendererId)
+                                .postMessage({ type: IPyWidgetMessages.IPyWidgets_DoNotRenderWidgets }, nbEditor)
+                                .then(noop, noop);
+                            return;
+                        }
+                        if (!webview.isReady) {
+                            traceInfoIfCI('Web view is not ready to receive widget messages');
+                            return;
+                        }
+                        traceInfoIfCI('Web view is ready to receive widget messages');
+                        this.readyMessageReceived = true;
+                        this.sendPendingWebViewMessages();
+                        // At this point, we know the kernels are ready, and the webview is ready to receive messages.
+                        // Its possible the webview was initially unable to render some widgets, but now that everything is ready we
+                        // should be able to render them now.
+                        // E.g assume we're dealing with remote kernel specs,
+                        // Then we start a kernel, next we create a controller that points to the live kernel & change the controller.
+                        // What happens now is the, webview is re-loaded (as theres a change in the controllers).
+                        // However if the user were to run a cell, then the output would get displayed and it could end up,
+                        // being displayed in the old webview & when the webview is re-loaded, it would be displayed in the new webview.
+                        // However its possble the kernel is not yet ready at that point in time.
+                        // We could solve this issue easily by not executing cells, until the webview is ready,
+                        // however that would have significant performance implications.
+                        // Hence for ipywidgets, once the webview has completely been initialized, we can attempt to re-render the widgets.
+                        const nbEditor = this.serviceContainer
+                            .get<IVSCodeNotebook>(IVSCodeNotebook)
+                            .notebookEditors.find((item) => item.notebook === this.document);
+                        if (nbEditor) {
+                            traceInfoIfCI('Re-rendering widgets');
+                            notebooks
+                                .createRendererMessaging(IPyWidgetRendererId)
+                                .postMessage({ type: IPyWidgetMessages.IPyWidgets_ReRenderWidgets }, nbEditor)
+                                .then(noop, noop);
+                        }
                         promise.resolve();
                     }
                 },
@@ -155,6 +209,14 @@ export class CommonMessageCoordinator {
         this.getIPyWidgetScriptSource().onMessage(message, payload);
     }
 
+    private sendPendingWebViewMessages() {
+        if (!this.webview || !this.webview.isReady || !this.readyMessageReceived) {
+            return;
+        }
+        while (this.queuedMessages.length) {
+            this.webview.postMessage(this.queuedMessages.shift()!).then(noop, noop);
+        }
+    }
     private initialize() {
         traceVerbose('initialize CommonMessageCoordinator');
         // First hook up the widget script source that will listen to messages even before we start sending messages.