From 17426d7a08ed9a5ca6b1ceef620802ffd840cac2 Mon Sep 17 00:00:00 2001 From: Nathaniel H Date: Sun, 3 May 2026 11:52:51 -0700 Subject: [PATCH 1/2] Add APIUsage tracking to extChatEndpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(ext-endpoint): propagate usage from DataPart stream in ExtensionContributedChatEndpoint The extension-contributed chat endpoint previously returned a hardcoded zeroed-out usage object on every successful response, causing the Copilot Chat context-window meter to never update when routing through third-party language model providers. Add a `usageFromDataPart()` helper that checks incoming `LanguageModelDataPart` chunks for a native usage shape (`prompt_tokens`, `completion_tokens`, `total_tokens` — all numbers) without keying on any MIME type. This mirrors the shape-based check already present in the non-provider (internal) streaming path. In `makeChatRequest2()`: - Declare `usage: APIUsage | undefined` before the stream loop. - After handling known MIME types (StatefulMarker, ContextManagement), call `usageFromDataPart()` on every DataPart and capture the result. - Return `usage ?? ` so existing behaviour is preserved when no usage chunk arrives. No new MIME branch is introduced; any provider that emits a DataPart whose JSON payload matches the OpenAI usage shape will automatically populate the meter. --- .../endpoint/vscode-node/extChatEndpoint.ts | 41 ++++++++++++++++++- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/src/platform/endpoint/vscode-node/extChatEndpoint.ts b/src/platform/endpoint/vscode-node/extChatEndpoint.ts index 7f77e8202d..11cd595061 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,37 @@ 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; + }; + + if ( + typeof parsed.prompt_tokens === 'number' && + typeof parsed.completion_tokens === 'number' && + typeof parsed.total_tokens === 'number' + ) { + return { + 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 }, + }; + } + } 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 +235,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 +263,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 +287,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 }; From 02930736407c2bb2d0f35934b0c9fdd79780c004 Mon Sep 17 00:00:00 2001 From: Nate Corcoran Date: Sun, 3 May 2026 12:33:39 -0700 Subject: [PATCH 2/2] fix(endpoint): preserve completion_tokens_details in ext chat usage parsing - Extend usageFromDataPart to include completion_tokens_details when present in the payload, preventing reasoning/prediction token breakdown from being silently dropped for extension-contributed models. - Add extChatEndpoint.spec.ts with provider-specific regression tests covering shape-based usage chunk detection and completion_tokens_details passthrough, mirroring the coverage pattern of the regular SSE stream. --- .../endpoint/vscode-node/extChatEndpoint.ts | 9 +- .../vscode-node/test/extChatEndpoint.spec.ts | 150 ++++++++++++++++++ 2 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 src/platform/endpoint/vscode-node/test/extChatEndpoint.spec.ts diff --git a/src/platform/endpoint/vscode-node/extChatEndpoint.ts b/src/platform/endpoint/vscode-node/extChatEndpoint.ts index 11cd595061..11212b44d6 100644 --- a/src/platform/endpoint/vscode-node/extChatEndpoint.ts +++ b/src/platform/endpoint/vscode-node/extChatEndpoint.ts @@ -44,6 +44,7 @@ function usageFromDataPart(part: vscode.LanguageModelDataPart): APIUsage | undef completion_tokens?: unknown; total_tokens?: unknown; prompt_tokens_details?: unknown; + completion_tokens_details?: unknown; }; if ( @@ -51,7 +52,7 @@ function usageFromDataPart(part: vscode.LanguageModelDataPart): APIUsage | undef typeof parsed.completion_tokens === 'number' && typeof parsed.total_tokens === 'number' ) { - return { + const usage: APIUsage = { prompt_tokens: parsed.prompt_tokens, completion_tokens: parsed.completion_tokens, total_tokens: parsed.total_tokens, @@ -60,6 +61,12 @@ function usageFromDataPart(part: vscode.LanguageModelDataPart): APIUsage | undef ? (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. 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(); + }); +});