diff --git a/src/platform/endpoint/vscode-node/extChatEndpoint.ts b/src/platform/endpoint/vscode-node/extChatEndpoint.ts index 7f77e8202d..11212b44d6 100644 --- a/src/platform/endpoint/vscode-node/extChatEndpoint.ts +++ b/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 }; diff --git a/src/platform/endpoint/vscode-node/test/extChatEndpoint.spec.ts b/src/platform/endpoint/vscode-node/test/extChatEndpoint.spec.ts new file mode 100644 index 0000000000..c9c9faca4d --- /dev/null +++ b/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 { + 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 { + 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 { + 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(); + }); +});