Start mocking out token use in response

Author: lramos15

src/extension/intents/node/toolCallingLoop.ts

@@ -517,7 +517,8 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 					toolCalls,
 					toolInputRetry,
 					statefulMarker,
-					thinking: thinkingItem
+					thinking: thinkingItem,
+					usage: fetchResult.usage
 				}),
 				chatResult,
 				hadIgnoredFiles: buildPromptResult.hasIgnoredFiles,

src/extension/prompt/common/intents.ts

@@ -6,6 +6,7 @@
 import type * as vscode from 'vscode';
 import { NotebookDocumentSnapshot } from '../../../platform/editing/common/notebookDocumentSnapshot';
 import { TextDocumentSnapshot } from '../../../platform/editing/common/textDocumentSnapshot';
+import { APIUsage } from '../../../platform/networking/common/openai';
 import { ThinkingData } from '../../../platform/thinking/common/thinking';
 import { ResourceMap } from '../../../util/vs/base/common/map';
 import { generateUuid } from '../../../util/vs/base/common/uuid';
@@ -30,6 +31,7 @@ export interface IToolCallRound {
 	toolCalls: IToolCall[];
 	thinking?: ThinkingData;
 	statefulMarker?: string;
+	usage?: APIUsage;
 }
 
 export interface InternalToolReference extends vscode.ChatLanguageModelToolReference {

src/extension/prompt/common/toolCallRound.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 import { FetchSuccess } from '../../../platform/chat/common/commonTypes';
+import { APIUsage } from '../../../platform/networking/common/openai';
 import { isEncryptedThinkingDelta, ThinkingData, ThinkingDelta } from '../../../platform/thinking/common/thinking';
 import { generateUuid } from '../../../util/vs/base/common/uuid';
 import { IToolCall, IToolCallRound } from './intents';
@@ -27,7 +28,8 @@ export class ToolCallRound implements IToolCallRound {
 			params.toolInputRetry,
 			params.id,
 			params.statefulMarker,
-			params.thinking
+			params.thinking,
+			params.usage
 		);
 		round.summary = params.summary;
 		return round;
@@ -39,14 +41,17 @@ export class ToolCallRound implements IToolCallRound {
 	 * @param toolInputRetry The number of times this round has been retried due to tool input validation failures
 	 * @param id A stable identifier for this round
 	 * @param statefulMarker Optional stateful marker used with the responses API
+	 * @param thinking Optional thinking data from the model
+	 * @param usage Optional API usage data for this round
 	 */
 	constructor(
 		public readonly response: string,
 		public readonly toolCalls: IToolCall[] = [],
 		public readonly toolInputRetry: number = 0,
 		public readonly id: string = ToolCallRound.generateID(),
 		public readonly statefulMarker?: string,
-		public readonly thinking?: ThinkingData
+		public readonly thinking?: ThinkingData,
+		public readonly usage?: APIUsage
 	) { }
 
 	private static generateID(): string {

src/extension/prompt/node/chatParticipantRequestHandler.ts

@@ -12,7 +12,7 @@ import { CanceledMessage, ChatLocation } from '../../../platform/chat/common/com
 import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider';
 import { IIgnoreService } from '../../../platform/ignore/common/ignoreService';
 import { ILogService } from '../../../platform/log/common/logService';
-import { FilterReason } from '../../../platform/networking/common/openai';
+import { APIUsage, FilterReason } from '../../../platform/networking/common/openai';
 import { ITabsAndEditorsService } from '../../../platform/tabs/common/tabsAndEditorsService';
 import { getWorkspaceFileDisplayPath, IWorkspaceService } from '../../../platform/workspace/common/workspaceService';
 import { ChatResponseStreamImpl } from '../../../util/common/chatResponseStreamImpl';
@@ -36,7 +36,7 @@ import { UnknownIntent } from '../../intents/node/unknownIntent';
 import { ContributedToolName } from '../../tools/common/toolNames';
 import { ChatVariablesCollection } from '../common/chatVariablesCollection';
 import { Conversation, getGlobalContextCacheKey, GlobalContextMessageMetadata, ICopilotChatResult, ICopilotChatResultIn, normalizeSummariesOnRounds, RenderedUserMessageMetadata, Turn, TurnStatus } from '../common/conversation';
-import { InternalToolReference } from '../common/intents';
+import { InternalToolReference, IToolCallRound } from '../common/intents';
 import { ChatTelemetryBuilder } from './chatParticipantTelemetry';
 import { DefaultIntentRequestHandler } from './defaultIntentRequestHandler';
 import { IDocumentContext } from './documentContext';
@@ -254,9 +254,14 @@ export class ChatParticipantRequestHandler {
 
 				result = await chatResult;
 				const endpoint = await this._endpointProvider.getChatEndpoint(this.request);
+				const resultMetadata = (result as ICopilotChatResultIn).metadata;
+				const totalUsage = getTotalUsage(resultMetadata?.toolCallRounds);
+				const usageStr = totalUsage
+					? ` • ↑${totalUsage.prompt_tokens} ↓${totalUsage.completion_tokens} tokens`
+					: '';
 				result.details = this._authService.copilotToken?.isNoAuthUser ?
-					`${endpoint.name}` :
-					`${endpoint.name} • ${endpoint.multiplier ?? 0}x`;
+					`${endpoint.name}${usageStr}` :
+					`${endpoint.name} • ${endpoint.multiplier ?? 0}x${usageStr}`;
 			}
 
 			this._conversationStore.addConversation(this.turn.id, this.conversation);
@@ -463,3 +468,36 @@ function anchorPartToMarkdown(workspaceService: IWorkspaceService, anchor: ChatR
 
 	return `[${text}](${path} ${anchor.title ? `"${anchor.title}"` : ''})`;
 }
+
+/**
+ * Calculates the total API usage by summing usage from all tool call rounds.
+ */
+function getTotalUsage(toolCallRounds: readonly IToolCallRound[] | undefined): APIUsage | undefined {
+	if (!toolCallRounds?.length) {
+		return undefined;
+	}
+
+	const initial: APIUsage = {
+		completion_tokens: 0,
+		prompt_tokens: 0,
+		total_tokens: 0,
+		prompt_tokens_details: { cached_tokens: 0 }
+	};
+
+	const hasAnyUsage = toolCallRounds.some(round => round.usage);
+	if (!hasAnyUsage) {
+		return undefined;
+	}
+
+	return toolCallRounds.reduce((acc, round): APIUsage => {
+		const usage = round.usage || initial;
+		return {
+			completion_tokens: acc.completion_tokens + usage.completion_tokens,
+			prompt_tokens: acc.prompt_tokens + usage.prompt_tokens,
+			total_tokens: acc.total_tokens + usage.total_tokens,
+			prompt_tokens_details: {
+				cached_tokens: (acc.prompt_tokens_details?.cached_tokens ?? 0) + (usage.prompt_tokens_details?.cached_tokens ?? 0),
+			}
+		};
+	}, initial);
+}