chat: allow tools to report progress

src/vs/platform/extensions/common/extensionsApiProposals.ts

@@ -395,6 +395,9 @@ const _allApiProposals = {
 	tokenInformation: {
 		proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.tokenInformation.d.ts',
 	},
+	toolProgress: {
+		proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.toolProgress.d.ts',
+	},
 	treeViewActiveItem: {
 		proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.treeViewActiveItem.d.ts',
 	},

src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts

@@ -6,7 +6,7 @@
 import { CancellationToken } from '../../../base/common/cancellation.js';
 import { Disposable, DisposableMap } from '../../../base/common/lifecycle.js';
 import { revive } from '../../../base/common/marshalling.js';
-import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolInvocation, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
+import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolInvocation, IToolProgressStep, IToolResult, ToolProgress } from '../../contrib/chat/common/languageModelToolsService.js';
 import { IExtHostContext, extHostNamedCustomer } from '../../services/extensions/common/extHostCustomers.js';
 import { Dto } from '../../services/extensions/common/proxyIdentifier.js';
 import { ExtHostContext, ExtHostLanguageModelToolsShape, MainContext, MainThreadLanguageModelToolsShape } from '../common/extHost.protocol.js';
@@ -16,7 +16,10 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
 
 	private readonly _proxy: ExtHostLanguageModelToolsShape;
 	private readonly _tools = this._register(new DisposableMap<string>());
-	private readonly _countTokenCallbacks = new Map</* call ID */string, CountTokensCallback>();
+	private readonly _runningToolCalls = new Map</* call ID */string, {
+		countTokens: CountTokensCallback;
+		progress: ToolProgress;
+	}>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -45,26 +48,30 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
 		};
 	}
 
