Estimate token usage for vscode.lm providers that report none

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.

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<unknown> in VS Code API
 		return this.convertResponseStream(
@@ -175,6 +192,7 @@ export class VscodeLmClient implements ModelClient {
 				| vscode.LanguageModelToolCallPart
 				| vscode.LanguageModelDataPart
 			>,
+			estimateInput,
 		);
 	}
 
@@ -185,11 +203,18 @@ export class VscodeLmClient implements ModelClient {
 		vscodeStream: AsyncIterable<
 			vscode.LanguageModelTextPart | vscode.LanguageModelToolCallPart | vscode.LanguageModelDataPart
 		>,
+		estimateInput: () => Promise<number | undefined>,
 	): AsyncIterable<LMStreamPart> {
+		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",
 					toolCallId: part.callId,
@@ -266,10 +291,12 @@ 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",
@@ -287,6 +314,47 @@ 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;
+			}
+			const finishReason = sawToolCall ? "tool-calls" : "stop";
+			yield {
+				type: "finish-step",
+				finishReason,
+				rawFinishReason: finishReason,
+				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 } },
+				},
+			};
+		}
 	}
 }
 

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 () => [],
+};