Implement Session support through the extension host (#3228)

Full plumb through of Session support. Also fixed some test issues

- Load session and get necessary information in kernels list
- Run Cell button now works as expected
- Added a ToggleAction base class which can be used for anything that switches icons. I'd still prefer to have this be dynamic and as clean as the extension classes
- Fixed account test unhandled promise rejections (caused by incorrect / invalid tests) that made it hard to see all the test run output.

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

@@ -6,24 +6,25 @@
 
 import * as sqlops from 'sqlops';
 import * as vscode from 'vscode';
+
 import { TPromise } from 'vs/base/common/winjs.base';
 import { IMainContext } from 'vs/workbench/api/node/extHost.protocol';
 import { Disposable } from 'vs/workbench/api/node/extHostTypes';
 import { localize } from 'vs/nls';
-
+import URI, { UriComponents } from 'vs/base/common/uri';
 
 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';
+import { INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, INotebookFutureDetails, FutureMessageType } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+type Adapter = sqlops.nb.NotebookProvider | sqlops.nb.NotebookManager | sqlops.nb.ISession | sqlops.nb.IKernel | sqlops.nb.IFuture;
 
 export class ExtHostNotebook implements ExtHostNotebookShape {
 	private static _handlePool: number = 0;
 
 	private readonly _proxy: MainThreadNotebookShape;
-	private _providers = new Map<number, sqlops.nb.NotebookProvider>();
+	private _adapters = new Map<number, Adapter>();
 	// Notebook URI to manager lookup.
-	private _managers = new Map<number, NotebookManagerAdapter>();
-	constructor(private _mainContext: IMainContext) {
+	constructor(_mainContext: IMainContext) {
 		this._proxy = _mainContext.getProxy(SqlMainContext.MainThreadNotebook);
 	}
 
@@ -39,7 +40,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		}
 
 		return {
-			handle: adapter.managerHandle,
+			handle: adapter.handle,
 			hasContentManager: !!adapter.contentManager,
 			hasServerManager: !!adapter.serverManager
 		};
@@ -50,7 +51,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		let manager = this.findManagerForUri(uriString);
 		if (manager) {
 			manager.provider.handleNotebookClosed(uri);
-			this._managers.delete(manager.managerHandle);
+			// Note: deliberately not removing handle.
 		}
 	}
 
@@ -70,6 +71,124 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		return this._withContentManager(managerHandle, (contentManager) => contentManager.save(URI.revive(notebookUri), notebook));
 	}
 