+	$acceptToolProgress(callId: string, progress: IToolProgressStep): void {
+		this._runningToolCalls.get(callId)?.progress.report(progress);
+	}
+
 	$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number> {
-		const fn = this._countTokenCallbacks.get(callId);
+		const fn = this._runningToolCalls.get(callId);
 		if (!fn) {
 			throw new Error(`Tool invocation call ${callId} not found`);
 		}
 
-		return fn(input, token);
+		return fn.countTokens(input, token);
 	}
 
 	$registerTool(id: string): void {
 		const disposable = this._languageModelToolsService.registerToolImplementation(
 			id,
 			{
-				invoke: async (dto, countTokens, token) => {
+				invoke: async (dto, countTokens, progress, token) => {
 					try {
-						this._countTokenCallbacks.set(dto.callId, countTokens);
+						this._runningToolCalls.set(dto.callId, { countTokens, progress });
 						const resultDto = await this._proxy.$invokeTool(dto, token);
 						return revive(resultDto) as IToolResult;
 					} finally {
-						this._countTokenCallbacks.delete(dto.callId);
+						this._runningToolCalls.delete(dto.callId);
 					}
 				},
 				prepareToolInvocation: (parameters, token) => this._proxy.$prepareToolInvocation(id, parameters, token),

src/vs/workbench/api/common/extHost.protocol.ts

@@ -59,7 +59,7 @@ import { IChatContentInlineReference, IChatFollowup, IChatNotebookEdit, IChatPro
 import { IChatRequestVariableValue } from '../../contrib/chat/common/chatVariables.js';
 import { ChatAgentLocation } from '../../contrib/chat/common/constants.js';
 import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata, ILanguageModelChatSelector, ILanguageModelsChangeEvent } from '../../contrib/chat/common/languageModels.js';
-import { IPreparedToolInvocation, IToolData, IToolInvocation, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
+import { IPreparedToolInvocation, IToolData, IToolInvocation, IToolProgressStep, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
 import { DebugConfigurationProviderTriggerKind, IAdapterDescriptor, IConfig, IDebugSessionReplMode, IDebugTestRunReference, IDebugVisualization, IDebugVisualizationContext, IDebugVisualizationTreeItem, MainThreadDebugVisualization } from '../../contrib/debug/common/debug.js';
 import { McpCollectionDefinition, McpConnectionState, McpServerDefinition, McpServerLaunch } from '../../contrib/mcp/common/mcpTypes.js';
 import * as notebookCommon from '../../contrib/notebook/common/notebookCommon.js';
@@ -1380,6 +1380,7 @@ export type IToolDataDto = Omit<IToolData, 'when'>;
 
 export interface MainThreadLanguageModelToolsShape extends IDisposable {
 	$getTools(): Promise<Dto<IToolDataDto>[]>;
+	$acceptToolProgress(callId: string, progress: IToolProgressStep): void;
 	$invokeTool(dto: IToolInvocation, token?: CancellationToken): Promise<Dto<IToolResult>>;
 	$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number>;
 	$registerTool(id: string): void;

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

@@ -138,7 +138,21 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 			};
 		}
 
-		const extensionResult = await raceCancellation(Promise.resolve(item.tool.invoke(options, token)), token);
+		let progress: vscode.Progress<{ message?: string | vscode.MarkdownString; increment?: number }> | undefined;
+		if (isProposedApiEnabled(item.extension, 'toolProgress')) {
+			progress = {
+				report: value => {
+					this._proxy.$acceptToolProgress(dto.callId, {
+						message: typeConvert.MarkdownString.fromStrict(value.message),
+						increment: value.increment,
+						total: 100,
+					});
+				}
+			};
+		}
+
+		// todo: 'any' cast because TS can't handle the overloads
+		const extensionResult = await raceCancellation(Promise.resolve((item.tool.invoke as any)(options, token, progress!)), token);
 		if (!extensionResult) {
 			throw new CancellationError();
 		}

src/vs/workbench/contrib/chat/browser/chatContentParts/chatToolInvocationPart.ts

@@ -9,6 +9,7 @@ import { Codicon } from '../../../../../base/common/codicons.js';
 import { Emitter } from '../../../../../base/common/event.js';
 import { IMarkdownString, MarkdownString } from '../../../../../base/common/htmlContent.js';
 import { Disposable, DisposableStore, IDisposable, thenIfNotDisposed, toDisposable } from '../../../../../base/common/lifecycle.js';
+import { autorunWithStore } from '../../../../../base/common/observable.js';
 import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { URI } from '../../../../../base/common/uri.js';
 import { generateUuid } from '../../../../../base/common/uuid.js';
@@ -454,27 +455,37 @@ class ChatToolInvocationSubPart extends Disposable {
 	}
 
 	private createProgressPart(): HTMLElement {
-		let content: IMarkdownString;
 		if (this.toolInvocation.isComplete && this.toolInvocation.isConfirmed !== false && this.toolInvocation.pastTenseMessage) {
-			content = typeof this.toolInvocation.pastTenseMessage === 'string' ?
-				new MarkdownString().appendText(this.toolInvocation.pastTenseMessage) :
-				this.toolInvocation.pastTenseMessage;
+			const part = this.renderProgressContent(this.toolInvocation.pastTenseMessage);
+			this._register(part);
+			return part.domNode;
 		} else {
-			content = typeof this.toolInvocation.invocationMessage === 'string' ?
-				new MarkdownString().appendText(this.toolInvocation.invocationMessage + '…') :
-				MarkdownString.lift(this.toolInvocation.invocationMessage).appendText('…');
+			const container = document.createElement('div');
+			const progressObservable = this.toolInvocation.kind === 'toolInvocation' ? this.toolInvocation.progress : undefined;
+			this._register(autorunWithStore((reader, store) => {
+				const progress = progressObservable?.read(reader);
+				const part = store.add(this.renderProgressContent(progress?.message || this.toolInvocation.invocationMessage));
+				dom.reset(container, part.domNode);
+			}));
+			return container;
+		}
+	}
+
+	private renderProgressContent(content: IMarkdownString | string) {
+		if (typeof content === 'string') {
+			content = new MarkdownString().appendText(content);
 		}
 
 		const progressMessage: IChatProgressMessage = {
 			kind: 'progressMessage',
 			content
 		};
+
 		const iconOverride = !this.toolInvocation.isConfirmed ?
 			Codicon.error :
 			this.toolInvocation.isComplete ?
 				Codicon.check : undefined;
-		const progressPart = this._register(this.instantiationService.createInstance(ChatProgressContentPart, progressMessage, this.renderer, this.context, undefined, true, iconOverride));
-		return progressPart.domNode;
+		return this.instantiationService.createInstance(ChatProgressContentPart, progressMessage, this.renderer, this.context, undefined, true, iconOverride);
 	}
 
 	private createTerminalMarkdownProgressPart(toolInvocation: IChatToolInvocation | IChatToolInvocationSerialized, terminalData: IChatTerminalToolInvocationData): HTMLElement {

src/vs/workbench/contrib/chat/browser/chatSetup.ts

@@ -7,6 +7,7 @@ import { $ } from '../../../../base/browser/dom.js';
 import { Dialog } from '../../../../base/browser/ui/dialog/dialog.js';
 import { toAction, WorkbenchActionExecutedClassification, WorkbenchActionExecutedEvent } from '../../../../base/common/actions.js';
 import { timeout } from '../../../../base/common/async.js';
+import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { Codicon } from '../../../../base/common/codicons.js';
 import { toErrorMessage } from '../../../../base/common/errorMessage.js';
 import { isCancellationError } from '../../../../base/common/errors.js';
@@ -52,21 +53,20 @@ import { IHostService } from '../../../services/host/browser/host.js';
 import { IWorkbenchLayoutService, Parts } from '../../../services/layout/browser/layoutService.js';
 import { ILifecycleService } from '../../../services/lifecycle/common/lifecycle.js';
 import { IViewsService } from '../../../services/views/common/viewsService.js';
+import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolProgress } from '../../chat/common/languageModelToolsService.js';
 import { IExtensionsWorkbenchService } from '../../extensions/common/extensions.js';
 import { IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../common/chatAgents.js';
 import { ChatContextKeys } from '../common/chatContextKeys.js';
 import { ChatEntitlement, ChatEntitlementContext, ChatEntitlementRequests, ChatEntitlementService, IChatEntitlementService } from '../common/chatEntitlementService.js';
-import { IChatRequestModel, ChatRequestModel, ChatModel, IChatRequestVariableData, IChatRequestToolEntry } from '../common/chatModel.js';
+import { ChatModel, ChatRequestModel, IChatRequestModel, IChatRequestToolEntry, IChatRequestVariableData } from '../common/chatModel.js';
+import { ChatRequestAgentPart, ChatRequestToolPart } from '../common/chatParserTypes.js';
 import { IChatProgress, IChatService } from '../common/chatService.js';
 import { ChatAgentLocation, ChatConfiguration, ChatMode, validateChatMode } from '../common/constants.js';
 import { ILanguageModelsService } from '../common/languageModels.js';
 import { CHAT_CATEGORY, CHAT_OPEN_ACTION_ID, CHAT_SETUP_ACTION_ID } from './actions/chatActions.js';
 import { ChatViewId, IChatWidgetService, showCopilotView } from './chat.js';
 import { CHAT_SIDEBAR_PANEL_ID } from './chatViewPane.js';
 import './media/chatSetup.css';
-import { ChatRequestAgentPart, ChatRequestToolPart } from '../common/chatParserTypes.js';
-import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../chat/common/languageModelToolsService.js';
-import { CancellationToken } from '../../../../base/common/cancellation.js';
 
 const defaultChat = {
 	extensionId: product.defaultChatAgent?.extensionId ?? '',
@@ -518,7 +518,7 @@ class SetupTool extends Disposable implements IToolImpl {
 		super();
 	}
 
-	invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
+	invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken): Promise<IToolResult> {
 		const result: IToolResult = {
 			content: [
 				{

src/vs/workbench/contrib/chat/browser/languageModelToolsService.ts

@@ -12,7 +12,7 @@ import { Emitter } from '../../../../base/common/event.js';
 import { MarkdownString } from '../../../../base/common/htmlContent.js';
 import { Iterable } from '../../../../base/common/iterator.js';
 import { Lazy } from '../../../../base/common/lazy.js';
-import { Disposable, DisposableStore, dispose, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
 import { LRUCache } from '../../../../base/common/map.js';
 import { localize } from '../../../../nls.js';
 import { IAccessibilityService } from '../../../../platform/accessibility/common/accessibility.js';
@@ -40,6 +40,11 @@ interface IToolEntry {
 	impl?: IToolImpl;
 }
 
+interface ITrackedCall {
+	invocation?: ChatToolInvocation;
+	store: IDisposable;
+}
+
 export class LanguageModelToolsService extends Disposable implements ILanguageModelToolsService {
 	_serviceBrand: undefined;
 
@@ -53,7 +58,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 	private _toolContextKeys = new Set<string>();
 	private readonly _ctxToolsCount: IContextKey<number>;
 
-	private _callsByRequestId = new Map<string, IDisposable[]>();
+	private _callsByRequestId = new Map<string, ITrackedCall[]>();
 
 	private _workspaceToolConfirmStore: Lazy<ToolConfirmStore>;
 	private _profileToolConfirmStore: Lazy<ToolConfirmStore>;
@@ -235,7 +240,8 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 				if (!this._callsByRequestId.has(requestId)) {
 					this._callsByRequestId.set(requestId, []);
 				}
-				this._callsByRequestId.get(requestId)!.push(store);
+				const trackedCall: ITrackedCall = { store };
+				this._callsByRequestId.get(requestId)!.push(trackedCall);
 
 				const source = new CancellationTokenSource();
 				store.add(toDisposable(() => {
@@ -252,6 +258,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 
 				const prepared = await this.prepareToolInvocation(tool, dto, token);
 				toolInvocation = new ChatToolInvocation(prepared, tool.data, dto.callId);
+				trackedCall.invocation = toolInvocation;
 				if (this.shouldAutoConfirm(tool.data.id, tool.data.runsInWorkspace)) {
 					toolInvocation.confirmed.complete(true);
 				}
@@ -285,7 +292,11 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 				throw new CancellationError();
 			}
 
-			toolResult = await tool.impl.invoke(dto, countTokens, token);
+			toolResult = await tool.impl.invoke(dto, countTokens, {
+				report: step => {
+					toolInvocation?.acceptProgress(step);
+				}
+			}, token);
 			this.ensureToolDetails(dto, toolResult, tool.data);
 
 			this._telemetryService.publicLog2<LanguageModelToolInvokedEvent, LanguageModelToolInvokedClassification>(
@@ -403,7 +414,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 	private cleanupCallDisposables(requestId: string, store: DisposableStore): void {
 		const disposables = this._callsByRequestId.get(requestId);
 		if (disposables) {
-			const index = disposables.indexOf(store);
+			const index = disposables.findIndex(d => d.store === store);
 			if (index > -1) {
 				disposables.splice(index, 1);
 			}
@@ -417,15 +428,15 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 	cancelToolCallsForRequest(requestId: string): void {
 		const calls = this._callsByRequestId.get(requestId);
 		if (calls) {
-			calls.forEach(call => call.dispose());
+			calls.forEach(call => call.store.dispose());
 			this._callsByRequestId.delete(requestId);
 		}
 	}
 
 	public override dispose(): void {
 		super.dispose();
 
-		this._callsByRequestId.forEach(calls => dispose(calls));
+		this._callsByRequestId.forEach(calls => calls.forEach(call => call.store.dispose()));
 		this._ctxToolsCount.reset();
 	}
 }

src/vs/workbench/contrib/chat/common/chatProgressTypes/chatToolInvocation.ts

@@ -5,9 +5,10 @@
 
 import { DeferredPromise } from '../../../../../base/common/async.js';
 import { IMarkdownString } from '../../../../../base/common/htmlContent.js';
+import { observableValue } from '../../../../../base/common/observable.js';
 import { localize } from '../../../../../nls.js';
 import { IChatTerminalToolInvocationData, IChatToolInputInvocationData, IChatToolInvocation, IChatToolInvocationSerialized } from '../chatService.js';
-import { IPreparedToolInvocation, IToolConfirmationMessages, IToolData, IToolResult } from '../languageModelToolsService.js';
+import { IPreparedToolInvocation, IToolConfirmationMessages, IToolData, IToolProgressStep, IToolResult } from '../languageModelToolsService.js';
 
 export class ChatToolInvocation implements IChatToolInvocation {
 	public readonly kind: 'toolInvocation' = 'toolInvocation';
@@ -45,6 +46,8 @@ export class ChatToolInvocation implements IChatToolInvocation {
 
 	public readonly toolSpecificData?: IChatTerminalToolInvocationData | IChatToolInputInvocationData;
 
+	public readonly progress = observableValue<{ message?: string | IMarkdownString; progress: number }>(this, { progress: 0 });
+
 	constructor(preparedInvocation: IPreparedToolInvocation | undefined, toolData: IToolData, public readonly toolCallId: string) {
 		const defaultMessage = localize('toolInvocationMessage', "Using {0}", `"${toolData.displayName}"`);
 		const invocationMessage = preparedInvocation?.invocationMessage ?? defaultMessage;
@@ -84,6 +87,14 @@ export class ChatToolInvocation implements IChatToolInvocation {
 		return this._confirmationMessages;
 	}
 
+	public acceptProgress(step: IToolProgressStep) {
+		const prev = this.progress.get();
+		this.progress.set({
+			progress: step.increment ? (prev.progress + step.increment) : prev.progress,
+			message: step.message,
+		}, undefined);
+	}
+
 	public toJSON(): IChatToolInvocationSerialized {
 		return {
 			kind: 'toolInvocationSerialized',

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

@@ -7,6 +7,7 @@ import { DeferredPromise } from '../../../../base/common/async.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { Event } from '../../../../base/common/event.js';
 import { IMarkdownString } from '../../../../base/common/htmlContent.js';
+import { IObservable } from '../../../../base/common/observable.js';
 import { ThemeIcon } from '../../../../base/common/themables.js';
 import { URI } from '../../../../base/common/uri.js';
 import { IRange, Range } from '../../../../editor/common/core/range.js';
@@ -234,6 +235,7 @@ export interface IChatToolInvocation {
 	invocationMessage: string | IMarkdownString;
 	pastTenseMessage: string | IMarkdownString | undefined;
 	resultDetails: IToolResult['toolResultDetails'];
+	progress: IObservable<{ message?: string | IMarkdownString; progress: number }>;
 	readonly toolId: string;
 	readonly toolCallId: string;