From 87dc506eca183dc167207677a47b25654dc532bd Mon Sep 17 00:00:00 2001
From: Daniel Imms <2193314+Tyriar@users.noreply.github.com>
Date: Mon, 24 Feb 2025 07:30:40 -0800
Subject: [PATCH] Ensure data events are flushed before execution end fires

Fixes #241592
---
 .../terminal.shellIntegration.test.ts         |  6 ++---
 .../common/extHostTerminalShellIntegration.ts | 23 +++++++++++++++++--
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/extensions/vscode-api-tests/src/singlefolder-tests/terminal.shellIntegration.test.ts b/extensions/vscode-api-tests/src/singlefolder-tests/terminal.shellIntegration.test.ts
index 3a63841d600..d5d84c58f4d 100644
--- a/extensions/vscode-api-tests/src/singlefolder-tests/terminal.shellIntegration.test.ts
+++ b/extensions/vscode-api-tests/src/singlefolder-tests/terminal.shellIntegration.test.ts
@@ -142,7 +142,7 @@ import { assertNoRpc } from '../utils';
 		await closeTerminalAsync(terminal);
 	});
 
-	test.skip('TerminalShellExecution.read iterables should be available between the start and end execution events', async () => {
+	test('TerminalShellExecution.read iterables should be available between the start and end execution events', async () => {
 		const { terminal, shellIntegration } = await createTerminalAndWaitForShellIntegration();
 		const events: string[] = [];
 		disposables.push(window.onDidStartTerminalShellExecution(() => events.push('start')));
@@ -164,7 +164,7 @@ import { assertNoRpc } from '../utils';
 		await closeTerminalAsync(terminal);
 	});
 
-	test.skip('TerminalShellExecution.read events should fire with contents of command', async () => {
+	test('TerminalShellExecution.read events should fire with contents of command', async () => {
 		const { terminal, shellIntegration } = await createTerminalAndWaitForShellIntegration();
 		const events: string[] = [];
 
@@ -179,7 +179,7 @@ import { assertNoRpc } from '../utils';
 		await closeTerminalAsync(terminal);
 	});
 
-	test.skip('TerminalShellExecution.read events should give separate iterables per call', async () => {
+	test('TerminalShellExecution.read events should give separate iterables per call', async () => {
 		const { terminal, shellIntegration } = await createTerminalAndWaitForShellIntegration();
 
 		const { execution, endEvent } = executeCommandAsync(shellIntegration, 'echo hello');
diff --git a/src/vs/workbench/api/common/extHostTerminalShellIntegration.ts b/src/vs/workbench/api/common/extHostTerminalShellIntegration.ts
index 684782ee896..c2d7b9e3172 100644
--- a/src/vs/workbench/api/common/extHostTerminalShellIntegration.ts
+++ b/src/vs/workbench/api/common/extHostTerminalShellIntegration.ts
@@ -235,8 +235,17 @@ class InternalTerminalShellIntegration extends Disposable {
 	endShellExecution(commandLine: vscode.TerminalShellExecutionCommandLine | undefined, exitCode: number | undefined): void {
 		if (this._currentExecution) {
 			this._currentExecution.endExecution(commandLine);
-			this._onDidRequestEndExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: this._currentExecution.value, exitCode });
-			this._currentExecution = undefined;
+			const currentExecution = this._currentExecution;
+			// 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) {
+					this._onDidRequestEndExecution.fire({ terminal: this._terminal, shellIntegration: this.value, execution: currentExecution.value, exitCode });
+					this._currentExecution = undefined;
+				}
+			});
 		}
 	}
 
@@ -314,10 +323,15 @@ class InternalTerminalShellExecution {
 		this._dataStream = undefined;
 		this._ended = true;
 	}
+
+	async flush(): Promise<void> {
+		await this._dataStream?.flush();
+	}
 }
 
 class ShellExecutionDataStream extends Disposable {
 	private _barrier: Barrier | undefined;
+	private _iterables: AsyncIterableObject<string>[] = [];
 	private _emitters: AsyncIterableEmitter<string>[] = [];
 
 	createIterable(): AsyncIterable<string> {
@@ -329,6 +343,7 @@ class ShellExecutionDataStream extends Disposable {
 			this._emitters.push(emitter);
 			await barrier.wait();
 		});
+		this._iterables.push(iterable);
 		return iterable;
 	}
 
@@ -342,4 +357,8 @@ class ShellExecutionDataStream extends Disposable {
 		this._barrier?.open();
 		this._barrier = undefined;
 	}
+
+	async flush(): Promise<void> {
+		await Promise.all(this._iterables.map(e => e.toPromise()));
+	}
 }