Handle delayed Notebook provider registration (#3526)

* Handle delayed Notebook provider registration
- Fixes #3197 Notebooks: builtin provider always used on reopen with notebook file visible
- Fixes #3414 Can't refresh kernel after connect to big data cluster

There are 3 parts to this fix:
- If no notebook provider other than the default is installed, we warn users and prompt to install the SQL2019 extension
- We wait on the extension host registration to complete before determining which provider to use
- We know that the extension registration of the provider instance will be after package.json is read, so if we wait after registration for 10 seconds to give this a chance to happen before returning a provider to the front end

* Remove launch.json change that was added accidentally

* Fix timeout not being the expected value

* Removed console log left in during debugging

* Remove unnecessary whitespace

* Fix unit test failure

* Name the registration better, and remove outdated comments

src/sql/parts/common/customInputConverter.ts

@@ -5,9 +5,8 @@
 
 import * as path from 'path';
 
-import { Registry } from 'vs/platform/registry/common/platform';
 import { EditorInput, IEditorInput } from 'vs/workbench/common/editor';
-import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { IInstantiationService, ServiceIdentifier } from 'vs/platform/instantiation/common/instantiation';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
 import { FileEditorInput } from 'vs/workbench/parts/files/common/editors/fileEditorInput';
 import URI from 'vs/base/common/uri';
@@ -17,8 +16,7 @@ import { QueryInput } from 'sql/parts/query/common/queryInput';
 import { IQueryEditorOptions } from 'sql/parts/query/common/queryEditorService';
 import { QueryPlanInput } from 'sql/parts/queryPlan/queryPlanInput';
 import { NotebookInput, NotebookInputModel, NotebookInputValidator } from 'sql/parts/notebook/notebookInput';
-import { Extensions, INotebookProviderRegistry } from 'sql/services/notebook/notebookRegistry';
+import { DEFAULT_NOTEBOOK_PROVIDER, INotebookService } from 'sql/services/notebook/notebookService';
-import { DEFAULT_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
 import { getProviderForFileName } from 'sql/parts/notebook/notebookUtils';
 
 const fs = require('fs');
@@ -59,20 +57,20 @@ export function convertEditorInput(input: EditorInput, options: IQueryEditorOpti
 
 		//Notebook
 		let notebookValidator = instantiationService.createInstance(NotebookInputValidator);
-		uri = getNotebookEditorUri(input);
+		uri = getNotebookEditorUri(input, instantiationService);
 		if(uri && notebookValidator.isNotebookEnabled()){
-			//TODO: We need to pass in notebook data either through notebook input or notebook service
+			return withService<INotebookService, NotebookInput>(instantiationService, INotebookService, notebookService => {
 				let fileName: string = 'untitled';
 				let providerId: string = DEFAULT_NOTEBOOK_PROVIDER;
 				if (input) {
 					fileName = input.getName();
-				providerId = getProviderForFileName(fileName);
+					providerId = getProviderForFileName(fileName, notebookService);
 				}
 				let notebookInputModel = new NotebookInputModel(uri, undefined, false, undefined);
 				notebookInputModel.providerId = providerId;
-			//TO DO: Second parameter has to be the content.
 				let notebookInput: NotebookInput = instantiationService.createInstance(NotebookInput, fileName, notebookInputModel);
 				return notebookInput;
+			});
 		}
 	}
 	return input;
@@ -159,7 +157,7 @@ function getQueryPlanEditorUri(input: EditorInput): URI {
  * If input is a supported notebook editor file (.ipynb), return it's URI. Otherwise return undefined.
  * @param input The EditorInput to get the URI of.
  */
-function getNotebookEditorUri(input: EditorInput): URI {
+function getNotebookEditorUri(input: EditorInput, instantiationService: IInstantiationService): URI {
 	if (!input || !input.getName()) {
 		return undefined;
 	}
@@ -170,7 +168,7 @@ function getNotebookEditorUri(input: EditorInput): URI {
 	if (!(input instanceof NotebookInput)) {
 		let uri: URI = getSupportedInputResource(input);
 		if (uri) {
-			if (hasFileExtension(getNotebookFileExtensions(), input, false)) {
+			if (hasFileExtension(getNotebookFileExtensions(instantiationService), input, false)) {
 				return uri;
 			}
 		}
@@ -179,9 +177,17 @@ function getNotebookEditorUri(input: EditorInput): URI {
 	return undefined;
 }
 
-function getNotebookFileExtensions() {
+function getNotebookFileExtensions(instantiationService: IInstantiationService): string[] {
-	let notebookRegistry = Registry.as<INotebookProviderRegistry>(Extensions.NotebookProviderContribution);
+	return withService<INotebookService, string[]>(instantiationService, INotebookService, notebookService => {
-	return notebookRegistry.getSupportedFileExtensions();
+		return notebookService.getSupportedFileExtensions();
+	});
+}
+
+function withService<TService, TResult>(instantiationService: IInstantiationService, serviceId: ServiceIdentifier<TService>, action: (service: TService) => TResult, ): TResult {
+	return instantiationService.invokeFunction(accessor => {
+		let service = accessor.get(serviceId);
+		return action(service);
+	});
 }
 
 /**

src/sql/parts/notebook/notebook.component.ts

@@ -9,41 +9,42 @@ import { OnInit, Component, Inject, forwardRef, ElementRef, ChangeDetectorRef, V
 
 import { IColorTheme, IWorkbenchThemeService } from 'vs/workbench/services/themes/common/workbenchThemeService';
 import * as themeColors from 'vs/workbench/common/theme';
-import { INotificationService, INotification } from 'vs/platform/notification/common/notification';
+import { INotificationService, INotification, Severity } from 'vs/platform/notification/common/notification';
 import { localize } from 'vs/nls';
-
-import { CommonServiceInterface } from 'sql/services/common/commonServiceInterface.service';
-import { AngularDisposable } from 'sql/base/common/lifecycle';
-
-import { CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
-import { ICellModel, IModelFactory, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
-import { IConnectionManagementService, IConnectionDialogService } from 'sql/parts/connection/common/connectionManagement';
-import { INotebookService, INotebookParams, INotebookManager, INotebookEditor, DEFAULT_NOTEBOOK_FILETYPE } from 'sql/services/notebook/notebookService';
-import { IBootstrapParams } from 'sql/services/bootstrap/bootstrapService';
-import { NotebookModel, NotebookContentChange } from 'sql/parts/notebook/models/notebookModel';
-import { ModelFactory } from 'sql/parts/notebook/models/modelFactory';
-import * as notebookUtils from './notebookUtils';
-import { Deferred } from 'sql/base/common/promise';
-import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
-import { Taskbar } from 'sql/base/browser/ui/taskbar/taskbar';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IContextMenuService, IContextViewService } from 'vs/platform/contextview/browser/contextView';
-import { KernelsDropdown, AttachToDropdown, AddCellAction, TrustedAction, SaveNotebookAction } from 'sql/parts/notebook/notebookActions';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { attachSelectBoxStyler } from 'vs/platform/theme/common/styler';
 import { MenuId, IMenuService, MenuItemAction } from 'vs/platform/actions/common/actions';
 import { IAction, Action, IActionItem } from 'vs/base/common/actions';
 import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
 import { fillInActions, LabeledMenuItemActionItem } from 'vs/platform/actions/browser/menuItemActionItem';
-import { IObjectExplorerService } from 'sql/parts/objectExplorer/common/objectExplorerService';
-import * as TaskUtilities from 'sql/workbench/common/taskUtilities';
-import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { Schemas } from 'vs/base/common/network';
 import URI from 'vs/base/common/uri';
 import { IHistoryService } from 'vs/workbench/services/history/common/history';
 import * as paths from 'vs/base/common/paths';
 import { IWindowService } from 'vs/platform/windows/common/windows';
 import { TPromise } from 'vs/base/common/winjs.base';
+import { IViewletService } from 'vs/workbench/services/viewlet/browser/viewlet';
+import { VIEWLET_ID, IExtensionsViewlet } from 'vs/workbench/parts/extensions/common/extensions';
+
+import { CommonServiceInterface } from 'sql/services/common/commonServiceInterface.service';
+import { AngularDisposable } from 'sql/base/common/lifecycle';
+import { CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
+import { ICellModel, IModelFactory, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
+import { IConnectionManagementService, IConnectionDialogService } from 'sql/parts/connection/common/connectionManagement';
+import { INotebookService, INotebookParams, INotebookManager, INotebookEditor, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
+import { IBootstrapParams } from 'sql/services/bootstrap/bootstrapService';
+import { NotebookModel, NotebookContentChange } from 'sql/parts/notebook/models/notebookModel';
+import { ModelFactory } from 'sql/parts/notebook/models/modelFactory';
+import * as notebookUtils from 'sql/parts/notebook/notebookUtils';
+import { Deferred } from 'sql/base/common/promise';
+import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
+import { Taskbar } from 'sql/base/browser/ui/taskbar/taskbar';
+import { KernelsDropdown, AttachToDropdown, AddCellAction, TrustedAction, SaveNotebookAction } from 'sql/parts/notebook/notebookActions';
+import { IObjectExplorerService } from 'sql/parts/objectExplorer/common/objectExplorerService';
+import * as TaskUtilities from 'sql/workbench/common/taskUtilities';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export const NOTEBOOK_SELECTOR: string = 'notebook-component';
@@ -88,6 +89,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		@Inject(IKeybindingService) private keybindingService: IKeybindingService,
 		@Inject(IHistoryService) private historyService: IHistoryService,
 		@Inject(IWindowService) private windowService: IWindowService,
+		@Inject(IViewletService) private viewletService: IViewletService
 	) {
 		super();
 		this.updateProfile();
@@ -228,6 +230,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	}
 
 	private async loadModel(): Promise<void> {
+		await this.awaitNonDefaultProvider();
 		this.notebookManager = await this.notebookService.getOrCreateNotebookManager(this._notebookParams.providerId, this._notebookParams.notebookUri);
 		let model = new NotebookModel({
 			factory: this.modelFactory,
@@ -247,6 +250,34 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._changeRef.detectChanges();
 	}
 
+	private async awaitNonDefaultProvider(): Promise<void> {
+		// Wait on registration for now. Long-term would be good to cache and refresh
+		await this.notebookService.registrationComplete;
+		// Refresh the provider if we had been using default
+		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
+			this._notebookParams.providerId = notebookUtils.getProviderForFileName(this._notebookParams.notebookUri.fsPath, this.notebookService);
+		}
+		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
+			// If it's still the default, warn them they should install an extension
+			this.notificationService.prompt(Severity.Warning,
+				localize('noKernelInstalled', 'Please install the SQL Server 2019 extension to run cells'),
+				[{
+					label: localize('installSql2019Extension', 'Install Extension'),
+					run: () => this.openExtensionGallery()
+				}]);
+		}
+	}
+
+	private async openExtensionGallery(): Promise<void> {
+		try {
+			let viewlet = await this.viewletService.openViewlet(VIEWLET_ID, true) as IExtensionsViewlet;
+			viewlet.search('sql-vnext');
+			viewlet.focus();
+		} catch (error) {
+			this.notificationService.error(error.message);
+		}
+	}
+
 	// Updates toolbar components
 	private updateToolbarComponents(isTrusted: boolean)
 	{

src/sql/parts/notebook/notebookUtils.ts

@@ -11,9 +11,7 @@ import * as os from 'os';
 import * as pfs from 'vs/base/node/pfs';
 import { localize } from 'vs/nls';
 import { IOutputChannel } from 'vs/workbench/parts/output/common/output';
-import { Registry } from 'vs/platform/registry/common/platform';
+import { DEFAULT_NOTEBOOK_PROVIDER, DEFAULT_NOTEBOOK_FILETYPE, INotebookService } from 'sql/services/notebook/notebookService';
-import { INotebookProviderRegistry, Extensions } from 'sql/services/notebook/notebookRegistry';
-import { DEFAULT_NOTEBOOK_PROVIDER, DEFAULT_NOTEBOOK_FILETYPE } from 'sql/services/notebook/notebookService';
 
 
 /**
@@ -41,18 +39,17 @@ export async function mkDir(dirPath: string, outputChannel?: IOutputChannel): Pr
 	}
 }
 
-export function getProviderForFileName(fileName: string): string {
+export function getProviderForFileName(fileName: string, notebookService: INotebookService): string {
 	let fileExt = path.extname(fileName);
 	let provider: string;
-	let notebookRegistry = Registry.as<INotebookProviderRegistry>(Extensions.NotebookProviderContribution);
 	// First try to get provider for actual file type
 	if (fileExt && fileExt.startsWith('.')) {
 		fileExt = fileExt.slice(1,fileExt.length);
-		provider = notebookRegistry.getProviderForFileType(fileExt);
+		provider = notebookService.getProviderForFileType(fileExt);
 	}
 	// Fallback to provider for default file type (assume this is a global handler)
 	if (!provider) {
-		provider = notebookRegistry.getProviderForFileType(DEFAULT_NOTEBOOK_FILETYPE);
+		provider = notebookService.getProviderForFileType(DEFAULT_NOTEBOOK_FILETYPE);
 	}
 	// Finally if all else fails, use the built-in handler
 	if (!provider) {

src/sql/services/notebook/notebookRegistry.ts

@@ -9,12 +9,13 @@ import { IJSONSchema } from 'vs/base/common/jsonSchema';
 import { ExtensionsRegistry, IExtensionPointUser } from 'vs/workbench/services/extensions/common/extensionsRegistry';
 import { localize } from 'vs/nls';
 import * as platform from 'vs/platform/registry/common/platform';
+import { Event, Emitter } from 'vs/base/common/event';
 
 export const Extensions = {
 	NotebookProviderContribution: 'notebook.providers'
 };
 
-export interface NotebookProviderDescription {
+export interface NotebookProviderRegistration {
 	provider: string;
 	fileExtensions: string | string[];
 }
@@ -54,42 +55,28 @@ let notebookContrib: IJSONSchema = {
 };
 
 export interface INotebookProviderRegistry {
-	registerNotebookProvider(provider: NotebookProviderDescription): void;
+	readonly registrations: NotebookProviderRegistration[];
-	getSupportedFileExtensions(): string[];
+	readonly onNewRegistration: Event<{ id: string, registration: NotebookProviderRegistration }>;
-	getProviderForFileType(fileType: string): string;
+
+	registerNotebookProvider(registration: NotebookProviderRegistration): void;
 }
 
 class NotebookProviderRegistry implements INotebookProviderRegistry {
-	private providerIdToProviders = new Map<string, NotebookProviderDescription>();
+	private providerIdToRegistration = new Map<string, NotebookProviderRegistration>();
-	private fileToProviders = new Map<string, NotebookProviderDescription>();
+	private _onNewRegistration = new Emitter<{ id: string, registration: NotebookProviderRegistration }>();
+	public readonly onNewRegistration: Event<{ id: string, registration: NotebookProviderRegistration }> = this._onNewRegistration.event;
 
-	registerNotebookProvider(provider: NotebookProviderDescription): void {
+	registerNotebookProvider(registration: NotebookProviderRegistration): void {
 		// Note: this method intentionally overrides default provider for a file type.
 		// This means that any built-in provider will be overridden by registered extensions
-		this.providerIdToProviders.set(provider.provider, provider);
+		this.providerIdToRegistration.set(registration.provider, registration);
-		if (provider.fileExtensions) {
+		this._onNewRegistration.fire( { id: registration.provider, registration: registration });
-			if (Array.isArray<string>(provider.fileExtensions)) {
-				for (let fileType of provider.fileExtensions) {
-					this.addFileProvider(fileType, provider);
-				}
-			} else {
-				this.addFileProvider(provider.fileExtensions, provider);
-			}
-		}
 	}
 
-	private addFileProvider(fileType: string, provider: NotebookProviderDescription) {
+	public get registrations(): NotebookProviderRegistration[] {
-		this.fileToProviders.set(fileType.toUpperCase(), provider);
+		let registrationArray: NotebookProviderRegistration[] = [];
-	}
+		this.providerIdToRegistration.forEach(p => registrationArray.push(p));
-
+		return registrationArray;
-	getSupportedFileExtensions(): string[] {
-		return Array.from(this.fileToProviders.keys());
-	}
-
-	getProviderForFileType(fileType: string): string {
-		fileType = fileType.toUpperCase();
-		let provider = this.fileToProviders.get(fileType);
-		return provider ? provider.provider : undefined;
 	}
 }
 
@@ -97,15 +84,15 @@ const notebookProviderRegistry = new NotebookProviderRegistry();
 platform.Registry.add(Extensions.NotebookProviderContribution, notebookProviderRegistry);
 
 
-ExtensionsRegistry.registerExtensionPoint<NotebookProviderDescription | NotebookProviderDescription[]>(Extensions.NotebookProviderContribution, [], notebookContrib).setHandler(extensions => {
+ExtensionsRegistry.registerExtensionPoint<NotebookProviderRegistration | NotebookProviderRegistration[]>(Extensions.NotebookProviderContribution, [], notebookContrib).setHandler(extensions => {
 
-	function handleExtension(contrib: NotebookProviderDescription, extension: IExtensionPointUser<any>) {
+	function handleExtension(contrib: NotebookProviderRegistration, extension: IExtensionPointUser<any>) {
 		notebookProviderRegistry.registerNotebookProvider(contrib);
 	}
 
 	for (let extension of extensions) {
 		const { value } = extension;
-		if (Array.isArray<NotebookProviderDescription>(value)) {
+		if (Array.isArray<NotebookProviderRegistration>(value)) {
 			for (let command of value) {
 				handleExtension(command, extension);
 			}

src/sql/services/notebook/notebookService.ts

@@ -26,10 +26,12 @@ export const DEFAULT_NOTEBOOK_FILETYPE = 'IPYNB';
 export interface INotebookService {
 	_serviceBrand: any;
 
-	onNotebookEditorAdd: Event<INotebookEditor>;
+	readonly onNotebookEditorAdd: Event<INotebookEditor>;
-	onNotebookEditorRemove: Event<INotebookEditor>;
+	readonly onNotebookEditorRemove: Event<INotebookEditor>;
 	onNotebookEditorRename: Event<INotebookEditor>;
 
+	readonly isRegistrationComplete: boolean;
+	readonly registrationComplete: Promise<void>;
 	/**
 	 * Register a metadata provider
 	 */
@@ -40,6 +42,10 @@ export interface INotebookService {
 	 */
 	unregisterProvider(providerId: string): void;
 
+	getSupportedFileExtensions(): string[];
+
+	getProviderForFileType(fileType: string): string;
+
 	/**
 	 * Initializes and returns a Notebook manager that can handle all important calls to open, display, and
 	 * run cells in a notebook.

src/sql/services/notebook/notebookServiceImpl.ts

@@ -18,42 +18,142 @@ import { RenderMimeRegistry } from 'sql/parts/notebook/outputs/registry';
 import { standardRendererFactories } from 'sql/parts/notebook/outputs/factories';
 import { LocalContentManager } from 'sql/services/notebook/localContentManager';
 import { SessionManager } from 'sql/services/notebook/sessionManager';
-import { Extensions, INotebookProviderRegistry } from 'sql/services/notebook/notebookRegistry';
+import { Extensions, INotebookProviderRegistry, NotebookProviderRegistration } from 'sql/services/notebook/notebookRegistry';
 import { Emitter, Event } from 'vs/base/common/event';
+import { Memento } from 'vs/workbench/common/memento';
+import { IStorageService } from 'vs/platform/storage/common/storage';
+import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
+import { IExtensionManagementService, IExtensionIdentifier } from 'vs/platform/extensionManagement/common/extensionManagement';
+import { Disposable } from 'vs/base/common/lifecycle';
+import { getIdFromLocalExtensionId } from 'vs/platform/extensionManagement/common/extensionManagementUtil';
+import { Deferred } from 'sql/base/common/promise';
 
+export interface NotebookProviderProperties {
+	provider: string;
+	fileExtensions: string[];
+}
 
-export class NotebookService implements INotebookService {
+interface NotebookProviderCache {
+	[id: string]: NotebookProviderProperties;
+}
+
+interface NotebookProvidersMemento {
+	notebookProviderCache: NotebookProviderCache;
+}
+
+const notebookRegistry = Registry.as<INotebookProviderRegistry>(Extensions.NotebookProviderContribution);
+
+class ProviderDescriptor {
+	private _instanceReady = new Deferred<INotebookProvider>();
+	constructor(private providerId: string, private _instance?: INotebookProvider) {
+		if (_instance) {
+			this._instanceReady.resolve(_instance);
+		}
+	}
+
+	public get instanceReady(): Promise<INotebookProvider> {
+		return this._instanceReady.promise;
+	}
+
+	public get instance(): INotebookProvider {
+		return this._instance;
+	}
+	public set instance(value: INotebookProvider) {
+		this._instance = value;
+		this._instanceReady.resolve(value);
+	}
+}
+
+export class NotebookService extends Disposable implements INotebookService {
 	_serviceBrand: any;
+
+	private _memento = new Memento('notebookProviders');
 	private _mimeRegistry: RenderMimeRegistry;
-	private _providers: Map<string, INotebookProvider> = new Map();
+	private _providers: Map<string, ProviderDescriptor> = new Map();
 	private _managers: Map<string, INotebookManager> = new Map();
 	private _onNotebookEditorAdd = new Emitter<INotebookEditor>();
 	private _onNotebookEditorRemove = new Emitter<INotebookEditor>();
 	private _onNotebookEditorRename = new Emitter<INotebookEditor>();
 	private _editors = new Map<string, INotebookEditor>();
+	private _fileToProviders = new Map<string, NotebookProviderRegistration>();
+	private _registrationComplete = new Deferred<void>();
+	private _isRegistrationComplete = false;
 
-	constructor() {
+	constructor(
+		@IStorageService private _storageService: IStorageService,
+		@IExtensionService extensionService: IExtensionService,
+		@IExtensionManagementService extensionManagementService: IExtensionManagementService
+	) {
+		super();
+		this._register(notebookRegistry.onNewRegistration(this.updateRegisteredProviders, this));
 		this.registerDefaultProvider();
-	}
 
-	private registerDefaultProvider() {
+		if (extensionService) {
-		let defaultProvider = new BuiltinProvider();
+				extensionService.whenInstalledExtensionsRegistered().then(() => {
-		this.registerProvider(defaultProvider.providerId, defaultProvider);
+				this.cleanupProviders();
-		let registry = Registry.as<INotebookProviderRegistry>(Extensions.NotebookProviderContribution);
+				this._isRegistrationComplete = true;
-		registry.registerNotebookProvider({
+				this._registrationComplete.resolve();
-			provider: defaultProvider.providerId,
-			fileExtensions: DEFAULT_NOTEBOOK_FILETYPE
 			});
 		}
+		if (extensionManagementService) {
+			this._register(extensionManagementService.onDidUninstallExtension(({ identifier }) => this.removeContributedProvidersFromCache(identifier, extensionService)));
+		}
+	}
 
-	registerProvider(providerId: string, provider: INotebookProvider): void {
+	private updateRegisteredProviders(p: { id: string; registration: NotebookProviderRegistration; }) {
-		this._providers.set(providerId, provider);
+		let registration = p.registration;
+
+		if (!this._providers.has(p.id)) {
+			this._providers.set(p.id, new ProviderDescriptor(p.id));
+		}
+		if (registration.fileExtensions) {
+			if (Array.isArray<string>(registration.fileExtensions)) {
+				for (let fileType of registration.fileExtensions) {
+					this.addFileProvider(fileType, registration);
+				}
+			}
+			else {
+				this.addFileProvider(registration.fileExtensions, registration);
+			}
+		}
+	}
+
+	registerProvider(providerId: string, instance: INotebookProvider): void {
+		let providerDescriptor = this._providers.get(providerId);
+		if (providerDescriptor) {
+			// Update, which will resolve the promise for anyone waiting on the instance to be registered
+			providerDescriptor.instance = instance;
+		} else {
+			this._providers.set(providerId, new ProviderDescriptor(providerId, instance));
+		}
 	}
 
 	unregisterProvider(providerId: string): void {
 		this._providers.delete(providerId);
 	}
 
+	get isRegistrationComplete(): boolean {
+		return this._isRegistrationComplete;
+	}
+
+	get registrationComplete(): Promise<void> {
+		return this._registrationComplete.promise;
+	}
+
+	private addFileProvider(fileType: string, provider: NotebookProviderRegistration) {
+		this._fileToProviders.set(fileType.toUpperCase(), provider);
+	}
+
+	getSupportedFileExtensions(): string[] {
+		return Array.from(this._fileToProviders.keys());
+	}
+
+	getProviderForFileType(fileType: string): string {
+		fileType = fileType.toUpperCase();
+		let provider = this._fileToProviders.get(fileType);
+		return provider ? provider.provider : undefined;
+	}
+
 	public shutdown(): void {
 		this._managers.forEach(manager => {
 			if (manager.serverManager) {
@@ -119,32 +219,59 @@ export class NotebookService implements INotebookService {
 		}
 	}
 
-	private sendNotebookCloseToProvider(editor: INotebookEditor) {
+	private sendNotebookCloseToProvider(editor: INotebookEditor): void {
 		let notebookUri = editor.notebookParams.notebookUri;
 		let uriString = notebookUri.toString();
 		let manager = this._managers.get(uriString);
 		if (manager) {
+			// As we have a manager, we can assume provider is ready
 			this._managers.delete(uriString);
 			let provider = this._providers.get(manager.providerId);
-			provider.handleNotebookClosed(notebookUri);
+			provider.instance.handleNotebookClosed(notebookUri);
 		}
 	}
 
 	// PRIVATE HELPERS /////////////////////////////////////////////////////
-	private doWithProvider<T>(providerId: string, op: (provider: INotebookProvider) => Thenable<T>): Thenable<T> {
+	private async doWithProvider<T>(providerId: string, op: (provider: INotebookProvider) => Thenable<T>): Promise<T> {
 		// Make sure the provider exists before attempting to retrieve accounts
-		let provider: INotebookProvider;
+		let provider: INotebookProvider = await this.getProviderInstance(providerId);
-		if (this._providers.has(providerId)) {
+		return op(provider);
-			provider = this._providers.get(providerId);
-		}
-		else {
-			provider = this._providers.get(DEFAULT_NOTEBOOK_PROVIDER);
 	}
 
-		if (!provider) {
+	private async getProviderInstance(providerId: string, timeout?: number): Promise<INotebookProvider> {
-			return Promise.reject(new Error(localize('notebookServiceNoProvider', 'Notebook provider does not exist'))).then();
+		let providerDescriptor = this._providers.get(providerId);
+		let instance: INotebookProvider;
+
+		// Try get from actual provider, waiting on its registration
+		if (providerDescriptor) {
+			if (!providerDescriptor.instance) {
+				instance = await this.waitOnProviderAvailability(providerDescriptor);
+			} else {
+				instance = providerDescriptor.instance;
 			}
-		return op(provider);
+		}
+
+		// Fall back to default if this failed
+		if (!instance) {
+			providerDescriptor = this._providers.get(DEFAULT_NOTEBOOK_PROVIDER);
+			instance = providerDescriptor ? providerDescriptor.instance : undefined;
+		}
+
+		// Should never happen, but if default wasn't registered we should throw
+		if (!instance) {
+			throw new Error(localize('notebookServiceNoProvider', 'Notebook provider does not exist'));
+		}
+		return instance;
+	}
+
+	private waitOnProviderAvailability(providerDescriptor: ProviderDescriptor, timeout?: number): Promise<INotebookProvider> {
+		// Wait up to 10 seconds for the provider to be registered
+		timeout = timeout || 10000;
+		let promises: Promise<INotebookProvider>[] = [
+			providerDescriptor.instanceReady,
+			new Promise<INotebookProvider>((resolve, reject) => setTimeout(() => resolve(), timeout))
+		];
+		return Promise.race(promises);
 	}
 
 	//Returns an instantiation of RenderMimeRegistry class
@@ -156,6 +283,41 @@ export class NotebookService implements INotebookService {
 		}
 		return this._mimeRegistry;
 	}
+
+	private get providersMemento(): NotebookProvidersMemento {
+		return this._memento.getMemento(this._storageService) as NotebookProvidersMemento;
+	}
+
+	private cleanupProviders(): void {
+		let knownProviders = Object.keys(notebookRegistry.registrations);
+		let cache = this.providersMemento.notebookProviderCache;
+		for (let key in cache) {
+			if (!knownProviders.includes(key)) {
+				this._providers.delete(key);
+				delete cache[key];
+			}
+		}
+	}
+
+	private registerDefaultProvider() {
+		let defaultProvider = new BuiltinProvider();
+		this.registerProvider(defaultProvider.providerId, defaultProvider);
+		notebookRegistry.registerNotebookProvider({
+			provider: defaultProvider.providerId,
+			fileExtensions: DEFAULT_NOTEBOOK_FILETYPE
+		});
+	}
+
+	private removeContributedProvidersFromCache(identifier: IExtensionIdentifier, extensionService: IExtensionService) {
+		let extensionid = getIdFromLocalExtensionId(identifier.id);
+		extensionService.getExtensions().then(i => {
+			let extension = i.find(c => c.id === extensionid);
+			if (extension && extension.contributes['notebookProvider']) {
+				let id = extension.contributes['notebookProvider'].providerId;
+				delete this.providersMemento.notebookProviderCache[id];
+			}
+		});
+	}
 }
 
 export class BuiltinProvider implements INotebookProvider {

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -258,7 +258,8 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		extHostContext: IExtHostContext,
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IEditorService private _editorService: IEditorService,
-		@IEditorGroupsService private _editorGroupService: IEditorGroupsService
+		@IEditorGroupsService private _editorGroupService: IEditorGroupsService,
+		@INotebookService private readonly _notebookService: INotebookService
 	) {
 		super();
 		if (extHostContext) {
@@ -306,7 +307,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		if(!providerId)
 		{
 			// Ensure there is always a sensible provider ID for this file type
-			providerId = getProviderForFileName(uri.fsPath);
+			providerId = getProviderForFileName(uri.fsPath, this._notebookService);
 		}
 
 		model.providerId = providerId;