fix: command execution event listener race condition (e22f167e) · Commits · GitLab.org / gitlab-vscode-extension

src/common/duo_workflow/terminal_manager.test.ts

+30 −0

Original line number	Diff line number	Diff line
		@@ -228,6 +228,36 @@ describe('Terminal Manager', () => {
		expect(exitCode).toBe(0);
		});

		it('registers the execution-end listener before starting the command', async () => {
		mockShellIntegrationGetter.mockReturnValue(mockShellIntegration);

		jest.mocked(mockExecution.read).mockImplementation(async function* read() {
		yield '';
		});

		const result = eventFn({ workflowId, command, args, silent });

		await jest.advanceTimersByTimeAsync(3000);

		shellExecutionFn({
		execution: mockExecution,
		exitCode: 0,
		terminal: mockTerminal,
		shellIntegration: mockShellIntegration,
		});

		await result;

		// invocationCallOrder is a monotonically increasing call-order number recorded
		// per mock call, so a smaller value proves the listener was registered first.
		const listenerOrder = jest.mocked(window.onDidEndTerminalShellExecution).mock
		.invocationCallOrder[0];
		const executeCommandOrder = jest.mocked(mockShellIntegration.executeCommand).mock
		.invocationCallOrder[0];

		expect(listenerOrder).toBeLessThan(executeCommandOrder);
		});

		it('reuses a terminal when the same workflow ID and silent flag are used', async () => {
		mockShellIntegrationGetter.mockReturnValue(mockShellIntegration);

src/common/duo_workflow/terminal_manager.ts

+127 −40

Original line number	Diff line number	Diff line
		import { Terminal, window, Disposable, TerminalShellExecution } from 'vscode';
		import {
		Terminal,
		window,
		Disposable,
		TerminalShellExecution,
		TerminalShellIntegration,
		} from 'vscode';
		import { BaseLanguageClient } from 'vscode-languageclient';
		import { createInterfaceId } from '@gitlab/needle';
		import { log } from '../log';
		@@ -16,6 +22,14 @@ type RunningExecution = {
		resolver: (value: { exitCode: number; output: string }) => void;
		};

		type StartedExecution = {
		execution: TerminalShellExecution;
		/** Undefined when the output stream could not be opened; completion is still tracked. */
		stream: AsyncIterable<string> \| undefined;
		executionEnded: Promise<{ exitCode: number }>;
		dispose: () => void;
		};

		export interface TerminalManager extends Disposable {
		setupRequests(client: BaseLanguageClient): void;
		}
		@@ -37,7 +51,9 @@ export class TerminalManagerImpl implements TerminalManager {
		setupRequests(client: BaseLanguageClient) {
		this.#disposables.push(
		client.onRequest('$/gitlab/runCommand', async ({ workflowId, command, args, silent }) => {
		log.debug(`Running command: ${command} ${args?.join(' ')}`);
		log.debug(
		`[TerminalManager] Received runCommand request for workflow ${workflowId}: ${command} ${args?.join(' ') ?? ''} (silent: ${silent})`,
		);

		return this.#executeCommand(workflowId, command, args, silent);
		}),
		@@ -45,16 +61,17 @@ export class TerminalManagerImpl implements TerminalManager {

		this.#disposables.push(
		client.onNotification('$/gitlab/cancelRunningCommand', ({ workflowId }) => {
		log.debug(`Received cancel request for workflow: ${workflowId}`);
		log.debug(`[TerminalManager] Received cancel request for workflow ${workflowId}`);
		this.#cancelRunningCommand(workflowId);
		}),
		);

		this.#disposables.push(
		window.onDidCloseTerminal(closedTerminal => {
		for (const [workflowId, terminal] of this.#terminals.entries()) {
		for (const [key, terminal] of this.#terminals.entries()) {
		if (terminal === closedTerminal) {
		this.#terminals.delete(workflowId);
		log.debug(`[TerminalManager] Terminal closed, removing from cache (key: ${key})`);
		this.#terminals.delete(key);
		break;
		}
		}
		@@ -71,7 +88,9 @@ export class TerminalManagerImpl implements TerminalManager {
		const terminal = await this.#getOrCreateTerminal(workflowId, silent);

		if (!terminal.shellIntegration) {
		log.warn('Terminal lost shell integration, recreating...');
		log.warn(
		`[TerminalManager] Terminal for workflow ${workflowId} lost shell integration, recreating...`,
		);
		this.#terminals.delete(lookupKey(workflowId, silent));
		terminal.dispose();
		return this.#executeCommand(workflowId, command, args, silent);
		@@ -81,80 +100,144 @@ export class TerminalManagerImpl implements TerminalManager {
		terminal.show(true);
		}

		const exec = args
		? terminal.shellIntegration.executeCommand(command, args)
		: terminal.shellIntegration.executeCommand(command);

		const output: string[] = [];
		let executionEndDisposable: Disposable \| undefined;
		const startTime = Date.now();

		// Create promises for both normal completion and cancellation
		const executionEndPromise = new Promise<{ exitCode: number; cancelled: false }>(resolve => {
		executionEndDisposable = window.onDidEndTerminalShellExecution(({ execution, exitCode }) => {
		if (exec === execution) {
		executionEndDisposable?.dispose();
		resolve({ exitCode: exitCode ?? 0, cancelled: false });
		}
		});
		});
		const { execution, stream, executionEnded, dispose } = this.#startExecution(
		terminal.shellIntegration,
		command,
		args,
		workflowId,
		);

		const cancellationPromise = new Promise<{ exitCode: number; cancelled: true }>(resolve => {
		// Store the running execution so it can be cancelled
		const cancelled = new Promise<{ cancelled: true; exitCode: number }>(resolve => {
		this.#runningExecutions.set(workflowId, {
		execution: exec,
		execution,
		terminal,
		outputCollector: output,
		resolver: (value: { exitCode: number; output: string }) => {
		resolve({ exitCode: value.exitCode, cancelled: true });
		},
		resolver: ({ exitCode }) => resolve({ cancelled: true, exitCode }),
		});
		});

		const streamPromise = (async () => {
		const streamDrained = (async () => {
		if (!stream) {
		return;
		}
		try {
		const stream = exec.read();
		for await (const data of stream) {
		output.push(data);
		}
		} catch (error) {
		log.error('Error reading stream:', error);
		log.error(`[TerminalManager] Error reading stream for workflow ${workflowId}:`, error);
		}
		})();

		// Wait for either normal completion or cancellation
		const result = await Promise.race([executionEndPromise, cancellationPromise]);
		const result = await Promise.race([
		executionEnded.then(({ exitCode }) => ({ cancelled: false as const, exitCode })),
		cancelled,
		]);

		executionEndDisposable?.dispose();
		dispose();
		this.#runningExecutions.delete(workflowId);

		// If cancelled, return immediately with partial output
		const durationMs = Date.now() - startTime;

		if (result.cancelled) {
		log.debug(
		`[TerminalManager] Command for workflow ${workflowId} cancelled after ${durationMs}ms (${output.length} output chunks)`,
		);
		return { exitCode: result.exitCode, output: output.join('') };
		}

		// Otherwise wait for stream to finish and return full output
		await streamPromise;
		await streamDrained;
		log.debug(
		`[TerminalManager] Command for workflow ${workflowId} completed in ${durationMs}ms with exit code ${result.exitCode} (${output.length} output chunks)`,
		);
		return { exitCode: result.exitCode, output: output.join('') };
		}

		async #getOrCreateTerminal(workflowId: string, silent: boolean) {
		return (
		this.#terminals.get(lookupKey(workflowId, silent)) ?? this.#createTerminal(workflowId, silent)
		/**
		* Starts a command and wires up completion tracking in a single synchronous step.
		*
		* This method is intentionally not `async`: the end-execution listener must be registered
		* before the command starts, and the output stream must be opened immediately, because VS
		* Code delivers terminal events on a later tick, does not replay missed events, and `read()`
		* drops output written before its first call.
		*/
		#startExecution(
		shellIntegration: TerminalShellIntegration,
		command: string,
		args: string[] \| undefined,
		workflowId: string,
		): StartedExecution {
		let execution: TerminalShellExecution \| undefined;
		let executionEndDisposable: Disposable \| undefined;

		const executionEnded = new Promise<{ exitCode: number }>(resolve => {
		executionEndDisposable = window.onDidEndTerminalShellExecution(
		({ execution: ended, exitCode }) => {
		if (ended === execution) {
		log.debug(
		`[TerminalManager] Shell execution ended for workflow ${workflowId} with exit code ${exitCode ?? 'undefined'}`,
		);
		executionEndDisposable?.dispose();
		resolve({ exitCode: exitCode ?? 0 });
		}
		},
		);
		});

		execution = args
		? shellIntegration.executeCommand(command, args)
		: shellIntegration.executeCommand(command);

		let stream: AsyncIterable<string> \| undefined;
		try {
		stream = execution.read();
		} catch (error) {
		log.error(`[TerminalManager] Error starting stream read for workflow ${workflowId}:`, error);
		}

		log.debug(`[TerminalManager] Started shell execution for workflow ${workflowId}`);

		return {
		execution,
		stream,
		executionEnded,
		dispose: () => executionEndDisposable?.dispose(),
		};
		}

		async #getOrCreateTerminal(workflowId: string, silent: boolean) {
		const existingTerminal = this.#terminals.get(lookupKey(workflowId, silent));
		if (existingTerminal) {
		log.debug(`[TerminalManager] Reusing existing terminal for workflow ${workflowId}`);
		return existingTerminal;
		}

		return this.#createTerminal(workflowId, silent);
		}

		async #createTerminal(workflowId: string, silent: boolean) {
		log.debug(`[TerminalManager] Creating new terminal for workflow ${workflowId}`);
		const terminal = window.createTerminal({
		name: `GitLab Duo Agent Platform ${workflowId}`,
		isTransient: true,
		hideFromUser: silent,
		});

		const createStartTime = Date.now();
		try {
		await this.#listenForShellIntegration(terminal);
		log.debug(
		`[TerminalManager] Terminal ready with shell integration for workflow ${workflowId} (took ${Date.now() - createStartTime}ms)`,
		);
		this.#terminals.set(lookupKey(workflowId, silent), terminal);
		return terminal;
		} catch (error) {
		log.warn(
		`[TerminalManager] Failed to establish shell integration for workflow ${workflowId}: ${error instanceof Error ? error.message : String(error)}`,
		);
		terminal.dispose();
		throw error;
		}
		@@ -164,11 +247,15 @@ export class TerminalManagerImpl implements TerminalManager {
		const runningExec = this.#runningExecutions.get(workflowId);

		if (!runningExec) {
		log.warn(`No running command found for workflow ${workflowId}, may have already completed`);
		log.warn(
		`[TerminalManager] No running command found for workflow ${workflowId}, may have already completed`,
		);
		return;
		}

		log.info(`Cancelling running command for workflow ${workflowId}`);
		log.info(
		`[TerminalManager] Cancelling running command for workflow ${workflowId}, sending SIGINT (${runningExec.outputCollector.length} output chunks captured so far)`,
		);

		// Remove entry first to prevent race with #executeCommand cleanup
		this.#runningExecutions.delete(workflowId);