Select Kernel Source Per Document Experiment

package.json

@@ -835,6 +835,10 @@
             {
                 "command": "jupyter.switchToAnotherRemoteKernels",
                 "title": "%DataScience.switchToAnotherRemoteKernelsTitle%"
+            },
+            {
+                "command": "jupyter.pickDocumentKernelSource",
+                "title": "%DataScience.pickDocumentKernelSourceTitle%"
             }
         ],
         "menus": {
@@ -997,6 +1001,10 @@
                 {
                     "command": "jupyter.switchToAnotherRemoteKernels",
                     "when": "jupyter.showingRemoteKernels"
+                },
+                {
+                    "command": "jupyter.pickDocumentKernelSource",
+                    "when": "jupyter.pickDocumentKernelSourceContext"
                 }
             ],
             "interactive/toolbar": [
@@ -1491,6 +1499,10 @@
                 {
                     "command": "jupyter.interactive.clearAllCells",
                     "when": "activeEditor == 'workbench.editor.interactive' && isWorkspaceTrusted"
+                },
+                {
+                    "command": "jupyter.pickDocumentKernelSource",
+                    "when": "false"
                 }
             ],
             "debug/variables/context": [
@@ -2094,7 +2106,8 @@
         "portsAttributes",
         "quickPickSortByLabel",
         "notebookKernelSource",
-        "interactiveWindow"
+        "interactiveWindow",
+        "notebookControllerAffinityHidden"
     ],
     "scripts": {
         "package": "gulp clean && gulp prePublishBundle && vsce package -o ms-toolsai-jupyter-insiders.vsix",

package.nls.json

@@ -1118,6 +1118,7 @@
     "DataScience.switchToLocalKernelsTitle": "Connect to Local Kernels",
     "DataScience.switchToRemoteKernelsTitle": "Connect to a Jupyter Server",
     "DataScience.switchToAnotherRemoteKernelsTitle": "Connect to Another Jupyter Server",
+    "DataScience.pickDocumentKernelSourceTitle": "Select Kernel Source",
     "DataScience.failedToInstallPythonExtension": "Failed to install the Python Extension.",
     "DataScience.kernelPrefixForRemote": "(Remote)",
     "DataScience.pythonFileOverridesPythonPackage": "This file could potentially override an existing Python package and interfere with kernel execution, consider renaming it.",

src/commands.ts

@@ -205,4 +205,5 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     [DSCommands.SwitchToLocalKernels]: [];
     [DSCommands.SwitchToRemoteKernels]: [];
     [DSCommands.SwitchToAnotherRemoteKernels]: [];
+    [DSCommands.PickDocumentKernelSource]: [NotebookDocument | undefined];
 }

src/kernels/kernelFinder.ts

@@ -6,7 +6,7 @@ import { CancellationToken, Event, EventEmitter } from 'vscode';
 import { IDisposable, IDisposableRegistry, Resource } from '../platform/common/types';
 import { StopWatch } from '../platform/common/utils/stopWatch';
 import { traceInfoIfCI } from '../platform/logging';
