Implement await terminal tool

Fixes #288565

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/terminal.chatAgentTools.contribution.ts

@@ -21,6 +21,7 @@ import { TerminalContextMenuGroup } from '../../../terminal/browser/terminalMenu
 import { TerminalContextKeys } from '../../../terminal/common/terminalContextKey.js';
 import { TerminalChatAgentToolsCommandId } from '../common/terminal.chatAgentTools.js';
 import { TerminalChatAgentToolsSettingId } from '../common/terminalChatAgentToolsConfiguration.js';
+import { AwaitTerminalTool, AwaitTerminalToolData } from './tools/awaitTerminalTool.js';
 import { GetTerminalLastCommandTool, GetTerminalLastCommandToolData } from './tools/getTerminalLastCommandTool.js';
 import { GetTerminalOutputTool, GetTerminalOutputToolData } from './tools/getTerminalOutputTool.js';
 import { GetTerminalSelectionTool, GetTerminalSelectionToolData } from './tools/getTerminalSelectionTool.js';
@@ -91,6 +92,10 @@ class ChatAgentToolsContribution extends Disposable implements IWorkbenchContrib
 		this._register(toolsService.registerTool(GetTerminalOutputToolData, getTerminalOutputTool));
 		this._register(toolsService.executeToolSet.addTool(GetTerminalOutputToolData));
 
