Add APIUsage tracking to extChatEndpoint

src/platform/endpoint/vscode-node/extChatEndpoint.ts

@@ -18,7 +18,7 @@ import { ContextManagementResponse } from '../../networking/common/anthropic';
 import { FinishedCallback, OpenAiFunctionTool, OptionalChatRequestParams } from '../../networking/common/fetch';
 import { Response } from '../../networking/common/fetcherService';
 import { IChatEndpoint, ICreateEndpointBodyOptions, IEndpointBody, IMakeChatRequestOptions } from '../../networking/common/networking';
-import { ChatCompletion } from '../../networking/common/openai';
+import { APIUsage, ChatCompletion } from '../../networking/common/openai';
 import { IOTelService } from '../../otel/common/otelService';
 import { retrieveCapturingTokenByCorrelation, storeCapturingTokenForCorrelation } from '../../requestLogger/node/requestLogger';
 import { ITelemetryService } from '../../telemetry/common/telemetry';
@@ -37,6 +37,44 @@ enum ChatImageMimeType {
 	BMP = 'image/bmp',
 }
 
+function usageFromDataPart(part: vscode.LanguageModelDataPart): APIUsage | undefined {
+	try {
+		const parsed = JSON.parse(new TextDecoder().decode(part.data)) as {
+			prompt_tokens?: unknown;
+			completion_tokens?: unknown;
+			total_tokens?: unknown;
+			prompt_tokens_details?: unknown;
+			completion_tokens_details?: unknown;
+		};
+
+		if (
+			typeof parsed.prompt_tokens === 'number' &&
+			typeof parsed.completion_tokens === 'number' &&
+			typeof parsed.total_tokens === 'number'
+		) {
+			const usage: APIUsage = {
+				prompt_tokens: parsed.prompt_tokens,
+				completion_tokens: parsed.completion_tokens,
+				total_tokens: parsed.total_tokens,
+				prompt_tokens_details:
+					parsed.prompt_tokens_details && typeof parsed.prompt_tokens_details === 'object'
+						? (parsed.prompt_tokens_details as APIUsage['prompt_tokens_details'])
+						: { cached_tokens: 0 },
+			};
+
+			if (parsed.completion_tokens_details && typeof parsed.completion_tokens_details === 'object') {
+				usage.completion_tokens_details = parsed.completion_tokens_details as APIUsage['completion_tokens_details'];
+			}
+
+			return usage;
+		}
+	} catch {
+		// Ignore non-JSON or non-usage DataParts.
+	}
+
+	return undefined;
+}
+
 export class ExtensionContributedChatEndpoint implements IChatEndpoint {
 	private readonly _maxTokens: number;
 	public readonly isDefault: boolean = false;
@@ -204,6 +242,7 @@ export class ExtensionContributedChatEndpoint implements IChatEndpoint {
 			const response = await this.languageModel.sendRequest(vscodeMessages, vscodeOptions, token);
 			let text = '';
 			let numToolsCalled = 0;
+			let usage: APIUsage | undefined;
 			const requestId = ourRequestId;
 
 			// consume stream
@@ -231,6 +270,11 @@ export class ExtensionContributedChatEndpoint implements IChatEndpoint {
 						const contextManagement = JSON.parse(new TextDecoder().decode(chunk.data)) as ContextManagementResponse;
 						await streamRecorder.callback?.(text, 0, { text: '', contextManagement });
 					}
+					// Mirror native usage-only chunk handling by shape, not MIME.
+					const maybeUsage = usageFromDataPart(chunk);
+					if (maybeUsage) {
+						usage = maybeUsage;
+					}
 				} else if (chunk instanceof vscode.LanguageModelThinkingPart) {
 					if (streamRecorder.callback) {
 						await streamRecorder.callback(text, 0, {
@@ -250,7 +294,7 @@ export class ExtensionContributedChatEndpoint implements IChatEndpoint {
 					type: ChatFetchResponseType.Success,
 					requestId,
 					serverRequestId: requestId,
-					usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, prompt_tokens_details: { cached_tokens: 0 } },
+					usage: usage ?? { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, prompt_tokens_details: { cached_tokens: 0 } },
 					value: text,
 					resolvedModel: this.languageModel.id
 				};

src/platform/endpoint/vscode-node/test/extChatEndpoint.spec.ts

@@ -0,0 +1,150 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Raw } from '@vscode/prompt-tsx';
+import { describe, expect, it, vi } from 'vitest';
+import * as vscode from 'vscode';
+import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
+import { ChatFetchResponseType, ChatLocation } from '../../../chat/common/commonTypes';
+import { IOTelService } from '../../../otel/common/otelService';
+import { ExtensionContributedChatEndpoint } from '../extChatEndpoint';
+
+vi.mock('../../../requestLogger/node/requestLogger', () => ({
+	storeCapturingTokenForCorrelation: vi.fn(),
+	retrieveCapturingTokenByCorrelation: vi.fn(),
+}));
+
+class MockLanguageModelChat implements Partial<vscode.LanguageModelChat> {
+	public readonly id = 'test.model';
+	public readonly name = 'Test Model';
+	public readonly version = '1.0.0';
+	public readonly family = 'test';
+	public readonly vendor = 'test-vendor';
+	public readonly maxInputTokens = 8192;
+	public readonly capabilities = {
+		supportsToolCalling: false,
+		supportsImageToText: false,
+	};
+
+	constructor(private readonly parts: readonly vscode.LanguageModelResponsePart[]) {
+	}
+
+	sendRequest(): Thenable<vscode.LanguageModelChatResponse> {
+		const parts = this.parts;
+		return Promise.resolve({
+			stream: (async function* () {
+				for (const part of parts) {
+					yield part;
+				}
+			})(),
+			text: (async function* () {
+				for (const part of parts) {
+					if (part instanceof vscode.LanguageModelTextPart) {
+						yield part.value;
+					}
+				}
+			})(),
+		} as unknown as vscode.LanguageModelChatResponse);
+	}
+
+	countTokens(_input: string | vscode.LanguageModelChatMessage | vscode.LanguageModelChatMessage2, _token?: vscode.CancellationToken): Thenable<number> {
+		return Promise.resolve(0);
+	}
+}
+
+function createEndpoint(parts: readonly vscode.LanguageModelResponsePart[]): ExtensionContributedChatEndpoint {
+	const languageModel = new MockLanguageModelChat(parts);
+	const instantiationService = {} as unknown as IInstantiationService;
+	const otelService = {
+		getActiveTraceContext: () => undefined,
+	} as unknown as IOTelService;
+
+	return new ExtensionContributedChatEndpoint(languageModel as unknown as vscode.LanguageModelChat, instantiationService, otelService);
+}
+
+function createUserMessage(text: string): Raw.ChatMessage {
+	return {
+		role: Raw.ChatRole.User,
+		content: [{ type: Raw.ChatCompletionContentPartKind.Text, text }],
+	};
+}
+
+describe('ExtensionContributedChatEndpoint', () => {
+	it('parses usage data by shape even when mime type is provider-specific', async () => {
+		const usagePart = new vscode.LanguageModelDataPart(
+			new TextEncoder().encode(JSON.stringify({
+				prompt_tokens: 11,
+				completion_tokens: 7,
+				total_tokens: 18,
+				prompt_tokens_details: { cached_tokens: 4 },
+				completion_tokens_details: {
+					reasoning_tokens: 2,
+					accepted_prediction_tokens: 3,
+					rejected_prediction_tokens: 1,
+				},
+			})),
+			'application/vnd.provider.usage+json'
+		);
+
+		const endpoint = createEndpoint([
+			new vscode.LanguageModelTextPart('hello'),
+			usagePart,
+		]);
+
+		const result = await endpoint.makeChatRequest2({
+			debugName: 'test',
+			messages: [createUserMessage('hi')],
+			requestOptions: undefined,
+			finishedCb: undefined,
+			location: ChatLocation.Panel,
+			source: undefined,
+		}, new vscode.CancellationTokenSource().token);
+
+		expect(result.type).toBe(ChatFetchResponseType.Success);
+		if (result.type !== ChatFetchResponseType.Success) {
+			return;
+		}
+
+		expect(result.usage?.prompt_tokens).toBe(11);
+		expect(result.usage?.completion_tokens).toBe(7);
+		expect(result.usage?.total_tokens).toBe(18);
+		expect(result.usage?.prompt_tokens_details?.cached_tokens).toBe(4);
+		expect(result.usage?.completion_tokens_details?.reasoning_tokens).toBe(2);
+		expect(result.usage?.completion_tokens_details?.accepted_prediction_tokens).toBe(3);
+		expect(result.usage?.completion_tokens_details?.rejected_prediction_tokens).toBe(1);
+	});
+
+	it('falls back to default usage when data chunks are not usage-shaped', async () => {
+		const nonUsagePart = new vscode.LanguageModelDataPart(
+			new TextEncoder().encode(JSON.stringify({ context: { id: 'abc' } })),
+			'application/json'
+		);
+
+		const endpoint = createEndpoint([
+			new vscode.LanguageModelTextPart('hello'),
+			nonUsagePart,
+		]);
+
+		const result = await endpoint.makeChatRequest2({
+			debugName: 'test',
+			messages: [createUserMessage('hi')],
+			requestOptions: undefined,
+			finishedCb: undefined,
+			location: ChatLocation.Panel,
+			source: undefined,
+		}, new vscode.CancellationTokenSource().token);
+
+		expect(result.type).toBe(ChatFetchResponseType.Success);
+		if (result.type !== ChatFetchResponseType.Success) {
+			return;
+		}
+
+		expect(result.usage?.prompt_tokens).toBe(0);
+		expect(result.usage?.completion_tokens).toBe(0);
+		expect(result.usage?.total_tokens).toBe(0);
+		expect(result.usage?.prompt_tokens_details?.cached_tokens).toBe(0);
+		expect(result.usage?.completion_tokens_details).toBeUndefined();
+	});
+});