-import { IContributedKernelFinder } from './internalTypes';
+import { IContributedKernelFinder, IContributedKernelFinderInfo } from './internalTypes';
 import { IKernelFinder, KernelConnectionMetadata } from './types';
 
 /**
@@ -16,6 +16,10 @@ import { IKernelFinder, KernelConnectionMetadata } from './types';
 export class KernelFinder implements IKernelFinder {
     private startTimeForFetching?: StopWatch;
     private _finders: IContributedKernelFinder[] = [];
+    private connectionFinderMapping: Map<string, IContributedKernelFinderInfo> = new Map<
+        string,
+        IContributedKernelFinderInfo
+    >();
 
     private _onDidChangeKernels = new EventEmitter<void>();
     onDidChangeKernels: Event<void> = this._onDidChangeKernels.event;
@@ -60,8 +64,17 @@ export class KernelFinder implements IKernelFinder {
 
         const kernels: KernelConnectionMetadata[] = [];
 
+        // List kernels might be called after finders or connections are removed so just clear out and regenerate
+        this.connectionFinderMapping.clear();
+
         for (const finder of this._finders) {
             const contributedKernels = finder.listContributedKernels(resource);
+
+            // Add our connection => finder mapping
+            contributedKernels.forEach((connection) => {
+                this.connectionFinderMapping.set(connection.id, finder);
+            });
+
             kernels.push(...contributedKernels);
         }
 
@@ -73,4 +86,15 @@ export class KernelFinder implements IKernelFinder {
 
         return kernels;
     }
+
+    // Check our mappings to see what connection supplies this metadata, since metadatas can be created outside of finders
+    // allow for undefined as a return value
+    public getFinderForConnection(kernelMetadata: KernelConnectionMetadata): IContributedKernelFinderInfo | undefined {
+        return this.connectionFinderMapping.get(kernelMetadata.id);
+    }
+
+    // Give the info for what kernel finders are currently registered
+    public get registered(): IContributedKernelFinderInfo[] {
+        return this._finders as IContributedKernelFinderInfo[];
+    }
 }

src/kernels/types.ts

@@ -20,6 +20,7 @@ import { PythonEnvironment } from '../platform/pythonEnvironments/info';
 import { IAsyncDisposable, IDisplayOptions, Resource } from '../platform/common/types';
 import { IBackupFile, IJupyterKernel } from './jupyter/types';
 import { PythonEnvironment_PythonApi } from '../platform/api/types';
+import { IContributedKernelFinderInfo } from './internalTypes';
 
 export type WebSocketData = string | Buffer | ArrayBuffer | Buffer[];
 
@@ -95,7 +96,6 @@ export type PythonKernelConnectionMetadata = Readonly<{
  * Unexpected as connections are defined once & not changed, if we need to change then user needs to create a new connection.
  */
 export type KernelConnectionMetadata = RemoteKernelConnectionMetadata | LocalKernelConnectionMetadata;
-
 /**
  * Connection metadata for local kernels. Makes it easier to not have to check for the live connection type.
  */
@@ -612,6 +612,14 @@ export const IKernelFinder = Symbol('IKernelFinder');
 export interface IKernelFinder {
     onDidChangeKernels: Event<void>;
     listKernels(resource: Resource, cancelToken?: CancellationToken): Promise<KernelConnectionMetadata[]>;
+    /*
+     * For a given kernel connection metadata return what kernel finder found it
+     */
+    getFinderForConnection(kernelMetadata: KernelConnectionMetadata): IContributedKernelFinderInfo | undefined;
+    /*
+     * Return basic info on all currently registered kernel finders
+     */
+    registered: IContributedKernelFinderInfo[];
 }
 
 export type KernelAction = 'start' | 'interrupt' | 'restart' | 'execution';

src/notebooks/controllers/commands/pickDocumentKernelSourceCommand.ts

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+'use strict';
+import { inject, injectable } from 'inversify';
+import { NotebookDocument, window } from 'vscode';
+import { IExtensionSingleActivationService } from '../../../platform/activation/types';
+import { ICommandManager } from '../../../platform/common/application/types';
+import { Commands } from '../../../platform/common/constants';
+import { ContextKey } from '../../../platform/common/contextKey';
+import { IConfigurationService, IDisposableRegistry } from '../../../platform/common/types';
+import { noop } from '../../../platform/common/utils/misc';
+import { INotebookKernelSourceSelector } from '../types';
+
+// Command that we will place into the kernel picker to determine what the controller source is for this document
+@injectable()
+export class PickDocumentKernelSourceCommand implements IExtensionSingleActivationService {
+    private showPickDocumentKernelSourceContext: ContextKey;
+    constructor(
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+        @inject(ICommandManager) private readonly commandManager: ICommandManager,
+        @inject(IConfigurationService) private readonly configService: IConfigurationService,
+        @inject(INotebookKernelSourceSelector) private readonly kernelSourceSelector: INotebookKernelSourceSelector
+    ) {
+        // Context keys to control when these commands are shown
+        this.showPickDocumentKernelSourceContext = new ContextKey(
+            'jupyter.pickDocumentKernelSourceContext',
+            this.commandManager
+        );
+    }
+
+    public async activate(): Promise<void> {
+        // Register for config changes
+        this.disposables.push(this.configService.getSettings().onDidChange(this.updateVisibility));
+
+        // Register our command to execute
+        this.disposables.push(
+            this.commandManager.registerCommand(Commands.PickDocumentKernelSource, this.pickDocumentKernelSource, this)
+        );
+
+        this.updateVisibility();
+    }
+
+    private async pickDocumentKernelSource(_notebook?: NotebookDocument) {
+        // We want to get the context here from the command, but that needs a fix on the core side: https://github.com/microsoft/vscode/issues/161445
+        if (window.activeNotebookEditor && this.configService.getSettings().kernelPickerType === 'Insiders') {
+            await this.kernelSourceSelector.selectKernelSource(window.activeNotebookEditor.notebook);
+        }
+    }
+
+    // Only show this command if we are in our Insiders picker type
+    private updateVisibility() {
+        if (this.configService.getSettings().kernelPickerType === 'Insiders') {
+            this.showPickDocumentKernelSourceContext.set(true).then(noop, noop);
+        } else {
+            this.showPickDocumentKernelSourceContext.set(false).then(noop, noop);
+        }
+    }
+}