+	$refreshSpecs(managerHandle: number): Thenable<sqlops.nb.IAllKernels> {
+		return this._withSessionManager(managerHandle, async (sessionManager) => {
+			await sessionManager.ready;
+			return sessionManager.specs;
+		});
+	}
+
+	$startNewSession(managerHandle: number, options: sqlops.nb.ISessionOptions): Thenable<INotebookSessionDetails> {
+		return this._withSessionManager(managerHandle, async (sessionManager) => {
+			let session = await sessionManager.startNew(options);
+			let sessionId = this._addNewAdapter(session);
+			let kernelDetails: INotebookKernelDetails = undefined;
+			if (session.kernel) {
+				kernelDetails = this.saveKernel(session.kernel);
+			}
+			let details: INotebookSessionDetails = {
+				sessionId: sessionId,
+				id: session.id,
+				path: session.path,
+				name: session.name,
+				type: session.type,
+				status: session.status,
+				canChangeKernels: session.canChangeKernels,
+				kernelDetails: kernelDetails
+			};
+			return details;
+		});
+	}
+
+	private saveKernel(kernel: sqlops.nb.IKernel): INotebookKernelDetails {
+		let kernelId = this._addNewAdapter(kernel);
+		let kernelDetails: INotebookKernelDetails = {
+			kernelId: kernelId,
+			id: kernel.id,
+			info: kernel.info,
+			name: kernel.name,
+			supportsIntellisense: kernel.supportsIntellisense
+		};
+		return kernelDetails;
+	}
+
+	$shutdownSession(managerHandle: number, sessionId: string): Thenable<void> {
+		return this._withSessionManager(managerHandle, async (sessionManager) => {
+			return sessionManager.shutdown(sessionId);
+		});
+	}
+
+	$changeKernel(sessionId: number, kernelInfo: sqlops.nb.IKernelSpec): Thenable<INotebookKernelDetails> {
+		let session = this._getAdapter<sqlops.nb.ISession>(sessionId);
+		return session.changeKernel(kernelInfo).then(kernel => this.saveKernel(kernel));
+	}
+
+	$getKernelReadyStatus(kernelId: number): Thenable<sqlops.nb.IInfoReply> {
+		let kernel = this._getAdapter<sqlops.nb.IKernel>(kernelId);
+		return kernel.ready.then(success => kernel.info);
+	}
+
+	$getKernelSpec(kernelId: number): Thenable<sqlops.nb.IKernelSpec> {
+		let kernel = this._getAdapter<sqlops.nb.IKernel>(kernelId);
+		return kernel.getSpec();
+	}
+
+	$requestComplete(kernelId: number, content: sqlops.nb.ICompleteRequest): Thenable<sqlops.nb.ICompleteReplyMsg> {
+		let kernel = this._getAdapter<sqlops.nb.IKernel>(kernelId);
+		return kernel.requestComplete(content);
+	}
+
+	$requestExecute(kernelId: number, content: sqlops.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable<INotebookFutureDetails> {
+		let kernel = this._getAdapter<sqlops.nb.IKernel>(kernelId);
+		let future = kernel.requestExecute(content, disposeOnDone);
+		let futureId = this._addNewAdapter(future);
+		this.hookFutureDone(futureId, future);
+		this.hookFutureMessages(futureId, future);
+		return Promise.resolve({
+			futureId: futureId,
+			msg: future.msg
+		});
+	}
+
+	private hookFutureDone(futureId: number, future: sqlops.nb.IFuture): void {
+		future.done.then(success => {
+			return this._proxy.$onFutureDone(futureId, { succeeded: true, rejectReason: undefined });
+		}, err => {
+			let rejectReason: string;
+			if (typeof err === 'string') {
+				rejectReason = err;
+			}
+			else if (err instanceof Error && typeof err.message === 'string') {
+				rejectReason = err.message;
+			}
+			else {
+				rejectReason = err;
+			}
+			return this._proxy.$onFutureDone(futureId, { succeeded: false, rejectReason: rejectReason });
+		});
+	}
+
+	private hookFutureMessages(futureId: number, future: sqlops.nb.IFuture): void {
+		future.setReplyHandler({ handle: (msg) => this._proxy.$onFutureMessage(futureId, FutureMessageType.Reply, msg) });
+		future.setStdInHandler({ handle: (msg) => this._proxy.$onFutureMessage(futureId, FutureMessageType.StdIn, msg) });
+		future.setIOPubHandler({ handle: (msg) => this._proxy.$onFutureMessage(futureId, FutureMessageType.IOPub, msg) });
+	}
+
+	$interruptKernel(kernelId: number): Thenable<void> {
+		let kernel = this._getAdapter<sqlops.nb.IKernel>(kernelId);
+		return kernel.interrupt();
+	}
+
+	$sendInputReply(futureId: number, content: sqlops.nb.IInputReply): void {
+		let future = this._getAdapter<sqlops.nb.IFuture>(futureId);
+		return future.sendInputReply(content);
+	}
+
+	$disposeFuture(futureId: number): void {
+		let future = this._getAdapter<sqlops.nb.IFuture>(futureId);
+		future.dispose();
+	}
+
 	//#endregion
 
 	//#region APIs called by extensions
@@ -77,7 +196,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		if (!provider || !provider.providerId) {
 			throw new Error(localize('providerRequired', 'A NotebookProvider with valid providerId must be passed to this method'));
 		}
-		const handle = this._addNewProvider(provider);
+		const handle = this._addNewAdapter(provider);
 		this._proxy.$registerNotebookProvider(provider.providerId, handle);
 		return this._createDisposable(handle);
 	}
@@ -86,8 +205,18 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 
 	//#region private methods
 
