feat(copilotcli): enhance usage reporting in CopilotCLISession with detailed token metrics

extensions/copilot/src/extension/chatSessions/copilotcli/node/copilotcliSession.ts

@@ -13,6 +13,7 @@
 import { CopilotChatAttr, GenAiAttr, GenAiOperationName, IOTelService, ISpanHandle, SpanKind, SpanStatusCode, truncateForOTel } from '../../../../platform/otel/common/index';
 import { CapturingToken } from '../../../../platform/requestLogger/common/capturingToken';
 import { IRequestLogger, LoggedRequestKind } from '../../../../platform/requestLogger/common/requestLogger';
+import { PromptTokenCategory, PromptTokenLabel } from '../../../../platform/tokenizer/node/promptTokenDetails';
 import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
 import { raceCancellation } from '../../../../util/vs/base/common/async';
 import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
@@ -403,6 +404,22 @@
 
 		const chunkMessageIds = new Set<string>();
 		const assistantMessageChunks: string[] = [];
+		let lastUsageInfo: UsageInfoData | undefined;
+		const reportUsage = (promptTokens: number, completionTokens: number) => {
+			if (token.isCancellationRequested || !this._stream) {
+				return;
+			}
+			this._stream.usage({
+				promptTokens,
+				completionTokens,
+				promptTokenDetails: buildPromptTokenDetails(lastUsageInfo),
+			});
+		};
+		const updateUsageInfo = (async () => {
+			const metrics = await this._sdkSession.usage.getMetrics();
+			const promptTokens = lastUsageInfo?.currentTokens || metrics.lastCallInputTokens;
+			reportUsage(promptTokens, metrics.lastCallOutputTokens);
+		})();
 		try {
 			const shouldHandleExitPlanModeRequests = this.configurationService.getConfig(ConfigKey.Advanced.CLIPlanExitModeEnabled);
 			disposables.add(toDisposable(this._sdkSession.on('*', (event) => {
@@ -558,12 +575,19 @@
 			})));
 			disposables.add(toDisposable(this._sdkSession.on('assistant.usage', (event) => {
 				if (this._stream && typeof event.data.outputTokens === 'number' && typeof event.data.inputTokens === 'number') {
-					this._stream.usage({
-						completionTokens: event.data.outputTokens,
-						promptTokens: event.data.inputTokens,
-					});
+					reportUsage(event.data.inputTokens, event.data.outputTokens);
 				}
 			})));
+			disposables.add(toDisposable(this._sdkSession.on('session.usage_info', (event) => {
+				lastUsageInfo = {
+					currentTokens: event.data.currentTokens,
+					systemTokens: event.data.systemTokens,
+					conversationTokens: event.data.conversationTokens,
+					toolDefinitionsTokens: event.data.toolDefinitionsTokens,
+					tokenLimit: event.data.tokenLimit,
+				};
+				reportUsage(lastUsageInfo.currentTokens, 0);
+			})));
 			disposables.add(toDisposable(this._sdkSession.on('assistant.message_delta', (event) => {
 				// Support for streaming delta messages.
 				if (typeof event.data.deltaContent === 'string' && event.data.deltaContent.length) {
@@ -723,7 +747,6 @@
 				await this.sendRequestInternal(input, attachments, false, logStartTime);
 			}
 			this.logService.trace(`[CopilotCLISession] Invoking session (completed) ${this.sessionId}`);
-
 			const resolvedToolIdEditMap: Record<string, string> = {};
 			await Promise.all(Array.from(toolIdEditMap.entries()).map(async ([toolId, editFilePromise]) => {
 				const editId = await editFilePromise.catch(() => undefined);
@@ -741,6 +764,9 @@
 					this.logService.error(`[CopilotCLISession] Failed to update chat session metadata store for request ${request.id}`, error);
 				});
 			}
+			await updateUsageInfo.catch(error => {
+				this.logService.error(`[CopilotCLISession] Failed to update usage info after request ${request.id}`, error);
+			});
 			this._status = ChatSessionStatus.Completed;
 			this._statusChange.fire(this._status);
 
@@ -1270,3 +1296,41 @@
 	}
 }
 
+interface UsageInfoData {
+	readonly currentTokens: number;
+	readonly systemTokens?: number;
+	readonly conversationTokens?: number;
+	readonly toolDefinitionsTokens?: number;
+	readonly tokenLimit?: number;
+}
+
+function buildPromptTokenDetails(usageInfo: UsageInfoData | undefined): { category: string; label: string; percentageOfPrompt: number }[] | undefined {
+	if (!usageInfo || usageInfo.currentTokens <= 0) {
+		return undefined;
+	}
+	const details: { category: string; label: string; percentageOfPrompt: number }[] = [];
+	const total = usageInfo.currentTokens;
+	if (usageInfo.systemTokens && usageInfo.systemTokens > 0) {
+		details.push({
+			category: PromptTokenCategory.System,
+			label: PromptTokenLabel.SystemInstructions,
+			percentageOfPrompt: Math.round((usageInfo.systemTokens / total) * 100),
+		});
+	}
+	if (usageInfo.toolDefinitionsTokens && usageInfo.toolDefinitionsTokens > 0) {
+		details.push({
+			category: PromptTokenCategory.System,
+			label: PromptTokenLabel.Tools,
+			percentageOfPrompt: Math.round((usageInfo.toolDefinitionsTokens / total) * 100),
+		});
+	}
+	if (usageInfo.conversationTokens && usageInfo.conversationTokens > 0) {
+		details.push({
+			category: PromptTokenCategory.UserContext,
+			label: PromptTokenLabel.Messages,
+			percentageOfPrompt: Math.round((usageInfo.conversationTokens / total) * 100),
+		});
+	}
+	return details.length > 0 ? details : undefined;
+}
+