+		const awaitTerminalTool = instantiationService.createInstance(AwaitTerminalTool);
+		this._register(toolsService.registerTool(AwaitTerminalToolData, awaitTerminalTool));
+		this._register(toolsService.executeToolSet.addTool(AwaitTerminalToolData));
+
 		instantiationService.invokeFunction(createRunInTerminalToolData).then(runInTerminalToolData => {
 			const runInTerminalTool = instantiationService.createInstance(RunInTerminalTool);
 			this._register(toolsService.registerTool(runInTerminalToolData, runInTerminalTool));

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/awaitTerminalTool.ts

@@ -0,0 +1,132 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { CancellationError } from '../../../../../../base/common/errors.js';
+import type { CancellationToken } from '../../../../../../base/common/cancellation.js';
+import { Codicon } from '../../../../../../base/common/codicons.js';
+import { Disposable } from '../../../../../../base/common/lifecycle.js';
+import { localize } from '../../../../../../nls.js';
+import { ToolDataSource, type CountTokensCallback, type IPreparedToolInvocation, type IToolData, type IToolImpl, type IToolInvocation, type IToolInvocationPreparationContext, type IToolResult, type ToolProgress } from '../../../../chat/common/tools/languageModelToolsService.js';
+import { RunInTerminalTool } from './runInTerminalTool.js';
+import { timeout } from '../../../../../../base/common/async.js';
+
+export const AwaitTerminalToolData: IToolData = {
+	id: 'await_terminal',
+	toolReferenceName: 'awaitTerminal',
+	displayName: localize('awaitTerminalTool.displayName', 'Await Terminal'),
+	modelDescription: 'Wait for a background terminal command to complete. Returns the output, exit code, or timeout status.',
+	icon: Codicon.terminal,
+	source: ToolDataSource.Internal,
+	inputSchema: {
+		type: 'object',
+		properties: {
+			id: {
+				type: 'string',
+				description: 'The ID of the terminal to await (returned by run_in_terminal when isBackground=true).'
+			},
+			timeout: {
+				type: 'number',
+				description: 'Timeout in milliseconds. If the command does not complete within this time, returns the output collected so far with a timeout indicator. Use 0 for no timeout.'
+			},
+		},
+		required: [
+			'id',
+			'timeout',
+		]
+	}
+};
+
+export interface IAwaitTerminalInputParams {
+	id: string;
+	timeout: number;
+}
+
+export class AwaitTerminalTool extends Disposable implements IToolImpl {
+	async prepareToolInvocation(_context: IToolInvocationPreparationContext, _token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
+		return {
+			invocationMessage: localize('await.progressive', "Awaiting terminal completion"),
+			pastTenseMessage: localize('await.past', "Awaited terminal completion"),
+		};
+	}
+
+	async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _progress: ToolProgress, token: CancellationToken): Promise<IToolResult> {
+		const args = invocation.parameters as IAwaitTerminalInputParams;
+
+		const execution = RunInTerminalTool.getExecution(args.id);
+		if (!execution) {
+			return {
+				content: [{
+					kind: 'text',
+					value: `Error: No active terminal execution found with ID ${args.id}. The terminal may have already completed or the ID is invalid.`
+				}]
+			};
+		}
+
+		try {
+			let result: { output?: string; exitCode?: number; error?: string; didEnterAltBuffer?: boolean };
+			const hasTimeout = args.timeout > 0;
+
+			if (hasTimeout) {
+				// Race completion against timeout
+				const timeoutPromise = timeout(args.timeout).then(() => ({ type: 'timeout' as const }));
+				const completionPromise = execution.completionPromise.then(r => ({ type: 'completed' as const, result: r }));
+
+				const raceResult = await Promise.race([completionPromise, timeoutPromise]);
+
+				if (token.isCancellationRequested) {
+					throw new CancellationError();
+				}
+
+				if (raceResult.type === 'timeout') {
+					// Timeout reached - return partial output
+					const partialOutput = execution.getOutput();
+					return {
+						toolMetadata: {
+							exitCode: undefined,
+							timedOut: true
+						},
+						content: [{
+							kind: 'text',
+							value: `Terminal ${args.id} timed out after ${args.timeout}ms. Output collected so far:\n${partialOutput}`
+						}]
+					};
+				}
+
+				result = raceResult.result;
+			} else {
+				// No timeout - await completion directly
+				result = await execution.completionPromise;
+
+				if (token.isCancellationRequested) {
+					throw new CancellationError();
+				}
+			}
+
+			// Command completed
+			const output = execution.getOutput();
+			const exitCodeText = result.exitCode !== undefined ? ` (exit code: ${result.exitCode})` : '';
+
+			return {
+				toolMetadata: {
+					exitCode: result.exitCode
+				},
+				content: [{
+					kind: 'text',
+					value: `Terminal ${args.id} completed${exitCodeText}:\n${output}`
+				}]
+			};
+		} catch (e) {
+			if (e instanceof CancellationError) {
+				throw e;
+			}
+			return {
+				content: [{
+					kind: 'text',
+					value: `Error awaiting terminal ${args.id}: ${e instanceof Error ? e.message : String(e)}`
+				}]
+			};
+		}
+	}
+}

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts

@@ -270,6 +270,22 @@ export interface IRunInTerminalInputParams {
 	timeout?: number;
 }
 
+/**
+ * Interface for accessing a running terminal execution.
+ * Used by tools that need to await or interact with background terminal commands.
+ */
+export interface IActiveTerminalExecution {
+	/**
+	 * Promise that resolves when the terminal command completes.
+	 */
+	readonly completionPromise: Promise<ITerminalExecuteStrategyResult>;
+
+	/**
+	 * Gets the current output from the terminal.
+	 */
+	getOutput(): string;
+}
+
 /**
  * A set of characters to ignore when reporting telemetry
  */
@@ -307,6 +323,14 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 		return execution.getOutput();
 	}
 
+	/**
+	 * Gets an active terminal execution by ID. Returns undefined if not found.
+	 * Can be used to await the completion of a background terminal command.
+	 */
+	public static getExecution(id: string): IActiveTerminalExecution | undefined {
+		return RunInTerminalTool._activeExecutions.get(id);
+	}
+
 	constructor(
 		@IChatService protected readonly _chatService: IChatService,
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
@@ -1144,7 +1168,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
  * mode. This unified class replaces the previous split between foreground strategy execution and
  * BackgroundTerminalExecution, allowing seamless switching between modes.
  */
-class ActiveTerminalExecution extends Disposable {
+class ActiveTerminalExecution extends Disposable implements IActiveTerminalExecution {
 	private _startMarker: IXtermMarker | undefined;
 	private _isBackground: boolean;
 	private readonly _completionDeferred: DeferredPromise<ITerminalExecuteStrategyResult>;