Handle partial tool streaming

src/extension/linkify/common/responseStreamWithLinkification.ts

@@ -154,8 +154,8 @@ export class ResponseStreamWithLinkification implements FinalizableChatResponseS
 		return this;
 	}
 
-	prepareToolInvocation(toolName: string): ChatResponseStream {
-		this.enqueue(() => this._progress.prepareToolInvocation(toolName), false);
+	prepareToolInvocation(toolCallId: string, toolName: string): ChatResponseStream {
+		this.enqueue(() => this._progress.prepareToolInvocation(toolCallId, toolName), false);
 		return this;
 	}
 

src/extension/prompt/node/pseudoStartStopConversationCallback.ts

@@ -69,7 +69,18 @@ export class PseudoStopStartResponseProcessor implements IResponseProcessor {
 		}
 
 		if (delta.beginToolCalls?.length) {
-			progress.prepareToolInvocation(getContributedToolName(delta.beginToolCalls[0].name));
+			for (const beginCall of delta.beginToolCalls) {
+				progress.prepareToolInvocation(beginCall.id ?? '', getContributedToolName(beginCall.name), {});
+			}
+		}
+
+		if (delta.copilotToolCallStreamUpdates?.length) {
+			for (const update of delta.copilotToolCallStreamUpdates) {
+				if (!update.name) {
+					continue;
+				}
+				progress.prepareToolInvocation(update.id ?? '', getContributedToolName(update.name), { partialInput: tryParsePartialToolInput(update.arguments) });
+			}
 		}
 	}
 
@@ -211,3 +222,15 @@ export function reportCitations(delta: IResponseDelta, progress: ChatResponseStr
 		});
 	}
 }
+
+function tryParsePartialToolInput(raw: string | undefined): unknown {
+	if (!raw) {
+		return raw;
+	}
+
+	try {
+		return JSON.parse(raw);
+	} catch {
+		return raw;
+	}
+}

src/extension/tools/node/createFileTool.tsx

@@ -180,6 +180,37 @@ export class CreateFileTool implements ICopilotTool<ICreateFileParams> {
 		};
 	}
 
