Begin defining Extension-based Notebook Provider (#3172)

Implements provider contribution in the MainThreadNotebook, with matching function calls in the ExtHostNotebook class. This will allow us to proxy through notebook providers (specifically, creation of a notebook manager with required content, server managers) from an extension up through to the main process.

Implemented in this PR:
- Callthroughs for content and server manager APIs
- Very basic unit tests covering provider & manager registration

Not implemented:
- Fuller unit tests on the specific callthrough methods for content & server manager.
- Contribution point needed to test this (so we can actually pass through the extension's existing Notebook implementation)

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

@@ -34,7 +34,7 @@ export class ClientSession implements IClientSession {
 	private _iopubMessageEmitter = new Emitter<nb.IMessage>();
 	private _unhandledMessageEmitter = new Emitter<nb.IMessage>();
 	private _propertyChangedEmitter = new Emitter<'path' | 'name' | 'type'>();
-	private _path: string;
+	private _notebookUri: URI;
 	private _type: string;
 	private _name: string;
 	private _isReady: boolean;
@@ -53,7 +53,7 @@ export class ClientSession implements IClientSession {
 	private _kernelConfigActions: ((kernelName: string) => Promise<any>)[] = [];
 
 	constructor(private options: IClientSessionOptions) {
-		this._path = options.path;
+		this._notebookUri = options.notebookUri;
 		this.notebookManager = options.notebookManager;
 		this._isReady = false;
 		this._ready = new Deferred<void>();
@@ -104,8 +104,9 @@ export class ClientSession implements IClientSession {
 	private async startSessionInstance(kernelName: string): Promise<void> {
 		let session: nb.ISession;
 		try {
+			// TODO #3164 should use URI instead of path for startNew
 			session = await this.notebookManager.sessionManager.startNew({
-				path: this.path,
+				path: this.notebookUri.fsPath,
 				kernelName: kernelName
 				// TODO add kernel name if saved in the document
 			});
@@ -115,7 +116,7 @@ export class ClientSession implements IClientSession {
 			if (err && err.response && err.response.status === 501) {
 				this.options.notificationService.warn(nls.localize('sparkKernelRequiresConnection', 'Kernel {0} was not found. The default kernel will be used instead.', kernelName));
 				session = await this.notebookManager.sessionManager.startNew({
-					path: this.path,
+					path: this.notebookUri.fsPath,
 					kernelName: undefined
 				});
 			} else {
@@ -169,8 +170,8 @@ export class ClientSession implements IClientSession {
 	public get kernel(): nb.IKernel | null {
 		return this._session ? this._session.kernel : undefined;
 	}
-	public get path(): string {
+	public get notebookUri(): URI {
-		return this._path;
+		return this._notebookUri;
 	}
 	public get name(): string {
 		return this._name;

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

@@ -21,7 +21,7 @@ import { NotebookConnection } from 'sql/parts/notebook/models/notebookConnection
 import { IConnectionManagementService } from 'sql/parts/connection/common/connectionManagement';
 
 export interface IClientSessionOptions {
-	path: string;
+	notebookUri: URI;
 	notebookManager: INotebookManager;
 	notificationService: INotificationService;
 }
@@ -73,7 +73,7 @@ export interface IClientSession extends IDisposable {
 	/**
 	 * The current path associated with the client session.
 	 */
-	readonly path: string;
+	readonly notebookUri: URI;
 
 	/**
 	 * The current name associated with the client session.
@@ -354,7 +354,7 @@ export interface INotebookModelOptions {
 	/**
 	 * Path to the local or remote notebook
 	 */
-	path: string;
+	notebookUri: URI;
 
 	/**
 	 * Factory for creating cells and client sessions

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

@@ -81,7 +81,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	constructor(private notebookOptions: INotebookModelOptions, startSessionImmediately?: boolean, private connectionProfile?: IConnectionProfile) {
 		super();
-		if (!notebookOptions || !notebookOptions.path || !notebookOptions.notebookManager) {
+		if (!notebookOptions || !notebookOptions.notebookUri || !notebookOptions.notebookManager) {
 			throw new Error('path or notebook service not defined');
 		}
 		if (startSessionImmediately) {
@@ -183,7 +183,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	public async requestModelLoad(isTrusted: boolean = false): Promise<void> {
 		try {
 			this._trustedMode = isTrusted;
-			let contents = await this.notebookManager.contentManager.getNotebookContents(this.notebookOptions.path);
+			let contents = await this.notebookManager.contentManager.getNotebookContents(this.notebookOptions.notebookUri);
 			let factory = this.notebookOptions.factory;
 			// if cells already exist, create them with language info (if it is saved)
 			this._cells = undefined;
@@ -248,7 +248,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public backgroundStartSession(): void {
 		this._clientSession = this.notebookOptions.factory.createClientSession({
-			path: this.notebookOptions.path,
+			notebookUri: this.notebookOptions.notebookUri,
 			notebookManager: this.notebookManager,
 			notificationService: this.notebookOptions.notificationService
 		});
@@ -416,7 +416,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		if (!notebook) {
 			return false;
 		}
-		await this.notebookManager.contentManager.save(this.notebookOptions.path, notebook);
+		await this.notebookManager.contentManager.save(this.notebookOptions.notebookUri, notebook);
 		this._contentChangedEmitter.fire({
 			changeType: NotebookChangeType.DirtyStateChanged,
 			isDirty: false

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

@@ -129,7 +129,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit {
 		this.notebookManager = await this.notebookService.getOrCreateNotebookManager(this.notebookParams.providerId, this.notebookParams.notebookUri);
 		let model = new NotebookModel({
 			factory: this.modelFactory,
-			path: this.notebookParams.notebookUri.fsPath,
+			notebookUri: this.notebookParams.notebookUri,
 			connectionService: this.connectionManagementService,
 			notificationService: this.notificationService,
 			notebookManager: this.notebookManager

src/sql/services/notebook/localContentManager.ts

@@ -4,29 +4,33 @@
  *--------------------------------------------------------------------------------------------*/
 
 'use strict';
+
 import { nb } from 'sqlops';
 
 import * as json from 'vs/base/common/json';
 import * as pfs from 'vs/base/node/pfs';
+import URI from 'vs/base/common/uri';
 
 import ContentManager = nb.ContentManager;
 import INotebook = nb.INotebook;
 
 export class LocalContentManager implements ContentManager {
-    public async getNotebookContents(path: string): Promise<INotebook> {
+    public async getNotebookContents(notebookUri: URI): Promise<INotebook> {
-        if (!path) {
+        if (!notebookUri) {
             return undefined;
         }
+        // TODO validate this is an actual file URI, and error if not
+        let path = notebookUri.fsPath;
         // Note: intentionally letting caller handle exceptions
         let notebookFileBuffer = await pfs.readFile(path);
         return <INotebook>json.parse(notebookFileBuffer.toString());
     }
 
-    public async save(path: string, notebook: INotebook): Promise<INotebook> {
+    public async save(notebookUri: URI, notebook: INotebook): Promise<INotebook> {
         // Convert to JSON with pretty-print functionality
         let contents = JSON.stringify(notebook, undefined, '    ');
+        let path = notebookUri.fsPath;
         await pfs.writeFile(path, contents);
         return notebook;
     }
-
 }

src/sql/sqlops.proposed.d.ts

@@ -1430,19 +1430,19 @@ declare module 'sqlops' {
 			/* Reads contents from a Uri representing a local or remote notebook and returns a
 			 * JSON object containing the cells and metadata about the notebook
 			 */
-			getNotebookContents(path: string): Thenable<INotebook>;
+			getNotebookContents(notebookUri: vscode.Uri): Thenable<INotebook>;
 
 			/**
 			 * Save a file.
 			 *
-			 * @param path - The desired file path.
+			 * @param notebookUri - The desired file path.
 			 *
 			 * @param notebook - notebook to be saved.
 			 *
 			 * @returns A thenable which resolves with the file content model when the
 			 *   file is saved.
 			 */
-			save(path: string, notebook: INotebook): Thenable<INotebook>;
+			save(notebookUri: vscode.Uri, notebook: INotebook): Thenable<INotebook>;
 		}
 
 		export interface INotebook {

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -411,3 +411,9 @@ export class SqlThemeIcon {
 		this.id = id;
 	}
 }
+
+export interface INotebookManagerDetails {
+	handle: number;
+	hasContentManager: boolean;
+	hasServerManager: boolean;
+}

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

@@ -13,23 +13,61 @@ import { localize } from 'vs/nls';
 
 
 import { ExtHostNotebookShape, MainThreadNotebookShape, SqlMainContext } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import URI, { UriComponents } from 'vs/base/common/uri';
+import { INotebookManagerDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export class ExtHostNotebook implements ExtHostNotebookShape {
 	private static _handlePool: number = 0;
 
 	private readonly _proxy: MainThreadNotebookShape;
 	private _providers = new Map<number, sqlops.nb.NotebookProvider>();
-
+	// Notebook URI to manager lookup.
+	private _managers = new Map<number, NotebookManagerAdapter>();
 	constructor(private _mainContext: IMainContext) {
 		this._proxy = _mainContext.getProxy(SqlMainContext.MainThreadNotebook);
 	}
 
 	//#region APIs called by main thread
-	getNotebookManager(notebookUri: vscode.Uri): Thenable<number> {
+	async $getNotebookManager(providerHandle: number, notebookUri: UriComponents): Promise<INotebookManagerDetails> {
-		throw new Error('Not implemented');
+		let uri = URI.revive(notebookUri);
+		let uriString = uri.toString();
+		let adapter = this.findManagerForUri(uriString);
+		if (!adapter) {
+			adapter = await this._withProvider(providerHandle, (provider) => {
+				return this.createManager(provider, uri);
+			});
+		}
+
+		return {
+			handle: adapter.managerHandle,
+			hasContentManager: !!adapter.contentManager,
+			hasServerManager: !!adapter.serverManager
+		};
 	}
-	handleNotebookClosed(notebookUri: vscode.Uri): void {
+	$handleNotebookClosed(notebookUri: UriComponents): void {
-		throw new Error('Not implemented');
+		let uri = URI.revive(notebookUri);
+		let uriString = uri.toString();
+		let manager = this.findManagerForUri(uriString);
+		if (manager) {
+			manager.provider.handleNotebookClosed(uri);
+			this._managers.delete(manager.managerHandle);
+		}
+	}
+
+	$doStartServer(managerHandle: number): Thenable<void> {
+		return this._withServerManager(managerHandle, (serverManager) => serverManager.startServer());
+	}
+
+	$doStopServer(managerHandle: number): Thenable<void> {
+		return this._withServerManager(managerHandle, (serverManager) => serverManager.stopServer());
+	}
+
+	$getNotebookContents(managerHandle: number, notebookUri: UriComponents): Thenable<sqlops.nb.INotebook> {
+		return this._withContentManager(managerHandle, (contentManager) => contentManager.getNotebookContents(URI.revive(notebookUri)));
+	}
+
+	$save(managerHandle: number, notebookUri: UriComponents, notebook: sqlops.nb.INotebook): Thenable<sqlops.nb.INotebook> {
+		return this._withContentManager(managerHandle, (contentManager) => contentManager.save(URI.revive(notebookUri), notebook));
 	}
 
 	//#endregion
@@ -47,6 +85,25 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 
 
 	//#region private methods
+
+	private findManagerForUri(uriString: string): NotebookManagerAdapter {
+		for(let manager of Array.from(this._managers.values())) {
+			if (manager.uriString === uriString) {
+				return manager;
+			}
+		}
+		return undefined;
+	}
+
+	private async createManager(provider: sqlops.nb.NotebookProvider, notebookUri: URI): Promise<NotebookManagerAdapter> {
+		let manager = await provider.getNotebookManager(notebookUri);
+		let uriString = notebookUri.toString();
+		let handle = this._nextHandle();
+		let adapter = new NotebookManagerAdapter(provider, handle, manager, uriString);
+		this._managers.set(handle, adapter);
+		return adapter;
+	}
+
 	private _createDisposable(handle: number): Disposable {
 		return new Disposable(() => {
 			this._providers.delete(handle);
@@ -58,12 +115,50 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		return ExtHostNotebook._handlePool++;
 	}
 
-	private _withProvider<A, R>(handle: number, ctor: { new(...args: any[]): A }, callback: (adapter: A) => TPromise<R>): TPromise<R> {
+	private _withProvider<R>(handle: number, callback: (provider: sqlops.nb.NotebookProvider) => R | PromiseLike<R>): TPromise<R> {
 		let provider = this._providers.get(handle);
-		if (!(provider instanceof ctor)) {
+		if (provider === undefined) {
-			return TPromise.wrapError<R>(new Error('no adapter found'));
+			return TPromise.wrapError<R>(new Error(localize('errNoProvider', 'no notebook provider found')));
 		}
-		return callback(<any>provider);
+		return TPromise.wrap(callback(provider));
+	}
+
+	private _withNotebookManager<R>(handle: number, callback: (manager: NotebookManagerAdapter) => R | PromiseLike<R>): TPromise<R> {
+		let manager = this._managers.get(handle);
+		if (manager === undefined) {
+			return TPromise.wrapError<R>(new Error(localize('errNoManager', 'No Manager found')));
+		}
+		return TPromise.wrap(callback(manager));
+	}
+
+	private _withServerManager<R>(handle: number, callback: (manager: sqlops.nb.ServerManager) => R | PromiseLike<R>): TPromise<R> {
+		return this._withNotebookManager(handle, (notebookManager) => {
+			let serverManager = notebookManager.serverManager;
+			if (!serverManager) {
+				return TPromise.wrapError(new Error(localize('noServerManager', 'Notebook Manager for notebook {0} does not have a server manager. Cannot perform operations on it', notebookManager.uriString)));
+			}
+			return callback(serverManager);
+		});
+	}
+
+	private _withContentManager<R>(handle: number, callback: (manager: sqlops.nb.ContentManager) => R | PromiseLike<R>): TPromise<R> {
+		return this._withNotebookManager(handle, (notebookManager) => {
+			let contentManager = notebookManager.contentManager;
+			if (!contentManager) {
+				return TPromise.wrapError(new Error(localize('noContentManager', 'Notebook Manager for notebook {0} does not have a content manager. Cannot perform operations on it', notebookManager.uriString)));
+			}
+			return callback(contentManager);
+		});
+	}
+
+	private _withSessionManager<R>(handle: number, callback: (manager: sqlops.nb.SessionManager) => R | PromiseLike<R>): TPromise<R> {
+		return this._withNotebookManager(handle, (notebookManager) => {
+			let sessionManager = notebookManager.sessionManager;
+			if (!sessionManager) {
+				return TPromise.wrapError(new Error(localize('noSessionManager', 'Notebook Manager for notebook {0} does not have a session manager. Cannot perform operations on it', notebookManager.uriString)));
+			}
+			return callback(sessionManager);
+		});
 	}
 
 	private _addNewProvider(adapter: sqlops.nb.NotebookProvider): number {
@@ -74,3 +169,26 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	//#endregion
 }
 
+
+class NotebookManagerAdapter implements sqlops.nb.NotebookManager {
+	constructor(
+		public readonly provider: sqlops.nb.NotebookProvider,
+		public readonly managerHandle: number,
+		private manager: sqlops.nb.NotebookManager,
+		public readonly uriString: string
+	) {
+	}
+
+	public get contentManager(): sqlops.nb.ContentManager {
+		return this.manager.contentManager;
+	}
+
+	public get sessionManager(): sqlops.nb.SessionManager {
+		return this.manager.sessionManager;
+	}
+
+	public get serverManager(): sqlops.nb.ServerManager {
+		return this.manager.serverManager;
+	}
+
+}

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

@@ -9,14 +9,18 @@ import { SqlExtHostContext, SqlMainContext, ExtHostNotebookShape, MainThreadNote
 import { extHostNamedCustomer } from 'vs/workbench/api/electron-browser/extHostCustomers';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IExtHostContext } from 'vs/workbench/api/node/extHost.protocol';
-import { INotebookService, INotebookProvider, INotebookManager } from 'sql/services/notebook/notebookService';
+import { Event, Emitter } from 'vs/base/common/event';
 import URI from 'vs/base/common/uri';
 
+import { INotebookService, INotebookProvider, INotebookManager } from 'sql/services/notebook/notebookService';
+import { INotebookManagerDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { LocalContentManager } from 'sql/services/notebook/localContentManager';
+
 @extHostNamedCustomer(SqlMainContext.MainThreadNotebook)
 export class MainThreadNotebook extends Disposable implements MainThreadNotebookShape {
 
 	private _proxy: ExtHostNotebookShape;
-	private _registrations: { [handle: number]: NotebookProviderWrapper } = Object.create(null);
+	private _providers = new Map<number, NotebookProviderWrapper>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -30,17 +34,17 @@ export class MainThreadNotebook extends Disposable implements MainThreadNotebook
 
 	//#region Extension host callable methods
 	public $registerNotebookProvider(providerId: string, handle: number): void {
-		let notebookProvider = new NotebookProviderWrapper(providerId, handle);
+		let notebookProvider = new NotebookProviderWrapper(this._proxy, providerId, handle);
-		this._registrations[providerId] = notebookProvider;
+		this._providers.set(handle, notebookProvider);
 		this.notebookService.registerProvider(providerId, notebookProvider);
 	}
 
 	public $unregisterNotebookProvider(handle: number): void {
-		let registration = this._registrations[handle];
+		let registration = this._providers.get(handle);
 		if (registration) {
 			this.notebookService.unregisterProvider(registration.providerId);
 			registration.dispose();
-			delete this._registrations[handle];
+			this._providers.delete(handle);
 		}
 	}
 
@@ -49,29 +53,147 @@ export class MainThreadNotebook extends Disposable implements MainThreadNotebook
 }
 
 class NotebookProviderWrapper extends Disposable implements INotebookProvider {
+	private _managers = new Map<string, NotebookManagerWrapper>();
 
-	constructor(public readonly providerId, public readonly handle: number) {
+	constructor(private _proxy: ExtHostNotebookShape, public readonly providerId, public readonly providerHandle: number) {
 		super();
 	}
 
 	getNotebookManager(notebookUri: URI): Thenable<INotebookManager> {
 		// TODO must call through to setup in the extension host
-		return Promise.resolve(new NotebookManagerWrapper(this.providerId));
+		return this.doGetNotebookManager(notebookUri);
+	}
+
+	private async doGetNotebookManager(notebookUri: URI): Promise<INotebookManager> {
+		let uriString = notebookUri.toString();
+		let manager = this._managers.get(uriString);
+		if (!manager) {
+			manager = new NotebookManagerWrapper(this._proxy, this.providerId, notebookUri);
+			await manager.initialize(this.providerHandle);
+			this._managers.set(uriString, manager);
+		}
+		return manager;
 	}
 
 	handleNotebookClosed(notebookUri: URI): void {
-		// TODO implement call through to extension host
+		this._proxy.$handleNotebookClosed(notebookUri);
 	}
 
-
 }
 
 class NotebookManagerWrapper implements INotebookManager {
-	constructor(public readonly providerId) {
+	private _sessionManager: sqlops.nb.SessionManager;
+	private _contentManager: sqlops.nb.ContentManager;
+	private _serverManager: sqlops.nb.ServerManager;
+	private managerDetails: INotebookManagerDetails;
 
+	constructor(private _proxy: ExtHostNotebookShape,
+		public readonly providerId,
+		private notebookUri: URI
+	) { }
+
+	public async initialize(providerHandle: number): Promise<NotebookManagerWrapper> {
+		this.managerDetails = await this._proxy.$getNotebookManager(providerHandle, this.notebookUri);
+		let managerHandle = this.managerDetails.handle;
+		this._contentManager = this.managerDetails.hasContentManager ? new ContentManagerWrapper(managerHandle, this._proxy) : new LocalContentManager();
+		this._serverManager = this.managerDetails.hasServerManager ? new ServerManagerWrapper(managerHandle, this._proxy) : undefined;
+		this._sessionManager = new SessionManagerWrapper(managerHandle, this._proxy);
+		return this;
+	}
+
+	public get sessionManager(): sqlops.nb.SessionManager {
+		return this._sessionManager;
+	}
+	public get contentManager(): sqlops.nb.ContentManager {
+		return this._contentManager;
+	}
+	public get serverManager(): sqlops.nb.ServerManager {
+		return this._serverManager;
+	}
+
+	public get managerHandle(): number {
+		return this.managerDetails.handle;
 	}
-	sessionManager: sqlops.nb.SessionManager;
+
-	contentManager: sqlops.nb.ContentManager;
+}
-	serverManager: sqlops.nb.ServerManager;
+
+class ContentManagerWrapper implements sqlops.nb.ContentManager {
+
+	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+	}
+	getNotebookContents(notebookUri: URI): Thenable<sqlops.nb.INotebook> {
+		return this._proxy.$getNotebookContents(this.handle, notebookUri);
+	}
+
+	save(path: URI, notebook: sqlops.nb.INotebook): Thenable<sqlops.nb.INotebook> {
+		return this._proxy.$save(this.handle, path, notebook);
+	}
+}
+
+class ServerManagerWrapper implements sqlops.nb.ServerManager {
+	private onServerStartedEmitter: Emitter<void>;
+	private _isStarted: boolean;
+	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+		this._isStarted = false;
+	}
+
+	get isStarted(): boolean {
+		return this._isStarted;
+	}
+
+	get onServerStarted(): Event<void> {
+		return this.onServerStartedEmitter.event;
+	}
+
+	startServer(): Thenable<void> {
+		return this.doStartServer();
+	}
+
+	private async doStartServer(): Promise<void> {
+		await this._proxy.$doStartServer(this.handle);
+		this._isStarted = true;
+		this.onServerStartedEmitter.fire();
+	}
+
+	stopServer(): Thenable<void> {
+		return this.doStopServer();
+	}
+
+	private async doStopServer(): Promise<void> {
+		try {
+			await this._proxy.$doStopServer(this.handle);
+		} finally {
+			// Always consider this a stopping event, even if a failure occurred.
+			this._isStarted = false;
+		}
+	}
+}
+
+class SessionManagerWrapper implements sqlops.nb.SessionManager {
+	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+	}
+
+	get isReady(): boolean {
+		throw new Error('Method not implemented.');
+
+	}
+
+	get ready(): Thenable<void> {
+		throw new Error('Method not implemented.');
+
+	}
+	get specs(): sqlops.nb.IAllKernels {
+		throw new Error('Method not implemented.');
+
+	}
+
+	startNew(options: sqlops.nb.ISessionOptions): Thenable<sqlops.nb.ISession> {
+		throw new Error('Method not implemented.');
+	}
+
+	shutdown(id: string): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
+
 
 }

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -21,7 +21,7 @@ import { ITreeComponentItem } from 'sql/workbench/common/views';
 import { ITaskHandlerDescription } from 'sql/platform/tasks/common/tasks';
 import {
 	IItemConfig, ModelComponentTypes, IComponentShape, IModelViewDialogDetails, IModelViewTabDetails, IModelViewButtonDetails,
-	IModelViewWizardDetails, IModelViewWizardPageDetails
+	IModelViewWizardDetails, IModelViewWizardPageDetails, INotebookManagerDetails
 } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export abstract class ExtHostAccountManagementShape {
@@ -711,11 +711,16 @@ export interface ExtHostNotebookShape {
 
 	/**
 	 * Looks up a notebook manager for a given notebook URI
+	 * @param {number} providerHandle
 	 * @param {vscode.Uri} notebookUri
 	 * @returns {Thenable<string>} handle of the manager to be used when sending
 	 */
-	getNotebookManager(notebookUri: vscode.Uri): Thenable<number>;
+	$getNotebookManager(providerHandle: number, notebookUri: UriComponents): Thenable<INotebookManagerDetails>;
-	handleNotebookClosed(notebookUri: vscode.Uri): void;
+	$handleNotebookClosed(notebookUri: UriComponents): void;
+	$doStartServer(managerHandle: number): Thenable<void>;
+	$doStopServer(managerHandle: number): Thenable<void>;
+	$getNotebookContents(managerHandle: number, notebookUri: UriComponents): Thenable<sqlops.nb.INotebook>;
+	$save(managerHandle: number, notebookUri: UriComponents, notebook: sqlops.nb.INotebook): Thenable<sqlops.nb.INotebook>;
 
 }
 
@@ -723,3 +728,4 @@ export interface MainThreadNotebookShape extends IDisposable {
 	$registerNotebookProvider(providerId: string, handle: number): void;
 	$unregisterNotebookProvider(handle: number): void;
 }
+

src/sqltest/utils/testUtils.ts

@@ -0,0 +1,22 @@
+
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+'use strict';
+
+import * as assert from 'assert';
+
+export async function assertThrowsAsync(fn, regExp?: string): Promise<void> {
+    let f = () => {
+        // Empty
+    };
+    try {
+      await fn();
+    } catch (e) {
+      f = () => { throw e; };
+    } finally {
+      assert.throws(f, regExp);
+    }
+}

src/sqltest/workbench/api/extHostBackgroundTaskManagement.test.ts

@@ -63,7 +63,7 @@ suite('ExtHostBackgroundTaskManagement Tests', () => {
 			operation: (op: sqlops.BackgroundOperation) => {
 				op.onCanceled(() => {
 					op.updateStatus(TaskStatus.Canceled);
-				})
+				});
 			},
 			operationId: operationId
 		};

src/sqltest/workbench/api/exthostNotebook.test.ts

@@ -0,0 +1,146 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+'use strict';
+
+import * as sqlops from 'sqlops';
+import * as vscode from 'vscode';
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+
+import URI from 'vs/base/common/uri';
+import { IMainContext } from 'vs/workbench/api/node/extHost.protocol';
+
+import { ExtHostNotebook } from 'sql/workbench/api/node/extHostNotebook';
+import { MainThreadNotebookShape } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import * as testUtils from '../../utils/testUtils';
+import { INotebookManagerDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+suite('ExtHostNotebook Tests', () => {
+
+	let extHostNotebook: ExtHostNotebook;
+	let mockProxy: TypeMoq.Mock<MainThreadNotebookShape>;
+	let notebookUri: URI;
+	let notebookProviderMock: TypeMoq.Mock<NotebookProviderStub>;
+	setup(() => {
+		mockProxy = TypeMoq.Mock.ofInstance(<MainThreadNotebookShape> {
+			$registerNotebookProvider: (providerId, handle) => undefined,
+			$unregisterNotebookProvider: (handle) => undefined,
+			dispose: () => undefined
+		});
+		let mainContext = <IMainContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+		extHostNotebook = new ExtHostNotebook(mainContext);
+		notebookUri = URI.parse('file:/user/default/my.ipynb');
+		notebookProviderMock = TypeMoq.Mock.ofType(NotebookProviderStub, TypeMoq.MockBehavior.Loose);
+		notebookProviderMock.callBase = true;
+	});
+
+	suite('getNotebookManager', () => {
+		test('Should throw if no matching provider is defined', async () => {
+			await testUtils.assertThrowsAsync(() => extHostNotebook.$getNotebookManager(-1, notebookUri));
+		});
+		suite('with provider', () => {
+			let providerHandle: number = -1;
+
+			setup(() => {
+				mockProxy.setup(p =>
+					p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId), TypeMoq.It.isAnyNumber()))
+					.returns((providerId, handle) => {
+						providerHandle = handle;
+						return undefined;
+					});
+
+				// Register the provider so we can test behavior with this present
+				extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			});
+
+			test('Should return a notebook manager with correct info on content and server manager existence', async () => {
+				// Given the provider returns a manager with no
+				let expectedManager = new NotebookManagerStub();
+				notebookProviderMock.setup(p => p.getNotebookManager(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedManager));
+
+				// When I call through using the handle provided during registration
+				let managerDetails: INotebookManagerDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+
+				// Then I expect the same manager to be returned
+				assert.ok(managerDetails.hasContentManager === false, 'Expect no content manager defined');
+				assert.ok(managerDetails.hasServerManager === false, 'Expect no server manager defined');
+				assert.ok(managerDetails.handle > 0, 'Expect a valid handle defined');
+			});
+
+			test('Should have a unique handle for each notebook URI', async () => {
+				// Given the we request 2 URIs
+				let expectedManager = new NotebookManagerStub();
+				notebookProviderMock.setup(p => p.getNotebookManager(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedManager));
+
+				// When I call through using the handle provided during registration
+				let originalManagerDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+				let differentDetails = await extHostNotebook.$getNotebookManager(providerHandle, URI.parse('file://other/file.ipynb'));
+				let sameDetails = await extHostNotebook.$getNotebookManager(providerHandle, notebookUri);
+
+				// Then I expect the 2 different handles in the managers returned.
+				// This is because we can't easily track identity of the managers, so just track which one is assigned to
+				// a notebook by the handle ID
+				assert.notEqual(originalManagerDetails.handle, differentDetails.handle, 'Should have unique handle for each manager');
+				assert.equal(originalManagerDetails.handle, sameDetails.handle, 'Should have same handle when same URI is passed in');
+
+			});
+		});
+	});
+
+	suite('registerNotebookProvider', () => {
+		let savedHandle: number = -1;
+		setup(() => {
+			mockProxy.setup(p =>
+				p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId), TypeMoq.It.isAnyNumber()))
+				.returns((providerId, handle) => {
+					savedHandle = handle;
+					return undefined;
+				});
+		});
+
+		test('Should register with a new handle to the proxy', () => {
+			extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			mockProxy.verify(p =>
+				p.$registerNotebookProvider(TypeMoq.It.isValue(notebookProviderMock.object.providerId),
+				TypeMoq.It.isAnyNumber()), TypeMoq.Times.once());
+				// It shouldn't unregister until requested
+				mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.never());
+
+		});
+
+		test('Should call unregister on disposing', () => {
+			let disposable = extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
+			disposable.dispose();
+			mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.once());
+		});
+	});
+});
+
+class NotebookProviderStub implements sqlops.nb.NotebookProvider {
+	providerId: string = 'TestProvider';
+
+	getNotebookManager(notebookUri: vscode.Uri): Thenable<sqlops.nb.NotebookManager> {
+		throw new Error('Method not implemented.');
+	}
+	handleNotebookClosed(notebookUri: vscode.Uri): void {
+		throw new Error('Method not implemented.');
+	}
+}
+
+class NotebookManagerStub implements sqlops.nb.NotebookManager {
+	get contentManager(): sqlops.nb.ContentManager {
+		return undefined;
+	}
+
+	get sessionManager(): sqlops.nb.SessionManager {
+		return undefined;
+	}
+
+	get serverManager(): sqlops.nb.ServerManager {
+		return undefined;
+	}
+}

src/sqltest/workbench/api/mainThreadNotebook.test.ts

@@ -0,0 +1,123 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+'use strict';
+
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+
+import URI from 'vs/base/common/uri';
+import { IExtHostContext } from 'vs/workbench/api/node/extHost.protocol';
+
+import { ExtHostNotebookShape } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import { MainThreadNotebook } from 'sql/workbench/api/node/mainThreadNotebook';
+import { NotebookService } from 'sql/services/notebook/notebookServiceImpl';
+import { INotebookProvider } from 'sql/services/notebook/notebookService';
+import { INotebookManagerDetails } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { LocalContentManager } from 'sql/services/notebook/localContentManager';
+
+
+suite('MainThreadNotebook Tests', () => {
+
+	let mainThreadNotebook: MainThreadNotebook;
+	let mockProxy: TypeMoq.Mock<ExtHostNotebookShape>;
+	let notebookUri: URI;
+	let mockNotebookService: TypeMoq.Mock<NotebookService>;
+	let providerId = 'TestProvider';
+	setup(() => {
+		mockProxy = TypeMoq.Mock.ofInstance(<ExtHostNotebookShape> {
+			$getNotebookManager: (handle, uri) => undefined,
+			$handleNotebookClosed: (uri) => undefined,
+			$getNotebookContents: (handle, uri) => undefined,
+			$save: (handle, uri, notebook) => undefined,
+			$doStartServer: (handle) => undefined,
+			$doStopServer: (handle) => undefined
+		});
+		let extContext = <IExtHostContext>{
+			getProxy: proxyType => mockProxy.object
+		};
+		mockNotebookService = TypeMoq.Mock.ofType(NotebookService);
+		notebookUri = URI.parse('file:/user/default/my.ipynb');
+		mainThreadNotebook = new MainThreadNotebook(extContext, mockNotebookService.object);
+	});
+
+	suite('On registering a provider', () => {
+		let provider: INotebookProvider;
+		let registeredProviderId: string;
+		setup(() => {
+			mockNotebookService.setup(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns((id, providerImpl) => {
+				registeredProviderId = id;
+				provider = providerImpl;
+			});
+		});
+
+		test('should call through to notebook service', () => {
+			// When I register a provider
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+			// Then I expect a provider implementation to be passed to the service
+			mockNotebookService.verify(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny()), TypeMoq.Times.once());
+			assert.equal(provider.providerId, providerId);
+		});
+		test('should unregister in service', () => {
+			// Given we have a provider
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+			// When I unregister a provider twice
+			mainThreadNotebook.$unregisterNotebookProvider(1);
+			mainThreadNotebook.$unregisterNotebookProvider(1);
+			// Then I expect it to be unregistered in the service just 1 time
+			mockNotebookService.verify(s => s.unregisterProvider(TypeMoq.It.isValue(providerId)), TypeMoq.Times.once());
+		});
+	});
+
+	suite('getNotebookManager', () => {
+		let managerWithAllFeatures: INotebookManagerDetails;
+		let provider: INotebookProvider;
+
+		setup(() => {
+			managerWithAllFeatures = {
+				handle: 2,
+				hasContentManager: true,
+				hasServerManager: true
+			};
+			mockNotebookService.setup(s => s.registerProvider(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns((id, providerImpl) => {
+				provider = providerImpl;
+			});
+			mainThreadNotebook.$registerNotebookProvider(providerId, 1);
+
+		});
+
+		test('should return manager with default content manager & undefined server manager if extension host has none', async () => {
+			// Given the extension provider doesn't have acontent or server manager
+			let details: INotebookManagerDetails = {
+				handle: 2,
+				hasContentManager: false,
+				hasServerManager: false
+			};
+			mockProxy.setup(p => p.$getNotebookManager(TypeMoq.It.isAnyNumber(), TypeMoq.It.isValue(notebookUri)))
+			.returns(() => Promise.resolve(details));
+
+			// When I get the notebook manager
+			let manager = await provider.getNotebookManager(notebookUri);
+
+			// Then it should use the built-in content manager
+			assert.ok(manager.contentManager instanceof LocalContentManager);
+			// And it should not define a server manager
+			assert.equal(manager.serverManager, undefined);
+		});
+
+		test('should return manager with a content & server manager if extension host has these', async () => {
+			// Given the extension provider doesn't have acontent or server manager
+			mockProxy.setup(p => p.$getNotebookManager(TypeMoq.It.isAnyNumber(), TypeMoq.It.isValue(notebookUri)))
+			.returns(() => Promise.resolve(managerWithAllFeatures));
+
+			// When I get the notebook manager
+			let manager = await provider.getNotebookManager(notebookUri);
+
+			// Then it shouldn't have wrappers for the content or server manager
+			assert.ok(!(manager.contentManager instanceof LocalContentManager));
+			assert.notEqual(manager.serverManager, undefined);
+		});
+	});
+
+});