microsoft · DonJayamanne · Mar 7, 2024 · Mar 7, 2024
diff --git a/src/kernels/jupyter/connection/preferredRemoteKernelIdProvider.ts b/src/kernels/jupyter/connection/preferredRemoteKernelIdProvider.ts
@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import { inject, injectable, named } from 'inversify';
-import { Memento, Uri } from 'vscode';
+import { Memento, NotebookDocument, Uri } from 'vscode';
 import { traceVerbose } from '../../../platform/logging';
 import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
 import { IMemento, GLOBAL_MEMENTO, ICryptoUtils } from '../../../platform/common/types';
@@ -26,10 +26,11 @@ export class PreferredRemoteKernelIdProvider {
         @inject(ICryptoUtils) private crypto: ICryptoUtils
     ) {}
 
-    public async getPreferredRemoteKernelId(uri: Uri): Promise<string | undefined> {
+    public async getPreferredRemoteKernelId(notebook: NotebookDocument): Promise<string | undefined> {
         // Stored as a list so we don't take up too much space
         const list: KernelIdListEntry[] = this.globalMemento.get<KernelIdListEntry[]>(ActiveKernelIdList, []);
         if (list.length) {
+            const uri = notebook.uri;
             // Not using a map as we're only going to store the last 40 items.
             const fileHash = await this.crypto.createHash(uri.toString());
             const entry = list.find((l) => l.fileHash === fileHash);

diff --git a/src/notebooks/controllers/liveKernelSwitcher.ts b/src/notebooks/controllers/liveKernelSwitcher.ts
@@ -42,7 +42,7 @@ export class LiveKernelSwitcher implements IExtensionSyncActivationService {
         if (!isJupyterNotebook(notebook)) {
             return;
         }
-        const preferredRemote = await this.preferredRemoteKernelIdProvider.getPreferredRemoteKernelId(notebook.uri);
+        const preferredRemote = await this.preferredRemoteKernelIdProvider.getPreferredRemoteKernelId(notebook);
         if (!preferredRemote) {
             return;
         }

diff --git a/src/notebooks/controllers/preferredKernelConnectionService.ts b/src/notebooks/controllers/preferredKernelConnectionService.ts
@@ -53,7 +53,7 @@ export class PreferredKernelConnectionService {
     ): Promise<RemoteKernelConnectionMetadata | undefined> {
         const preferredRemoteKernelId = await ServiceContainer.instance
             .get<PreferredRemoteKernelIdProvider>(PreferredRemoteKernelIdProvider)
-            .getPreferredRemoteKernelId(notebook.uri);
+            .getPreferredRemoteKernelId(notebook);
 
         const findLiveKernelConnection = async () => {
             let liveKernelMatchingIdFromCurrentKernels = kernelFinder.kernels.find(

diff --git a/src/notebooks/controllers/preferredKernelConnectionService.unit.test.ts b/src/notebooks/controllers/preferredKernelConnectionService.unit.test.ts
@@ -23,7 +23,6 @@ import { IDisposable } from '../../platform/common/types';
 import { NotebookMetadata } from '../../platform/common/utils';
 import { IInterpreterService } from '../../platform/interpreter/contracts';
 import { ServiceContainer } from '../../platform/ioc/container';
-import { uriEquals } from '../../test/datascience/helpers';
 import { TestNotebookDocument } from '../../test/datascience/notebook/executionHelper';
 import { PreferredKernelConnectionService } from './preferredKernelConnectionService';
 import { JupyterConnection } from '../../kernels/jupyter/connection/jupyterConnection';
@@ -185,7 +184,7 @@ suite('Preferred Kernel Connection', () => {
     teardown(() => (disposables = dispose(disposables)));
     suite('Live Remote Kernels (preferred match)', () => {
         test('Find preferred kernel spec if there is no exact match for the live kernel connection (match kernel spec name)', async () => {
-            when(preferredRemoteKernelProvider.getPreferredRemoteKernelId(uriEquals(notebook.uri))).thenResolve(
+            when(preferredRemoteKernelProvider.getPreferredRemoteKernelId(notebook)).thenResolve(
                 remoteLiveKernelConnection2.id
             );
             when(remoteKernelFinder.status).thenReturn('idle');
@@ -201,7 +200,7 @@ suite('Preferred Kernel Connection', () => {
             assert.strictEqual(preferredKernel, remoteJavaKernelSpec);
         });
         test('Find preferred kernel spec if there is no exact match for the live kernel connection (match kernel spec language)', async () => {
-            when(preferredRemoteKernelProvider.getPreferredRemoteKernelId(uriEquals(notebook.uri))).thenResolve(
+            when(preferredRemoteKernelProvider.getPreferredRemoteKernelId(notebook)).thenResolve(
                 remoteLiveKernelConnection2.id
             );
             when(remoteKernelFinder.status).thenReturn('idle');

diff --git a/src/test/datascience/notebook/controllerDefaultService.ts b/src/test/datascience/notebook/controllerDefaultService.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { NotebookDocument, workspace } from 'vscode';
+import { NotebookDocument } from 'vscode';
 import { isPythonNotebook } from '../../../kernels/helpers';
 import { PreferredRemoteKernelIdProvider } from '../../../kernels/jupyter/connection/preferredRemoteKernelIdProvider';
 import {
@@ -10,11 +10,10 @@ import {
     PYTHON_LANGUAGE,
     isWebExtension
 } from '../../../platform/common/constants';
-import { IDisposableRegistry, Resource } from '../../../platform/common/types';
+import { IDisposableRegistry } from '../../../platform/common/types';
 import { getNotebookMetadata } from '../../../platform/common/utils';
 import { IInterpreterService } from '../../../platform/interpreter/contracts';
 import { traceInfoIfCI, traceDecoratorVerbose, traceError } from '../../../platform/logging';
-import { isEqual } from '../../../platform/vscode-path/resources';
 import { createActiveInterpreterController } from './helpers';
 import { IControllerRegistration, IVSCodeNotebookController } from '../../../notebooks/controllers/types';
 import { IServiceContainer } from '../../../platform/ioc/types';
@@ -43,18 +42,14 @@ export class ControllerDefaultService {
         return ControllerDefaultService._instance;
     }
     public async computeDefaultController(
-        resource: Resource,
+        notebook: NotebookDocument | undefined,
         viewType: typeof JupyterNotebookView | typeof InteractiveWindowView
     ): Promise<IVSCodeNotebookController | undefined> {
         if (!IS_REMOTE_NATIVE_TEST() && this.interpreters) {
             traceInfoIfCI('CreateActiveInterpreterController');
-            return createActiveInterpreterController(viewType, resource, this.interpreters, this.registration);
+            return createActiveInterpreterController(viewType, notebook?.uri, this.interpreters, this.registration);
         } else {
             traceInfoIfCI('CreateDefaultRemoteController');
-            const notebook =
-                viewType === JupyterNotebookView
-                    ? workspace.notebookDocuments.find((item) => isEqual(item.uri, resource, true))
-                    : undefined;
             const controller = await this.createDefaultRemoteController(viewType, notebook);
             // If we're running on web, there is no active interpreter to fall back to
             if (controller || isWebExtension()) {
@@ -78,7 +73,7 @@ export class ControllerDefaultService {
         const kernelName = metadata ? metadata.kernelspec?.name : undefined;
         const preferredRemoteKernelId =
             notebook && this.preferredRemoteFinder
-                ? await this.preferredRemoteFinder.getPreferredRemoteKernelId(notebook.uri)
+                ? await this.preferredRemoteFinder.getPreferredRemoteKernelId(notebook)
                 : undefined;
 
         if (preferredRemoteKernelId) {

diff --git a/src/test/datascience/notebook/controllerPreferredService.ts b/src/test/datascience/notebook/controllerPreferredService.ts
@@ -151,7 +151,7 @@ export class ControllerPreferredService {
                 isPythonNbOrInteractiveWindow
             ) {
                 const defaultPythonController = await this.defaultService.computeDefaultController(
-                    document.uri,
+                    document,
                     document.notebookType
                 );
                 if (preferredSearchToken.token.isCancellationRequested) {
@@ -268,7 +268,7 @@ export class ControllerPreferredService {
 
                 // For interactive set the preferred controller as the interpreter or default
                 const defaultInteractiveController = await this.defaultService.computeDefaultController(
-                    document.uri,
+                    document,
                     'interactive'
                 );
                 preferredConnection = defaultInteractiveController?.connection;
@@ -411,10 +411,9 @@ export class ControllerPreferredService {
         cancelToken: CancellationToken,
         preferredInterpreter: PythonEnvironment | undefined
     ): Promise<KernelConnectionMetadata | undefined> {
-        const uri = notebook.uri;
         let preferredConnection: KernelConnectionMetadata | undefined;
         const rankedConnections = await this.kernelRankHelper.rankKernels(
-            uri,
+            notebook,
             this.registration.all,
             notebookMetadata,
             preferredInterpreter,
@@ -438,7 +437,7 @@ export class ControllerPreferredService {
             }
 
             // Are we an exact match based on metadata hash / name / ect...?
-            const isExactMatch = await this.kernelRankHelper.isExactMatch(uri, potentialMatch, notebookMetadata);
+            const isExactMatch = await this.kernelRankHelper.isExactMatch(notebook, potentialMatch, notebookMetadata);
             if (cancelToken.isCancellationRequested) {
                 return;
             }

diff --git a/src/test/datascience/notebook/kernelRankingHelper.ts b/src/test/datascience/notebook/kernelRankingHelper.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { CancellationToken } from 'vscode';
+import { CancellationToken, NotebookDocument } from 'vscode';
 import { PreferredRemoteKernelIdProvider } from '../../../kernels/jupyter/connection/preferredRemoteKernelIdProvider';
 import type * as nbformat from '@jupyterlab/nbformat';
 import {
@@ -14,10 +14,9 @@ import {
     isPythonNotebook
 } from '../../../kernels/helpers';
 import { IJupyterKernelSpec, KernelConnectionMetadata, PythonKernelConnectionMetadata } from '../../../kernels/types';
-import { isCI, PYTHON_LANGUAGE } from '../../../platform/common/constants';
+import { InteractiveWindowView, isCI, PYTHON_LANGUAGE } from '../../../platform/common/constants';
 import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
-import { Resource } from '../../../platform/common/types';
-import { getResourceType, NotebookMetadata } from '../../../platform/common/utils';
+import { NotebookMetadata } from '../../../platform/common/utils';
 import { traceError, traceInfo, traceInfoIfCI } from '../../../platform/logging';
 import { PythonEnvironment } from '../../../platform/pythonEnvironments/info';
 import { getInterpreterHash } from '../../../platform/pythonEnvironments/info/interpreter';
@@ -93,14 +92,14 @@ function getVSCodeInfoInInMetadata(
 }
 export async function rankKernels(
     kernels: KernelConnectionMetadata[],
-    resource: Resource,
+    notebook: NotebookDocument,
     notebookMetadata: nbformat.INotebookMetadata | undefined,
     preferredInterpreter: PythonEnvironment | undefined,
     preferredRemoteKernelId: string | undefined,
     cancelToken?: CancellationToken
 ): Promise<KernelConnectionMetadata[] | undefined> {
     traceInfo(
-        `Find preferred kernel for ${getDisplayPath(resource)} with metadata ${JSON.stringify(
+        `Find preferred kernel for ${getDisplayPath(notebook.uri)} with metadata ${JSON.stringify(
             notebookMetadata || {}
         )} & preferred interpreter ${
             preferredInterpreter?.uri ? getDisplayPath(preferredInterpreter?.uri) : '<undefined>'
@@ -162,7 +161,7 @@ export async function rankKernels(
 
     // Now perform our big comparison on the kernel list
     // Interactive window always defaults to Python kernels.
-    if (getResourceType(resource) === 'interactive') {
+    if (notebook.notebookType === InteractiveWindowView) {
         // TODO: Based on the resource, we should be able to find the language.
         possibleNbMetadataLanguage = PYTHON_LANGUAGE;
     } else {
@@ -186,7 +185,7 @@ export async function rankKernels(
     );
     kernels.sort((a, b) =>
         compareKernels(
-            resource,
+            notebook,
             possibleNbMetadataLanguage,
             actualNbMetadataLanguage,
             notebookMetadata,
@@ -283,7 +282,7 @@ function isKernelSpecExactMatch(
 }
 
 export function compareKernels(
-    _resource: Resource,
+    _resource: NotebookDocument,
     possibleNbMetadataLanguage: string | undefined,
     actualNbMetadataLanguage: string | undefined,
     notebookMetadata: nbformat.INotebookMetadata | undefined,
@@ -978,23 +977,20 @@ export class KernelRankingHelper {
     constructor(private readonly preferredRemoteFinder: PreferredRemoteKernelIdProvider) {}
 
     public async rankKernels(
-        resource: Resource,
+        notebook: NotebookDocument,
         kernels: KernelConnectionMetadata[],
         notebookMetadata?: nbformat.INotebookMetadata | undefined,
         preferredInterpreter?: PythonEnvironment,
         cancelToken?: CancellationToken
     ): Promise<KernelConnectionMetadata[] | undefined> {
         try {
-            const preferredRemoteKernelId =
-                resource && this.preferredRemoteFinder
-                    ? await this.preferredRemoteFinder.getPreferredRemoteKernelId(resource)
-                    : undefined;
+            const preferredRemoteKernelId = await this.preferredRemoteFinder.getPreferredRemoteKernelId(notebook);
             if (cancelToken?.isCancellationRequested) {
                 return;
             }
             let rankedKernels = await rankKernels(
                 kernels,
-                resource,
+                notebook,
                 notebookMetadata,
                 preferredInterpreter,
                 preferredRemoteKernelId,
@@ -1009,14 +1005,11 @@ export class KernelRankingHelper {
     }
 
     public async isExactMatch(
-        resource: Resource,
+        notebook: NotebookDocument,
         kernelConnection: KernelConnectionMetadata,
         notebookMetadata: nbformat.INotebookMetadata | undefined
     ): Promise<boolean> {
-        const preferredRemoteKernelId =
-            resource && this.preferredRemoteFinder
-                ? await this.preferredRemoteFinder.getPreferredRemoteKernelId(resource)
-                : undefined;
+        const preferredRemoteKernelId = await this.preferredRemoteFinder.getPreferredRemoteKernelId(notebook);
 
         return isExactMatch(kernelConnection, notebookMetadata, preferredRemoteKernelId);
     }

diff --git a/src/test/datascience/notebook/remoteNotebookEditor.vscode.common.test.ts b/src/test/datascience/notebook/remoteNotebookEditor.vscode.common.test.ts
@@ -213,7 +213,7 @@ export async function runCellAndVerifyUpdateOfPreferredRemoteKernelId(
     // If we nb it as soon as output appears, its possible the kernel id hasn't been saved yet & we mess that up.
     // Optionally we could wait for 100ms.
     await waitForCondition(
-        async () => !!(await remoteKernelIdProvider.getPreferredRemoteKernelId(nbEditor.notebook.uri)),
+        async () => !!(await remoteKernelIdProvider.getPreferredRemoteKernelId(nbEditor.notebook)),
         5_000,
         'Remote Kernel id not saved'
     );