Support returning a preferred remote controller

.vscode/settings.json

@@ -1,7 +1,7 @@
 // Place your settings in this file to overwrite default and user settings.
 {
     "files.exclude": {
-        "out": true, // set this to true to hide the "out" folder with the compiled JS files
+        "out": false, // set this to true to hide the "out" folder with the compiled JS files
         "**/*.pyc": true,
         ".nyc_output": true,
         "obj": true,

package.nls.json

@@ -288,8 +288,8 @@
     "DataScience.jupyterSelectPasswordPrompt": "Enter your password",
     "DataScience.removeRemoteJupyterServerEntryInQuickPick": "Remove",
     "DataScience.jupyterNotebookFailure": "Jupyter notebook failed to launch. \r\n{0}",
-    "DataScience.remoteJupyterConnectionFailedWithServer": "Failed to connect to the remote Jupyter Server {0}. Verify the server is running and reachable.",
-    "DataScience.remoteJupyterConnectionFailedWithServerWithError":"Failed to connect to the remote Jupyter Server {0}. Verify the server is running and reachable. ({1}).",
+    "DataScience.remoteJupyterConnectionFailedWithServer": "Failed to connect to the remote Jupyter Server '{0}'. Verify the server is running and reachable.",
+    "DataScience.remoteJupyterConnectionFailedWithServerWithError":"Failed to connect to the remote Jupyter Server '{0}'. Verify the server is running and reachable. ({1}).",
     "DataScience.remoteJupyterConnectionFailedWithoutServerWithError": "Connection failure. Verify the server is running and reachable. ({0}).",
     "DataScience.jupyterNotebookRemoteConnectFailedWeb": "Failed to connect to remote Jupyter server.\r\nCheck that the Jupyter Server URI can be reached from a browser.\r\n{0}. Click [here](https://aka.ms/vscjremoteweb) for more information.",
     "DataScience.jupyterNotebookRemoteConnectSelfCertsFailed": "Failed to connect to remote Jupyter notebook.\r\nSpecified server is using self signed certs. Enable Allow Unauthorized Remote Connection setting to connect anyways\r\n{0}\r\n{1}",

src/extension.node.ts

@@ -126,7 +126,8 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             registerPythonApi: noop,
             registerRemoteServerProvider: noop,
             showDataViewer: () => Promise.resolve(),
-            getKernelService: () => Promise.resolve(undefined)
+            getKernelService: () => Promise.resolve(undefined),
+            getSuggestedController: () => Promise.resolve(undefined)
         };
     }
 }

src/extension.web.ts

@@ -117,7 +117,8 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             registerPythonApi: noop,
             registerRemoteServerProvider: noop,
             showDataViewer: () => Promise.resolve(),
-            getKernelService: () => Promise.resolve(undefined)
+            getKernelService: () => Promise.resolve(undefined),
+            getSuggestedController: () => Promise.resolve(undefined)
         };
     }
 }

src/kernels/jupyter/commands/serverSelector.ts

