Fix #4047 Redesign notebook model to handle single client session (#4371)

* Start single client session based on the default kernel or saved kernel in NB.

* Added kernel displayName to standardKernel.
Modified name to allign wtih Juptyer Kernel.name.
So we can show the displayName during startup and use the name to start the session.

* Change session.OnSessionReady event in KernelDropDown

* Added model.KernelChnaged for switching kernel in the same provider

* Fixed session.Ready sequence

* Fixed merge issues

* Solve merged issue

* Fixed wrong kernel name in saved NB

* Added new event in Model to notify kernel change.
Toolbar depends on ModelReady to load

* Change attachTo to wait for ModelReady like KenelDropDown

* sanitizeSavedKernelInfo to fix invalid kernel and display_name. For example: PySpark1111 and PySpark 1111


* Added _contextsChangingEmitter to change loadContext msg when changing kernel

* Resolve PR comments

src/sql/azdata.proposed.d.ts

@@ -4098,6 +4098,7 @@ declare module 'azdata' {
 
 		export interface IStandardKernel {
 			readonly name: string;
+			readonly displayName: string;
 			readonly connectionProviderIds: string[];
 		}
 
@@ -4206,7 +4207,7 @@ declare module 'azdata' {
 
 		export interface ILanguageInfo {
 			name: string;
-			version: string;
+			version?: string;
 			mimetype?: string;
 			codemirror_mode?: string | ICodeMirrorMode;
 		}

src/sql/parts/notebook/models/clientSession.ts

@@ -44,6 +44,7 @@ export class ClientSession implements IClientSession {
 	private _errorMessage: string;
 	private _cachedKernelSpec: nb.IKernelSpec;
 	private _kernelChangeHandlers: KernelChangeHandler[] = [];
+	private _defaultKernel: nb.IKernelSpec;
 
 	//#endregion
 
@@ -59,6 +60,7 @@ export class ClientSession implements IClientSession {
 		this._isReady = false;
 		this._ready = new Deferred<void>();
 		this._kernelChangeCompleted = new Deferred<void>();
+		this._defaultKernel = options.kernelSpec;
 	}
 
 	public async initialize(): Promise<void> {
@@ -95,8 +97,8 @@ export class ClientSession implements IClientSession {
 			if (!this.notebookManager.sessionManager.isReady) {
 				await this.notebookManager.sessionManager.ready;
 			}
-			if (this._kernelPreference && this._kernelPreference.shouldStart) {
-				await this.startSessionInstance(this._kernelPreference.name);
+			if (this._defaultKernel) {
+				await this.startSessionInstance(this._defaultKernel.name);
 			}
 		}
 	}
@@ -228,7 +230,7 @@ export class ClientSession implements IClientSession {
 	async changeKernel(options: nb.IKernelSpec, oldValue?: nb.IKernel): Promise<nb.IKernel> {
 		this._kernelChangeCompleted = new Deferred<void>();
 		this._isReady = false;
-		let oldKernel = oldValue? oldValue: this.kernel;
+		let oldKernel = oldValue ? oldValue : this.kernel;
 		let newKernel = this.kernel;
 
 		let kernel = await this.doChangeKernel(options);

src/sql/parts/notebook/models/modelInterfaces.ts

@@ -22,11 +22,13 @@ import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHos
 import { IStandardKernelWithProvider } from 'sql/parts/notebook/notebookUtils';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
+import { localize } from 'vs/nls';
 
 export interface IClientSessionOptions {
 	notebookUri: URI;
 	notebookManager: INotebookManager;
 	notificationService: INotificationService;
+	kernelSpec: nb.IKernelSpec;
 }
 
 /**
@@ -309,6 +311,11 @@ export interface INotebookModel {
 	 */
 	readonly contextsChanged: Event<void>;
 
+	/**
+	 * Event fired on when switching kernel and should show loading context
+	 */
+	readonly contextsLoading: Event<void>;
+
 	/**
 	 * The specs for available kernels, or undefined if these have
 	 * not been loaded yet
@@ -389,6 +396,12 @@ export interface INotebookModel {
 
 	getApplicableConnectionProviderIds(kernelName: string): string[];
 
+	/**
+	 * Get the standardKernelWithProvider by name
+	 * @param name The kernel name
+	 */
+	getStandardKernelFromName(name: string): IStandardKernelWithProvider;
+
 	/** Event fired once we get call back from ConfigureConnection method in sqlops extension */
 	readonly onValidConnectionSelected: Event<boolean>;
 }
@@ -510,4 +523,11 @@ export interface ICellMagicMapper {
 
 export namespace notebookConstants {
 	export const SQL = 'SQL';
+	export const SQL_CONNECTION_PROVIDER = 'MSSQL';
+	export const sqlKernel: string = localize('sqlKernel', 'SQL');
+	export const sqlKernelSpec: nb.IKernelSpec = ({
+		name: sqlKernel,
+		language: 'sql',
+		display_name: sqlKernel
+	});
 }

src/sql/parts/notebook/models/notebookContexts.ts

@@ -123,15 +123,14 @@ export class NotebookContexts {
 	/**
 	 *
 	 * @param specs kernel specs (comes from session manager)
-	 * @param connectionInfo connection profile
-	 * @param savedKernelInfo kernel info loaded from
+	 * @param displayName kernel info loaded from
 	 */
-	public static getDefaultKernel(specs: nb.IAllKernels, connectionInfo: IConnectionProfile, savedKernelInfo: nb.IKernelInfo): nb.IKernelSpec {
+	public static getDefaultKernel(specs: nb.IAllKernels, displayName: string): nb.IKernelSpec {
 		let defaultKernel: nb.IKernelSpec;
 		if (specs) {
 			// find the saved kernel (if it exists)
-			if (savedKernelInfo) {
-				defaultKernel = specs.kernels.find((kernel) => kernel.name === savedKernelInfo.name);
+			if (displayName) {
+				defaultKernel = specs.kernels.find((kernel) => kernel.display_name === displayName);
 			}
 			// if no saved kernel exists, use the default KernelSpec
 			if (!defaultKernel) {
@@ -144,10 +143,7 @@ export class NotebookContexts {
 
 		// If no default kernel specified (should never happen), default to SQL
 		if (!defaultKernel) {
-			defaultKernel = {
-				name: notebookConstants.SQL,
-				display_name: notebookConstants.SQL
-			};
+			defaultKernel = notebookConstants.sqlKernelSpec;
 		}
 		return defaultKernel;
 	}

src/sql/parts/notebook/models/notebookModel.ts

@@ -9,10 +9,10 @@ import { nb, connection } from 'azdata';
 
 import { localize } from 'vs/nls';
 import { Event, Emitter } from 'vs/base/common/event';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 
 import { CellModel } from './cell';
-import { IClientSession, INotebookModel, IDefaultConnection, INotebookModelOptions, ICellModel, NotebookContentChange } from './modelInterfaces';
+import { IClientSession, INotebookModel, IDefaultConnection, INotebookModelOptions, ICellModel, NotebookContentChange, notebookConstants } from './modelInterfaces';
 import { NotebookChangeType, CellType } from 'sql/parts/notebook/models/contracts';
 import { nbversion } from '../notebookConstants';
 import * as notebookUtils from '../notebookUtils';
@@ -41,13 +41,13 @@ export class ErrorInfo {
 
 export class NotebookModel extends Disposable implements INotebookModel {
 	private _contextsChangedEmitter = new Emitter<void>();
+	private _contextsLoadingEmitter = new Emitter<void>();
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
 	private _kernelsChangedEmitter = new Emitter<nb.IKernelSpec>();
+	private _kernelChangedEmitter = new Emitter<nb.IKernelChangedArgs>();
 	private _layoutChanged = new Emitter<void>();
 	private _inErrorState: boolean = false;
-	private _clientSessions: IClientSession[] = [];
 	private _activeClientSession: IClientSession;
-	private _oldClientSession: IClientSession;
 	private _sessionLoadFinished: Promise<void>;
 	private _onClientSessionReady = new Emitter<IClientSession>();
 	private _onProviderIdChanged = new Emitter<string>();
@@ -69,21 +69,24 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _kernelDisplayNameToConnectionProviderIds: Map<string, string[]> = new Map<string, string[]>();
 	private _kernelDisplayNameToNotebookProviderIds: Map<string, string> = new Map<string, string>();
 	private _onValidConnectionSelected = new Emitter<boolean>();
+	private _oldKernel: nb.IKernel;
+	private _clientSessionListeners: IDisposable[] = [];
 
 	constructor(private _notebookOptions: INotebookModelOptions, startSessionImmediately?: boolean, private connectionProfile?: IConnectionProfile) {
 		super();
 		if (!_notebookOptions || !_notebookOptions.notebookUri || !_notebookOptions.notebookManagers) {
 			throw new Error('path or notebook service not defined');
 		}
-		if (startSessionImmediately) {
-			this.backgroundStartSession();
-		}
 		this._trustedMode = false;
 		this._providerId = _notebookOptions.providerId;
 		this._onProviderIdChanged.fire(this._providerId);
 		this._notebookOptions.standardKernels.forEach(kernel => {
-			this._kernelDisplayNameToConnectionProviderIds.set(kernel.name, kernel.connectionProviderIds);
-			this._kernelDisplayNameToNotebookProviderIds.set(kernel.name, kernel.notebookProvider);
+			let displayName = kernel.displayName;
+			if (!displayName) {
+				displayName = kernel.name;
+			}
+			this._kernelDisplayNameToConnectionProviderIds.set(displayName, kernel.connectionProviderIds);
+			this._kernelDisplayNameToNotebookProviderIds.set(displayName, kernel.notebookProvider);
 		});
 		if (this._notebookOptions.layoutChanged) {
 			this._notebookOptions.layoutChanged(() => this._layoutChanged.fire());
@@ -109,6 +112,13 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return manager;
 	}
 
+	public getNotebookManager(providerId: string): INotebookManager {
+		if (providerId) {
+			return this.notebookManagers.find(manager => manager.providerId === providerId);
+		}
+		return undefined;
+	}
+
 	public get notebookOptions(): INotebookModelOptions {
 		return this._notebookOptions;
 	}
@@ -144,7 +154,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public get kernelChanged(): Event<nb.IKernelChangedArgs> {
-		return this._activeClientSession.kernelChanged;
+		return this._kernelChangedEmitter.event;
 	}
 
 	public get kernelsChanged(): Event<nb.IKernelSpec> {
@@ -163,6 +173,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._contextsChangedEmitter.event;
 	}
 
+	public get contextsLoading(): Event<void> {
+		return this._contextsLoadingEmitter.event;
+	}
+
 	public get cells(): ICellModel[] {
 		return this._cells;
 	}
@@ -189,6 +203,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return specs;
 	}
 
+	public standardKernelsDisplayName(): string[] {
+		return Array.from(this._kernelDisplayNameToNotebookProviderIds.keys());
+	}
+
 	public get inErrorState(): boolean {
 		return this._inErrorState;
 	}
@@ -257,25 +275,17 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			if (this._notebookOptions && this._notebookOptions.contentManager) {
 				contents = await this._notebookOptions.contentManager.loadContent();
 			}
-
 			let factory = this._notebookOptions.factory;
 			// if cells already exist, create them with language info (if it is saved)
 			this._cells = [];
-			this._defaultLanguageInfo = {
-				name: this._providerId === SQL_NOTEBOOK_PROVIDER ? 'sql' : 'python',
-				version: ''
-			};
 			if (contents) {
 				this._defaultLanguageInfo = this.getDefaultLanguageInfo(contents);
 				this._savedKernelInfo = this.getSavedKernelInfo(contents);
-				this.setProviderIdForKernel(this._savedKernelInfo);
-				if (this._savedKernelInfo) {
-					this._defaultKernel = this._savedKernelInfo;
-				}
 				if (contents.cells && contents.cells.length > 0) {
 					this._cells = contents.cells.map(c => factory.createCell(c, { notebook: this, isTrusted: isTrusted }));
 				}
 			}
+			this.setDefaultKernelAndProviderId();
 			this.trySetLanguageFromLangInfo();
 		} catch (error) {
 			this._inErrorState = true;
@@ -380,17 +390,21 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		this._onErrorEmitter.fire({ message: error, severity: Severity.Error });
 	}
 
-	public backgroundStartSession(): void {
-		// TODO: only one session should be active at a time, depending on the current provider
-		this.notebookManagers.forEach(manager => {
+	public async startSession(manager: INotebookManager, displayName?: string): Promise<void> {
+		if (displayName) {
+			let standardKernel = this._notebookOptions.standardKernels.find(kernel => kernel.displayName === displayName);
+			this._defaultKernel = displayName ? { name: standardKernel.name, display_name: standardKernel.displayName } : this._defaultKernel;
+		}
+		if (this._defaultKernel) {
 			let clientSession = this._notebookOptions.factory.createClientSession({
 				notebookUri: this._notebookOptions.notebookUri,
 				notebookManager: manager,
-				notificationService: this._notebookOptions.notificationService
+				notificationService: this._notebookOptions.notificationService,
+				kernelSpec: this._defaultKernel
 			});
-			this._clientSessions.push(clientSession);
 			if (!this._activeClientSession) {
-				this._activeClientSession = clientSession;
+				this.updateActiveClientSession(clientSession);
+
 			}
 			let profile = new ConnectionProfile(this._notebookOptions.capabilitiesService, this.connectionProfile);
 
@@ -400,25 +414,84 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				this._activeConnection = undefined;
 			}
 
-			clientSession.initialize();
-			this._sessionLoadFinished = clientSession.ready.then(async () => {
-				if (clientSession.isInErrorState) {
-					this.setErrorState(clientSession.errorMessage);
-				} else {
-					this._onClientSessionReady.fire(clientSession);
-					// Once session is loaded, can use the session manager to retrieve useful info
-					this.loadKernelInfo(clientSession);
+			await clientSession.initialize();
+			// By somehow we have to wait for ready, otherwise may not be called for some cases.
+			await clientSession.ready;
+			if (clientSession.isInErrorState) {
+				this.setErrorState(clientSession.errorMessage);
+			} else {
+				this._onClientSessionReady.fire(clientSession);
+				// Once session is loaded, can use the session manager to retrieve useful info
+				this.loadKernelInfo(clientSession, this.defaultKernel.display_name);
+			}
+		}
+	}
+
+	// When changing kernel, update the active session and register the kernel change event
+	// So KernelDropDown could get the event fired when added listerner on Model.KernelChange
+	private updateActiveClientSession(clientSession: IClientSession) {
+		this.clearClientSessionListeners();
+		this._activeClientSession = clientSession;
+		this._clientSessionListeners.push(this._activeClientSession.kernelChanged(e => this._kernelChangedEmitter.fire(e)));
+	}
+
+	private clearClientSessionListeners() {
+		this._clientSessionListeners.forEach(listener => listener.dispose());
+		this._clientSessionListeners = [];
+	}
+
+	public setDefaultKernelAndProviderId() {
+		if (this._savedKernelInfo) {
+			this.sanitizeSavedKernelInfo();
+			let provider = this._kernelDisplayNameToNotebookProviderIds.get(this._savedKernelInfo.display_name);
+			if (provider && provider !== this._providerId) {
+				this._providerId = provider;
+			}
+			this._defaultKernel = this._savedKernelInfo;
+		} else if (this._defaultKernel) {
+			let providerId = this._kernelDisplayNameToNotebookProviderIds.get(this._defaultKernel.display_name);
+			if (providerId) {
+				if (this._providerId !== providerId) {
+					this._providerId = providerId;
 				}
-			});
-		});
+			} else {
+				this._defaultKernel = notebookConstants.sqlKernelSpec;
+				this._providerId = SQL_NOTEBOOK_PROVIDER;
+			}
+		} else {
+			this._defaultKernel = notebookConstants.sqlKernelSpec;
+			this._providerId = SQL_NOTEBOOK_PROVIDER;
+		}
+		// update default language
+		this._defaultLanguageInfo = {
+			name: this._providerId === SQL_NOTEBOOK_PROVIDER ? 'sql' : 'python',
+			version: ''
+		};
 	}
 
 	private isValidConnection(profile: IConnectionProfile | connection.Connection) {
-		let standardKernels = this._notebookOptions.standardKernels.find(kernel => this._savedKernelInfo && kernel.name === this._savedKernelInfo.display_name);
+		let standardKernels = this._notebookOptions.standardKernels.find(kernel => this._savedKernelInfo && kernel.displayName === this._savedKernelInfo.display_name);
 		let connectionProviderIds = standardKernels ? standardKernels.connectionProviderIds : undefined;
 		return profile && connectionProviderIds && connectionProviderIds.find(provider => provider === profile.providerName) !== undefined;
 	}
 
+	public getStandardKernelFromName(name: string): notebookUtils.IStandardKernelWithProvider {
+		if (name) {
+			let kernel = this._notebookOptions.standardKernels.find(kernel => kernel.name.toLowerCase() === name.toLowerCase());
+			return kernel;
+		}
+		return undefined;
+	}
+
+	public getStandardKernelFromDisplayName(displayName: string): notebookUtils.IStandardKernelWithProvider {
+		if (displayName) {
+			let kernel = this._notebookOptions.standardKernels.find(kernel => kernel.displayName.toLowerCase() === displayName.toLowerCase());
+			return kernel;
+		}
+		return undefined;
+	}
+
+
 	public get languageInfo(): nb.ILanguageInfo {
 		return this._defaultLanguageInfo;
 	}
@@ -469,20 +542,26 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public changeKernel(displayName: string): void {
-		let spec = this.getKernelSpecFromDisplayName(displayName);
-		this.doChangeKernel(spec);
+		this._contextsLoadingEmitter.fire();
+		this.doChangeKernel(displayName, true);
 	}
 
-	public doChangeKernel(kernelSpec: nb.IKernelSpec): Promise<void> {
-		this.setProviderIdForKernel(kernelSpec);
-		// Ensure that the kernel we try to switch to is a valid kernel; if not, use the default
-		let kernelSpecs = this.getKernelSpecs();
-		if (kernelSpecs && kernelSpecs.length > 0 && kernelSpecs.findIndex(k => k.name === kernelSpec.name) < 0) {
-			kernelSpec = kernelSpecs.find(spec => spec.name === this.notebookManager.sessionManager.specs.defaultKernel);
+	public async doChangeKernel(displayName: string, mustSetProvider: boolean = true): Promise<void> {
+		if (mustSetProvider) {
+			await this.setProviderIdAndStartSession(displayName);
+		}
+		let spec = this.getKernelSpecFromDisplayName(displayName);
+		if (spec) {
+			// Ensure that the kernel we try to switch to is a valid kernel; if not, use the default
+			let kernelSpecs = this.getKernelSpecs();
+			if (kernelSpecs && kernelSpecs.length > 0 && kernelSpecs.findIndex(k => k.display_name === spec.display_name) < 0) {
+				spec = kernelSpecs.find(spec => spec.name === this.notebookManager.sessionManager.specs.defaultKernel);
+			}
+		} else {
+			spec = notebookConstants.sqlKernelSpec;
 		}
 		if (this._activeClientSession && this._activeClientSession.isReady) {
-			let oldKernel = this._oldClientSession ? this._oldClientSession.kernel : null;
-			return this._activeClientSession.changeKernel(kernelSpec, oldKernel)
+			return this._activeClientSession.changeKernel(spec, this._oldKernel)
 				.then((kernel) => {
 					this.updateKernelInfo(kernel);
 					kernel.ready.then(() => {
@@ -546,32 +625,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
-	private loadKernelInfo(clientSession: IClientSession): void {
+	private loadKernelInfo(clientSession: IClientSession, displayName: string): void {
 		clientSession.onKernelChanging(async (e) => {
 			await this.loadActiveContexts(e);
 		});
 		clientSession.statusChanged(async (session) => {
 			this._kernelsChangedEmitter.fire(session.kernel);
 		});
-		if (!this.notebookManager) {
-			return;
-		}
-		try {
-			let sessionManager = this.notebookManager.sessionManager;
-			if (sessionManager) {
-				if (!this._defaultKernel) {
-					this._defaultKernel = NotebookContexts.getDefaultKernel(sessionManager.specs, this.connectionProfile, this._savedKernelInfo);
-				}
-				let spec = this.getKernelSpecFromDisplayName(this._defaultKernel.display_name);
-				if (spec) {
-					this._defaultKernel = spec;
-				}
-				this.doChangeKernel(this._defaultKernel);
-			}
-		} catch (err) {
-			let msg = notebookUtils.getErrorMessage(err);
-			this.notifyError(localize('loadKernelFailed', 'Loading kernel info failed: {0}', msg));
-		}
+
+		this.doChangeKernel(displayName, false);
 	}
 
 	// Get default language if saved in notebook file
@@ -600,7 +662,23 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return kernel;
 	}
 
-	private getDisplayNameFromSpecName(kernel: nb.IKernel): string {
+	private sanitizeSavedKernelInfo() {
+		if (this._savedKernelInfo) {
+			let displayName = this.sanitizeDisplayName(this._savedKernelInfo.display_name);
+
+			if (this._savedKernelInfo.display_name !== displayName) {
+				this._savedKernelInfo.display_name = displayName;
+			}
+
+			let standardKernel = this._notebookOptions.standardKernels.find(kernel => kernel.displayName === displayName || displayName.startsWith(kernel.displayName));
+			if (standardKernel && this._savedKernelInfo.name && this._savedKernelInfo.name !== standardKernel.name) {
+				this._savedKernelInfo.name = standardKernel.name;
+				this._savedKernelInfo.display_name = standardKernel.displayName;
+			}
+		}
+	}
+
+	public getDisplayNameFromSpecName(kernel: nb.IKernel): string {
 		let specs = this.notebookManager.sessionManager.specs;
 		if (!specs || !specs.kernels) {
 			return kernel.name;
@@ -638,22 +716,26 @@ export class NotebookModel extends Disposable implements INotebookModel {
 					this._activeConnection = undefined;
 				}
 			}
-			if (this._activeClientSession) {
-				try {
-					await this._activeClientSession.ready;
-				} catch (err) {
-					this.notifyError(localize('shutdownClientSessionError', 'A client session error occurred when closing the notebook: {0}', err));
-				}
-				await this._activeClientSession.shutdown();
-				this._clientSessions = undefined;
-				this._activeClientSession = undefined;
-
-			}
+			await this.shutdownActiveSession();
 		} catch (err) {
 			this.notifyError(localize('shutdownError', 'An error occurred when closing the notebook: {0}', err));
 		}
 	}
 
+	private async shutdownActiveSession() {
+		if (this._activeClientSession) {
+			try {
+				await this._activeClientSession.ready;
+			}
+			catch (err) {
+				this.notifyError(localize('shutdownClientSessionError', 'A client session error occurred when closing the notebook: {0}', err));
+			}
+			await this._activeClientSession.shutdown();
+			this.clearClientSessionListeners();
+			this._activeClientSession = undefined;
+		}
+	}
+
 	private async loadActiveContexts(kernelChangedArgs: nb.IKernelChangedArgs): Promise<void> {
 		if (kernelChangedArgs && kernelChangedArgs.newValue && kernelChangedArgs.newValue.name) {
 			let kernelDisplayName = this.getDisplayNameFromSpecName(kernelChangedArgs.newValue);
@@ -711,48 +793,52 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	/**
-	 * Set _providerId and _activeClientSession based on a kernelSpec representing new kernel
-	 * @param kernelSpec KernelSpec for new kernel
+	 * Set _providerId and start session if it is new provider
+	 * @param displayName Kernel dispay name
 	 */
-	private setProviderIdForKernel(kernelSpec: nb.IKernelSpec): void {
-		if (!kernelSpec) {
-			// Just use the 1st non-default provider, we don't have a better heuristic
-			let notebookManagers = this._notebookOptions.notebookManagers.filter(manager => manager.providerId !== DEFAULT_NOTEBOOK_PROVIDER);
-			if (!notebookManagers.length) {
-				notebookManagers = this._notebookOptions.notebookManagers;
-			}
-			if (notebookManagers.length > 0) {
-				this._providerId = notebookManagers[0].providerId;
-			}
-		} else {
-			let sessionManagerFound: boolean = false;
-			for (let i = 0; i < this.notebookManagers.length; i++) {
-				if (this.notebookManagers[i].sessionManager && this.notebookManagers[i].sessionManager.specs && this.notebookManagers[i].sessionManager.specs.kernels) {
-					let index = this.notebookManagers[i].sessionManager.specs.kernels.findIndex(kernel => kernel.name === kernelSpec.name);
-					if (index >= 0 && this._clientSessions && this._clientSessions.length > 0) {
-						if (this._activeClientSession) {
-							this._oldClientSession = this._activeClientSession;
+	private async setProviderIdAndStartSession(displayName: string): Promise<void> {
+		if (displayName) {
+			if (this._activeClientSession && this._activeClientSession.isReady) {
+				this._oldKernel = this._activeClientSession.kernel;
+				let providerId = this.tryFindProviderForKernel(displayName);
+
+				if (providerId) {
+					if (providerId !== this._providerId) {
+						this._providerId = providerId;
+						this._onProviderIdChanged.fire(this._providerId);
+
+						await this.shutdownActiveSession();
+
+						try {
+							let manager = this.getNotebookManager(providerId);
+							if (manager) {
+								await this.startSession(manager, displayName);
+							} else {
+								throw new Error(localize('ProviderNoManager', "Can't find notebook manager for provider {0}", providerId));
+							}
 						}
-						this._activeClientSession = this._clientSessions[i];
-						if (this.notebookManagers[i].providerId !== this._providerId) {
-							this._providerId = this.notebookManagers[i].providerId;
-							this._onProviderIdChanged.fire(this._providerId);
+						catch (err) {
+							console.log(err);
 						}
-						sessionManagerFound = true;
-						break;
+					} else {
+						console.log(`No provider found supporting the kernel: ${displayName}`);
 					}
 				}
 			}
+		}
+	}
 
-			// If no SessionManager exists, utilize passed in StandardKernels to see if we can intelligently set _providerId
-			if (!sessionManagerFound) {
-				let provider = this._kernelDisplayNameToNotebookProviderIds.get(kernelSpec.display_name);
-				if (provider && provider !== this._providerId) {
-					this._providerId = provider;
-					this._onProviderIdChanged.fire(this._providerId);
-				}
+	private tryFindProviderForKernel(displayName: string) {
+		if (!displayName) {
+			return undefined;
+		}
+		let standardKernel = this.getStandardKernelFromDisplayName(displayName);
+		if (standardKernel && this._oldKernel && this._oldKernel.name !== standardKernel.name) {
+			if (this._kernelDisplayNameToNotebookProviderIds.has(displayName)) {
+				return this._kernelDisplayNameToNotebookProviderIds.get(displayName);
 			}
 		}
+		return undefined;
 	}
 
 	// Get kernel specs from current sessionManager

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

@@ -220,6 +220,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	private async doLoad(): Promise<void> {
 		try {
+			await this.setNotebookManager();
 			await this.loadModel();
 			this.setLoading(false);
 			this._modelReadyDeferred.resolve(this._model);
@@ -243,10 +244,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		let providerId = 'sql'; // this is tricky; really should also depend on the connection profile
 		this.setContextKeyServiceWithProviderId(providerId);
 		this.fillInActionsForCurrentContext();
-		for (let providerId of this._notebookParams.providers) {
-			let notebookManager = await this.notebookService.getOrCreateNotebookManager(providerId, this._notebookParams.notebookUri);
-			this.notebookManagers.push(notebookManager);
-		}
+
 		let model = new NotebookModel({
 			factory: this.modelFactory,
 			notebookUri: this._notebookParams.notebookUri,
@@ -268,10 +266,17 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._model = this._register(model);
 		this.updateToolbarComponents(this._model.trustedMode);
 		this._modelRegisteredDeferred.resolve(this._model);
-		model.backgroundStartSession();
+		await model.startSession(this.model.notebookManager);
 		this.detectChanges();
 	}
 
+	private async setNotebookManager() {
+		for (let providerId of this._notebookParams.providers) {
+			let notebookManager = await this.notebookService.getOrCreateNotebookManager(providerId, this._notebookParams.notebookUri);
+			this.notebookManagers.push(notebookManager);
+		}
+	}
+
 	private async awaitNonDefaultProvider(): Promise<void> {
 		// Wait on registration for now. Long-term would be good to cache and refresh
 		await this.notebookService.registrationComplete;
@@ -348,12 +353,12 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	protected initActionBar() {
 		let kernelContainer = document.createElement('div');
-		let kernelDropdown = new KernelsDropdown(kernelContainer, this.contextViewService, this.modelRegistered);
+		let kernelDropdown = new KernelsDropdown(kernelContainer, this.contextViewService, this.modelReady);
 		kernelDropdown.render(kernelContainer);
 		attachSelectBoxStyler(kernelDropdown, this.themeService);
 
 		let attachToContainer = document.createElement('div');
-		let attachToDropdown = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelRegistered,
+		let attachToDropdown = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelReady,
 			this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService);
 		attachToDropdown.render(attachToContainer);
 		attachSelectBoxStyler(attachToDropdown, this.themeService);

src/sql/parts/notebook/notebookActions.ts

@@ -12,7 +12,7 @@ import { IContextViewProvider } from 'vs/base/browser/ui/contextview/contextview
 import { INotificationService, Severity, INotificationActions } from 'vs/platform/notification/common/notification';
 
 import { SelectBox, ISelectBoxOptionsWithLabel } from 'sql/base/browser/ui/selectBox/selectBox';
-import { INotebookModel, IDefaultConnection } from 'sql/parts/notebook/models/modelInterfaces';
+import { INotebookModel } from 'sql/parts/notebook/models/modelInterfaces';
 import { CellType } from 'sql/parts/notebook/models/contracts';
 import { NotebookComponent } from 'sql/parts/notebook/notebook.component';
 import { getErrorMessage, formatServerNameWithDatabaseNameForAttachTo, getServerFromFormattedAttachToName, getDatabaseFromFormattedAttachToName } from 'sql/parts/notebook/notebookUtils';
@@ -21,6 +21,7 @@ import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilit
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { noKernel } from 'sql/workbench/services/notebook/common/sessionManager';
 import { IConnectionDialogService } from 'sql/workbench/services/connection/common/connectionDialogService';
+import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
 
 const msgLoading = localize('loading', 'Loading kernels...');
 const kernelLabel: string = localize('Kernel', 'Kernel: ');
@@ -233,16 +234,12 @@ export class TrustedAction extends ToggleableAction {
 }
 
 export class KernelsDropdown extends SelectBox {
-	private model: INotebookModel;
-	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelRegistered: Promise<INotebookModel>
-	) {
-		let selectBoxOptionsWithLabel: ISelectBoxOptionsWithLabel = {
-			labelText: kernelLabel,
-			labelOnTop: false
-		};
-		super([msgLoading], msgLoading, contextViewProvider, container, selectBoxOptionsWithLabel);
-		if (modelRegistered) {
-			modelRegistered
+	private model: NotebookModel;
+	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelReady: Promise<INotebookModel>) {
+		super([msgLoading], msgLoading, contextViewProvider, container, { labelText: kernelLabel, labelOnTop: false } as ISelectBoxOptionsWithLabel);
+
+		if (modelReady) {
+			modelReady
 				.then((model) => this.updateModel(model))
 				.catch((err) => {
 					// No-op for now
@@ -253,44 +250,42 @@ export class KernelsDropdown extends SelectBox {
 	}
 
 	updateModel(model: INotebookModel): void {
-		this.model = model;
-		this._register(model.kernelsChanged((defaultKernel) => {
-			this.updateKernel(defaultKernel);
+		this.model = model as NotebookModel;
+		this._register(this.model.kernelChanged((changedArgs: azdata.nb.IKernelChangedArgs) => {
+			this.updateKernel(changedArgs.newValue);
 		}));
-		if (model.clientSession) {
-			this._register(model.clientSession.kernelChanged((changedArgs: azdata.nb.IKernelChangedArgs) => {
-				if (changedArgs.newValue) {
-					this.updateKernel(changedArgs.newValue);
-				}
-			}));
-		}
 	}
 
 	// Update SelectBox values
-	private updateKernel(defaultKernel: azdata.nb.IKernelSpec) {
-		let specs = this.model.specs;
-		if (specs && specs.kernels) {
-			let index = specs.kernels.findIndex((kernel => kernel.name === defaultKernel.name));
-			this.setOptions(specs.kernels.map(kernel => kernel.display_name), index);
+	public updateKernel(kernel: azdata.nb.IKernel) {
+		if (kernel) {
+			let standardKernel = this.model.getStandardKernelFromName(kernel.name);
+
+			let kernels: string[] = this.model.standardKernelsDisplayName();
+			if (kernels && standardKernel) {
+				let index = kernels.findIndex((kernel => kernel === standardKernel.displayName));
+				this.setOptions(kernels, index);
+			}
 		}
 	}
 
 	public doChangeKernel(displayName: string): void {
+		this.setOptions([msgLoading], 0);
 		this.model.changeKernel(displayName);
 	}
 }
 
 export class AttachToDropdown extends SelectBox {
-	private model: INotebookModel;
+	private model: NotebookModel;
 
-	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelRegistered: Promise<INotebookModel>,
+	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelReady: Promise<INotebookModel>,
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IConnectionDialogService private _connectionDialogService: IConnectionDialogService,
 		@INotificationService private _notificationService: INotificationService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService) {
 		super([msgLoadingContexts], msgLoadingContexts, contextViewProvider, container, { labelText: attachToLabel, labelOnTop: false } as ISelectBoxOptionsWithLabel);
-		if (modelRegistered) {
-			modelRegistered
+		if (modelReady) {
+			modelReady
 				.then(model => {
 					this.updateModel(model);
 					this.updateAttachToDropdown(model);
@@ -305,17 +300,19 @@ export class AttachToDropdown extends SelectBox {
 	}
 
 	public updateModel(model: INotebookModel): void {
-		this.model = model;
+		this.model = model as NotebookModel;
 		this._register(model.contextsChanged(() => {
 			let kernelDisplayName: string = this.getKernelDisplayName();
 			if (kernelDisplayName) {
 				this.loadAttachToDropdown(this.model, kernelDisplayName);
 			}
 		}));
+		this._register(this.model.contextsLoading(() => {
+			this.setOptions([msgLoadingContexts], 0);
+		}));
 	}
 
 	private updateAttachToDropdown(model: INotebookModel): void {
-		this.model = model;
 		model.onValidConnectionSelected(validConnection => {
 			let kernelDisplayName: string = this.getKernelDisplayName();
 			if (kernelDisplayName) {

src/sql/parts/notebook/notebookInput.ts

@@ -222,6 +222,7 @@ export class NotebookInput extends EditorInput {
 			this._standardKernels.push({
 				connectionProviderIds: kernel.connectionProviderIds,
 				name: kernel.name,
+				displayName: kernel.displayName,
 				notebookProvider: kernel.notebookProvider
 			});
 		});

src/sql/parts/notebook/notebookUtils.ts

@@ -12,7 +12,6 @@ import * as pfs from 'vs/base/node/pfs';
 import { localize } from 'vs/nls';
 import { IOutputChannel } from 'vs/workbench/parts/output/common/output';
 import { DEFAULT_NOTEBOOK_PROVIDER, DEFAULT_NOTEBOOK_FILETYPE, INotebookService } from 'sql/workbench/services/notebook/common/notebookService';
-import { ContextKeyExpr, IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 
 
@@ -97,6 +96,7 @@ export function getDatabaseFromFormattedAttachToName(name: string): string {
 
 export interface IStandardKernelWithProvider {
 	readonly name: string;
+	readonly displayName: string;
 	readonly connectionProviderIds: string[];
 	readonly notebookProvider: string;
 }

src/sql/sqlops.proposed.d.ts

@@ -2050,6 +2050,7 @@ declare module 'sqlops' {
 
 		export interface IStandardKernel {
 			readonly name: string;
+			readonly displayName: string;
 			readonly connectionProviderIds: string[];
 		}
 

src/sql/workbench/services/notebook/common/notebookRegistry.ts

@@ -52,6 +52,9 @@ let notebookProviderType: IJSONSchema = {
 						name: {
 							type: 'string',
 						},
+						displayName: {
+							type: 'string',
+						},
 						connectionProviderIds: {
 							type: 'array',
 							items: {

src/sql/workbench/services/notebook/common/notebookServiceImpl.ts

@@ -175,6 +175,7 @@ export class NotebookService extends Disposable implements INotebookService {
 			if (provider) {
 				this._providerToStandardKernels.set(notebookConstants.SQL, [{
 					name: notebookConstants.SQL,
+					displayName: notebookConstants.SQL,
 					connectionProviderIds: sqlConnectionTypes
 				}]);
 			}
@@ -451,7 +452,7 @@ export class NotebookService extends Disposable implements INotebookService {
 		notebookRegistry.registerNotebookProvider({
 			provider: sqlProvider.providerId,
 			fileExtensions: DEFAULT_NOTEBOOK_FILETYPE,
-			standardKernels: { name: 'SQL', connectionProviderIds: ['MSSQL'] }
+			standardKernels: { name: notebookConstants.SQL, displayName: notebookConstants.SQL, connectionProviderIds: [notebookConstants.SQL_CONNECTION_PROVIDER] }
 		});
 	}
 

src/sql/workbench/services/notebook/sql/sqlNotebookManager.ts

@@ -6,13 +6,14 @@
 'use strict';
 
 import { nb } from 'azdata';
+import * as vscode from 'vscode';
 
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
-import { INotebookManager, SQL_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
+import { SQL_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
 import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
 import { SqlSessionManager } from 'sql/workbench/services/notebook/sql/sqlSessionManager';
 
-export class SqlNotebookManager implements INotebookManager {
+export class SqlNotebookManager implements nb.NotebookProvider {
 	private _contentManager: nb.ContentManager;
 	private _sessionManager: nb.SessionManager;
 
@@ -36,4 +37,16 @@ export class SqlNotebookManager implements INotebookManager {
 	public get sessionManager(): nb.SessionManager {
 		return this._sessionManager;
 	}
+
+	getNotebookManager(notebookUri: vscode.Uri): Thenable<nb.NotebookManager> {
+		throw new Error('Method not implemented.');
+	}
+
+	handleNotebookClosed(notebookUri: vscode.Uri): void {
+		throw new Error('Method not implemented.');
+	}
+
+	public get standardKernels(): nb.IStandardKernel[] {
+		return [];
+	}
 }

src/sql/workbench/services/notebook/sql/sqlSessionManager.ts

@@ -8,7 +8,7 @@ import * as os from 'os';
 import { nb, QueryExecuteSubsetResult, IDbColumn, BatchSummary, IResultMessage, ResultSetSummary } from 'azdata';
 import { localize } from 'vs/nls';
 import * as strings from 'vs/base/common/strings';
-import { FutureInternal, ILanguageMagic } from 'sql/parts/notebook/models/modelInterfaces';
+import { FutureInternal, ILanguageMagic, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
 import QueryRunner, { EventType } from 'sql/platform/query/common/queryRunner';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
@@ -24,18 +24,11 @@ import { elapsedTimeLabel } from 'sql/parts/query/common/localizedConstants';
 import * as notebookUtils from 'sql/parts/notebook/notebookUtils';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 
-export const sqlKernel: string = localize('sqlKernel', "SQL");
-export const sqlKernelError: string = localize('sqlKernelError', "SQL kernel error");
+export const sqlKernelError: string = localize("sqlKernelError", "SQL kernel error");
 export const MAX_ROWS = 5000;
 export const NotebookConfigSectionName = 'notebook';
 export const MaxTableRowsConfigName = 'maxTableRows';
 
-const sqlKernelSpec: nb.IKernelSpec = ({
-	name: sqlKernel,
-	language: 'sql',
-	display_name: sqlKernel
-});
-
 const languageMagics: ILanguageMagic[] = [{
 	language: 'Python',
 	magic: 'lang_python'
@@ -65,8 +58,8 @@ export class SqlSessionManager implements nb.SessionManager {
 
 	public get specs(): nb.IAllKernels {
 		let allKernels: nb.IAllKernels = {
-			defaultKernel: sqlKernel,
-			kernels: [sqlKernelSpec]
+			defaultKernel: notebookConstants.sqlKernel,
+			kernels: [notebookConstants.sqlKernelSpec]
 		};
 		return allKernels;
 	}
@@ -176,7 +169,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	}
 
 	public get name(): string {
-		return sqlKernel;
+		return notebookConstants.sqlKernel;
 	}
 
 	public get supportsIntellisense(): boolean {
@@ -217,7 +210,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	}
 
 	getSpec(): Thenable<nb.IKernelSpec> {
-		return Promise.resolve(sqlKernelSpec);
+		return Promise.resolve(notebookConstants.sqlKernelSpec);
 	}
 
 	requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean): nb.IFuture {