Merge pull request microsoft#243442 from microsoft/tyriar/237208_2

Support tracking sub-executions within a single executeCommand

Author: Tyriar

src/vs/workbench/api/common/extHostTerminalShellIntegration.ts

@@ -149,12 +149,20 @@ export class ExtHostTerminalShellIntegration extends Disposable implements IExtH
 	}
 }
 
-class InternalTerminalShellIntegration extends Disposable {
+interface IExecutionProperties {
+	isMultiLine: boolean;
+	unresolvedCommandLines: string[] | undefined;
+}
+
+export class InternalTerminalShellIntegration extends Disposable {
 	private _pendingExecutions: InternalTerminalShellExecution[] = [];
+	private _pendingEndingExecution: InternalTerminalShellExecution | undefined;
 
+	private _currentExecutionProperties: IExecutionProperties | undefined;
 	private _currentExecution: InternalTerminalShellExecution | undefined;
 	get currentExecution(): InternalTerminalShellExecution | undefined { return this._currentExecution; }
 
+
 	private _env: vscode.TerminalShellIntegrationEnvironment | undefined;
 	private _cwd: URI | undefined;
 	private _hasRichCommandDetection: boolean = false;
@@ -226,27 +234,70 @@ class InternalTerminalShellIntegration extends Disposable {
 
 	requestNewShellExecution(commandLine: vscode.TerminalShellExecutionCommandLine, cwd: URI | undefined) {
 		const execution = new InternalTerminalShellExecution(commandLine, cwd ?? this._cwd);
+		const unresolvedCommandLines = splitAndSanitizeCommandLine(commandLine.value);
+		if (unresolvedCommandLines.length > 1) {
+			this._currentExecutionProperties = {
+				isMultiLine: true,
+				unresolvedCommandLines: splitAndSanitizeCommandLine(commandLine.value),
+			};
+		}
 		this._pendingExecutions.push(execution);
 		this._onDidRequestNewExecution.fire(commandLine.value);
 		return execution;
 	}
 
-	startShellExecution(commandLine: vscode.TerminalShellExecutionCommandLine, cwd: URI | undefined): InternalTerminalShellExecution {
+	startShellExecution(commandLine: vscode.TerminalShellExecutionCommandLine, cwd: URI | undefined): undefined {
+		// Since an execution is starting, fire the end event for any execution that is awaiting to
+		// end. When this happens it means that the data stream may not be flushed and therefore may
+		// fire events after the end event.
+		if (this._pendingEndingExecution) {
+			this._onDidRequestEndExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: this._pendingEndingExecution.value, exitCode: undefined });
+			this._pendingEndingExecution = undefined;
+		}
+
 		if (this._currentExecution) {
+			// If the current execution is multi-line, check if this command line is part of it.
+			if (this._currentExecutionProperties?.isMultiLine && this._currentExecutionProperties.unresolvedCommandLines) {
+				const subExecutionResult = isSubExecution(this._currentExecutionProperties.unresolvedCommandLines, commandLine);
+				if (subExecutionResult) {
+					this._currentExecutionProperties.unresolvedCommandLines = subExecutionResult.unresolvedCommandLines;
+					return;
+				}
+			}
 			if (this._hasRichCommandDetection) {
 				console.warn('Rich command detection is enabled but an execution started before the last ended');
 			}
 			this._currentExecution.endExecution(undefined);
+			this._currentExecution.flush();
 			this._onDidRequestEndExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: this._currentExecution.value, exitCode: undefined });
 		}
 
 		// Get the matching pending execution, how strict this is depends on the confidence of the
 		// command line
 		let currentExecution: InternalTerminalShellExecution | undefined;
 		if (commandLine.confidence === TerminalShellExecutionCommandLineConfidence.High) {
-			const index = this._pendingExecutions.findIndex(e => e.value.commandLine.value === commandLine.value);
-			if (index !== -1) {
-				currentExecution = this._pendingExecutions.splice(index, 1)[0];
+			for (const [i, execution] of this._pendingExecutions.entries()) {
+				if (execution.value.commandLine.value === commandLine.value) {
+					currentExecution = execution;
+					this._currentExecutionProperties = {
+						isMultiLine: false,
+						unresolvedCommandLines: undefined,
+					};
+					currentExecution = execution;
+					this._pendingExecutions.splice(i, 1);
+					break;
+				} else {
+					const subExecutionResult = isSubExecution(splitAndSanitizeCommandLine(execution.value.commandLine.value), commandLine);
+					if (subExecutionResult) {
+						this._currentExecutionProperties = {
+							isMultiLine: true,
+							unresolvedCommandLines: subExecutionResult.unresolvedCommandLines,
+						};
+						currentExecution = execution;
+						this._pendingExecutions.splice(i, 1);
+						break;
+					}
+				}
 			}
 		} else {
 			currentExecution = this._pendingExecutions.shift();
@@ -259,27 +310,36 @@ class InternalTerminalShellIntegration extends Disposable {
 		}
 
 		this._currentExecution = currentExecution;
-
 		this._onDidStartTerminalShellExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: this._currentExecution.value });
-		return this._currentExecution;
 	}
 
 	emitData(data: string): void {
 		this.currentExecution?.emitData(data);
 	}
 
 	endShellExecution(commandLine: vscode.TerminalShellExecutionCommandLine | undefined, exitCode: number | undefined): void {
+		// If the current execution is multi-line, don't end it until the next command line is
+		// confirmed to not be a part of it.
+		if (this._currentExecutionProperties?.isMultiLine) {
+			if (this._currentExecutionProperties.unresolvedCommandLines && this._currentExecutionProperties.unresolvedCommandLines.length > 0) {
+				return;
+			}
+		}
+
 		if (this._currentExecution) {
-			this._currentExecution.endExecution(commandLine);
+			const commandLineForEvent = this._currentExecutionProperties?.isMultiLine ? this._currentExecution.value.commandLine : commandLine;
+			this._currentExecution.endExecution(commandLineForEvent);
 			const currentExecution = this._currentExecution;
+			this._pendingEndingExecution = currentExecution;
+			this._currentExecution = undefined;
 			// IMPORTANT: Ensure the current execution's data events are flushed in order to
 			// prevent data events firing after the end event fires.
 			currentExecution.flush().then(() => {
 				// Only fire if it's still the same execution, if it's changed it would have already
 				// been fired.
-				if (this._currentExecution === currentExecution) {
+				if (this._pendingEndingExecution === currentExecution) {
 					this._onDidRequestEndExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: currentExecution.value, exitCode });
-					this._currentExecution = undefined;
+					this._pendingEndingExecution = undefined;
 				}
 			});
 		}
