Adopt new Jupyter API internally (#14084)

* Move display names into Jupyter Server Uri storage

* fix formatting

* Avoid validation of Jupyter Server providers

* Stop using display name from Uri storage

* WIP

* Misc

* Adopt the new Jupyter Provider API internally

* Fix test

* fix formatting

* comments

src/api.internal.d.ts

@@ -1,9 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-// The following is required to make sure the types are merged correctly.
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
-import { CancellationToken } from 'vscode';
+import { Uri } from 'vscode';
 
 // These types are only used internally within the extension.
 // Never to be exposed to other extensions.
@@ -27,6 +25,20 @@ declare module './api' {
  }
 
  export interface IJupyterUriProvider {
+ /**
+ * Link to documentation for this provider.
+ * Used internally to display the link for `Existing Jupyter Servers` in the quick pick.
+ */
+ documentation?: Uri;
+ /**
+ * Ability to retrieve the displayName without having to get the auth information.
+ * Only used internally when we need the displayName.
+ * The getServerUri could end up prompting for username/password when connecting to the remote servers.
+ */
  getServerUriWithoutAuthInfo?(handle: string): Promise<IJupyterServerUri>;
+ /**
+ * Internally used by Jupyter extension to track the extension that owns this provider.
+ */
+ readonly extensionId: string;
  }
 }

src/commands.ts

