Support streaming of partial tool data

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

@@ -93,6 +93,7 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
 					}
 				},
 				prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token),
+				handleToolStream: (context, token) => this._proxy.$handleToolStream(id, context, token),
 			});
 		this._tools.set(id, disposable);
 	}

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

@@ -64,7 +64,7 @@ import { IChatSessionItem, IChatSessionProviderOptionGroup, IChatSessionProvider
 import { IChatRequestVariableValue } from '../../contrib/chat/common/chatVariables.js';
 import { ChatAgentLocation } from '../../contrib/chat/common/constants.js';
 import { IChatMessage, IChatResponsePart, ILanguageModelChatMetadataAndIdentifier, ILanguageModelChatSelector } from '../../contrib/chat/common/languageModels.js';
-import { IPreparedToolInvocation, IToolInvocation, IToolInvocationPreparationContext, IToolProgressStep, IToolResult, ToolDataSource } from '../../contrib/chat/common/languageModelToolsService.js';
+import { IPreparedToolInvocation, IStreamedToolInvocation, IToolInvocation, IToolInvocationPreparationContext, IToolInvocationStreamContext, IToolProgressStep, IToolResult, ToolDataSource } 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';
@@ -1498,6 +1498,7 @@ export interface ExtHostLanguageModelToolsShape {
 	$invokeTool(dto: Dto<IToolInvocation>, token: CancellationToken): Promise<Dto<IToolResult> | SerializableObjectWithBuffers<Dto<IToolResult>>>;
 	$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number>;
 
+	$handleToolStream(toolId: string, context: IToolInvocationStreamContext, token: CancellationToken): Promise<IStreamedToolInvocation | undefined>;
 	$prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
 }
 

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

@@ -296,11 +296,11 @@ export class ChatAgentResponseStream {
 					_report(dto);
 					return this;
 				},
-				prepareToolInvocation(toolName) {
+				prepareToolInvocation(toolCallId, toolName, streamData) {
 					throwIfDone(this.prepareToolInvocation);
 					checkProposedApiEnabled(that._extension, 'chatParticipantAdditions');
 
-					const part = new extHostTypes.ChatPrepareToolInvocationPart(toolName);
+					const part = new extHostTypes.ChatPrepareToolInvocationPart(toolCallId, toolName, streamData);
 					const dto = typeConvert.ChatPrepareToolInvocationPart.from(part);
 					_report(dto);
 					return this;

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

@@ -12,7 +12,7 @@ import { IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
 import { revive } from '../../../base/common/marshalling.js';
 import { generateUuid } from '../../../base/common/uuid.js';
 import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
-import { IPreparedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolInvocationPreparationContext, IToolResult, ToolInvocationPresentation } from '../../contrib/chat/common/languageModelToolsService.js';
+import { IPreparedToolInvocation, IStreamedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolInvocationPreparationContext, IToolInvocationStreamContext, IToolResult, ToolInvocationPresentation } from '../../contrib/chat/common/languageModelToolsService.js';
 import { ExtensionEditToolId, InternalEditToolId } from '../../contrib/chat/common/tools/editFileTool.js';
 import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/tools.js';
 import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
@@ -242,6 +242,37 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
 		return model;
 	}
 
+	async $handleToolStream(toolId: string, context: IToolInvocationStreamContext, token: CancellationToken): Promise<IStreamedToolInvocation | undefined> {
+		const item = this._registeredTools.get(toolId);
+		if (!item) {
+			throw new Error(`Unknown tool ${toolId}`);
+		}
+
+		// Only call handleToolStream if it's defined on the tool
+		if (!item.tool.handleToolStream) {
+			return undefined;
+		}
+
+		// Ensure the chatParticipantAdditions API is enabled
+		checkProposedApiEnabled(item.extension, 'chatParticipantAdditions');
+
+		const options: vscode.LanguageModelToolInvocationStreamOptions<any> = {
+			rawInput: context.rawInput,
+			chatRequestId: context.chatRequestId,
+			chatSessionId: context.chatSessionId,
+			chatInteractionId: context.chatInteractionId
+		};
+
+		const result = await item.tool.handleToolStream(options, token);
+		if (!result) {
+			return undefined;
+		}
+
+		return {
+			invocationMessage: typeConvert.MarkdownString.fromStrict(result.invocationMessage)
+		};
+	}
+
 	async $prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
 		const item = this._registeredTools.get(toolId);
 		if (!item) {

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

@@ -2815,12 +2815,16 @@ export namespace ChatPrepareToolInvocationPart {
 	export function from(part: vscode.ChatPrepareToolInvocationPart): IChatPrepareToolInvocationPart {
 		return {
 			kind: 'prepareToolInvocation',
+			toolCallId: part.toolCallId,
 			toolName: part.toolName,
+			streamData: part.streamData ? {
+				partialInput: part.streamData.partialInput
+			} : undefined
 		};
 	}
 
 	export function to(part: IChatPrepareToolInvocationPart): vscode.ChatPrepareToolInvocationPart {
-		return new types.ChatPrepareToolInvocationPart(part.toolName);
+		return new types.ChatPrepareToolInvocationPart(part.toolCallId, part.toolName, part.streamData);
 	}
 }
 

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

@@ -3338,12 +3338,20 @@ export class ChatResponseNotebookEditPart implements vscode.ChatResponseNotebook
 }
 
 export class ChatPrepareToolInvocationPart {
+	toolCallId: string;
 	toolName: string;
+	streamData?: {
+		readonly partialInput?: unknown;
+	};
 	/**
+	 * @param toolCallId Unique identifier for this tool call.
 	 * @param toolName The name of the tool being prepared for invocation.
+	 * @param streamData Optional streaming data with partial arguments.
 	 */
-	constructor(toolName: string) {
+	constructor(toolCallId: string, toolName: string, streamData?: { readonly partialInput?: unknown }) {
+		this.toolCallId = toolCallId;
 		this.toolName = toolName;
+		this.streamData = streamData;
 	}
 }
 

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

@@ -0,0 +1,117 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { CancellationToken } from '../../../../../base/common/cancellation.js';
+import { MarkdownString } from '../../../../../base/common/htmlContent.js';
+import { IMarkdownRenderer } from '../../../../../platform/markdown/browser/markdownRenderer.js';
+import { localize } from '../../../../../nls.js';
+import { IChatPrepareToolInvocationPart, IChatProgressMessage } from '../../common/chatService.js';
+import { IChatRendererContent, isResponseVM } from '../../common/chatViewModel.js';
+import { ChatTreeItem } from '../chat.js';
+import { IChatContentPart, IChatContentPartRenderContext } from './chatContentParts.js';
+import { ChatProgressContentPart } from './chatProgressContentPart.js';
+import { ILanguageModelToolsService, IToolInvocationStreamContext } from '../../common/languageModelToolsService.js';
+import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
+import { IChatMarkdownAnchorService } from './chatMarkdownAnchorService.js';
+import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
+import { ILogService } from '../../../../../platform/log/common/log.js';
+
+/**
+ * Content part for rendering prepareToolInvocation progress.
+ * This handles calling handleToolStream at the view layer and rendering the result.
+ * Similar to progressMessage, this hides when other content types arrive after it.
+ */
+export class ChatPrepareToolInvocationPart extends ChatProgressContentPart implements IChatContentPart {
+	private readonly element: ChatTreeItem;
+	private readonly isHiddenByFollowingContent: boolean;
+
+	constructor(
+		private prepareToolInvocation: IChatPrepareToolInvocationPart,
+		chatContentMarkdownRenderer: IMarkdownRenderer,
+		context: IChatContentPartRenderContext,
+		@IInstantiationService instantiationService: IInstantiationService,
+		@IChatMarkdownAnchorService chatMarkdownAnchorService: IChatMarkdownAnchorService,
+		@IConfigurationService configurationService: IConfigurationService,
+		@ILanguageModelToolsService private readonly languageModelToolsService: ILanguageModelToolsService,
+		@ILogService private readonly logService: ILogService
+	) {
+		// Create initial progress message - will be updated async
+		const initialProgressMessage = ChatPrepareToolInvocationPart.createInitialProgressMessage(prepareToolInvocation, languageModelToolsService);
+		super(initialProgressMessage, chatContentMarkdownRenderer, context, undefined, undefined, undefined, undefined, instantiationService, chatMarkdownAnchorService, configurationService);
+
+		this.element = context.element;
+
+		// Hide when following content contains parts that are not prepareToolInvocation or progressMessage
+		// This is similar to how progressMessage parts hide when other content arrives
+		const followingContent = context.content.slice(context.contentIndex + 1);
+		this.isHiddenByFollowingContent = followingContent.some(part => part.kind !== 'prepareToolInvocation' && part.kind !== 'progressMessage');
+
+		// Asynchronously call handleToolStream and update the message
+		if (!this.isHiddenByFollowingContent) {
+			this.fetchAndUpdateMessage(prepareToolInvocation);
+		}
+	}
+
+	private static createInitialProgressMessage(prepareToolInvocation: IChatPrepareToolInvocationPart, languageModelToolsService: ILanguageModelToolsService): IChatProgressMessage {
+		const toolData = languageModelToolsService.getTool(prepareToolInvocation.toolName);
+		const displayName = toolData?.displayName ?? prepareToolInvocation.toolName;
+		return {
+			kind: 'progressMessage',
+			content: new MarkdownString(localize('invokingTool', "Invoking tool: {0}", displayName))
+		};
+	}
+
+	private async fetchAndUpdateMessage(prepareToolInvocation: IChatPrepareToolInvocationPart): Promise<void> {
+		const toolData = this.languageModelToolsService.getTool(prepareToolInvocation.toolName);
+		if (!toolData) {
+			return;
+		}
+
+		const streamContext: IToolInvocationStreamContext = {
+			toolCallId: prepareToolInvocation.toolCallId,
+			rawInput: prepareToolInvocation.streamData?.partialInput,
+			chatRequestId: isResponseVM(this.element) ? this.element.requestId : undefined,
+			chatSessionId: this.element.sessionId
+		};
+
+		try {
+			const streamResult = await this.languageModelToolsService.handleToolStream(toolData.id, streamContext, CancellationToken.None);
+			if (streamResult?.invocationMessage) {
+				const progressContent = typeof streamResult.invocationMessage === 'string'
+					? new MarkdownString(streamResult.invocationMessage)
+					: new MarkdownString(streamResult.invocationMessage.value, { isTrusted: streamResult.invocationMessage.isTrusted, supportThemeIcons: streamResult.invocationMessage.supportThemeIcons, supportHtml: streamResult.invocationMessage.supportHtml });
+				this.updateMessage(progressContent);
+			}
+		} catch (error) {
+			this.logService.warn(`ChatPrepareToolInvocationPart: Error calling handleToolStream for tool ${toolData.id}`, error);
+		}
+	}
+
+	override hasSameContent(other: IChatRendererContent, followingContent: IChatRendererContent[], element: ChatTreeItem): boolean {
+		if (other.kind !== 'prepareToolInvocation') {
+			return false;
+		}
+
+		// If following content contains parts that are not prepareToolInvocation or progressMessage,
+		// we need to re-render to hide this part (similar to progressMessage behavior)
+		const shouldBeHidden = followingContent.some(part => part.kind !== 'prepareToolInvocation' && part.kind !== 'progressMessage');
+		if (shouldBeHidden && !this.isHiddenByFollowingContent) {
+			return false;
+		}
+
+		// Same toolCallId means this is an update to the same tool invocation
+		// We should reuse this part and update its content
+		if (other.toolCallId === this.prepareToolInvocation.toolCallId) {
+			// Update with new stream data
+			this.prepareToolInvocation = other;
+			if (!this.isHiddenByFollowingContent) {
+				this.fetchAndUpdateMessage(other);
+			}
+			return true;
+		}
+
+		return false;
+	}
+}

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

@@ -31,6 +31,7 @@ export class ChatProgressContentPart extends Disposable implements IChatContentP
 	private readonly showSpinner: boolean;
 	private readonly isHidden: boolean;
 	private readonly renderedMessage = this._register(new MutableDisposable<IRenderedMarkdown>());
+	private currentContent: IMarkdownString;
 
 	constructor(
 		progress: IChatProgressMessage | IChatTask | IChatTaskSerialized,
@@ -45,6 +46,7 @@ export class ChatProgressContentPart extends Disposable implements IChatContentP
 		@IConfigurationService private readonly configurationService: IConfigurationService
 	) {
 		super();
+		this.currentContent = progress.content;
 
 		const followingContent = context.content.slice(context.contentIndex + 1);
 		this.showSpinner = forceShowSpinner ?? shouldShowSpinner(followingContent, context.element);
@@ -100,6 +102,12 @@ export class ChatProgressContentPart extends Disposable implements IChatContentP
 
 		// Needs rerender when spinner state changes
 		const showSpinner = shouldShowSpinner(followingContent, element);
+
+		// Needs rerender when content changes
+		if (other.kind === 'progressMessage' && other.content.value !== this.currentContent.value) {
+			return false;
+		}
+
 		return other.kind === 'progressMessage' && this.showSpinner === showSpinner;
 	}
 

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

@@ -76,6 +76,7 @@ import { ChatExtensionsContentPart } from './chatContentParts/chatExtensionsCont
 import { ChatMarkdownContentPart } from './chatContentParts/chatMarkdownContentPart.js';
 import { ChatMcpServersInteractionContentPart } from './chatContentParts/chatMcpServersInteractionContentPart.js';
 import { ChatMultiDiffContentPart } from './chatContentParts/chatMultiDiffContentPart.js';
+import { ChatPrepareToolInvocationPart } from './chatContentParts/chatPrepareToolInvocationPart.js';
 import { ChatProgressContentPart, ChatWorkingProgressContentPart } from './chatContentParts/chatProgressContentPart.js';
 import { ChatPullRequestContentPart } from './chatContentParts/chatPullRequestContentPart.js';
 import { ChatQuotaExceededPart } from './chatContentParts/chatQuotaExceededPart.js';
@@ -771,6 +772,11 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 			return false;
 		}
 
+		// Don't show working if prepareToolInvocation is already present
+		if (partsToRender.some(part => part.kind === 'prepareToolInvocation')) {
+			return false;
+		}
+
 		// Show if no content, only "used references", ends with a complete tool call, or ends with complete text edits and there is no incomplete tool call (edits are still being applied some time after they are all generated)
 		const lastPart = findLast(partsToRender, part => part.kind !== 'markdownContent' || part.content.value.trim().length > 0);
 
@@ -790,7 +796,7 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 			((lastPart.kind === 'toolInvocation' || lastPart.kind === 'toolInvocationSerialized') && (IChatToolInvocation.isComplete(lastPart) || lastPart.presentation === 'hidden')) ||
 			((lastPart.kind === 'textEditGroup' || lastPart.kind === 'notebookEditGroup') && lastPart.done && !partsToRender.some(part => part.kind === 'toolInvocation' && !IChatToolInvocation.isComplete(part))) ||
 			(lastPart.kind === 'progressTask' && lastPart.deferred.isSettled) ||
-			lastPart.kind === 'prepareToolInvocation' || lastPart.kind === 'mcpServersStarting'
+			lastPart.kind === 'mcpServersStarting'
 		) {
 			return true;
 		}
@@ -1370,6 +1376,8 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 				return this.renderMcpServersInteractionRequired(content, context, templateData);
 			} else if (content.kind === 'thinking') {
 				return this.renderThinkingPart(content, context, templateData);
+			} else if (content.kind === 'prepareToolInvocation') {
+				return this.instantiationService.createInstance(ChatPrepareToolInvocationPart, content, this.chatContentMarkdownRenderer, context);
 			}
 
 			return this.renderNoContent(other => content.kind === other.kind);

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

@@ -40,7 +40,7 @@ import { ConfirmedReason, IChatService, IChatToolInvocation, ToolConfirmKind } f
 import { ChatRequestToolReferenceEntry, toToolSetVariableEntry, toToolVariableEntry } from '../common/chatVariableEntries.js';
 import { ChatConfiguration } from '../common/constants.js';
 import { ILanguageModelToolsConfirmationService } from '../common/languageModelToolsConfirmationService.js';
-import { CountTokensCallback, createToolSchemaUri, GithubCopilotToolReference, ILanguageModelToolsService, IPreparedToolInvocation, IToolAndToolSetEnablementMap, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultInputOutputDetails, stringifyPromptTsxPart, ToolDataSource, ToolSet, VSCodeToolReference } from '../common/languageModelToolsService.js';
+import { CountTokensCallback, createToolSchemaUri, GithubCopilotToolReference, ILanguageModelToolsService, IPreparedToolInvocation, IStreamedToolInvocation, IToolAndToolSetEnablementMap, IToolData, IToolImpl, IToolInvocation, IToolInvocationStreamContext, IToolResult, IToolResultInputOutputDetails, stringifyPromptTsxPart, ToolDataSource, ToolSet, VSCodeToolReference } from '../common/languageModelToolsService.js';
 import { getToolConfirmationAlert } from './chatAccessibilityProvider.js';
 
 const jsonSchemaRegistry = Registry.as<JSONContributionRegistry.IJSONContributionRegistry>(JSONContributionRegistry.Extensions.JSONContribution);
@@ -521,6 +521,35 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 		return prepared;
 	}
 