+	private getAdapters<A>(ctor: { new(...args: any[]): A }): A[] {
+		let matchingAdapters = [];
+		this._adapters.forEach(a => {
+			if (a instanceof ctor) {
+				matchingAdapters.push(a);
+			}
+		});
+		return matchingAdapters;
+	}
+
 	private findManagerForUri(uriString: string): NotebookManagerAdapter {
-		for(let manager of Array.from(this._managers.values())) {
+		for(let manager of this.getAdapters(NotebookManagerAdapter)) {
 			if (manager.uriString === uriString) {
 				return manager;
 			}
@@ -98,15 +227,14 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	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);
+		let adapter = new NotebookManagerAdapter(provider, manager, uriString);
+		adapter.handle = this._addNewAdapter(adapter);
 		return adapter;
 	}
 
 	private _createDisposable(handle: number): Disposable {
 		return new Disposable(() => {
-			this._providers.delete(handle);
+			this._adapters.delete(handle);
 			this._proxy.$unregisterNotebookProvider(handle);
 		});
 	}
@@ -116,7 +244,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	}
 
 	private _withProvider<R>(handle: number, callback: (provider: sqlops.nb.NotebookProvider) => R | PromiseLike<R>): TPromise<R> {
-		let provider = this._providers.get(handle);
+		let provider = this._adapters.get(handle) as sqlops.nb.NotebookProvider;
 		if (provider === undefined) {
 			return TPromise.wrapError<R>(new Error(localize('errNoProvider', 'no notebook provider found')));
 		}
@@ -124,7 +252,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	}
 
 	private _withNotebookManager<R>(handle: number, callback: (manager: NotebookManagerAdapter) => R | PromiseLike<R>): TPromise<R> {
-		let manager = this._managers.get(handle);
+		let manager = this._adapters.get(handle) as NotebookManagerAdapter;
 		if (manager === undefined) {
 			return TPromise.wrapError<R>(new Error(localize('errNoManager', 'No Manager found')));
 		}
@@ -161,19 +289,28 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		});
 	}
 
-	private _addNewProvider(adapter: sqlops.nb.NotebookProvider): number {
+	private _addNewAdapter(adapter: Adapter): number {
 		const handle = this._nextHandle();
-		this._providers.set(handle, adapter);
+		this._adapters.set(handle, adapter);
 		return handle;
 	}
+
+	private _getAdapter<T>(id: number): T {
+		let adapter = <T><any>this._adapters.get(id);
+		if (adapter === undefined) {
+			throw new Error('No adapter found');
+		}
+		return adapter;
+	}
+
 	//#endregion
 }
 
 
 class NotebookManagerAdapter implements sqlops.nb.NotebookManager {
+	public handle: number;
 	constructor(
 		public readonly provider: sqlops.nb.NotebookProvider,
-		public readonly managerHandle: number,
 		private manager: sqlops.nb.NotebookManager,
 		public readonly uriString: string
 	) {
@@ -190,5 +327,4 @@ class NotebookManagerAdapter implements sqlops.nb.NotebookManager {
 	public get serverManager(): sqlops.nb.ServerManager {
 		return this.manager.serverManager;
 	}
-
-}
+}

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

@@ -13,14 +13,17 @@ 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 { INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, FutureMessageType, INotebookFutureDetails, INotebookFutureDone } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { LocalContentManager } from 'sql/services/notebook/localContentManager';
+import { Deferred } from 'sql/base/common/promise';
+import { FutureInternal } from 'sql/parts/notebook/models/modelInterfaces';
 
 @extHostNamedCustomer(SqlMainContext.MainThreadNotebook)
 export class MainThreadNotebook extends Disposable implements MainThreadNotebookShape {
 
 	private _proxy: ExtHostNotebookShape;
 	private _providers = new Map<number, NotebookProviderWrapper>();
+	private _futures = new Map<number, FutureWrapper>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -32,9 +35,21 @@ export class MainThreadNotebook extends Disposable implements MainThreadNotebook
 		}
 	}
 
+	public addFuture(futureId: number, future: FutureWrapper): void {
+		this._futures.set(futureId, future);
+	}
+
+	public disposeFuture(futureId: number): void {
+		this._futures.delete(futureId);
+	}
+
 	//#region Extension host callable methods
 	public $registerNotebookProvider(providerId: string, handle: number): void {
-		let notebookProvider = new NotebookProviderWrapper(this._proxy, providerId, handle);
+		let proxy: Proxies = {
+			main: this,
+			ext: this._proxy
+		};
+		let notebookProvider = new NotebookProviderWrapper(proxy, providerId, handle);
 		this._providers.set(handle, notebookProvider);
 		this.notebookService.registerProvider(providerId, notebookProvider);
 	}
