Wait for agent loop to finish in automation (#262370)

* add `workbench.action.chat.open::waitForComplete`

* s/waitForCompletion/blockOnResponse

* cleanup tests

* remove unused comment

* Apply suggestion from @connor4312

Co-authored-by: Connor Peet <connor@peet.io>

* fixup tests

* don't block test loop on invokeTool

* Revert "don't block test loop on invokeTool"

This reverts commit d8d16dbe79.

* fix tool confirmation test

* attempt to account for the flip of isPendingConfirmation

* [DEBUG] debug CI flake

* register tool so it exists in all test envs

* finish configuring custom tool

* run test in seperate chat windows

* revert debug changes

* remove timeout dep

* fix assertion

* cleaup tests by examining output of command directly

---------

Co-authored-by: Connor Peet <connor@peet.io>

extensions/vscode-api-tests/package.json

@@ -87,6 +87,17 @@
         "commands": []
       }
     ],
+    "languageModelTools": [
+			{
+				"name": "requires_confirmation_tool",
+				"toolReferenceName": "requires_confirmation_tool",
+				"displayName": "Requires Confirmation Tool",
+				"modelDescription": "A noop tool to trigger confirmation.",
+				"canBeReferencedInPrompt": true,
+				"icon": "$(files)",
+				"inputSchema": {}
+			}
+    ],
     "configuration": {
       "type": "object",
       "title": "Test Config",

extensions/vscode-api-tests/src/singlefolder-tests/chat.test.ts

@@ -124,6 +124,59 @@ suite('chat', () => {
 		assert.strictEqual(request3.context.history.length, 2); // request + response = 2
 	});
 
+	test('workbench.action.chat.open.blockOnResponse defaults to non-blocking for backwards compatibility', async () => {
+		const toolRegistration = lm.registerTool<void>('requires_confirmation_tool', {
+			invoke: async (_options, _token) => null, prepareInvocation: async (_options, _token) => {
+				return { invocationMessage: 'Invoking', pastTenseMessage: 'Invoked', confirmationMessages: { title: 'Confirm', message: 'Are you sure?' } };
+			}
+		});
+
+		const participant = chat.createChatParticipant('api-test.participant', async (_request, _context, _progress, _token) => {
+			await lm.invokeTool('requires_confirmation_tool', {
+				input: {},
+				toolInvocationToken: _request.toolInvocationToken,
+			});
+			return { metadata: { complete: true } };
+		});
+		disposables.push(participant, toolRegistration);
+
+		await commands.executeCommand('workbench.action.chat.newChat');
+		const result = await commands.executeCommand('workbench.action.chat.open', { query: 'hello' });
+		assert.strictEqual(result, undefined);
+	});
+
+	test('workbench.action.chat.open.blockOnResponse resolves when waiting for user confirmation to run a tool', async () => {
+		const toolRegistration = lm.registerTool<void>('requires_confirmation_tool', {
+			invoke: async (_options, _token) => null, prepareInvocation: async (_options, _token) => {
+				return { invocationMessage: 'Invoking', pastTenseMessage: 'Invoked', confirmationMessages: { title: 'Confirm', message: 'Are you sure?' } };
+			}
+		});
+
+		const participant = chat.createChatParticipant('api-test.participant', async (_request, _context, _progress, _token) => {
+			await lm.invokeTool('requires_confirmation_tool', {
+				input: {},
+				toolInvocationToken: _request.toolInvocationToken,
+			});
+			return { metadata: { complete: true } };
+		});
+		disposables.push(participant, toolRegistration);
+
+		await commands.executeCommand('workbench.action.chat.newChat');
+		const result: any = await commands.executeCommand('workbench.action.chat.open', { query: 'hello', blockOnResponse: true });
+		assert.strictEqual(result?.type, 'confirmation');
+	});
+
+	test('workbench.action.chat.open.blockOnResponse resolves when an error is hit', async () => {
+		const participant = chat.createChatParticipant('api-test.participant', async (_request, _context, _progress, _token) => {
+			return { errorDetails: { code: 'rate_limited', message: `You've been rate limited. Try again later!` } };
+		});
+		disposables.push(participant);
+
+		await commands.executeCommand('workbench.action.chat.newChat');
+		const result = await commands.executeCommand('workbench.action.chat.open', { query: 'hello', blockOnResponse: true });
+		assert.strictEqual((result as any).errorDetails.code, 'rate_limited');
+	});
+
 	test.skip('title provider is called for first request', async () => {
 		let calls = 0;
 		const deferred = new DeferredPromise<void>();

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

@@ -599,7 +599,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
 						responseIsIncomplete: true
 					};
 				}
-				if (errorDetails?.responseIsRedacted || errorDetails?.isQuotaExceeded || errorDetails?.confirmationButtons) {
+				if (errorDetails?.responseIsRedacted || errorDetails?.isQuotaExceeded || errorDetails?.confirmationButtons || errorDetails?.code) {
 					checkProposedApiEnabled(agent.extension, 'chatParticipantPrivate');
 				}
 

src/vs/workbench/contrib/chat/browser/actions/chatActions.ts

@@ -51,7 +51,7 @@ import { IWorkbenchLayoutService, Parts } from '../../../../services/layout/brow
 import { IViewsService } from '../../../../services/views/common/viewsService.js';
 import { IPreferencesService } from '../../../../services/preferences/common/preferences.js';
 import { EXTENSIONS_CATEGORY, IExtensionsWorkbenchService } from '../../../extensions/common/extensions.js';
-import { IChatAgentService } from '../../common/chatAgents.js';
+import { IChatAgentResult, IChatAgentService } from '../../common/chatAgents.js';
 import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { IChatEditingSession, ModifiedFileEntryState } from '../../common/chatEditingService.js';
 import { ChatEntitlement, IChatEntitlementService } from '../../common/chatEntitlementService.js';
@@ -72,6 +72,7 @@ import { VIEWLET_ID } from '../chatSessions.js';
 import { ChatViewPane } from '../chatViewPane.js';
 import { convertBufferToScreenshotVariable } from '../contrib/screenshot.js';
 import { clearChatEditor } from './chatClear.js';
+import { IChatResponseModel } from '../../common/chatModel.js';
 
 export const CHAT_CATEGORY = localize2('chat.category', 'Chat');
 
@@ -111,6 +112,10 @@ export interface IChatViewOpenOptions {
 	 * The mode ID or name to open the chat in.
 	 */
 	mode?: ChatModeKind | string;
+	/**
+	 * Wait to resolve the command until the chat response reaches a terminal state (complete, error, or pending user confirmation, etc.).
+	 */
+	blockOnResponse?: boolean;
 }
 
 export interface IChatViewOpenRequestEntry {
@@ -136,7 +141,7 @@ abstract class OpenChatGlobalAction extends Action2 {
 		});
 	}
 
-	override async run(accessor: ServicesAccessor, opts?: string | IChatViewOpenOptions): Promise<void> {
+	override async run(accessor: ServicesAccessor, opts?: string | IChatViewOpenOptions): Promise<IChatAgentResult & { type?: 'confirmation' } | undefined> {
 		opts = typeof opts === 'string' ? { query: opts } : opts;
 
 		const chatService = accessor.get(IChatService);
@@ -184,13 +189,16 @@ abstract class OpenChatGlobalAction extends Action2 {
 				}
 			}
 		}
+
+		let resp: Promise<IChatResponseModel | undefined> | undefined;
+
 		if (opts?.query) {
 			if (opts.isPartialQuery) {
 				chatWidget.setInput(opts.query);
 			} else {
 				await chatWidget.waitForReady();
 				await waitForDefaultAgent(chatAgentService, chatWidget.input.currentModeKind);
-				chatWidget.acceptInput(opts.query);
+				resp = chatWidget.acceptInput(opts.query);
 			}
 		}
 		if (opts?.toolIds && opts.toolIds.length > 0) {
@@ -210,6 +218,24 @@ abstract class OpenChatGlobalAction extends Action2 {
 		}
 
 		chatWidget.focusInput();
+
+		if (opts?.blockOnResponse) {
+			const response = await resp;
+			if (response) {
+				await new Promise<void>(resolve => {
+					const d = response.onDidChange(async () => {
+						if (response.isComplete || response.isPendingConfirmation.get()) {
+							d.dispose();
+							resolve();
+						}
+					});
+				});
+
+				return { ...response.result, type: response.isPendingConfirmation.get() ? 'confirmation' : undefined };
+			}
+		}
+
+		return undefined;
 	}
 
 	private async handleSwitchToMode(switchToMode: IChatMode, chatWidget: IChatWidget, instaService: IInstantiationService, commandService: ICommandService): Promise<void> {

src/vs/workbench/contrib/chat/common/chatService.ts

@@ -51,6 +51,7 @@ export interface IChatResponseErrorDetails {
 	isQuotaExceeded?: boolean;
 	level?: ChatErrorLevel;
 	confirmationButtons?: IChatResponseErrorDetailsConfirmationButton[];
+	code?: string;
 }
 
 export interface IChatResponseProgressFileTreeData {

src/vscode-dts/vscode.proposed.chatParticipantPrivate.d.ts

@@ -183,6 +183,8 @@ declare module 'vscode' {
 		isQuotaExceeded?: boolean;
 
 		level?: ChatErrorLevel;
+
+		code?: string;
 	}
 
 	export namespace chat {