+	async handleToolStream(toolId: string, context: IToolInvocationStreamContext, token: CancellationToken): Promise<IStreamedToolInvocation | undefined> {
+		const tool = this._tools.get(toolId);
+		if (!tool) {
+			this._logService.warn(`[LanguageModelToolsService#handleToolStream] Tool ${toolId} not found`);
+			return undefined;
+		}
+
+		if (!tool.impl) {
+			await this._extensionService.activateByEvent(`onLanguageModelTool:${toolId}`);
+			const activatedTool = this._tools.get(toolId);
+			if (!activatedTool?.impl) {
+				this._logService.warn(`[LanguageModelToolsService#handleToolStream] Tool ${toolId} has no implementation after activation`);
+				return undefined;
+			}
+		}
+
+		if (!tool.impl?.handleToolStream) {
+			// Tool doesn't implement handleToolStream, that's fine
+			return undefined;
+		}
+
+		try {
+			return await tool.impl.handleToolStream(context, token);
+		} catch (error) {
+			this._logService.error(`[LanguageModelToolsService#handleToolStream] Error calling handleToolStream for tool ${toolId}:`, error);
+			return undefined;
+		}
+	}
+
 	private playAccessibilitySignal(toolInvocations: ChatToolInvocation[]): void {
 		const autoApproved = this._configurationService.getValue(ChatConfiguration.GlobalAutoApprove);
 		if (autoApproved) {