@@ -48,14 +63,32 @@ export class MainThreadNotebook extends Disposable implements MainThreadNotebook
 		}
 	}
 
-	//#endregion
+	public $onFutureMessage(futureId: number, type: FutureMessageType, payload: sqlops.nb.IMessage): void {
+		let future = this._futures.get(futureId);
+		if (future) {
+			future.onMessage(type, payload);
+		}
+	}
 
+	public $onFutureDone(futureId: number, done: INotebookFutureDone): void {
+		let future = this._futures.get(futureId);
+		if (future) {
+			future.onDone(done);
+		}
+
+	}
+	//#endregion
+}
+
+interface Proxies {
+	main: MainThreadNotebook;
+	ext: ExtHostNotebookShape;
 }
 
 class NotebookProviderWrapper extends Disposable implements INotebookProvider {
-	private _managers = new Map<string, NotebookManagerWrapper>();
+	private _notebookUriToManagerMap = new Map<string, NotebookManagerWrapper>();
 
-	constructor(private _proxy: ExtHostNotebookShape, public readonly providerId, public readonly providerHandle: number) {
+	constructor(private _proxy: Proxies, public readonly providerId, public readonly providerHandle: number) {
 		super();
 	}
 
@@ -66,19 +99,19 @@ class NotebookProviderWrapper extends Disposable implements INotebookProvider {
 
 	private async doGetNotebookManager(notebookUri: URI): Promise<INotebookManager> {
 		let uriString = notebookUri.toString();
-		let manager = this._managers.get(uriString);
+		let manager = this._notebookUriToManagerMap.get(uriString);
 		if (!manager) {
 			manager = new NotebookManagerWrapper(this._proxy, this.providerId, notebookUri);
 			await manager.initialize(this.providerHandle);
-			this._managers.set(uriString, manager);
+			this._notebookUriToManagerMap.set(uriString, manager);
 		}
 		return manager;
 	}
 
 	handleNotebookClosed(notebookUri: URI): void {
-		this._proxy.$handleNotebookClosed(notebookUri);
+		this._notebookUriToManagerMap.delete(notebookUri.toString());
+		this._proxy.ext.$handleNotebookClosed(notebookUri);
 	}
-
 }
 
 class NotebookManagerWrapper implements INotebookManager {
@@ -87,13 +120,13 @@ class NotebookManagerWrapper implements INotebookManager {
 	private _serverManager: sqlops.nb.ServerManager;
 	private managerDetails: INotebookManagerDetails;
 
-	constructor(private _proxy: ExtHostNotebookShape,
+	constructor(private _proxy: Proxies,
 		public readonly providerId,
 		private notebookUri: URI
 	) { }
 
 	public async initialize(providerHandle: number): Promise<NotebookManagerWrapper> {
-		this.managerDetails = await this._proxy.$getNotebookManager(providerHandle, this.notebookUri);
+		this.managerDetails = await this._proxy.ext.$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;
@@ -114,26 +147,25 @@ class NotebookManagerWrapper implements INotebookManager {
 	public get managerHandle(): number {
 		return this.managerDetails.handle;
 	}
-
 }
 
 class ContentManagerWrapper implements sqlops.nb.ContentManager {
 
-	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+	constructor(private handle: number, private _proxy: Proxies) {
 	}
 	getNotebookContents(notebookUri: URI): Thenable<sqlops.nb.INotebook> {
-		return this._proxy.$getNotebookContents(this.handle, notebookUri);
+		return this._proxy.ext.$getNotebookContents(this.handle, notebookUri);
 	}
 
 	save(path: URI, notebook: sqlops.nb.INotebook): Thenable<sqlops.nb.INotebook> {
-		return this._proxy.$save(this.handle, path, notebook);
+		return this._proxy.ext.$save(this.handle, path, notebook);
 	}
 }
 
 class ServerManagerWrapper implements sqlops.nb.ServerManager {
-	private onServerStartedEmitter: Emitter<void>;
+	private onServerStartedEmitter = new Emitter<void>();
 	private _isStarted: boolean;
-	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+	constructor(private handle: number, private _proxy: Proxies) {
 		this._isStarted = false;
 	}
 
@@ -150,7 +182,7 @@ class ServerManagerWrapper implements sqlops.nb.ServerManager {
 	}
 
 	private async doStartServer(): Promise<void> {
-		await this._proxy.$doStartServer(this.handle);
+		await this._proxy.ext.$doStartServer(this.handle);
 		this._isStarted = true;
 		this.onServerStartedEmitter.fire();
 	}
@@ -161,7 +193,7 @@ class ServerManagerWrapper implements sqlops.nb.ServerManager {
 
 	private async doStopServer(): Promise<void> {
 		try {
-			await this._proxy.$doStopServer(this.handle);
+			await this._proxy.ext.$doStopServer(this.handle);
 		} finally {
 			// Always consider this a stopping event, even if a failure occurred.
 			this._isStarted = false;
@@ -170,30 +202,255 @@ class ServerManagerWrapper implements sqlops.nb.ServerManager {
 }
 
 class SessionManagerWrapper implements sqlops.nb.SessionManager {
-	constructor(private handle: number, private _proxy: ExtHostNotebookShape) {
+	private readyPromise: Promise<void>;
+	private _isReady: boolean;
+	private _specs: sqlops.nb.IAllKernels;
+	constructor(private managerHandle: number, private _proxy: Proxies) {
+		this._isReady = false;
+		this.readyPromise = this.initializeSessionManager();
 	}
 
 	get isReady(): boolean {
-		throw new Error('Method not implemented.');
-
+		return this._isReady;
 	}
 
 	get ready(): Thenable<void> {
-		throw new Error('Method not implemented.');
-
+		return this.readyPromise;
 	}
-	get specs(): sqlops.nb.IAllKernels {
-		throw new Error('Method not implemented.');
 
+	get specs(): sqlops.nb.IAllKernels {
+		return this._specs;
 	}
 
 	startNew(options: sqlops.nb.ISessionOptions): Thenable<sqlops.nb.ISession> {
-		throw new Error('Method not implemented.');
+		return this.doStartNew(options);
+	}
+
+	private async doStartNew(options: sqlops.nb.ISessionOptions): Promise<sqlops.nb.ISession> {
+		let sessionDetails = await this._proxy.ext.$startNewSession(this.managerHandle, options);
+		return new SessionWrapper(this._proxy, sessionDetails);
 	}
 
 	shutdown(id: string): Thenable<void> {
-		throw new Error('Method not implemented.');
+		return this._proxy.ext.$shutdownSession(this.managerHandle, id);
 	}
 
+	private async initializeSessionManager(): Promise<void> {
+		await this.refreshSpecs();
+		this._isReady = true;
+	}
 
-}
+	private async refreshSpecs(): Promise<void> {
+		let specs = await this._proxy.ext.$refreshSpecs(this.managerHandle);
+		if (specs) {
+			this._specs = specs;
+		}
+	}
+}
+
+class SessionWrapper implements sqlops.nb.ISession {
+	private _kernel: KernelWrapper;
+	constructor(private _proxy: Proxies, private sessionDetails: INotebookSessionDetails) {
+		if (sessionDetails && sessionDetails.kernelDetails) {
+			this._kernel = new KernelWrapper(_proxy, sessionDetails.kernelDetails);
+		}
+	}
+
+	get canChangeKernels(): boolean {
+		return this.sessionDetails.canChangeKernels;
+	}
+
+	get id(): string {
+		return this.sessionDetails.id;
+	}
+
+	get path(): string {
+		return this.sessionDetails.path;
+	}
+
+	get name(): string {
+		return this.sessionDetails.name;
+	}
+
+	get type(): string {
+		return this.sessionDetails.type;
+	}
+
+	get status(): sqlops.nb.KernelStatus {
+		return this.sessionDetails.status as sqlops.nb.KernelStatus;
+	}
+
+	get kernel(): sqlops.nb.IKernel {
+		return this._kernel;
+	}
+
+	changeKernel(kernelInfo: sqlops.nb.IKernelSpec): Thenable<sqlops.nb.IKernel> {
+		return this.doChangeKernel(kernelInfo);
+	}
+
+	private async doChangeKernel(kernelInfo: sqlops.nb.IKernelSpec): Promise<sqlops.nb.IKernel> {
+		let kernelDetails = await this._proxy.ext.$changeKernel(this.sessionDetails.sessionId, kernelInfo);
+		this._kernel = new KernelWrapper(this._proxy, kernelDetails);
+		return this._kernel;
+	}
+}
+
+class KernelWrapper implements sqlops.nb.IKernel {
+	private _isReady: boolean = false;
+	private _ready = new Deferred<void>();
+	private _info: sqlops.nb.IInfoReply;
+	constructor(private _proxy: Proxies, private kernelDetails: INotebookKernelDetails) {
+		this.initialize(kernelDetails);
+	}
+
+	private async initialize(kernelDetails: INotebookKernelDetails): Promise<void> {
+		try {
+			this._info = await this._proxy.ext.$getKernelReadyStatus(kernelDetails.kernelId);
+			this._isReady = true;
+			this._ready.resolve();
+		} catch (error) {
+			this._isReady = false;
+			this._ready.reject(error);
+		}
+	}
+
+	get isReady(): boolean {
+		return this._isReady;
+	}
+	get ready(): Thenable<void> {
+		return this._ready.promise;
+	}
+
+	get id(): string {
+		return this.kernelDetails.id;
+	}
+
+	get name(): string {
+		return this.kernelDetails.name;
+	}
+
+	get supportsIntellisense(): boolean {
+		return this.kernelDetails.supportsIntellisense;
+	}
+
+	get info(): sqlops.nb.IInfoReply {
+		return this._info;
+	}
+
+	getSpec(): Thenable<sqlops.nb.IKernelSpec> {
+		return this._proxy.ext.$getKernelSpec(this.kernelDetails.kernelId);
+	}
+
+	requestComplete(content: sqlops.nb.ICompleteRequest): Thenable<sqlops.nb.ICompleteReplyMsg> {
+		return this._proxy.ext.$requestComplete(this.kernelDetails.kernelId, content);
+	}
+
+	requestExecute(content: sqlops.nb.IExecuteRequest, disposeOnDone?: boolean): sqlops.nb.IFuture {
+		let future = new FutureWrapper(this._proxy);
+		this._proxy.ext.$requestExecute(this.kernelDetails.kernelId, content, disposeOnDone)
+		.then(details => {
+			future.setDetails(details);
+			// Save the future in the main thread notebook so extension can call through and reference it
+			this._proxy.main.addFuture(details.futureId, future);
+		}, error => future.setError(error));
+		return future;
+	}
+
+	interrupt(): Thenable<void> {
+		return this._proxy.ext.$interruptKernel(this.kernelDetails.kernelId);
+	}
+}
+
+
+class FutureWrapper implements FutureInternal {
+	private _futureId: number;
+	private _done = new Deferred<sqlops.nb.IShellMessage>();
+	private _messageHandlers = new Map<FutureMessageType, sqlops.nb.MessageHandler<sqlops.nb.IMessage>>();
+	private _msg: sqlops.nb.IMessage;
+	private _inProgress: boolean;
+
+	constructor(private _proxy: Proxies) {
+        this._inProgress = true;
+	}
+
+	public setDetails(details: INotebookFutureDetails): void {
+		this._futureId = details.futureId;
+		this._msg = details.msg;
+	}
+
+	public setError(error: Error | string): void {
+		this._done.reject(error);
+	}
+
+	public onMessage(type: FutureMessageType, payload: sqlops.nb.IMessage): void {
+		let handler = this._messageHandlers.get(type);
+		if (handler) {
+			try {
+				handler.handle(payload);
+			} catch (error) {
+				// TODO log errors from the handler
+			}
+		}
+	}
+
+	public onDone(done: INotebookFutureDone): void {
+		this._inProgress = false;
+		if (done.succeeded) {
+			this._done.resolve();
+		} else {
+			this._done.reject(new Error(done.rejectReason));
+		}
+	}
+
+	private addMessageHandler(type: FutureMessageType, handler: sqlops.nb.MessageHandler<sqlops.nb.IMessage>): void {
+		// Note: there can only be 1 message handler according to the Jupyter Notebook spec.
+		// You can use a message hook to override this / add additional side-processors
+		this._messageHandlers.set(type, handler);
+	}
+
+	//#region Public APIs
+	get inProgress(): boolean {
+		return this._inProgress;
+	}
+
+	set inProgress(value: boolean) {
+		this._inProgress = value;
+	}
+
+	get msg(): sqlops.nb.IMessage {
+		return this._msg;
+	}
+
+	get done(): Thenable<sqlops.nb.IShellMessage> {
+		return this._done.promise;
+	}
+
+	setReplyHandler(handler: sqlops.nb.MessageHandler<sqlops.nb.IShellMessage>): void {
+		this.addMessageHandler(FutureMessageType.Reply, handler);
+	}
+
+	setStdInHandler(handler: sqlops.nb.MessageHandler<sqlops.nb.IStdinMessage>): void {
+		this.addMessageHandler(FutureMessageType.StdIn, handler);
+	}
+
+	setIOPubHandler(handler: sqlops.nb.MessageHandler<sqlops.nb.IIOPubMessage>): void {
+		this.addMessageHandler(FutureMessageType.IOPub, handler);
+	}
+
+	sendInputReply(content: sqlops.nb.IInputReply): void {
+		this._proxy.ext.$sendInputReply(this._futureId, content);
+	}
+
+	dispose() {
+		this._proxy.main.disposeFuture(this._futureId);
+		this._proxy.ext.$disposeFuture(this._futureId);
+	}
+
+	registerMessageHook(hook: (msg: sqlops.nb.IIOPubMessage) => boolean | Thenable<boolean>): void {
+		throw new Error('Method not implemented.');
+	}
+	removeMessageHook(hook: (msg: sqlops.nb.IIOPubMessage) => boolean | Thenable<boolean>): void {
+		throw new Error('Method not implemented.');
+	}
+	//#endregion
+}

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, INotebookManagerDetails
+	IModelViewWizardDetails, IModelViewWizardPageDetails, INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, INotebookFutureDetails, FutureMessageType, INotebookFutureDone
 } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export abstract class ExtHostAccountManagementShape {
@@ -717,15 +717,39 @@ export interface ExtHostNotebookShape {
 	 */
 	$getNotebookManager(providerHandle: number, notebookUri: UriComponents): Thenable<INotebookManagerDetails>;
 	$handleNotebookClosed(notebookUri: UriComponents): void;
+
+	// Server Manager APIs
 	$doStartServer(managerHandle: number): Thenable<void>;
 	$doStopServer(managerHandle: number): Thenable<void>;
+
+	// Content Manager APIs
 	$getNotebookContents(managerHandle: number, notebookUri: UriComponents): Thenable<sqlops.nb.INotebook>;
 	$save(managerHandle: number, notebookUri: UriComponents, notebook: sqlops.nb.INotebook): Thenable<sqlops.nb.INotebook>;
 
+	// Session Manager APIs
+	$refreshSpecs(managerHandle: number): Thenable<sqlops.nb.IAllKernels>;
+	$startNewSession(managerHandle: number, options: sqlops.nb.ISessionOptions): Thenable<INotebookSessionDetails>;
+	$shutdownSession(managerHandle: number, sessionId: string): Thenable<void>;
+
+	// Session APIs
+	$changeKernel(sessionId: number, kernelInfo: sqlops.nb.IKernelSpec): Thenable<INotebookKernelDetails>;
+
+	// Kernel APIs
+	$getKernelReadyStatus(kernelId: number): Thenable<sqlops.nb.IInfoReply>;
+	$getKernelSpec(kernelId: number): Thenable<sqlops.nb.IKernelSpec>;
+	$requestComplete(kernelId: number, content: sqlops.nb.ICompleteRequest): Thenable<sqlops.nb.ICompleteReplyMsg>;
+	$requestExecute(kernelId: number, content: sqlops.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable<INotebookFutureDetails>;
+	$interruptKernel(kernelId: number): Thenable<void>;
+
+	// Future APIs
+	$sendInputReply(futureId: number, content: sqlops.nb.IInputReply): void;
+	$disposeFuture(futureId: number): void;
 }
 
 export interface MainThreadNotebookShape extends IDisposable {
 	$registerNotebookProvider(providerId: string, handle: number): void;
 	$unregisterNotebookProvider(handle: number): void;
+	$onFutureMessage(futureId: number, type: FutureMessageType, payload: sqlops.nb.IMessage): void;
+	$onFutureDone(futureId: number, done: INotebookFutureDone): void;
 }