Fix timing bug for back-to-back CustomExecution tasks

Fixes #131926

extensions/vscode-api-tests/src/singlefolder-tests/workspace.tasks.test.ts

@@ -363,6 +363,70 @@ import { assertNoRpc } from '../utils';
 					await tasksConfig.update('tasks', []);
 				});
 			});
+
+			test('Tasks can be run back to back', async () => {
+				class Pty implements Pseudoterminal {
+					writer = new EventEmitter<string>();
+					onDidWrite = this.writer.event;
+					closer = new EventEmitter<number | undefined>();
+					onDidClose = this.closer.event;
+
+					constructor(readonly num: number, readonly quick: boolean) { }
+
+					cleanup() {
+						this.writer.dispose();
+						this.closer.dispose();
+					}
+
+					open() {
+						this.writer.fire('starting\r\n');
+						setTimeout(() => {
+							this.closer.fire(this.num);
+							this.cleanup();
+						}, this.quick ? 1 : 200);
+					}
+
+					close() {
+						this.closer.fire(undefined);
+						this.cleanup();
+					}
+				}
+
+				async function runTask(num: number, quick: boolean) {
+					const pty = new Pty(num, quick);
+					const task = new Task(
+						{ type: 'task_bug', exampleProp: `hello world ${num}` },
+						TaskScope.Workspace, `task bug ${num}`, 'task bug',
+						new CustomExecution(
+							async () => {
+								return pty;
+							},
+						));
+					tasks.executeTask(task);
+					return new Promise<number | undefined>(resolve => {
+						pty.onDidClose(exitCode => {
+							resolve(exitCode);
+						});
+					});
+				}
+
+
+				const [r1, r2, r3, r4] = await Promise.all([
+					runTask(1, false), runTask(2, false), runTask(3, false), runTask(4, false)
+				]);
+				assert.strictEqual(r1, 1);
+				assert.strictEqual(r2, 2);
+				assert.strictEqual(r3, 3);
+				assert.strictEqual(r4, 4);
+
+				const [j1, j2, j3, j4] = await Promise.all([
+					runTask(5, true), runTask(6, true), runTask(7, true), runTask(8, true)
+				]);
+				assert.strictEqual(j1, 5);
+				assert.strictEqual(j2, 6);
+				assert.strictEqual(j3, 7);
+				assert.strictEqual(j4, 8);
+			});
 		});
 	});
 });

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

@@ -492,10 +492,6 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape, IExtHostTask
 	public async $onDidStartTask(execution: tasks.TaskExecutionDTO, terminalId: number, resolvedDefinition: tasks.TaskDefinitionDTO): Promise<void> {
 		const customExecution: types.CustomExecution | undefined = this._providedCustomExecutions2.get(execution.id);
 		if (customExecution) {
-			if (this._activeCustomExecutions2.get(execution.id) !== undefined) {
-				throw new Error('We should not be trying to start the same custom task executions twice.');
-			}
-
 			// Clone the custom execution to keep the original untouched. This is important for multiple runs of the same task.
 			this._activeCustomExecutions2.set(execution.id, customExecution);
 			this._terminalService.attachPtyToTerminal(terminalId, await customExecution.callback(resolvedDefinition));
@@ -625,6 +621,8 @@ export abstract class ExtHostTaskBase implements ExtHostTaskShape, IExtHostTask
 		const taskId = await this._proxy.$createTaskId(taskDTO);
 		if (!isProvided && !this._providedCustomExecutions2.has(taskId)) {
 			this._notProvidedCustomExecutions.add(taskId);
+			// Also add to active executions when not coming from a provider to prevent timing issue.
+			this._activeCustomExecutions2.set(taskId, <types.CustomExecution>task.execution);
 		}
 		this._providedCustomExecutions2.set(taskId, <types.CustomExecution>task.execution);
 	}