@@ -44,6 +44,7 @@ export interface ICommandNameArgumentTypeMapping {
  ['jupyter.selectJupyterInterpreter']: [];
  ['dataScience.ClearUserProviderJupyterServerCache']: [];
  ['jupyterViewVariables.focus']: [];
+ ['jupyter.selectLocalJupyterServer']: [];
  ['workbench.action.openSettings']: ['jupyter.kernels.excludePythonEnvironments'];
  [DSCommands.RunCurrentCell]: [];
  [DSCommands.RunCurrentCellAdvance]: [];

src/kernels/jupyter/connection/jupyterServerProviderRegistry.ts

@@ -13,7 +13,7 @@ import {
 import { Disposables } from '../../../platform/common/utils';
 import { IJupyterServerProviderRegistry, IJupyterUriProviderRegistration } from '../types';
 import { IDisposable, IDisposableRegistry } from '../../../platform/common/types';
-import { inject } from 'inversify';
+import { inject, injectable } from 'inversify';
 import { disposeAllDisposables } from '../../../platform/common/helpers';
 import { traceError } from '../../../platform/logging';
 import { JVSC_EXTENSION_ID } from '../../../platform/common/constants';
@@ -53,13 +53,19 @@ class JupyterUriProviderAdaptor extends Disposables implements IJupyterUriProvid
  public get displayName() {
  return this.provider.label;
  }
+ public get documentation() {
+ return this.provider.documentation;
+ }
  detail?: string | undefined;
  private _onDidChangeHandles = new EventEmitter<void>();
  onDidChangeHandles = this._onDidChangeHandles.event;
  private providerChanges: IDisposable[] = [];
  removeHandle?(handle: string): Promise<void>;
  getServerUriWithoutAuthInfo?(handle: string): Promise<IJupyterServerUri>;
- constructor(private readonly provider: JupyterServerCollection) {
+ constructor(
+ private readonly provider: JupyterServerCollection,
+ public readonly extensionId: string
+ ) {
  super();
  this.id = provider.id;
  this.hookupProviders();
@@ -221,6 +227,8 @@ class JupyterUriProviderAdaptor extends Disposables implements IJupyterUriProvid
  }
  }
 }
+
+@injectable()
 export class JupyterServerProviderRegistry extends Disposables implements IJupyterServerProviderRegistry {
  private readonly _onDidChangeProviders = new EventEmitter<void>();
  public get onDidChangeProviders() {
@@ -241,19 +249,33 @@ export class JupyterServerProviderRegistry extends Disposables implements IJupyt
  createJupyterServerCollection(extensionId: string, id: string, label: string): JupyterServerCollection {
  const extId = `${extensionId}#${id}`;
  if (this._serverProviders.has(extId)) {
- throw new Error(`Jupyter Server Provider with id ${extId} already exists`);
+ // When testing we might have a duplicate as we call the registration API in ctor of a test.
+ if (extensionId !== JVSC_EXTENSION_ID) {
+ throw new Error(`Jupyter Server Provider with id ${extId} already exists`);
+ }
  }
  const serverProvider = new JupyterServerCollectionImpl(extensionId, id, label);
  this._serverProviders.set(extId, serverProvider);
- const uriRegistration = this.jupyterUriProviderRegistration.registerProvider(
- new JupyterUriProviderAdaptor(serverProvider),
- extensionId
+ let uriRegistration: IDisposable | undefined;
+ serverProvider.onDidChangeProvider(
+ () => {
+ if (serverProvider.serverProvider) {
+ uriRegistration?.dispose();
+ uriRegistration = this.jupyterUriProviderRegistration.registerProvider(
+ new JupyterUriProviderAdaptor(serverProvider, extensionId),
+ extensionId
+ );
+ this.disposables.push(uriRegistration);
+ this._onDidChangeProviders.fire();
+ }
+ },
+ this,
+ this.disposables
  );
 
- this._onDidChangeProviders.fire();
  serverProvider.onDidDispose(
  () => {
- uriRegistration.dispose();
+ uriRegistration?.dispose();
  this._serverProviders.delete(extId);
  this._onDidChangeProviders.fire();
  },

src/kernels/jupyter/connection/jupyterUriProviderRegistration.ts

@@ -184,6 +184,9 @@ class JupyterUriProviderWrapper extends Disposables implements IInternalJupyterU
  public get detail() {
  return this.provider.detail;
  }
+ public get documentation() {
+ return this.provider.documentation;
+ }
 
  public get onDidChangeHandles() {
  return this.provider.onDidChangeHandles;

src/kernels/jupyter/connection/serverUriStorage.ts

@@ -273,15 +273,13 @@ class NewStorage {
  this._onDidAddUri.fire(item);
  }
  if (removedItems.length) {
- const removeJupyterUris = await Promise.all(
- removedItems.map(async (removedItem) => {
- return <IJupyterServerUriEntry>{
- provider: removedItem.serverHandle,
- time: removedItem.time,
- displayName: removedItem.displayName || ''
- };
- })
- );
+ const removeJupyterUris = removedItems.map((removedItem) => {
+ return <IJupyterServerUriEntry>{
+ provider: removedItem.serverHandle,
+ time: removedItem.time,
+ displayName: removedItem.displayName || ''
+ };
+ });
  this._onDidRemoveUris.fire(removeJupyterUris);
  }
  this._onDidChangeUri.fire();
@@ -348,17 +346,15 @@ class NewStorage {
  const data = await this.getAllRaw();
  const entries: IJupyterServerUriEntry[] = [];
 
- await Promise.all(
- data.map(async (item) => {
- const uri = generateIdFromRemoteProvider(item.serverHandle);
- const server: IJupyterServerUriEntry = {
- time: item.time,
- displayName: item.displayName || uri,
- provider: item.serverHandle
- };
- entries.push(server);
- })
- );
+ data.forEach(async (item) => {
+ const uri = generateIdFromRemoteProvider(item.serverHandle);
+ const server: IJupyterServerUriEntry = {
+ time: item.time,
+ displayName: item.displayName || uri,
+ provider: item.serverHandle
+ };
+ entries.push(server);
+ });
  return entries;
  }
  private async getAllRaw(): Promise<StorageMRUItem[]> {

src/kernels/jupyter/finder/remoteKernelFinderController.ts

@@ -22,6 +22,7 @@ import { IFileSystem } from '../../../platform/common/platform/types';
 import { ContributedKernelFinderKind } from '../../internalTypes';
 import { generateIdFromRemoteProvider } from '../jupyterUtils';
 import { swallowExceptions } from '../../../platform/common/utils/decorators';
+import { IJupyterUriProvider } from '../../../api';
 
 @injectable()
 export class RemoteKernelFinderController implements IExtensionSyncActivationService {
@@ -43,26 +44,46 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
  @inject(IJupyterUriProviderRegistration)
  private readonly jupyterPickerRegistration: IJupyterUriProviderRegistration
  ) {}
-
+ private readonly handledProviders = new WeakSet<IJupyterUriProvider>();
  activate() {
- // Check for when more URIs are added
  this.serverUriStorage.onDidAdd((server) => this.validateAndCreateFinder(server), this, this.disposables);
+ this.serverUriStorage.onDidChange(this.buildListOfFinders, this, this.disposables);
+ // Possible some extensions register their providers later.
+ // And also possible they load their old server later, hence we need to go through the
+ // MRU list again and try to build the finders, as the servers might now exist.
+ this.jupyterPickerRegistration.onDidChangeProviders(this.handleProviderHandleChanges, this, this.disposables);
 
  // Also check for when a URI is removed
  this.serverUriStorage.onDidRemove(this.urisRemoved, this, this.disposables);
 
+ // Add in the URIs that we already know about
+ this.buildListOfFinders();
+ }
+ private buildListOfFinders() {
  // Add in the URIs that we already know about
  this.serverUriStorage
  .getAll()
- .then(async (currentServers) => {
- await Promise.all(currentServers.map((server) => this.validateAndCreateFinder(server)));
- })
+ .then((currentServers) => currentServers.map((server) => this.validateAndCreateFinder(server).catch(noop)))
  .catch(noop);
  }
+ private handleProviderHandleChanges() {
+ this.jupyterPickerRegistration.providers.forEach((provider) => {
+ if (!this.handledProviders.has(provider)) {
+ this.handledProviders.add(provider);
+ if (provider.onDidChangeHandles) {
+ provider.onDidChangeHandles(this.buildListOfFinders, this, this.disposables);
+ }
+ }
+ });
+ this.buildListOfFinders();
+ }
  @swallowExceptions('Failed to create a Remote Kernel Finder')
  private async validateAndCreateFinder(serverUri: IJupyterServerUriEntry) {
- const info = await this.jupyterPickerRegistration.getJupyterServerUri(serverUri.provider, true);
- this.createRemoteKernelFinder(serverUri.provider, info.displayName);
+ const serverId = generateIdFromRemoteProvider(serverUri.provider);
+ if (!this.serverFinderMapping.has(serverId)) {
+ const info = await this.jupyterPickerRegistration.getJupyterServerUri(serverUri.provider, true);
+ this.createRemoteKernelFinder(serverUri.provider, info.displayName);
+ }
  }
 
  createRemoteKernelFinder(serverProviderHandle: JupyterServerProviderHandle, displayName: string) {

src/kernels/jupyter/serviceRegistry.node.ts

@@ -48,13 +48,15 @@ import {
  IJupyterRequestAgentCreator,
  ILiveRemoteKernelConnectionUsageTracker,
  IJupyterRemoteCachedKernelValidator,
- IJupyterServerHelper
+ IJupyterServerHelper,
+ IJupyterServerProviderRegistry
 } from './types';
 import { IJupyterCommandFactory, IJupyterSubCommandExecutionService } from './types.node';
 import { RemoteKernelFinderController } from './finder/remoteKernelFinderController';
 import { KernelSessionFactory } from '../common/kernelSessionFactory';
 import { OldJupyterKernelSessionFactory } from './session/oldJupyterKernelSessionFactory';
 import { JupyterKernelSessionFactory } from './session/jupyterKernelSessionFactory';
+import { JupyterServerProviderRegistry } from './connection/jupyterServerProviderRegistry';
 
 export function registerTypes(serviceManager: IServiceManager, _isDevMode: boolean) {
  serviceManager.add<IJupyterCommandFactory>(IJupyterCommandFactory, JupyterCommandFactory);
@@ -134,4 +136,8 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole
  IExtensionSyncActivationService,
  RemoteKernelFinderController
  );
+ serviceManager.addSingleton<IJupyterServerProviderRegistry>(
+ IJupyterServerProviderRegistry,
+ JupyterServerProviderRegistry
+ );
 }

src/kernels/jupyter/serviceRegistry.web.ts

@@ -26,12 +26,14 @@ import {
  IJupyterServerProvider,
  IJupyterRequestCreator,
  ILiveRemoteKernelConnectionUsageTracker,
- IJupyterRemoteCachedKernelValidator
+ IJupyterRemoteCachedKernelValidator,
+ IJupyterServerProviderRegistry
 } from './types';
 import { RemoteKernelFinderController } from './finder/remoteKernelFinderController';
 import { KernelSessionFactory } from '../common/kernelSessionFactory';
 import { OldJupyterKernelSessionFactory } from './session/oldJupyterKernelSessionFactory';
 import { JupyterKernelSessionFactory } from './session/jupyterKernelSessionFactory';
+import { JupyterServerProviderRegistry } from './connection/jupyterServerProviderRegistry';
 
 export function registerTypes(serviceManager: IServiceManager, _isDevMode: boolean) {
  serviceManager.addSingleton<IOldJupyterSessionManagerFactory>(
@@ -70,4 +72,8 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole
  IExtensionSyncActivationService,
  RemoteKernelFinderController
  );
+ serviceManager.addSingleton<IJupyterServerProviderRegistry>(
+ IJupyterServerProviderRegistry,
+ JupyterServerProviderRegistry
+ );
 }

src/notebooks/controllers/kernelSource/kernelSourceCommandHandler.ts

@@ -180,12 +180,19 @@ export class KernelSourceCommandHandler implements IExtensionSyncActivationServi
  provideNotebookKernelSourceActions: () => {
  return [
  {
- label:
- provider.displayName ??
- (provider.detail ? `${provider.detail} (${provider.id})` : provider.id),
- documentation: provider.id.startsWith('_builtin')
- ? Uri.parse('https://aka.ms/vscodeJuptyerExtKernelPickerExistingServer')
- : undefined,
+ get label() {
+ return (
+ provider.displayName ??
+ (provider.detail
+ ? `${provider.detail} (${provider.id}) ${Date.now()}`
+ : provider.id)
+ );
+ },
+ get documentation() {
+ return provider.id.startsWith('_builtin') && provider.documentation
+ ? provider.documentation
+ : undefined;
+ },
  command: {
  command: 'jupyter.kernel.selectJupyterServerKernel',
  arguments: [provider.extensionId, provider.id],
@@ -199,12 +206,19 @@ export class KernelSourceCommandHandler implements IExtensionSyncActivationServi
  provideNotebookKernelSourceActions: () => {
  return [
  {
- label:
- provider.displayName ??
- (provider.detail ? `${provider.detail} (${provider.id})` : provider.id),
- documentation: provider.id.startsWith('_builtin')
- ? Uri.parse('https://aka.ms/vscodeJuptyerExtKernelPickerExistingServer')
- : undefined,
+ get label() {
+ return (
+ provider.displayName ??
+ (provider.detail
+ ? `${provider.detail} (${provider.id}) ${Date.now()}`
+ : provider.id)
+ );
+ },
+ get documentation() {
+ return provider.id.startsWith('_builtin') && provider.documentation
+ ? provider.documentation
+ : undefined;
+ },
  command: {
  command: 'jupyter.kernel.selectJupyterServerKernel',
  arguments: [provider.extensionId, provider.id],

src/standalone/userJupyterServer/userServerUriProvider.unit.test.ts

@@ -5,6 +5,7 @@ import * as sinon from 'sinon';
 import { anything, instance, mock, verify, when } from 'ts-mockito';
 import {
  IJupyterRequestCreator,
+ IJupyterServerProviderRegistry,
  IJupyterServerUriStorage,
  IJupyterUriProviderRegistration
 } from '../../kernels/jupyter/types';
@@ -28,7 +29,8 @@ import {
  IClipboard,
  IApplicationShell,
  IEncryptedStorage,
- ICommandManager
+ ICommandManager,
+ IApplicationEnvironment
 } from '../../platform/common/application/types';
 import { noop, sleep } from '../../test/core';
 import { disposeAllDisposables } from '../../platform/common/helpers';
@@ -38,6 +40,7 @@ import { generateIdFromRemoteProvider } from '../../kernels/jupyter/jupyterUtils
 import { Common, DataScience } from '../../platform/common/utils/localize';
 import { IJupyterPasswordConnectInfo, JupyterPasswordConnect } from './jupyterPasswordConnect';
 import { IFileSystem } from '../../platform/common/platform/types';
+import { JupyterServerCollection } from '../../api';
 
 /* eslint-disable @typescript-eslint/no-explicit-any, , */
 suite('User Uri Provider', () => {
@@ -165,6 +168,14 @@ suite('User Uri Provider', () => {
 
  when(serverUriStorage.add(anything())).thenResolve();
  when(serverUriStorage.add(anything(), anything())).thenResolve();
+ const jupyterServerProviderRegistry = mock<IJupyterServerProviderRegistry>();
+ const collection = mock<JupyterServerCollection>();
+ when(collection.dispose()).thenReturn();
+ when(
+ jupyterServerProviderRegistry.createJupyterServerCollection(anything(), anything(), anything())
+ ).thenReturn(instance(collection));
+ const appEnv = mock<IApplicationEnvironment>();
+ when(appEnv.channel).thenReturn('stable');
  provider = new UserJupyterServerUrlProvider(
  instance(clipboard),
  instance(uriProviderRegistration),
@@ -182,7 +193,9 @@ suite('User Uri Provider', () => {
  undefined,
  instance(requestCreator),
  instance(mock<IExtensionContext>()),
- instance(mock<IFileSystem>())
+ instance(mock<IFileSystem>()),
+ instance(jupyterServerProviderRegistry),
+ instance(appEnv)
  );
  });
  teardown(async () => {