Consume new API (#3322)

* Consume new API

* Update snapshot

* Update api version

Author: lramos15

package.json

@@ -112,7 +112,7 @@
 		"mappedEditsProvider",
 		"aiRelatedInformation",
 		"aiSettingsSearch",
-		"chatParticipantAdditions",
+		"chatParticipantAdditions@1",
 		"defaultChatParticipant@4",
 		"contribSourceControlInputBoxMenu",
 		"authLearnMore",

src/extension/codeBlocks/node/codeBlockProcessor.ts

@@ -124,6 +124,7 @@ export class CodeBlockTrackingChatResponseStream implements ChatResponseStream {
 	externalEdit = this.forward(this._wrapped.externalEdit.bind(this._wrapped));
 	beginToolInvocation = this.forward(this._wrapped.beginToolInvocation.bind(this._wrapped));
 	updateToolInvocation = this.forward(this._wrapped.updateToolInvocation.bind(this._wrapped));
+	usage = this.forward(this._wrapped.usage.bind(this._wrapped));
 
 	questionCarousel(questions: ChatQuestion[], allowSkip?: boolean): Thenable<Record<string, unknown> | undefined> {
 		this._codeBlockProcessor.flush();

src/extension/intents/node/toolCallingLoop.ts

@@ -491,17 +491,16 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 			tools: availableTools,
 		});
 		fetchStreamSource?.resolve();
-		let chatResult = await processResponsePromise ?? undefined;
-
-		// hydrate the token usage into the chat result as this renders the context window widget
-		if (fetchResult.type === ChatFetchResponseType.Success && fetchResult.usage) {
-			chatResult = {
-				...chatResult, usage: {
-					completionTokens: fetchResult.usage.completion_tokens,
-					promptTokens: fetchResult.usage.prompt_tokens,
-					promptTokenDetails,
-				}
-			};
+		const chatResult = await processResponsePromise ?? undefined;
+
+		// Report token usage to the stream for rendering the context window widget
+		const stream = streamParticipants[streamParticipants.length - 1];
+		if (fetchResult.type === ChatFetchResponseType.Success && fetchResult.usage && stream) {
+			stream.usage({
+				completionTokens: fetchResult.usage.completion_tokens,
+				promptTokens: fetchResult.usage.prompt_tokens,
+				promptTokenDetails,
+			});
 		}
 
 		// Validate authentication session upgrade and handle accordingly

src/extension/linkify/common/responseStreamWithLinkification.ts

@@ -2,7 +2,7 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-import type { ChatQuestion, ChatResponseClearToPreviousToolInvocationReason, ChatResponseFileTree, ChatResponsePart, ChatResponseStream, ChatToolInvocationStreamData, ChatVulnerability, ChatWorkspaceFileEdit, Command, Location, NotebookEdit, TextEdit, ThinkingDelta, Uri } from 'vscode';
+import type { ChatQuestion, ChatResponseClearToPreviousToolInvocationReason, ChatResponseFileTree, ChatResponsePart, ChatResponseStream, ChatResultUsage, ChatToolInvocationStreamData, ChatVulnerability, ChatWorkspaceFileEdit, Command, Location, NotebookEdit, TextEdit, ThinkingDelta, Uri } from 'vscode';
 import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService';
 import { FinalizableChatResponseStream } from '../../../util/common/chatResponseStreamImpl';
 import { CancellationToken } from '../../../util/vs/base/common/cancellation';
@@ -171,6 +171,11 @@ export class ResponseStreamWithLinkification implements FinalizableChatResponseS
 		return this.enqueue(() => this._progress.questionCarousel(questions, allowSkip), true);
 	}
 
+	usage(usage: ChatResultUsage): ChatResponseStream {
+		this.enqueue(() => this._progress.usage(usage), false);
+		return this;
+	}
+
 	//#endregion
 
 	private sequencer: Promise<unknown> = Promise.resolve();

src/extension/prompt/node/test/__snapshots__/defaultIntentRequestHandler.spec.ts.snap

@@ -251,17 +251,6 @@ exports[`defaultIntentRequestHandler > confirms on max tool call iterations, and
       },
     ],
   },
-  "usage": {
-    "completionTokens": 0,
-    "promptTokenDetails": [
-      {
-        "category": "User Context",
-        "label": "Messages",
-        "percentageOfPrompt": 70,
-      },
-    ],
-    "promptTokens": 0,
-  },
 }
 `;
 
@@ -363,17 +352,6 @@ exports[`defaultIntentRequestHandler > confirms on max tool call iterations, and
       },
     ],
   },
-  "usage": {
-    "completionTokens": 0,
-    "promptTokenDetails": [
-      {
-        "category": "User Context",
-        "label": "Messages",
-        "percentageOfPrompt": 70,
-      },
-    ],
-    "promptTokens": 0,
-  },
 }
 `;
 
@@ -2473,17 +2451,6 @@ exports[`defaultIntentRequestHandler > makes a successful request with a single
       },
     ],
   },
-  "usage": {
-    "completionTokens": 0,
-    "promptTokenDetails": [
-      {
-        "category": "User Context",
-        "label": "Messages",
-        "percentageOfPrompt": 70,
-      },
-    ],
-    "promptTokens": 0,
-  },
 }
 `;
 
@@ -2822,17 +2789,6 @@ exports[`defaultIntentRequestHandler > makes a tool call turn 1`] = `
       },
     ],
   },
-  "usage": {
-    "completionTokens": 0,
-    "promptTokenDetails": [
-      {
-        "category": "User Context",
-        "label": "Messages",
-        "percentageOfPrompt": 70,
-      },
-    ],
-    "promptTokens": 0,
-  },
 }
 `;
 

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

@@ -3,6 +3,8 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+// version: 1
+
 declare module 'vscode' {
 
 	export interface ChatParticipant {
@@ -553,6 +555,13 @@ declare module 'vscode' {
 		push(part: ExtendedChatResponsePart): void;
 
 		clearToPreviousToolInvocation(reason: ChatResponseClearToPreviousToolInvocationReason): void;
+
+		/**
+		 * Report token usage information for this request.
+		 * This is typically called when the underlying language model provides usage statistics.
+		 * @param usage Token usage information including prompt and completion tokens
+		 */
+		usage(usage: ChatResultUsage): void;
 	}
 
 	export enum ChatResponseReferencePartStatusKind {
@@ -744,12 +753,6 @@ declare module 'vscode' {
 		 * An optional detail string that will be rendered at the end of the response in certain UI contexts.
 		 */
 		details?: string;
-
-		/**
-		 * Token usage information for this request, if available.
-		 * This is typically provided by the underlying language model.
-		 */
-		readonly usage?: ChatResultUsage;
 	}
 
 	export namespace chat {

src/util/common/chatResponseStreamImpl.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { ChatResponseReferencePartStatusKind } from '@vscode/prompt-tsx';
-import type { ChatQuestion, ChatResponseFileTree, ChatResponseStream, ChatToolInvocationStreamData, ChatVulnerability, ChatWorkspaceFileEdit, Command, ExtendedChatResponsePart, Location, NotebookEdit, Progress, ThinkingDelta, Uri } from 'vscode';
+import type { ChatQuestion, ChatResponseFileTree, ChatResponseStream, ChatResultUsage, ChatToolInvocationStreamData, ChatVulnerability, ChatWorkspaceFileEdit, Command, ExtendedChatResponsePart, Location, NotebookEdit, Progress, ThinkingDelta, Uri } from 'vscode';
 import { ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseExternalEditPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseWarningPart, ChatResponseWorkspaceEditPart, MarkdownString, TextEdit } from '../../vscodeTypes';
 import type { ThemeIcon } from '../vs/base/common/themables';
 
@@ -43,6 +43,9 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 			},
 			(questions, allowSkip) => {
 				return stream.questionCarousel(questions, allowSkip);
+			},
+			(usage) => {
+				stream.usage(usage);
 			}
 		);
 	}
@@ -66,6 +69,9 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 			},
 			(questions, allowSkip) => {
 				return stream.questionCarousel(questions, allowSkip);
+			},
+			(usage) => {
+				stream.usage(usage);
 			});
 	}
 
@@ -89,6 +95,9 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 			},
 			(questions, allowSkip) => {
 				return stream.questionCarousel(questions, allowSkip);
+			},
+			(usage) => {
+				stream.usage(usage);
 			});
 	}
 
@@ -99,6 +108,7 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 		private readonly _beginToolInvocation?: (toolCallId: string, toolName: string, streamData?: ChatToolInvocationStreamData) => void,
 		private readonly _updateToolInvocation?: (toolCallId: string, streamData: ChatToolInvocationStreamData) => void,
 		private readonly _questionCarousel?: (questions: ChatQuestion[], allowSkip?: boolean) => Thenable<Record<string, unknown> | undefined>,
+		private readonly _usage?: (usage: ChatResultUsage) => void,
 	) { }
 
 	async finalize(): Promise<void> {
@@ -219,4 +229,10 @@ export class ChatResponseStreamImpl implements FinalizableChatResponseStream {
 		}
 		return Promise.resolve(undefined);
 	}
+
+	usage(usage: ChatResultUsage): void {
+		if (this._usage) {
+			this._usage(usage);
+		}
+	}
 }