@@ -48,7 +48,7 @@ export class JupyterServerSelectorCommand implements IDisposable {
             // Find one that is the default for this remote
             if (notebook) {
                 // Recompute the preferred controller
-                await this.controllerManager.computePreferredNotebookController(notebook);
+                await this.controllerManager.initializePreferredNotebookController(notebook);
 
                 // That should have picked a preferred
                 const preferred = this.controllerManager.getPreferredNotebookController(notebook);

src/kernels/jupyter/jupyterConnection.ts

@@ -5,8 +5,9 @@ import { inject, injectable } from 'inversify';
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
 import { Identifiers } from '../../platform/common/constants';
 import { IDisposableRegistry } from '../../platform/common/types';
+import { RemoteJupyterServerUriProviderError } from '../../platform/errors/remoteJupyterServerUriProviderError';
 import { IJupyterConnection } from '../types';
-import { computeUriHash, createRemoteConnectionInfo } from './jupyterUtils';
+import { computeServerId, createRemoteConnectionInfo } from './jupyterUtils';
 import { ServerConnectionType } from './launcher/serverConnectionType';
 import {
     IJupyterServerUri,
@@ -65,7 +66,7 @@ export class JupyterConnection implements IExtensionSyncActivationService {
     private async getUriFromServerId(serverId: string) {
         // Since there's one server per session, don't use a resource to figure out these settings
         const savedList = await this.serverUriStorage.getSavedUriList();
-        return savedList.find((item) => computeUriHash(item.uri) === serverId)?.uri;
+        return savedList.find((item) => computeServerId(item.uri) === serverId)?.uri;
     }
     private async createConnectionInfoFromUri(uri: string) {
         // Prepare our map of server URIs
@@ -92,18 +93,24 @@ export class JupyterConnection implements IExtensionSyncActivationService {
     public async updateServerUri(uri: string): Promise<void> {
         const idAndHandle = this.extractJupyterServerHandleAndId(uri);
         if (idAndHandle) {
-            const serverUri = await this.jupyterPickerRegistration.getJupyterServerUri(
-                idAndHandle.id,
-                idAndHandle.handle
-            );
-            this.uriToJupyterServerUri.set(uri, serverUri);
-            // See if there's an expiration date
-            if (serverUri.expiration) {
-                const timeoutInMS = serverUri.expiration.getTime() - Date.now();
-                // Week seems long enough (in case the expiration is ridiculous)
-                if (timeoutInMS > 0 && timeoutInMS < 604800000) {
-                    this.pendingTimeouts.push(setTimeout(() => this.updateServerUri(uri).ignoreErrors(), timeoutInMS));
+            try {
+                const serverUri = await this.jupyterPickerRegistration.getJupyterServerUri(
+                    idAndHandle.id,
+                    idAndHandle.handle
+                );
+                this.uriToJupyterServerUri.set(uri, serverUri);
+                // See if there's an expiration date
+                if (serverUri.expiration) {
+                    const timeoutInMS = serverUri.expiration.getTime() - Date.now();
+                    // Week seems long enough (in case the expiration is ridiculous)
+                    if (timeoutInMS > 0 && timeoutInMS < 604800000) {
+                        this.pendingTimeouts.push(
+                            setTimeout(() => this.updateServerUri(uri).ignoreErrors(), timeoutInMS)
+                        );
+                    }
                 }
+            } catch (ex) {
+                throw new RemoteJupyterServerUriProviderError(idAndHandle.id, idAndHandle.handle, ex);
             }
         }
     }

src/kernels/jupyter/jupyterUriProviderRegistration.ts

@@ -1,9 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-import { inject, injectable } from 'inversify';
+import { inject, injectable, named } from 'inversify';
+import { Memento } from 'vscode';
 
-import { IExtensions } from '../../platform/common/types';
+import { GLOBAL_MEMENTO, IExtensions, IMemento } from '../../platform/common/types';
+import { swallowExceptions } from '../../platform/common/utils/decorators';
 import * as localize from '../../platform/common/utils/localize';
+import { noop } from '../../platform/common/utils/misc';
 import { JupyterUriProviderWrapper } from './jupyterUriProviderWrapper';
 import {
     IJupyterServerUri,
@@ -12,12 +15,17 @@ import {
     JupyterServerUriHandle
 } from './types';
 
+const REGISTRATION_ID_EXTENSION_OWNER_MEMENTO_KEY = 'REGISTRATION_ID_EXTENSION_OWNER_MEMENTO_KEY';
 @injectable()
 export class JupyterUriProviderRegistration implements IJupyterUriProviderRegistration {
     private loadedOtherExtensionsPromise: Promise<void> | undefined;
     private providers = new Map<string, Promise<IJupyterUriProvider>>();
+    private providerExtensionMapping = new Map<string, string>();
 
-    constructor(@inject(IExtensions) private readonly extensions: IExtensions) {}
+    constructor(
+        @inject(IExtensions) private readonly extensions: IExtensions,
+        @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly globalMemento: Memento
+    ) {}
 
     public async getProviders(): Promise<ReadonlyArray<IJupyterUriProvider>> {
         await this.checkOtherExtensions();
@@ -26,7 +34,7 @@ export class JupyterUriProviderRegistration implements IJupyterUriProviderRegist
         return Promise.all([...this.providers.values()]);
     }
 
-    public registerProvider(provider: IJupyterUriProvider) {
+    public async registerProvider(provider: IJupyterUriProvider) {
         if (!this.providers.has(provider.id)) {
             this.providers.set(provider.id, this.createProvider(provider));
         } else {
@@ -53,14 +61,35 @@ export class JupyterUriProviderRegistration implements IJupyterUriProviderRegist
     }
 
     private async loadOtherExtensions(): Promise<void> {
+        const extensionIds = new Set<string>();
+        this.globalMemento
+            .get<{ extensionId: string; providerId: string }[]>(REGISTRATION_ID_EXTENSION_OWNER_MEMENTO_KEY, [])
+            .forEach((item) => extensionIds.add(item.extensionId));
+
         const list = this.extensions.all
-            .filter((e) => e.packageJSON?.contributes?.pythonRemoteServerProvider)
-            .map((e) => (e.isActive ? Promise.resolve() : e.activate()));
+            .filter((e) => e.packageJSON?.contributes?.pythonRemoteServerProvider || extensionIds.has(e.id))
+            .map((e) => (e.isActive ? Promise.resolve() : e.activate().then(noop, noop)));
         await Promise.all(list);
     }
 
     private async createProvider(provider: IJupyterUriProvider): Promise<IJupyterUriProvider> {
         const info = await this.extensions.determineExtensionFromCallStack();
+        this.updateRegistrationInfo(provider.id, info.extensionId).catch(noop);
         return new JupyterUriProviderWrapper(provider, info.extensionId);
     }
+    @swallowExceptions()
+    private async updateRegistrationInfo(providerId: string, extensionId: string): Promise<void> {
+        const registeredList = this.globalMemento.get<{ extensionId: string; providerId: string }[]>(
+            REGISTRATION_ID_EXTENSION_OWNER_MEMENTO_KEY,
+            []
+        );
+        registeredList.forEach((item) => this.providerExtensionMapping.set(item.providerId, item.extensionId));
+        this.providerExtensionMapping.set(providerId, extensionId);
+
+        const newList: { extensionId: string; providerId: string }[] = [];
+        this.providerExtensionMapping.forEach((providerId, extensionId) => {
+            newList.push({ extensionId, providerId });
+        });
+        await this.globalMemento.update(REGISTRATION_ID_EXTENSION_OWNER_MEMENTO_KEY, newList);
+    }
 }

src/kernels/jupyter/jupyterUriProviderWrapper.ts

@@ -14,6 +14,9 @@ export class JupyterUriProviderWrapper implements IJupyterUriProvider {
         return this.provider.id;
     }
     public getQuickPickEntryItems(): vscode.QuickPickItem[] {
+        if (!this.provider.getQuickPickEntryItems) {
+            return [];
+        }
         return this.provider.getQuickPickEntryItems().map((q) => {
             return {
                 ...q,
@@ -26,10 +29,13 @@ export class JupyterUriProviderWrapper implements IJupyterUriProvider {
             };
         });
     }
-    public handleQuickPick(
+    public async handleQuickPick(
         item: vscode.QuickPickItem,
         back: boolean
     ): Promise<JupyterServerUriHandle | 'back' | undefined> {
+        if (!this.provider.handleQuickPick) {
+            return;
+        }
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         if ((item as any).original) {
             // eslint-disable-next-line @typescript-eslint/no-explicit-any

src/kernels/jupyter/jupyterUtils.ts

@@ -8,13 +8,13 @@ import { ConfigurationTarget, Uri } from 'vscode';
 import { IApplicationShell, IWorkspaceService } from '../../platform/common/application/types';
 import { noop } from '../../platform/common/utils/misc';
 import { IJupyterConnection } from '../types';
-import { IJupyterServerUri } from './types';
+import { IJupyterServerUri, JupyterServerUriHandle } from './types';
 import { getJupyterConnectionDisplayName } from './launcher/helpers';
 import { IConfigurationService, IWatchableJupyterSettings, Resource } from '../../platform/common/types';
 import { getFilePath } from '../../platform/common/platform/fs-paths';
 import { DataScience } from '../../platform/common/utils/localize';
 import { sendTelemetryEvent } from '../../telemetry';
-import { Telemetry } from '../../platform/common/constants';
+import { Identifiers, Telemetry } from '../../platform/common/constants';
 
 export function expandWorkingDir(
     workingDir: string | undefined,
@@ -133,6 +133,13 @@ export function createRemoteConnectionInfo(
     };
 }
 
-export function computeUriHash(uri: string) {
+export function computeServerId(uri: string) {
     return hashjs.sha256().update(uri).digest('hex');
 }
+
+export function generateUriFromRemoteProvider(id: string, result: JupyterServerUriHandle) {
+    // eslint-disable-next-line
+    return `${Identifiers.REMOTE_URI}?${Identifiers.REMOTE_URI_ID_PARAM}=${id}&${
+        Identifiers.REMOTE_URI_HANDLE_PARAM
+    }=${encodeURI(result)}`;
+}

src/kernels/jupyter/remoteKernelFinder.ts

@@ -21,7 +21,7 @@ import { IJupyterSessionManagerFactory, IJupyterSessionManager } from './types';
 import { sendKernelSpecTelemetry } from '../raw/finder/helper';
 import { traceError, traceInfoIfCI } from '../../platform/logging';
 import { IPythonExtensionChecker } from '../../platform/api/types';
-import { computeUriHash } from './jupyterUtils';
+import { computeServerId } from './jupyterUtils';
 
 // This class searches for a kernel that matches the given kernel name.
 // First it searches on a global persistent state, then on the installed python interpreters,
@@ -68,9 +68,9 @@ export class RemoteKernelFinder implements IRemoteKernelFinder {
                             kind: 'startUsingRemoteKernelSpec',
                             interpreter: await this.getInterpreter(s, connInfo.baseUrl),
                             kernelSpec: s,
-                            id: getKernelId(s, undefined, computeUriHash(connInfo.url)),
+                            id: getKernelId(s, undefined, computeServerId(connInfo.url)),
                             baseUrl: connInfo.baseUrl,
-                            serverId: computeUriHash(connInfo.url)
+                            serverId: computeServerId(connInfo.url)
                         };
                         return kernel;
                     })
@@ -101,7 +101,7 @@ export class RemoteKernelFinder implements IRemoteKernelFinder {
                         },
                         baseUrl: connInfo.baseUrl,
                         id: s.kernel?.id || '',
-                        serverId: computeUriHash(connInfo.url)
+                        serverId: computeServerId(connInfo.url)
                     };
                     return kernel;
                 });

src/kernels/jupyter/serverSelector.ts

@@ -18,16 +18,21 @@ import {
     InputFlowAction
 } from '../../platform/common/utils/multiStepInput';
 import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
-import { Telemetry, Identifiers } from '../../webviews/webview-side/common/constants';
+import { Telemetry } from '../../webviews/webview-side/common/constants';
 import {
     IJupyterUriProvider,
     IJupyterUriProviderRegistration,
     IJupyterServerUriStorage,
     JupyterServerUriHandle
 } from './types';
 import { IDataScienceErrorHandler } from '../../platform/errors/types';
-import { handleExpiredCertsError, handleSelfCertsError, computeUriHash } from './jupyterUtils';
 import { IConfigurationService, IsWebExtension } from '../../platform/common/types';
+import {
+    handleExpiredCertsError,
+    handleSelfCertsError,
+    computeServerId,
+    generateUriFromRemoteProvider
+} from './jupyterUtils';
 import { JupyterConnection } from './jupyterConnection';
 import { JupyterSelfCertsError } from '../../platform/errors/jupyterSelfCertsError';
 import { RemoteJupyterServerConnectionError } from '../../platform/errors/remoteJupyterServerConnectionError';
@@ -79,6 +84,32 @@ export class JupyterServerSelector {
         const multiStep = this.multiStepFactory.create<{}>();
         return multiStep.run(this.startSelectingURI.bind(this, allowLocal), {});
     }
+
+    @captureTelemetry(Telemetry.EnterJupyterURI)
+    @traceDecoratorError('Failed to enter Jupyter Uri')
+    public async enterJupyterURI(): Promise<string | undefined> {
+        let initialValue = defaultUri;
+        try {
+            const text = await this.clipboard.readText().catch(() => '');
+            const parsedUri = Uri.parse(text.trim(), true);
+            // Only display http/https uris.
+            initialValue = text && parsedUri && parsedUri.scheme.toLowerCase().startsWith('http') ? text : defaultUri;
+        } catch {
+            // We can ignore errors.
+        }
+        // Ask the user to enter a URI to connect to.
+        const uri = await this.applicationShell.showInputBox({
+            title: DataScience.jupyterSelectURIPrompt(),
+            value: initialValue || defaultUri,
+            validateInput: this.validateSelectJupyterURI,
+            prompt: ''
+        });
+
+        if (uri) {
+            await this.setJupyterURIToRemote(uri, true);
+            return computeServerId(uri);
+        }
+    }
     @captureTelemetry(Telemetry.SetJupyterURIToLocal)
     public async setJupyterURIToLocal(): Promise<void> {
         await this.serverUriStorage.setUri(Settings.JupyterServerLocalLaunch);
@@ -108,7 +139,7 @@ export class JupyterServerSelector {
                     sendTelemetryEvent(Telemetry.FetchError, undefined, { currentTask: 'connecting' });
                 }
                 await this.errorHandler.handleError(
-                    new RemoteJupyterServerConnectionError(userURI, computeUriHash(userURI), err)
+                    new RemoteJupyterServerConnectionError(userURI, computeServerId(userURI), err)
                 );
                 // Can't set the URI in this case.
                 return;
@@ -187,6 +218,9 @@ export class JupyterServerSelector {
         _input: IMultiStepInput<{}>,
         _state: {}
     ): Promise<InputStep<{}> | void> {
+        if (!provider.handleQuickPick) {
+            return;
+        }
         const result = await provider.handleQuickPick(item, true);
         if (result === 'back') {
             throw InputFlowAction.back;
@@ -197,18 +231,10 @@ export class JupyterServerSelector {
     }
     private async handleProviderQuickPick(id: string, result: JupyterServerUriHandle | undefined) {
         if (result) {
-            const uri = this.generateUriFromRemoteProvider(id, result);
+            const uri = generateUriFromRemoteProvider(id, result);
             await this.setJupyterURIToRemote(uri);
         }
     }
-
-    private generateUriFromRemoteProvider(id: string, result: JupyterServerUriHandle) {
-        // eslint-disable-next-line
-        return `${Identifiers.REMOTE_URI}?${Identifiers.REMOTE_URI_ID_PARAM}=${id}&${
-            Identifiers.REMOTE_URI_HANDLE_PARAM
-        }=${encodeURI(result)}`;
-    }
-
     private async selectRemoteURI(input: IMultiStepInput<{}>, _state: {}): Promise<InputStep<{}> | void> {
         let initialValue = defaultUri;
         try {
@@ -275,6 +301,9 @@ export class JupyterServerSelector {
         const providers = await this.extraUriProviders.getProviders();
         if (providers) {
             providers.forEach((p) => {
+                if (!p.getQuickPickEntryItems) {
+                    return;
+                }
                 const newProviderItems = p.getQuickPickEntryItems().map((i) => {
                     return { ...i, newChoice: false, provider: p };
                 });