+	async handleToolStream(options: vscode.LanguageModelToolInvocationStreamOptions<ICreateFileParams>, token: vscode.CancellationToken): Promise<vscode.LanguageModelToolStreamResult> {
+		let invocationMessage: MarkdownString;
+
+		if (options.rawInput && typeof options.rawInput === 'string') {
+			// Try to extract filePath from the raw JSON input
+			const filePathMatch = options.rawInput.match(/"filePath"\s*:\s*"([^"]+)"/);
+			const contentMatch = options.rawInput.match(/"content"\s*:\s*"((?:[^"\\]|\\.)*)/);
+
+			if (filePathMatch) {
+				const filePath = filePathMatch[1];
+				const uri = resolveToolInputPath(filePath, this.promptPathRepresentationService);
+
+				if (contentMatch) {
+					const content = contentMatch[1];
+					const charCount = content.length;
+					invocationMessage = new MarkdownString(l10n.t`Creating ${formatUriForFileWidget(uri)} (${charCount} characters)`);
+				} else {
+					invocationMessage = new MarkdownString(l10n.t`Creating ${formatUriForFileWidget(uri)}`);
+				}
+			} else {
+				invocationMessage = new MarkdownString(l10n.t`Creating file`);
+			}
+		} else {
+			invocationMessage = new MarkdownString(l10n.t`Creating file`);
+		}
+
+		return {
+			invocationMessage,
+		};
+	}
+
 	private sendTelemetry(requestId: string | undefined, model: string | undefined, fileExtension: string) {
 		/* __GDPR__
 			"createFileToolInvoked" : {

src/extension/vscode.proposed.chatParticipantAdditions.d.ts

@@ -80,8 +80,24 @@ declare module 'vscode' {
 	}
 
 	export class ChatPrepareToolInvocationPart {
+		/**
+		 * Unique identifier for this tool call, used to correlate streaming updates.
+		 */
+		toolCallId: string;
 		toolName: string;
-		constructor(toolName: string);
+		/**
+		 * Partial arguments that have streamed in for the tool invocation.
+		 */
+		streamData?: ChatToolInvocationStreamData;
+		constructor(toolCallId: string, toolName: string, streamData?: ChatToolInvocationStreamData);
+	}
+
+	export interface ChatToolInvocationStreamData {
+		/**
+		 * Partial or not-yet-validated arguments that have streamed from the language model.
+		 * Tools may use this to render interim UI while the full invocation input is collected.
+		 */
+		readonly partialInput?: unknown;
 	}
 
 	export interface ChatTerminalToolInvocationData {
@@ -347,7 +363,14 @@ declare module 'vscode' {
 
 		codeCitation(value: Uri, license: string, snippet: string): void;
 
-		prepareToolInvocation(toolName: string): void;
+		/**
+		 * Notifies the UI that a tool invocation is being prepared. Optional streaming data can be
+		 * provided to render partial arguments while the invocation input is still being generated.
+		 * @param toolCallId Unique identifier for this tool call, used to correlate streaming updates.
+		 * @param toolName The name of the tool being invoked.
+		 * @param streamData Optional streaming data with partial arguments.
+		 */
+		prepareToolInvocation(toolCallId: string, toolName: string, streamData?: ChatToolInvocationStreamData): void;
 
 		push(part: ExtendedChatResponsePart): void;
 
@@ -668,6 +691,36 @@ declare module 'vscode' {
 		model?: LanguageModelChat;
 	}
 
+	export interface LanguageModelToolInvocationStreamOptions<T> {
+		/**
+		 * Raw argument payload, such as the streamed JSON fragment from the language model.
+		 */
+		readonly rawInput?: unknown;
+
+		readonly chatRequestId?: string;
+		readonly chatSessionId?: string;
+		readonly chatInteractionId?: string;
+	}
+
+	export interface LanguageModelToolStreamResult {
+		/**
+		 * A customized progress message to show while the tool runs.
+		 */
+		invocationMessage?: string | MarkdownString;
+	}
+
+	export interface LanguageModelTool<T> {
+		/**
+		 * Called zero or more times before {@link LanguageModelTool.prepareInvocation} while the
+		 * language model streams argument data for the invocation. Use this to update progress
+		 * or UI with the partial arguments that have been generated so far.
+		 *
+		 * Implementations must be free of side-effects and should be resilient to receiving
+		 * malformed or incomplete input.
+		 */
+		handleToolStream?(options: LanguageModelToolInvocationStreamOptions<T>, token: CancellationToken): ProviderResult<LanguageModelToolStreamResult>;
+	}
+
 	export interface ChatRequest {
 		readonly modeInstructions?: string;
 		readonly modeInstructions2?: ChatRequestModeInstructions;

src/platform/chat/common/chatMLFetcher.ts

@@ -125,7 +125,7 @@ export class FetchStreamRecorder {
 		callback: FinishedCallback | undefined
 	) {
 		this.callback = async (text: string, index: number, delta: IResponseDelta): Promise<number | undefined> => {
-			if (this._firstTokenEmittedTime === undefined && (delta.text || delta.beginToolCalls || (typeof delta.thinking?.text === 'string' && delta.thinking?.text || delta.thinking?.text?.length) || delta.copilotToolCalls)) {
+			if (this._firstTokenEmittedTime === undefined && (delta.text || delta.beginToolCalls || (typeof delta.thinking?.text === 'string' && delta.thinking?.text || delta.thinking?.text?.length) || delta.copilotToolCalls || delta.copilotToolCallStreamUpdates)) {
 				this._firstTokenEmittedTime = Date.now();
 			}
 

src/platform/endpoint/node/messagesApi.ts

@@ -310,14 +310,15 @@ export class AnthropicMessagesProcessor {
 				return;
 			case 'content_block_start':
 				if (chunk.content_block?.type === 'tool_use' && chunk.index !== undefined) {
+					const toolCallId = chunk.content_block.id || generateUuid();
 					this.toolCallAccumulator.set(chunk.index, {
-						id: chunk.content_block.id || generateUuid(),
+						id: toolCallId,
 						name: chunk.content_block.name || '',
 						arguments: '',
 					});
 					onProgress({
 						text: '',
-						beginToolCalls: [{ name: chunk.content_block.name || '' }]
+						beginToolCalls: [{ name: chunk.content_block.name || '', id: toolCallId }]
 					});
 				} else if (chunk.content_block?.type === 'thinking' && chunk.index !== undefined) {
 					this.thinkingAccumulator.set(chunk.index, {

src/platform/endpoint/node/responsesApi.ts

@@ -411,7 +411,7 @@ export class OpenAIResponsesProcessor {
 				if (chunk.item.type === 'function_call') {
 					onProgress({
 						text: '',
-						beginToolCalls: [{ name: chunk.item.name }]
+						beginToolCalls: [{ name: chunk.item.name, id: chunk.item.call_id }]
 					});
 				}
 				return;

src/platform/networking/common/fetch.ts

@@ -84,8 +84,15 @@ export interface ICopilotToolCall {
 	id: string;
 }
 
+export interface ICopilotToolCallStreamUpdate {
+	name: string;
+	arguments: string;
+	id?: string;
+}
+
 export interface ICopilotBeginToolCall {
 	name: string;
+	id?: string;
 }
 
 /**
@@ -128,6 +135,7 @@ export interface IResponseDelta {
 	copilotReferences?: ICopilotReference[];
 	copilotErrors?: ICopilotError[];
 	copilotToolCalls?: ICopilotToolCall[];
+	copilotToolCallStreamUpdates?: ICopilotToolCallStreamUpdate[];
 	beginToolCalls?: ICopilotBeginToolCall[];
 	_deprecatedCopilotFunctionCalls?: ICopilotFunctionCall[];
 	copilotConfirmation?: ICopilotConfirmation;

src/platform/networking/node/stream.ts

@@ -10,7 +10,7 @@ import { ITelemetryService } from '../../telemetry/common/telemetry';
 import { TelemetryData } from '../../telemetry/common/telemetryData';
 import { RawThinkingDelta, ThinkingDelta } from '../../thinking/common/thinking';
 import { extractThinkingDeltaFromChoice, } from '../../thinking/common/thinkingUtils';
-import { FinishedCallback, getRequestId, ICodeVulnerabilityAnnotation, ICopilotBeginToolCall, ICopilotConfirmation, ICopilotError, ICopilotFunctionCall, ICopilotReference, ICopilotToolCall, IIPCodeCitation, isCodeCitationAnnotation, isCopilotAnnotation, RequestId } from '../common/fetch';
+import { FinishedCallback, getRequestId, ICodeVulnerabilityAnnotation, ICopilotBeginToolCall, ICopilotConfirmation, ICopilotError, ICopilotFunctionCall, ICopilotReference, ICopilotToolCall, ICopilotToolCallStreamUpdate, IIPCodeCitation, isCodeCitationAnnotation, isCopilotAnnotation, RequestId } from '../common/fetch';
 import { Response } from '../common/fetcherService';
 import { APIErrorResponse, APIJsonData, APIUsage, ChoiceLogProbs, FilterReason, FinishedCompletionReason, isApiUsage, IToolCall } from '../common/openai';
 
@@ -69,7 +69,9 @@ class StreamingToolCall {
 
 	constructor() { }
 
-	update(toolCall: IToolCall) {
+	update(toolCall: IToolCall): boolean {
+		let argumentsChanged = false;
+
 		if (toolCall.id) {
 			this.id = toolCall.id;
 		}
@@ -80,7 +82,10 @@ class StreamingToolCall {
 
 		if (toolCall.function?.arguments) {
 			this.arguments += toolCall.function.arguments;
+			argumentsChanged = true;
 		}
+
+		return argumentsChanged;
 	}
 }
 
@@ -103,16 +108,31 @@ class StreamingToolCalls {
 		return this.toolCalls.length > 0;
 	}
 
-	update(choice: ExtendedChoiceJSON) {
+	update(choice: ExtendedChoiceJSON): ICopilotToolCallStreamUpdate[] {
+		const updates: ICopilotToolCallStreamUpdate[] = [];
 		choice.delta?.tool_calls?.forEach(toolCall => {
-			let currentCall = this.toolCalls.at(-1);
-			if (!currentCall || (toolCall.id && currentCall.id !== toolCall.id)) {
+			let currentCall: StreamingToolCall | undefined;
+			if (toolCall.id) {
+				currentCall = this.toolCalls.find(call => call.id === toolCall.id);
+			}
+			if (!currentCall) {
+				currentCall = this.toolCalls.at(-1);
+			}
+			if (!currentCall || (toolCall.id && currentCall.id && currentCall.id !== toolCall.id)) {
 				currentCall = new StreamingToolCall();
 				this.toolCalls.push(currentCall);
 			}
 
-			currentCall.update(toolCall);
+			const argumentsChanged = currentCall.update(toolCall);
+			if (argumentsChanged && currentCall.name) {
+				updates.push({
+					name: currentCall.name,
+					arguments: currentCall.arguments,
+					id: currentCall.id,
+				});
+			}
 		});
+		return updates;
 	}
 }
 
@@ -300,7 +320,7 @@ export class SSEProcessor {
 				return;
 			}
 
-			// this.logService.public.debug(chunk.toString());
+			this.logService.debug(chunk.toString());
 			const [dataLines, remainder] = splitChunk(extraData + chunk.toString());
 			extraData = remainder;
 
@@ -418,7 +438,7 @@ export class SSEProcessor {
 
 					let finishOffset: number | undefined;
 
-					const emitSolution = async (delta?: { vulnAnnotations?: ICodeVulnerabilityAnnotation[]; ipCodeCitations?: IIPCodeCitation[]; references?: ICopilotReference[]; toolCalls?: ICopilotToolCall[]; functionCalls?: ICopilotFunctionCall[]; errors?: ICopilotError[]; beginToolCalls?: ICopilotBeginToolCall[]; thinking?: ThinkingDelta }) => {
+					const emitSolution = async (delta?: { vulnAnnotations?: ICodeVulnerabilityAnnotation[]; ipCodeCitations?: IIPCodeCitation[]; references?: ICopilotReference[]; toolCalls?: ICopilotToolCall[]; toolCallStreamUpdates?: ICopilotToolCallStreamUpdate[]; functionCalls?: ICopilotFunctionCall[]; errors?: ICopilotError[]; beginToolCalls?: ICopilotBeginToolCall[]; thinking?: ThinkingDelta }) => {
 						if (delta?.vulnAnnotations && (!Array.isArray(delta.vulnAnnotations) || !delta.vulnAnnotations.every(a => isCopilotAnnotation(a)))) {
 							delta.vulnAnnotations = undefined;
 						}
@@ -435,6 +455,7 @@ export class SSEProcessor {
 							ipCitations: delta?.ipCodeCitations,
 							copilotReferences: delta?.references,
 							copilotToolCalls: delta?.toolCalls,
+							copilotToolCallStreamUpdates: delta?.toolCallStreamUpdates,
 							_deprecatedCopilotFunctionCalls: delta?.functionCalls,
 							beginToolCalls: delta?.beginToolCalls,
 							copilotErrors: delta?.errors,
@@ -448,17 +469,26 @@ export class SSEProcessor {
 
 					let handled = true;
 					if (choice.delta?.tool_calls) {
-						if (!this.toolCalls.hasToolCalls()) {
-							const firstToolName = choice.delta.tool_calls.at(0)?.function?.name;
+						const hadExistingToolCalls = this.toolCalls.hasToolCalls();
+						if (!hadExistingToolCalls) {
+							const firstToolCall = choice.delta.tool_calls.at(0);
+							const firstToolName = firstToolCall?.function?.name;
 							if (firstToolName) {
 								if (solution.text.length) {
 									// Flush the linkifier stream. See #16465
 									solution.append({ index: 0, delta: { content: ' ' } });
 								}
-								await emitSolution({ beginToolCalls: [{ name: firstToolName }] });
+								if (await emitSolution({ beginToolCalls: [{ name: firstToolName, id: firstToolCall?.id }] })) {
+									continue;
+								}
+							}
+						}
+						const toolCallStreamUpdates = this.toolCalls.update(choice);
+						if (toolCallStreamUpdates.length) {
+							if (await emitSolution({ toolCallStreamUpdates })) {
+								continue;
 							}
 						}
-						this.toolCalls.update(choice);
 					} else if (choice.delta?.copilot_annotations?.CodeVulnerability || choice.delta?.copilot_annotations?.IPCodeCitations) {
 						if (await emitSolution()) {
 							continue;

src/util/common/chatResponseStreamImpl.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { ChatResponseReferencePartStatusKind } from '@vscode/prompt-tsx';
-import type { ChatResponseFileTree, ChatResponseStream, ChatVulnerability, Command, ExtendedChatResponsePart, Location, NotebookEdit, Progress, ThinkingDelta, Uri } from 'vscode';
+import type { ChatResponseFileTree, ChatResponseStream, ChatToolInvocationStreamData, ChatVulnerability, Command, ExtendedChatResponsePart, Location, NotebookEdit, Progress, ThinkingDelta, Uri } from 'vscode';
 import { ChatPrepareToolInvocationPart, ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseExternalEditPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseWarningPart, MarkdownString, TextEdit } from '../../vscodeTypes';
 import type { ThemeIcon } from '../vs/base/common/themables';
 
@@ -167,7 +167,7 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 		this._push(new ChatResponseWarningPart(value));
 	}
 
-	prepareToolInvocation(toolName: string): void {
-		this._push(new ChatPrepareToolInvocationPart(toolName));
+	prepareToolInvocation(toolCallId: string, toolName: string, streamData: ChatToolInvocationStreamData): void {
+		this._push(new ChatPrepareToolInvocationPart(toolCallId, toolName, streamData));
 	}
 }

src/util/common/test/shims/chatTypes.ts

@@ -213,12 +213,18 @@ export class ChatResponseConfirmationPart {
 }
 
 export class ChatPrepareToolInvocationPart {
+	toolCallId: string;
 	toolName: string;
+	streamData?: { partialInput?: unknown };
 	/**
+	 * @param toolCallId Unique identifier for this tool call.
 	 * @param toolName The name of the tool being prepared for invocation.
+	 * @param streamData Partial arguments that have streamed in for the tool invocation.
 	 */
-	constructor(toolName: string) {
+	constructor(toolCallId: string, toolName: string, streamData?: { partialInput?: unknown }) {
+		this.toolCallId = toolCallId;
 		this.toolName = toolName;
+		this.streamData = streamData;
 	}
 }