src/notebooks/controllers/kernelSource/notebookKernelSourceSelector.ts

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { NotebookDocument, QuickPickItem } from 'vscode';
+import { IContributedKernelFinderInfo } from '../../../kernels/internalTypes';
+import { IKernelFinder } from '../../../kernels/types';
+import { IApplicationShell } from '../../../platform/common/application/types';
+import { INotebookKernelSourceSelector, INotebookKernelSourceTracker } from '../types';
+
+interface KernelFinderQuickPickItem extends QuickPickItem {
+    kernelFinderInfo: IContributedKernelFinderInfo;
+}
+
+// Provides the UI to select a Kernel Source for a given notebook document
+@injectable()
+export class NotebookKernelSourceSelector implements INotebookKernelSourceSelector {
+    constructor(
+        @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
+        @inject(INotebookKernelSourceTracker) private readonly kernelSourceTracker: INotebookKernelSourceTracker,
+        @inject(IKernelFinder) private readonly kernelFinder: IKernelFinder
+    ) {}
+
+    public async selectKernelSource(notebook: NotebookDocument): Promise<void> {
+        const quickPickItems = this.kernelFinder.registered.map(this.toQuickPickItem);
+        const selectedItem = await this.applicationShell.showQuickPick(quickPickItems);
+
+        // If we selected something persist that value
+        if (selectedItem) {
+            this.kernelSourceTracker.setKernelSourceForNotebook(notebook, selectedItem.kernelFinderInfo);
+        }
+    }
+
+    toQuickPickItem(kernelFinderInfo: IContributedKernelFinderInfo): KernelFinderQuickPickItem {
+        return { kernelFinderInfo, label: kernelFinderInfo.displayName };
+    }
+}

src/notebooks/controllers/kernelSource/notebookKernelSourceTracker.ts

