From b186c6d168134c4fafee4382ceb0f070dd3b643b Mon Sep 17 00:00:00 2001 From: Winston Chang Date: Wed, 8 Jul 2026 02:29:10 -0500 Subject: [PATCH 1/2] Add estimated token usage for vscode.lm providers that report none When a vscode.lm stream (e.g. GitHub Copilot) ends without the Positron usage data part, VscodeLmClient now synthesizes a finish-step with locally estimated usage via LanguageModelChat.countTokens(), marked providerMetadata.positai.usage.isEstimated so consumers display it as approximate. Input is counted once over the final post-transform payload (lazy, never-rejecting, bounded LRU memo); output is counted from the accumulated stream. --- memory-bank/architecture.md | 67 +++- .../src/positron/VscodeLmClient.ts | 64 ++++ .../src/positron/__tests__/vscode-mock.ts | 101 ++++++ .../vscodelm-token-estimation.test.ts | 302 ++++++++++++++++++ .../src/positron/token-estimation.ts | 195 +++++++++++ 5 files changed, 728 insertions(+), 1 deletion(-) create mode 100644 packages/ai-provider-bridge/src/positron/__tests__/vscode-mock.ts create mode 100644 packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts create mode 100644 packages/ai-provider-bridge/src/positron/token-estimation.ts diff --git a/memory-bank/architecture.md b/memory-bank/architecture.md index 64c6ce5..1193368 100644 --- a/memory-bank/architecture.md +++ b/memory-bank/architecture.md @@ -1,6 +1,6 @@ --- title: ai-provider-bridge Architecture -description: Architecture of ai-provider-bridge -- entrypoints, invariants, code layout, credential system, and VS Code LM integration. +description: Architecture of ai-provider-bridge -- entrypoints, invariants, code layout, credential system, VS Code LM integration, and client-side token usage estimation for vscode.lm providers (e.g. Copilot) that report no usage. package: ai-provider-bridge --- @@ -102,6 +102,71 @@ The `/positron` entrypoint exposes `VscodeLmClient` and `listVscodeLmModels()` - - **No ProviderRegistry** -- `vscode.lm` models are host-authenticated; no `ProviderCredentials` needed. Consumers use `VscodeLmClient` + `listVscodeLmModels()` directly. - **Positron compat shim** -- `fromAiMessages2()` accepts `{ supportsToolResultImages: boolean }` option. Host extensions inject platform version checks at the call site. +### Client-side token usage estimation + +Neither the `vscode.lm` consumer API nor the provider-side `LanguageModelChatProvider` API has a +usage channel -- providers can only stream text, tool-call, data, and thinking parts. Positron +works around this for **its own** LM providers by emitting a side-channel +`LanguageModelDataPart` (`mimeType: "text/x-json"`, payload `{type: "usage", data: {...}}`) at +end of stream, which `VscodeLmClient.convertResponseStream()` parses into a `finish-step` with +real usage. The GitHub Copilot Chat extension's provider never emits that data part, so Copilot +requests (and any other third-party `vscode.lm` vendor) otherwise end with **no usage event at +all**. + +**Trigger rule**: rather than keying on vendor, `VscodeLmClient` synthesizes estimated usage +whenever the stream ends without having seen a usage data part. This covers Copilot today, +automatically steps aside if a provider starts reporting real usage, and needs no vendor +allowlist. Positron's own providers are unaffected since they always emit the data part first. + +**Mechanics** (`src/positron/token-estimation.ts`): + +- **Input**: counted once, in `chat()`, over the _final post-transform_ `vscodeMessages` array + (after system-prompt extraction, the Copilot system-as-user-message special case, and the + tool-result-image transform) plus serialized tool definitions and the `system` parameter where + it isn't already folded into messages -- counting "logical messages + system" separately would + double-count Copilot's prepended system message. Because `convertResponseStream()` only sees + the response stream (no messages/tools/system), `chat()` builds a lazy, **never-rejecting** + thunk (`.catch()` attached at creation, resolving to `number | undefined`) and passes it plus + the `LanguageModelChat` reference into the converter as a typed estimation context; the + converter only awaits it at stream end, so providers that report real usage never pay the + `countTokens()` cost. +- Per-message counts are memoized in a **module-level bounded LRU** (~500 entries) keyed by + `model.id` + a hash of the serialized final message (role + all content parts), plus a small + constant added per message for framing overhead. Module-level because `VscodeLmClient` + instances are per-request -- a per-instance cache would never hit. +- **Output**: accumulated from streamed text deltas and tool-call inputs during + `convertResponseStream()`; one `countTokens()` call on the accumulated text after the provider + stream ends. Steady state is ~2 `countTokens()` calls per turn. +- **Emission**: the synthesized `finish-step` uses the same shape as the real usage-data-part + path (`inputTokenDetails.noCacheTokens = inputTokens`, no cache fields) plus a marker, + `providerMetadata.positai.usage.isEstimated = true` -- see `messageArchitecture.md` in the main + monorepo for how core/UI consume that marker. +- **Failure policy**: estimation must never fail the request. Errors are caught and logged as a + warning; the stream ends with no usage, matching pre-estimation behavior. + +**Why `countTokens()` is cheap enough to call inline**: for Copilot models it's backed by a local +tiktoken BPE tokenizer (`cl100k_base`/`o200k_base` vocabularies bundled in the Copilot Chat +extension, tokenized in a worker thread) -- no network request, ~1-2ms of IPC per call (extension +host -> VS Code main -> Copilot Chat's provider -> its tokenizer worker). Counts are exact for +OpenAI-family models and an approximation for Anthropic/Gemini models served through Copilot +(same tokenizer, different real vocabulary). Message-framing and tool-schema overhead outside +`messages` is approximated with constants, and there is no cache-read/write breakdown to report +-- all input is counted as uncached. All of this is why results are estimates, not exact counts, +and why consumers must present them as such. + +**Why VS Code's own cost/usage UI can't be reused**: it isn't built on `vscode.lm` at all. The +Copilot Chat extension is simultaneously the LM provider, the chat participant, and the HTTP +client, so it reads usage from first-party channels unavailable to LM consumers: the raw +OpenAI-style `usage` block in the `api.githubcopilot.com` response body, `x-quota-snapshot-*` +response headers for credit/quota state, and (once adopted) the proposed +`vscode.proposed.languageModelPricing` API for per-model credit rates +(`inputCost`/`outputCost`/`cacheCost` per `LanguageModelChatInformation`). Even VS Code's generic +context-window indicator hardcodes "0 tokens used" for third-party LM providers today +([microsoft/vscode#309207](https://github.com/microsoft/vscode/issues/309207)) -- there is no +provider-to-consumer usage channel for anyone, which is why estimation is the only option until +that gap closes. `languageModelPricing` is the future hook for a real Copilot cost figure once +Positron's VS Code base includes it. + ## Provider Registration `register-all-providers.ts` registers every provider into a caller-owned `ProviderRegistry`, honoring `config.allowedProviders`. It relies on `src/provider-registration.ts`, which holds the `ProviderRegistrationConfig` interface, the `RegisterAllProviders` signature type, and the `isProviderAllowed` predicate. Consumers that want to restrict the available provider set pass `allowedProviders`; there is no build-time provider filtering. diff --git a/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts b/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts index 05bc3f8..b552b7c 100644 --- a/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts +++ b/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts @@ -20,6 +20,7 @@ import { } from "../tool-result-images"; import type { AiToolWithJsonSchema, CancellationToken, LMStreamPart, Logger } from "../types"; import { fromAiMessages2 } from "./message-formats"; +import { estimateInputTokens, estimateOutputTokens } from "./token-estimation"; import { ensureUint8Array } from "./utils"; /** @@ -167,6 +168,22 @@ export class VscodeLmClient implements ModelClient { params.cancellationToken, ); + // Lazy input-token estimation over the final post-transform payload. + // Invoked by the converter only if the stream ends without a usage data + // part (e.g. Copilot), so providers that report real usage never pay + // the countTokens cost. The returned promise never rejects. + // + // For Copilot the system prompt was prepended as a user message above, + // so it must not be counted again as a separate system parameter. + const estimateInput = () => + estimateInputTokens({ + model: this.model, + messages: vscodeMessages, + systemPrompt: this.model.vendor === "copilot" ? undefined : systemPrompt, + tools: options.tools, + logger: this.logger, + }); + // Convert response stream to platform-agnostic format // Note: chatResponse.stream is typed as AsyncIterable in VS Code API return this.convertResponseStream( @@ -175,6 +192,7 @@ export class VscodeLmClient implements ModelClient { | vscode.LanguageModelToolCallPart | vscode.LanguageModelDataPart >, + estimateInput, ); } @@ -185,11 +203,16 @@ export class VscodeLmClient implements ModelClient { vscodeStream: AsyncIterable< vscode.LanguageModelTextPart | vscode.LanguageModelToolCallPart | vscode.LanguageModelDataPart >, + estimateInput: () => Promise, ): AsyncIterable { + let sawUsage = false; + let outputText = ""; for await (const part of vscodeStream) { if (part instanceof vscode.LanguageModelTextPart) { + outputText += part.value; yield { type: "text-delta", id: "0", text: part.value }; } else if (part instanceof vscode.LanguageModelToolCallPart) { + outputText += `${part.name}\n${JSON.stringify(part.input)}\n`; yield { type: "tool-call", toolCallId: part.callId, @@ -266,6 +289,7 @@ export class VscodeLmClient implements ModelClient { } } + sawUsage = true; yield { type: "finish-step", finishReason: "stop", @@ -287,6 +311,46 @@ export class VscodeLmClient implements ModelClient { } } } + + // The stream ended without a usage data part (e.g. Copilot, whose + // vscode.lm provider reports no usage). Synthesize estimated usage from + // local token counts. Consumers detect the estimate via the + // providerMetadata marker and must display it as approximate. + if (!sawUsage) { + const inputTokens = await estimateInput(); + const outputTokens = await estimateOutputTokens(this.model, outputText, this.logger); + if (inputTokens === undefined || outputTokens === undefined) { + // Estimation failed; end the stream with no usage (previous behavior). + return; + } + yield { + type: "finish-step", + finishReason: "stop", + rawFinishReason: "stop", + usage: { + inputTokens, + inputTokenDetails: { + noCacheTokens: inputTokens, + cacheReadTokens: undefined, + cacheWriteTokens: undefined, + }, + outputTokens, + outputTokenDetails: { + textTokens: outputTokens, + reasoningTokens: undefined, + }, + totalTokens: inputTokens + outputTokens, + }, + response: { + id: "0", + modelId: "", + timestamp: new Date(), + }, + providerMetadata: { + positai: { usage: { isEstimated: true } }, + }, + }; + } } } diff --git a/packages/ai-provider-bridge/src/positron/__tests__/vscode-mock.ts b/packages/ai-provider-bridge/src/positron/__tests__/vscode-mock.ts new file mode 100644 index 0000000..21725d9 --- /dev/null +++ b/packages/ai-provider-bridge/src/positron/__tests__/vscode-mock.ts @@ -0,0 +1,101 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Minimal vscode module stub for testing the positron layer. + * + * Register in a test file with: + * vi.mock("vscode", () => import("./vscode-mock")); + * + * Only the runtime values (classes/enums) used by VscodeLmClient, + * message-formats, and token-estimation are stubbed; pure types need no + * runtime counterpart. + */ + +export class LanguageModelTextPart { + constructor(public value: string) {} +} + +export class LanguageModelToolCallPart { + constructor( + public callId: string, + public name: string, + public input: object, + ) {} +} + +export class LanguageModelDataPart { + constructor( + public data: Uint8Array, + public mimeType: string, + ) {} + + static json(value: unknown, mime = "text/x-json"): LanguageModelDataPart { + return new LanguageModelDataPart(new TextEncoder().encode(JSON.stringify(value)), mime); + } + + static text(value: string, mime = "text/plain"): LanguageModelDataPart { + return new LanguageModelDataPart(new TextEncoder().encode(value), mime); + } + + static image(data: Uint8Array, mimeType: string): LanguageModelDataPart { + return new LanguageModelDataPart(data, mimeType); + } +} + +export class LanguageModelToolResultPart { + constructor( + public callId: string, + public content: unknown[], + ) {} +} + +export class LanguageModelToolResultPart2 { + constructor( + public callId: string, + public content: unknown[], + ) {} +} + +export enum LanguageModelChatMessageRole { + User = 1, + Assistant = 2, +} + +export class LanguageModelChatMessage2 { + constructor( + public role: LanguageModelChatMessageRole, + public content: unknown[], + public name?: string, + ) {} + + static User(content: string | unknown[], name?: string): LanguageModelChatMessage2 { + const parts = typeof content === "string" ? [new LanguageModelTextPart(content)] : content; + return new LanguageModelChatMessage2(LanguageModelChatMessageRole.User, parts, name); + } + + static Assistant(content: string | unknown[], name?: string): LanguageModelChatMessage2 { + const parts = typeof content === "string" ? [new LanguageModelTextPart(content)] : content; + return new LanguageModelChatMessage2(LanguageModelChatMessageRole.Assistant, parts, name); + } +} + +export const LanguageModelChatMessage = LanguageModelChatMessage2; + +export enum LanguageModelChatToolMode { + Auto = 1, + Required = 2, +} + +export enum ChatImageMimeType { + PNG = "image/png", + JPEG = "image/jpeg", + GIF = "image/gif", + WEBP = "image/webp", + BMP = "image/bmp", +} + +export const lm = { + selectChatModels: async () => [], +}; diff --git a/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts b/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts new file mode 100644 index 0000000..14dae9c --- /dev/null +++ b/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts @@ -0,0 +1,302 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Tests for estimated token usage synthesis in VscodeLmClient and the + * token-estimation helpers. + */ + +import type * as ai from "ai"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type * as vscode from "vscode"; + +import type { Logger } from "../../types"; + +vi.mock("vscode", () => import("./vscode-mock")); + +import { + clearTokenEstimationCache, + estimateInputTokens, + serializeMessageForCounting, +} from "../token-estimation"; +import { VscodeLmClient } from "../VscodeLmClient"; +import * as vscodeMock from "./vscode-mock"; + +// ============================================================================ +// Fakes +// ============================================================================ + +const testLogger: Logger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + trace: vi.fn(), +}; + +/** Deterministic token count: 1 token per 4 characters, minimum 1. */ +function charCount(text: string): number { + return Math.max(1, Math.ceil(text.length / 4)); +} + +interface FakeModelOptions { + vendor: string; + id?: string; + /** Parts the response stream yields, in order. */ + streamParts: unknown[]; + /** Override countTokens behavior (e.g. to reject). */ + countTokens?: (text: string) => Promise; +} + +function makeFakeModel(options: FakeModelOptions) { + const countTokensCalls: string[] = []; + const fake = { + id: options.id ?? "fake-model-1", + vendor: options.vendor, + name: "Fake Model", + family: "fake", + version: "1.0", + maxInputTokens: 100_000, + countTokens: vi.fn(async (text: string) => { + countTokensCalls.push(text); + if (options.countTokens) { + return options.countTokens(text); + } + return charCount(text); + }), + sendRequest: vi.fn(async () => ({ + stream: (async function* () { + for (const part of options.streamParts) { + yield part; + } + })(), + text: (async function* () {})(), + })), + }; + // The fake implements the runtime surface VscodeLmClient touches; the + // declared vscode.LanguageModelChat type carries more members than any + // test double needs, hence the two-step cast. + return { + model: fake as unknown as vscode.LanguageModelChat, + countTokensCalls, + sendRequest: fake.sendRequest, + }; +} + +function userMessage(text: string): ai.ModelMessage { + return { role: "user", content: [{ type: "text", text }] }; +} + +async function collectParts(stream: AsyncIterable): Promise[]> { + const parts: Record[] = []; + for await (const part of stream) { + parts.push(part as Record); + } + return parts; +} + +const noopCancellation = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => {} }), +}; + +async function runChat( + model: vscode.LanguageModelChat, + params?: { messages?: ai.ModelMessage[]; systemPrompt?: string }, +): Promise[]> { + const client = new VscodeLmClient(model, testLogger); + const stream = await client.chat({ + model: "fake-model-1", + messages: params?.messages ?? [userMessage("Hello there")], + systemPrompt: params?.systemPrompt, + cancellationToken: noopCancellation, + }); + return collectParts(stream); +} + +function finishSteps(parts: Record[]): Record[] { + return parts.filter((p) => p.type === "finish-step"); +} + +// ============================================================================ +// Tests +// ============================================================================ + +beforeEach(() => { + clearTokenEstimationCache(); + vi.clearAllMocks(); +}); + +describe("VscodeLmClient estimated usage synthesis", () => { + it("synthesizes an estimated finish-step when the stream has no usage data part", async () => { + const { model } = makeFakeModel({ + vendor: "copilot", + streamParts: [ + new vscodeMock.LanguageModelTextPart("Hello "), + new vscodeMock.LanguageModelTextPart("world"), + ], + }); + + const parts = await runChat(model, { systemPrompt: "SYSTEM PROMPT" }); + const finishes = finishSteps(parts); + expect(finishes).toHaveLength(1); + + const finish = finishes[0]; + const usage = finish.usage as { + inputTokens: number; + outputTokens: number; + totalTokens: number; + }; + expect(usage.inputTokens).toBeGreaterThan(0); + expect(usage.outputTokens).toBe(charCount("Hello world")); + expect(usage.totalTokens).toBe(usage.inputTokens + usage.outputTokens); + + const metadata = finish.providerMetadata as { + positai?: { usage?: { isEstimated?: boolean } }; + }; + expect(metadata.positai?.usage?.isEstimated).toBe(true); + }); + + it("counts the Copilot system prompt exactly once (via the prepended user message)", async () => { + const { model, countTokensCalls } = makeFakeModel({ + vendor: "copilot", + streamParts: [new vscodeMock.LanguageModelTextPart("ok")], + }); + + await runChat(model, { systemPrompt: "UNIQUE_SYSTEM_MARKER" }); + + const callsWithMarker = countTokensCalls.filter((text) => + text.includes("UNIQUE_SYSTEM_MARKER"), + ); + // Counted once, inside the serialized prepended user message — never as + // a bare separate system-parameter count. + expect(callsWithMarker).toHaveLength(1); + expect(callsWithMarker[0]).not.toBe("UNIQUE_SYSTEM_MARKER"); + }); + + it("counts the system parameter separately for non-copilot vendors", async () => { + const { model, countTokensCalls } = makeFakeModel({ + vendor: "some-other-vendor", + streamParts: [new vscodeMock.LanguageModelTextPart("ok")], + }); + + await runChat(model, { systemPrompt: "UNIQUE_SYSTEM_MARKER" }); + + expect(countTokensCalls).toContain("UNIQUE_SYSTEM_MARKER"); + }); + + it("does not estimate when a real usage data part arrives", async () => { + const { model, countTokensCalls } = makeFakeModel({ + vendor: "anthropic-api", + streamParts: [ + new vscodeMock.LanguageModelTextPart("hi"), + vscodeMock.LanguageModelDataPart.json({ + type: "usage", + data: { inputTokens: 58, outputTokens: 7, cachedTokens: 30284 }, + }), + ], + }); + + const parts = await runChat(model); + const finishes = finishSteps(parts); + expect(finishes).toHaveLength(1); + + const usage = finishes[0].usage as { inputTokens: number; outputTokens: number }; + expect(usage.inputTokens).toBe(58); + expect(usage.outputTokens).toBe(7); + + const metadata = finishes[0].providerMetadata as + | { positai?: { usage?: { isEstimated?: boolean } } } + | undefined; + expect(metadata?.positai?.usage?.isEstimated).toBeUndefined(); + + // The lazy estimation thunk was never invoked. + expect(countTokensCalls).toHaveLength(0); + }); + + it("memoizes per-message input counts across requests", async () => { + const streamParts = () => [new vscodeMock.LanguageModelTextPart("out")]; + const first = makeFakeModel({ vendor: "copilot", streamParts: streamParts() }); + await runChat(first.model); + const firstCallCount = first.countTokensCalls.length; + expect(firstCallCount).toBeGreaterThan(1); + + // Same model id and messages: input counts come from the module-level + // cache; only the (uncached) output text is counted again. + const second = makeFakeModel({ vendor: "copilot", streamParts: streamParts() }); + await runChat(second.model); + expect(second.countTokensCalls).toHaveLength(1); + expect(second.countTokensCalls[0]).toBe("out"); + }); + + it("ends the stream without usage when countTokens fails, without failing the request", async () => { + const { model } = makeFakeModel({ + vendor: "copilot", + streamParts: [new vscodeMock.LanguageModelTextPart("hello")], + countTokens: async () => { + throw new Error("tokenizer unavailable"); + }, + }); + + const parts = await runChat(model); + expect(parts.some((p) => p.type === "text-delta")).toBe(true); + expect(finishSteps(parts)).toHaveLength(0); + expect(testLogger.warn).toHaveBeenCalled(); + }); +}); + +describe("token-estimation helpers", () => { + it("keys the cache on role, so identical text in different roles is counted separately", async () => { + const countTokens = vi.fn(async (text: string) => charCount(text)); + const model = { id: "m1", countTokens }; + + const asUser = vscodeMock.LanguageModelChatMessage2.User("same text"); + const asAssistant = vscodeMock.LanguageModelChatMessage2.Assistant("same text"); + + expect(serializeMessageForCounting(toVscodeMessage(asUser))).not.toBe( + serializeMessageForCounting(toVscodeMessage(asAssistant)), + ); + + await estimateInputTokens({ + model, + messages: [toVscodeMessage(asUser), toVscodeMessage(asAssistant)], + logger: testLogger, + }); + expect(countTokens).toHaveBeenCalledTimes(2); + }); + + it("includes tool definitions in the input estimate", async () => { + const countTokens = vi.fn(async (text: string) => charCount(text)); + const model = { id: "m1", countTokens }; + + const withoutTools = await estimateInputTokens({ + model, + messages: [toVscodeMessage(vscodeMock.LanguageModelChatMessage2.User("hi"))], + logger: testLogger, + }); + const withTools = await estimateInputTokens({ + model, + messages: [toVscodeMessage(vscodeMock.LanguageModelChatMessage2.User("hi"))], + tools: [ + { + name: "get_weather", + description: "Get the weather for a location", + inputSchema: { type: "object", properties: { location: { type: "string" } } }, + }, + ], + logger: testLogger, + }); + + expect(withoutTools).toBeDefined(); + expect(withTools).toBeDefined(); + expect(withTools!).toBeGreaterThan(withoutTools!); + }); +}); + +/** The mock message class satisfies the runtime shape the serializer walks. */ +function toVscodeMessage( + message: vscodeMock.LanguageModelChatMessage2, +): vscode.LanguageModelChatMessage2 { + return message as unknown as vscode.LanguageModelChatMessage2; +} diff --git a/packages/ai-provider-bridge/src/positron/token-estimation.ts b/packages/ai-provider-bridge/src/positron/token-estimation.ts new file mode 100644 index 0000000..6432491 --- /dev/null +++ b/packages/ai-provider-bridge/src/positron/token-estimation.ts @@ -0,0 +1,195 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Token usage estimation for vscode.lm providers that report no usage. + * + * Some vscode.lm providers (notably GitHub Copilot) never emit the Positron + * usage data part, so requests end with no token usage at all. This module + * estimates usage locally with `LanguageModelChat.countTokens()`, which for + * Copilot is backed by a bundled tiktoken tokenizer running in a worker + * thread — a few IPC hops per call, no network request. + * + * The counts are estimates: message framing overhead is approximated with a + * per-message constant, images/binary parts are represented by a short + * placeholder, and non-OpenAI models are tokenized with an approximating + * tokenizer on the provider side. Consumers must present these numbers as + * approximate. + */ + +import * as vscode from "vscode"; + +import type { Logger } from "../types"; + +/** + * Tokens added per message to approximate chat framing overhead (role + * markers and separators) that countTokens() on the bare content misses. + */ +const MESSAGE_FRAMING_TOKENS = 4; + +/** + * Minimal surface of vscode.LanguageModelChat needed for estimation. + * Narrow on purpose so tests can supply a fake without the full class. + */ +export interface EstimationModel { + readonly id: string; + countTokens(text: string, token?: vscode.CancellationToken): Thenable; +} + +// ============================================================================ +// Bounded LRU memo for per-message token counts +// ============================================================================ + +/** + * Module-level because VscodeLmClient instances are per-request — a + * per-instance cache would never hit. Entries are integers keyed by + * model id + content hash, so the memory bound is small. + */ +const MAX_CACHE_ENTRIES = 500; +const tokenCountCache = new Map(); + +/** FNV-1a 32-bit hash, hex-encoded. Length is included in the cache key to further reduce collision risk. */ +function fnv1a(text: string): string { + let hash = 0x811c9dc5; + for (let i = 0; i < text.length; i++) { + hash ^= text.charCodeAt(i); + hash = Math.imul(hash, 0x01000193); + } + return (hash >>> 0).toString(16); +} + +function cacheKey(modelId: string, text: string): string { + return `${modelId}:${text.length}:${fnv1a(text)}`; +} + +/** Exposed for tests. */ +export function clearTokenEstimationCache(): void { + tokenCountCache.clear(); +} + +async function countTextCached(model: EstimationModel, text: string): Promise { + const key = cacheKey(model.id, text); + const cached = tokenCountCache.get(key); + if (cached !== undefined) { + // Refresh recency (Map iteration order is insertion order). + tokenCountCache.delete(key); + tokenCountCache.set(key, cached); + return cached; + } + const count = await model.countTokens(text); + if (tokenCountCache.size >= MAX_CACHE_ENTRIES) { + const oldest = tokenCountCache.keys().next(); + if (!oldest.done) { + tokenCountCache.delete(oldest.value); + } + } + tokenCountCache.set(key, count); + return count; +} + +// ============================================================================ +// Message serialization for counting +// ============================================================================ + +/** + * Serialize a final (post-transform) VS Code message for token counting. + * Includes the role so identical text in different roles never shares a + * cached count. Binary/data parts contribute a short placeholder — their + * true token cost is provider-specific and unknowable here. + */ +export function serializeMessageForCounting( + message: vscode.LanguageModelChatMessage | vscode.LanguageModelChatMessage2, +): string { + const pieces: string[] = [`role:${message.role}`]; + for (const part of message.content) { + if (part instanceof vscode.LanguageModelTextPart) { + pieces.push(part.value); + } else if (part instanceof vscode.LanguageModelToolCallPart) { + pieces.push(part.name, JSON.stringify(part.input)); + } else if ( + part instanceof vscode.LanguageModelToolResultPart2 || + part instanceof vscode.LanguageModelToolResultPart + ) { + for (const inner of part.content) { + if (inner instanceof vscode.LanguageModelTextPart) { + pieces.push(inner.value); + } else { + pieces.push("[data]"); + } + } + } else if (part instanceof vscode.LanguageModelDataPart) { + pieces.push(`[data:${part.mimeType}]`); + } + } + return pieces.join("\n"); +} + +// ============================================================================ +// Input estimation +// ============================================================================ + +export interface EstimateInputTokensParams { + model: EstimationModel; + /** The final, post-transform payload passed to sendRequest(). */ + messages: ReadonlyArray; + /** + * The separate system parameter, only for vendors where it is NOT already + * folded into `messages` (for Copilot the system prompt is prepended as a + * user message, so passing it here too would double-count it). + */ + systemPrompt?: string; + tools?: ReadonlyArray; + logger: Logger; +} + +/** + * Estimate input tokens for the final request payload. + * + * Never rejects: estimation must never fail a request, and the resulting + * promise may be awaited long after creation (at stream end), so a rejection + * here could otherwise surface as an unhandled rejection mid-stream. + * + * @returns The estimate, or undefined if estimation failed. + */ +export async function estimateInputTokens( + params: EstimateInputTokensParams, +): Promise { + try { + const counts = await Promise.all([ + ...params.messages.map((m) => countTextCached(params.model, serializeMessageForCounting(m))), + params.systemPrompt ? countTextCached(params.model, params.systemPrompt) : 0, + params.tools && params.tools.length > 0 + ? countTextCached(params.model, JSON.stringify(params.tools)) + : 0, + ]); + const framing = params.messages.length * MESSAGE_FRAMING_TOKENS; + return counts.reduce((sum, c) => sum + c, framing); + } catch (e) { + params.logger.warn("Token usage estimation failed for input payload", e); + return undefined; + } +} + +/** + * Count tokens in accumulated output text. Never rejects. + * + * @returns The estimate, or undefined if counting failed. + */ +export async function estimateOutputTokens( + model: EstimationModel, + outputText: string, + logger: Logger, +): Promise { + if (outputText === "") { + return 0; + } + try { + // No memoization: output text is unique per response, so caching it + // would only churn the LRU. + return await model.countTokens(outputText); + } catch (e) { + logger.warn("Token usage estimation failed for output text", e); + return undefined; + } +} From ba60088c1211786e839f1928cebd72a191888ce5 Mon Sep 17 00:00:00 2001 From: Winston Chang Date: Fri, 10 Jul 2026 16:48:11 -0500 Subject: [PATCH 2/2] Preserve tool-call finish reason for vscode.lm usage --- .../src/positron/VscodeLmClient.ts | 12 +++-- .../vscodelm-token-estimation.test.ts | 47 +++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts b/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts index b552b7c..a70ca0d 100644 --- a/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts +++ b/packages/ai-provider-bridge/src/positron/VscodeLmClient.ts @@ -206,12 +206,14 @@ export class VscodeLmClient implements ModelClient { estimateInput: () => Promise, ): AsyncIterable { let sawUsage = false; + let sawToolCall = false; let outputText = ""; for await (const part of vscodeStream) { if (part instanceof vscode.LanguageModelTextPart) { outputText += part.value; yield { type: "text-delta", id: "0", text: part.value }; } else if (part instanceof vscode.LanguageModelToolCallPart) { + sawToolCall = true; outputText += `${part.name}\n${JSON.stringify(part.input)}\n`; yield { type: "tool-call", @@ -290,10 +292,11 @@ export class VscodeLmClient implements ModelClient { } sawUsage = true; + const finishReason = sawToolCall ? "tool-calls" : "stop"; yield { type: "finish-step", - finishReason: "stop", - rawFinishReason: "stop", + finishReason, + rawFinishReason: finishReason, usage: usageData, response: { id: "0", @@ -323,10 +326,11 @@ export class VscodeLmClient implements ModelClient { // Estimation failed; end the stream with no usage (previous behavior). return; } + const finishReason = sawToolCall ? "tool-calls" : "stop"; yield { type: "finish-step", - finishReason: "stop", - rawFinishReason: "stop", + finishReason, + rawFinishReason: finishReason, usage: { inputTokens, inputTokenDetails: { diff --git a/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts b/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts index 14dae9c..d06e9a3 100644 --- a/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts +++ b/packages/ai-provider-bridge/src/positron/__tests__/vscodelm-token-estimation.test.ts @@ -230,6 +230,53 @@ describe("VscodeLmClient estimated usage synthesis", () => { expect(second.countTokensCalls[0]).toBe("out"); }); + it("marks the synthesized finish-step as tool-calls when the stream emitted a tool call", async () => { + const { model } = makeFakeModel({ + vendor: "copilot", + streamParts: [ + new vscodeMock.LanguageModelTextPart("let me check"), + new vscodeMock.LanguageModelToolCallPart("call-1", "get_weather", { + location: "SF", + }), + ], + }); + + const finishes = finishSteps(await runChat(model)); + expect(finishes).toHaveLength(1); + expect(finishes[0].finishReason).toBe("tool-calls"); + expect(finishes[0].rawFinishReason).toBe("tool-calls"); + }); + + it("marks the synthesized finish-step as stop when no tool call was emitted", async () => { + const { model } = makeFakeModel({ + vendor: "copilot", + streamParts: [new vscodeMock.LanguageModelTextPart("all done")], + }); + + const finishes = finishSteps(await runChat(model)); + expect(finishes).toHaveLength(1); + expect(finishes[0].finishReason).toBe("stop"); + expect(finishes[0].rawFinishReason).toBe("stop"); + }); + + it("marks the real-usage finish-step as tool-calls when the stream emitted a tool call", async () => { + const { model } = makeFakeModel({ + vendor: "anthropic-api", + streamParts: [ + new vscodeMock.LanguageModelToolCallPart("call-1", "get_weather", { location: "SF" }), + vscodeMock.LanguageModelDataPart.json({ + type: "usage", + data: { inputTokens: 58, outputTokens: 7, cachedTokens: 0 }, + }), + ], + }); + + const finishes = finishSteps(await runChat(model)); + expect(finishes).toHaveLength(1); + expect(finishes[0].finishReason).toBe("tool-calls"); + expect(finishes[0].rawFinishReason).toBe("tool-calls"); + }); + it("ends the stream without usage when countTokens fails, without failing the request", async () => { const { model } = makeFakeModel({ vendor: "copilot",