Fix race condition in getting task execution

Fixes #96643

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

@@ -384,6 +384,7 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape {
 	protected _handleCounter: number;
 	protected _handlers: Map<number, HandlerData>;
 	protected _taskExecutions: Map<string, TaskExecutionImpl>;
+	protected _taskExecutionPromises: Map<string, Promise<TaskExecutionImpl>>;
 	protected _providedCustomExecutions2: Map<string, types.CustomExecution>;
 	private _notProvidedCustomExecutions: Set<string>; // Used for custom executions tasks that are created and run through executeTask.
 	protected _activeCustomExecutions2: Map<string, types.CustomExecution>;
@@ -412,6 +413,7 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape {
 		this._handleCounter = 0;
 		this._handlers = new Map<number, HandlerData>();
 		this._taskExecutions = new Map<string, TaskExecutionImpl>();
+		this._taskExecutionPromises = new Map<string, Promise<TaskExecutionImpl>>();
 		this._providedCustomExecutions2 = new Map<string, types.CustomExecution>();
 		this._notProvidedCustomExecutions = new Set<string>();
 		this._activeCustomExecutions2 = new Map<string, types.CustomExecution>();
@@ -496,6 +498,7 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape {
 
 	public async $OnDidEndTask(execution: tasks.TaskExecutionDTO): Promise<void> {
 		const _execution = await this.getTaskExecution(execution);
+		this._taskExecutionPromises.delete(execution.id);
 		this._taskExecutions.delete(execution.id);
 		this.customExecutionComplete(execution);
 		this._onDidTerminateTask.fire({
@@ -626,17 +629,24 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape {
 			return taskExecution;
 		}
 
-		let result: TaskExecutionImpl | undefined = this._taskExecutions.get(execution.id);
+		let result: Promise<TaskExecutionImpl> | undefined = this._taskExecutionPromises.get(execution.id);
 		if (result) {
 			return result;
 		}
-		const taskToCreate = task ? task : await TaskDTO.to(execution.task, this._workspaceProvider);
-		if (!taskToCreate) {
-			throw new Error('Unexpected: Task does not exist.');
-		}
-		const createdResult: TaskExecutionImpl = new TaskExecutionImpl(this, execution.id, taskToCreate);
-		this._taskExecutions.set(execution.id, createdResult);
-		return createdResult;
+		const createdResult: Promise<TaskExecutionImpl> = new Promise(async (resolve, reject) => {
+			const taskToCreate = task ? task : await TaskDTO.to(execution.task, this._workspaceProvider);
+			if (!taskToCreate) {
+				reject('Unexpected: Task does not exist.');
+			} else {
+				resolve(new TaskExecutionImpl(this, execution.id, taskToCreate));
+			}
+		});
+
+		this._taskExecutionPromises.set(execution.id, createdResult);
+		return createdResult.then(result => {
+			this._taskExecutions.set(execution.id, result);
+			return result;
+		});
 	}
 
 	protected checkDeprecation(task: vscode.Task, handler: HandlerData) {