@@ -0,0 +1,118 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { NotebookControllerAffinity2, NotebookDocument, workspace } from 'vscode';
+import { IContributedKernelFinderInfo } from '../../../kernels/internalTypes';
+import { IKernelFinder } from '../../../kernels/types';
+import { IExtensionSyncActivationService } from '../../../platform/activation/types';
+import { IDisposableRegistry } from '../../../platform/common/types';
+import { IControllerRegistration, INotebookKernelSourceTracker, IVSCodeNotebookController } from '../types';
+
+// Controls which kernel source is associated with each document, and controls hiding and showing kernel sources for them.
+@injectable()
+export class NotebookKernelSourceTracker implements INotebookKernelSourceTracker, IExtensionSyncActivationService {
+    private documentSourceMapping: Map<NotebookDocument, IContributedKernelFinderInfo | undefined> = new Map<
+        NotebookDocument,
+        IContributedKernelFinderInfo | undefined
+    >();
+
+    constructor(
+        @inject(IDisposableRegistry) private readonly disposableRegistry: IDisposableRegistry,
+        @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration,
+        @inject(IKernelFinder) private readonly kernelFinder: IKernelFinder
+    ) {}
+
+    activate(): void {
+        workspace.onDidOpenNotebookDocument(this.onDidOpenNotebookDocument, this, this.disposableRegistry);
+        workspace.onDidCloseNotebookDocument(this.onDidCloseNotebookDocument, this, this.disposableRegistry);
+        this.controllerRegistration.onCreated(this.onCreatedController, this, this.disposableRegistry);
+
+        // Tag all open documents
+        workspace.notebookDocuments.forEach(this.onDidOpenNotebookDocument.bind(this));
+    }
+
+    public getKernelSourceForNotebook(notebook: NotebookDocument): IContributedKernelFinderInfo | undefined {
+        return this.documentSourceMapping.get(notebook);
+    }
+    public setKernelSourceForNotebook(notebook: NotebookDocument, kernelSource: IContributedKernelFinderInfo): void {
+        this.documentSourceMapping.set(notebook, kernelSource);
+
+        // After setting the kernelsource we now need to change the affinity of the controllers to hide all controllers not from that finder
+        this.updateControllerAffinity(notebook, kernelSource);
+    }
+
+    // When a controller is created, see if it shows or hides for all open documents
+    private onCreatedController(controller: IVSCodeNotebookController) {
+        this.documentSourceMapping.forEach((finderInfo, notebook) => {
+            const controllerFinderInfo = this.kernelFinder.getFinderForConnection(controller.connection);
+            if (controllerFinderInfo && finderInfo && controllerFinderInfo.id === finderInfo.id) {
+                // Match, associate with controller
+                this.associateController(notebook, controller);
+            } else {
+                this.disassociateController(notebook, controller);
+            }
+        });
+    }
+
+    private updateControllerAffinity(notebook: NotebookDocument, kernelSource: IContributedKernelFinderInfo) {
+        // Find the controller associated with the given kernel source
+        const nonAssociatedControllers = this.controllerRegistration.registered.filter(
+            (controller) => !this.controllerMatchesKernelSource(controller, kernelSource)
+        );
+        const associatedControllers = this.controllerRegistration.registered.filter((controller) =>
+            this.controllerMatchesKernelSource(controller, kernelSource)
+        );
+
+        // At this point we need to pipe in our suggestion engine, right now everything will end up with default priority
+
+        // Change the visibility on our controllers for that document
+        nonAssociatedControllers.forEach((controller) => {
+            this.disassociateController(notebook, controller);
+        });
+        associatedControllers.forEach((controller) => {
+            this.associateController(notebook, controller);
+        });
+    }
+
+    // Matching function to filter if controllers match a specific source
+    private controllerMatchesKernelSource(
+        controller: IVSCodeNotebookController,
+        kernelSource: IContributedKernelFinderInfo
+    ): boolean {
+        const controllerFinderInfo = this.kernelFinder.getFinderForConnection(controller.connection);
+        if (controllerFinderInfo && controllerFinderInfo.id === kernelSource.id) {
+            return true;
+        }
+        return false;
+    }
+
+    private associateController(notebook: NotebookDocument, controller: IVSCodeNotebookController) {
+        controller.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity2.Default);
+    }
+
+    private disassociateController(notebook: NotebookDocument, controller: IVSCodeNotebookController) {
+        controller.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity2.Hidden);
+    }
+
+    private onDidOpenNotebookDocument(notebook: NotebookDocument) {
+        this.documentSourceMapping.set(notebook, undefined);
+
+        // We need persistance here moving forward, but for now, we just default to a fresh state of
+        // not having a kernel source selected when we first open a document.
+        this.controllerRegistration.registered.forEach((controller) => {
+            this.disassociateController(notebook, controller);
+        });
+    }
+
+    private onDidCloseNotebookDocument(notebook: NotebookDocument) {
+        // Associate controller back to default on close
+        this.controllerRegistration.registered.forEach((controller) => {
+            this.associateController(notebook, controller);
+        });
+
+        this.documentSourceMapping.delete(notebook);
+    }
+}