Notebook StdIn support to fix #5231 (#5232)

Fixes #5231 
- Add stdin handling. Has to be at UI level so add plumb through handling
- Add unit tests
- Add new StdIn component.

Testing:
Unit Tests and manual testing of following:
- Prompt for password using `getpass` in python.
   - Password prompt is hidden since "password" is true.
   - Hit enter, it completes
- prompt, stop cell running, StdIn disappears
- prompt, hit escape, stdIn disappears and stdIn request is handled.

Issues: focus isn't always set to the input even though we call focus.
Will investigate this further.

extensions/notebook/src/jupyter/jupyterKernel.ts

@@ -21,7 +21,7 @@ function toStdInMessage(msgImpl: KernelMessage.IStdinMessage): nb.IStdinMessage
 	return {
 		channel: msgImpl.channel,
 		type: msgImpl.channel,
-		content: msgImpl.content,
+		content: <any>msgImpl.content,
 		header: msgImpl.header,
 		parent_header: msgImpl.parent_header,
 		metadata: msgImpl.metadata

src/sql/azdata.proposed.d.ts

@@ -5110,6 +5110,10 @@ declare module 'azdata' {
 		 */
 		export interface IStdinMessage extends IMessage {
 			channel: 'stdin';
+			content: {
+				prompt: string;
+				password: boolean;
+			};
 		}
 
 		/**

src/sql/workbench/parts/notebook/cellViews/codeCell.component.html

@@ -11,5 +11,6 @@
 	<div style="flex: 0 0 auto; width: 100%; height: 100%; display: block">
 		<output-area-component *ngIf="cellModel.outputs && cellModel.outputs.length > 0" [cellModel]="cellModel" [activeCellId]="activeCellId">
 		</output-area-component>
+		<stdin-component *ngIf="isStdInVisible" [onSendInput]="inputDeferred" [stdIn]="stdIn" [cellModel]="cellModel"></stdin-component>
 	</div>
 </div>

src/sql/workbench/parts/notebook/cellViews/codeCell.component.ts

@@ -3,10 +3,12 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { nb } from 'azdata';
 import { OnInit, Component, Input, Inject, forwardRef, ChangeDetectorRef, SimpleChange, OnChanges, HostListener } from '@angular/core';
 import { CellView } from 'sql/workbench/parts/notebook/cellViews/interfaces';
 import { ICellModel } from 'sql/workbench/parts/notebook/models/modelInterfaces';
 import { NotebookModel } from 'sql/workbench/parts/notebook/models/notebookModel';
+import { Deferred } from 'sql/base/common/promise';
 
 
 export const CODE_SELECTOR: string = 'code-cell-component';
@@ -34,6 +36,9 @@ export class CodeCellComponent extends CellView implements OnInit, OnChanges {
 	private _model: NotebookModel;
 	private _activeCellId: string;
 
+	public inputDeferred: Deferred<string>;
+	public stdIn: nb.IStdinMessage;
+
 	constructor(
 		@Inject(forwardRef(() => ChangeDetectorRef)) private _changeRef: ChangeDetectorRef,
 	) {
@@ -45,6 +50,9 @@ export class CodeCellComponent extends CellView implements OnInit, OnChanges {
 			this._register(this.cellModel.onOutputsChanged(() => {
 				this._changeRef.detectChanges();
 			}));
+			// Register request handler, cleanup on dispose of this component
+			this.cellModel.setStdInHandler({ handle: (msg) => this.handleStdIn(msg) });
+			this._register({ dispose: () => this.cellModel.setStdInHandler(undefined) });
 		}
 	}
 
@@ -68,4 +76,32 @@ export class CodeCellComponent extends CellView implements OnInit, OnChanges {
 	public layout() {
 
 	}
+
+	handleStdIn(msg: nb.IStdinMessage): void | Thenable<void> {
+		if (msg) {
+			this.stdIn = msg;
+			this.inputDeferred = new Deferred();
+			this._changeRef.detectChanges();
+			return this.awaitStdIn();
+		}
+	}
+
+	private async awaitStdIn(): Promise<void> {
+		try {
+			let value = await this.inputDeferred.promise;
+			this.cellModel.future.sendInputReply({ value: value });
+		} catch (err) {
+			// Note: don't have a better way to handle completing input request. For now just canceling by sending empty string?
+			this.cellModel.future.sendInputReply({ value: '' });
+		} finally {
+			// Clean up so no matter what, the stdIn request goes away
+			this.stdIn = undefined;
+			this.inputDeferred = undefined;
+			this._changeRef.detectChanges();
+		}
+	}
+
+	get isStdInVisible(): boolean {
+		return !!(this.stdIn && this.inputDeferred);
+	}
 }

src/sql/workbench/parts/notebook/cellViews/outputArea.component.ts

@@ -22,8 +22,6 @@ export class OutputAreaComponent extends AngularDisposable implements OnInit {
 
 	private _activeCellId: string;
 
-	private readonly _minimumHeight = 30;
-
 	constructor(
 		@Inject(IWorkbenchThemeService) private themeService: IWorkbenchThemeService,
 		@Inject(forwardRef(() => ChangeDetectorRef)) private _changeRef: ChangeDetectorRef

src/sql/workbench/parts/notebook/cellViews/stdin.component.ts

@@ -0,0 +1,113 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import 'vs/css!./stdin';
+
+import {
+	Component, Input, Inject, ChangeDetectorRef, forwardRef,
+	ViewChild, ElementRef, AfterViewInit, HostListener
+} from '@angular/core';
+import { nb } from 'azdata';
+import { localize } from 'vs/nls';
+
+import { IInputOptions } from 'vs/base/browser/ui/inputbox/inputBox';
+import { IWorkbenchThemeService } from 'vs/workbench/services/themes/common/workbenchThemeService';
+import { IContextViewService } from 'vs/platform/contextview/browser/contextView';
+import { inputBackground, inputBorder } from 'vs/platform/theme/common/colorRegistry';
+import { StandardKeyboardEvent, IKeyboardEvent } from 'vs/base/browser/keyboardEvent';
+import { KeyCode } from 'vs/base/common/keyCodes';
+
+import { InputBox } from 'sql/base/browser/ui/inputBox/inputBox';
+import { attachInputBoxStyler } from 'sql/platform/theme/common/styler';
+import { AngularDisposable } from 'sql/base/node/lifecycle';
+import { Deferred } from 'sql/base/common/promise';
+import { ICellModel, CellExecutionState } from 'sql/workbench/parts/notebook/models/modelInterfaces';
+
+export const STDIN_SELECTOR: string = 'stdin-component';
+@Component({
+	selector: STDIN_SELECTOR,
+	template: `
+		<div class="prompt">{{prompt}}</div>
+		<div #input class="input"></div>
+	`
+})
+export class StdInComponent extends AngularDisposable implements AfterViewInit {
+	private _input: InputBox;
+	@ViewChild('input', { read: ElementRef }) private _inputContainer: ElementRef;
+
+	@Input() stdIn: nb.IStdinMessage;
+	@Input() onSendInput: Deferred<string>;
+	@Input() cellModel: ICellModel;
+
+
+	constructor(
+		@Inject(forwardRef(() => ChangeDetectorRef)) private changeRef: ChangeDetectorRef,
+		@Inject(IWorkbenchThemeService) private themeService: IWorkbenchThemeService,
+		@Inject(IContextViewService) private contextViewService: IContextViewService,
+		@Inject(forwardRef(() => ElementRef)) private el: ElementRef
+	) {
+		super();
+	}
+
+	ngAfterViewInit(): void {
+		let inputOptions: IInputOptions = {
+			placeholder: '',
+			ariaLabel: this.prompt
+		};
+		this._input = new InputBox(this._inputContainer.nativeElement, this.contextViewService, inputOptions);
+		if (this.password) {
+			this._input.inputElement.type = 'password';
+		}
+		this._register(this._input);
+		this._register(attachInputBoxStyler(this._input, this.themeService, {
+			inputValidationInfoBackground: inputBackground,
+			inputValidationInfoBorder: inputBorder,
+		}));
+		if (this.cellModel) {
+			this._register(this.cellModel.onExecutionStateChange((status) => this.handleExecutionChange(status)));
+		}
+		this._input.focus();
+	}
+
+	@HostListener('document:keydown', ['$event'])
+	public handleKeyboardInput(event: KeyboardEvent): void {
+		let e = new StandardKeyboardEvent(event);
+		switch (e.keyCode) {
+			case KeyCode.Enter:
+				// Indi
+				if (this.onSendInput) {
+					this.onSendInput.resolve(this._input.value);
+				}
+				e.stopPropagation();
+				break;
+			case KeyCode.Escape:
+				if (this.onSendInput) {
+					this.onSendInput.reject('');
+				}
+				e.stopPropagation();
+				break;
+			default:
+				// No-op
+				break;
+		}
+	}
+
+	handleExecutionChange(status: CellExecutionState): void {
+		if (status !== CellExecutionState.Running && this.onSendInput) {
+			this.onSendInput.reject('');
+		}
+	}
+
+	private get prompt(): string {
+		if (this.stdIn && this.stdIn.content && this.stdIn.content.prompt) {
+			return this.stdIn.content.prompt;
+		}
+		return localize('stdInLabel', "StdIn:");
+	}
+
+	private get password(): boolean {
+		return this.stdIn && this.stdIn.content && this.stdIn.content.password;
+	}
+}

src/sql/workbench/parts/notebook/cellViews/stdin.css

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+stdin-component {
+	display: flex;
+	flex-flow: row;
+	padding: 10px;
+	align-items: center;
+}
+
+stdin-component .prompt {
+	flex: 0 0 auto;
+}
+stdin-component .input {
+	flex: 1 1 auto;
+	padding-left: 10px;
+}

src/sql/workbench/parts/notebook/models/cell.ts

@@ -38,6 +38,7 @@ export class CellModel implements ICellModel {
 	private _cellUri: URI;
 	public id: string;
 	private _connectionManagementService: IConnectionManagementService;
+	private _stdInHandler: nb.MessageHandler<nb.IStdinMessage>;
 
 	constructor(private factory: IModelFactory, cellData?: nb.ICellContents, private _options?: ICellModelOptions) {
 		this.id = `${modelId++}`;
@@ -305,6 +306,7 @@ export class CellModel implements ICellModel {
 		this._future = future;
 		future.setReplyHandler({ handle: (msg) => this.handleReply(msg) });
 		future.setIOPubHandler({ handle: (msg) => this.handleIOPub(msg) });
+		future.setStdInHandler({ handle: (msg) => this.handleSdtIn(msg) });
 	}
 
 	public clearOutputs(): void {
@@ -427,6 +429,33 @@ export class CellModel implements ICellModel {
 		return transient['display_id'] as string;
 	}
 
+	public setStdInHandler(handler: nb.MessageHandler<nb.IStdinMessage>): void {
+		this._stdInHandler = handler;
+	}
+
+	/**
+	 * StdIn requires user interaction, so this is deferred to upstream UI
+	 * components. If one is registered the cell will call and wait on it, if not
+	 * it will immediately return to unblock error handling
+	 */
+	private handleSdtIn(msg: nb.IStdinMessage): void | Thenable<void> {
+		let handler = async () => {
+			if (!this._stdInHandler) {
+				// No-op
+				return;
+			}
+			try {
+				await this._stdInHandler.handle(msg);
+			} catch (err) {
+				if (this.future) {
+					// TODO should we error out in this case somehow? E.g. send Ctrl+C?
+					this.future.sendInputReply({ value: '' });
+				}
+			}
+		};
+		return handler();
+	}
+
 	public toJSON(): nb.ICellContents {
 		let cellJson: Partial<nb.ICellContents> = {
 			cell_type: this._cellType,

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

@@ -447,6 +447,7 @@ export interface ICellModel {
 	readonly executionState: CellExecutionState;
 	readonly notebookModel: NotebookModel;
 	setFuture(future: FutureInternal): void;
+	setStdInHandler(handler: nb.MessageHandler<nb.IStdinMessage>): void;
 	runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise<boolean>;
 	setOverrideLanguage(language: string);
 	equals(cellModel: ICellModel): boolean;

src/sql/workbench/parts/notebook/notebook.module.ts

@@ -26,6 +26,7 @@ import { CodeCellComponent } from 'sql/workbench/parts/notebook/cellViews/codeCe
 import { TextCellComponent } from 'sql/workbench/parts/notebook/cellViews/textCell.component';
 import { OutputAreaComponent } from 'sql/workbench/parts/notebook/cellViews/outputArea.component';
 import { OutputComponent } from 'sql/workbench/parts/notebook/cellViews/output.component';
+import { StdInComponent } from 'sql/workbench/parts/notebook/cellViews/stdin.component';
 import { PlaceholderCellComponent } from 'sql/workbench/parts/notebook/cellViews/placeholderCell.component';
 import LoadingSpinner from 'sql/workbench/electron-browser/modelComponents/loadingSpinner.component';
 
@@ -44,7 +45,8 @@ export const NotebookModule = (params, selector: string, instantiationService: I
 			NotebookComponent,
 			ComponentHostDirective,
 			OutputAreaComponent,
-			OutputComponent
+			OutputComponent,
+			StdInComponent
 		],
 		entryComponents: [NotebookComponent],
 		imports: [

src/sqltest/parts/notebook/model/cell.test.ts

@@ -16,6 +16,7 @@ import { ModelFactory } from 'sql/workbench/parts/notebook/models/modelFactory';
 import { NotebookModelStub } from '../common';
 import { EmptyFuture } from 'sql/workbench/services/notebook/common/sessionManager';
 import { ICellModel } from 'sql/workbench/parts/notebook/models/modelInterfaces';
+import { Deferred } from 'sql/base/common/promise';
 
 suite('Cell Model', function (): void {
 	let factory = new ModelFactory();
@@ -185,6 +186,19 @@ suite('Cell Model', function (): void {
 	suite('Model Future handling', function (): void {
 		let future: TypeMoq.Mock<EmptyFuture>;
 		let cell: ICellModel;
+		const stdInDefaultMessage: nb.IStdinMessage = {
+			channel: 'stdin',
+			type: 'stdin',
+			parent_header: undefined,
+			metadata: undefined,
+			header: <nb.IHeader>{
+				msg_type: 'stream'
+			},
+			content: {
+				prompt: 'Prompt',
+				password: false
+			}
+		};
 		setup(() => {
 			future = TypeMoq.Mock.ofType(EmptyFuture);
 			cell = factory.createCell({
@@ -239,8 +253,72 @@ suite('Cell Model', function (): void {
 			message.header.msg_type = 'display_data';
 			onIopub.handle(message);
 			should(outputs[1].output_type).equal('display_data');
+		});
 
-			// ... TODO: And when I sent a reply I expect it to be processed.
+		test('stdin should return void if no handler registered', async () => {
+			// Given stdIn does not have a request handler setup
+			let onStdIn: nb.MessageHandler<nb.IStdinMessage>;
+			future.setup(f => f.setStdInHandler(TypeMoq.It.isAny())).callback((handler) => onStdIn = handler);
+
+			// When I set it on the cell
+			cell.setFuture(future.object);
+
+			// Then I expect stdIn to have been hooked up
+			should(onStdIn).not.be.undefined();
+			// ... And when I send a stdIn request message
+			let result = onStdIn.handle(stdInDefaultMessage);
+			// Then I expect the promise to resolve
+			await result;
+			future.verify(f => f.sendInputReply(TypeMoq.It.isAny()), TypeMoq.Times.never());
+		});
+
+		test('stdin should wait on handler if handler registered', async () => {
+			// Given stdIn has a handler set up
+			let onStdIn: nb.MessageHandler<nb.IStdinMessage>;
+			future.setup(f => f.setStdInHandler(TypeMoq.It.isAny())).callback((handler) => onStdIn = handler);
+
+			let deferred = new Deferred<void>();
+			let stdInMessage: nb.IStdinMessage = undefined;
+			cell.setStdInHandler({
+				handle: (msg: nb.IStdinMessage) => {
+					stdInMessage = msg;
+					return deferred.promise;
+				}
+			});
+
+			// When I send a stdIn request message
+			cell.setFuture(future.object);
+			let result = onStdIn.handle(stdInDefaultMessage);
+			deferred.resolve();
+			// Then I expect promise to resolve since it should wait on upstream handling
+			await result;
+			// And I expect message to have been passed upstream and no message sent from the cell
+			should(stdInMessage).not.be.undefined();
+			should(stdInMessage.content.prompt).equal(stdInDefaultMessage.content.prompt);
+			should(stdInMessage.content.password).equal(stdInDefaultMessage.content.password);
+			future.verify(f => f.sendInputReply(TypeMoq.It.isAny()), TypeMoq.Times.never());
+		});
+		test('stdin should send default response if there is upstream error', async () => {
+			// Given stdIn has a handler set up
+			let onStdIn: nb.MessageHandler<nb.IStdinMessage>;
+			future.setup(f => f.setStdInHandler(TypeMoq.It.isAny())).callback((handler) => onStdIn = handler);
+
+			let deferred = new Deferred<void>();
+			let stdInMessage: nb.IStdinMessage = undefined;
+			cell.setStdInHandler({
+				handle: (msg: nb.IStdinMessage) => {
+					stdInMessage = msg;
+					return deferred.promise;
+				}
+			});
+
+			// When I send a stdIn request message
+			cell.setFuture(future.object);
+			let result = onStdIn.handle(stdInDefaultMessage);
+			deferred.reject('Something went wrong');
+			// Then I expect promise to resolve since it should wait on upstream handling
+			await result;
+			future.verify(f => f.sendInputReply(TypeMoq.It.isAny()), TypeMoq.Times.once());
 		});
 
 		test('should delete transient tag while handling incoming messages', async () => {