@@ -320,12 +380,11 @@ class InternalTerminalShellIntegration extends Disposable {
 }
 
 class InternalTerminalShellExecution {
-	private _dataStream: ShellExecutionDataStream | undefined;
-
-	private _ended: boolean = false;
-
 	readonly value: vscode.TerminalShellExecution;
 
+	private _dataStream: ShellExecutionDataStream | undefined;
+	private _isEnded: boolean = false;
+
 	constructor(
 		private _commandLine: vscode.TerminalShellExecutionCommandLine,
 		readonly cwd: URI | undefined,
@@ -346,7 +405,7 @@ class InternalTerminalShellExecution {
 
 	private _createDataStream(): AsyncIterable<string> {
 		if (!this._dataStream) {
-			if (this._ended) {
+			if (this._isEnded) {
 				return AsyncIterableObject.EMPTY;
 			}
 			this._dataStream = new ShellExecutionDataStream();
@@ -355,20 +414,25 @@ class InternalTerminalShellExecution {
 	}
 
 	emitData(data: string): void {
-		this._dataStream?.emitData(data);
+		if (!this._isEnded) {
+			this._dataStream?.emitData(data);
+		}
 	}
 
 	endExecution(commandLine: vscode.TerminalShellExecutionCommandLine | undefined): void {
 		if (commandLine) {
 			this._commandLine = commandLine;
 		}
 		this._dataStream?.endExecution();
-		this._dataStream = undefined;
-		this._ended = true;
+		this._isEnded = true;
 	}
 
 	async flush(): Promise<void> {
-		await this._dataStream?.flush();
+		if (this._dataStream) {
+			await this._dataStream.flush();
+			this._dataStream.dispose();
+			this._dataStream = undefined;
+		}
 	}
 }
 
@@ -405,3 +469,39 @@ class ShellExecutionDataStream extends Disposable {
 		await Promise.all(this._iterables.map(e => e.toPromise()));
 	}
 }
+
+function splitAndSanitizeCommandLine(commandLine: string): string[] {
+	return commandLine
+		.split('\n')
+		.map(line => line.trim())
+		.filter(line => line.length > 0);
+}
+
+/**
+ * When executing something that the shell considers multiple commands, such as
+ * a comment followed by a command, this needs to all be tracked under a single
+ * execution.
+ */
+function isSubExecution(unresolvedCommandLines: string[], commandLine: vscode.TerminalShellExecutionCommandLine): { unresolvedCommandLines: string[] } | false {
+	if (unresolvedCommandLines.length === 0) {
+		return false;
+	}
+	const newUnresolvedCommandLines = [...unresolvedCommandLines];
+	const subExecutionLines = splitAndSanitizeCommandLine(commandLine.value);
+	if (newUnresolvedCommandLines && newUnresolvedCommandLines.length > 0) {
+		// If all sub-execution lines are in the command line, this is part of the
+		// multi-line execution.
+		for (let i = 0; i < newUnresolvedCommandLines.length; i++) {
+			if (newUnresolvedCommandLines[i] !== subExecutionLines[i]) {
+				break;
+			}
+			newUnresolvedCommandLines.shift();
+			subExecutionLines.shift();
+		}
+
+		if (subExecutionLines.length === 0) {
+			return { unresolvedCommandLines: newUnresolvedCommandLines };
+		}
+	}
+	return false;
+}