feat: emit token usage as LanguageModelDataPart for context window indicator

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## 0.1.3
+
+### Changed
+
+- Set `stream_options.include_usage` on streaming requests so oMLX returns
+  token counts in the final SSE chunk.
+- Emit token usage as a `LanguageModelDataPart` (mime type `usage`) so the
+  Copilot Chat context window indicator shows real usage instead of 0%.
+  Requires VS Code 1.120 or newer (microsoft/vscode#315394).
+
 ## 0.1.2
 
 ### Changed

package.json

@@ -2,7 +2,7 @@
   "name": "omlx-copilot-chat",
   "displayName": "oMLX Copilot Chat",
   "description": "Use oMLX OpenAI-compatible models from the VS Code Copilot Chat model picker.",
-  "version": "0.1.2",
+  "version": "0.1.3",
   "publisher": "techopolis",
   "preview": true,
   "license": "MIT",
@@ -15,7 +15,7 @@
   },
   "homepage": "https://github.com/mikedoise/oMLX-Copilot-Chat#readme",
   "engines": {
-    "vscode": "^1.104.0"
+    "vscode": "^1.120.0"
   },
   "categories": [
     "AI",
@@ -118,7 +118,7 @@
   },
   "devDependencies": {
     "@types/node": "^22.0.0",
-    "@types/vscode": "^1.104.0",
+    "@types/vscode": "^1.120.0",
     "@vscode/vsce": "^3.6.0",
     "typescript": "^5.9.0"
   }

src/provider/OmlxClient.ts

@@ -8,12 +8,14 @@ import {
   OpenAIModelListResponse,
   OpenAITool,
   OpenAIToolCall,
-  OpenAIToolCallDelta
+  OpenAIToolCallDelta,
+  OpenAIUsage
 } from './openAICompatTypes';
 
 export type OmlxStreamPart =
   | { readonly type: 'text'; readonly value: string }
-  | { readonly type: 'toolCall'; readonly callId: string; readonly name: string; readonly input: object };
+  | { readonly type: 'toolCall'; readonly callId: string; readonly name: string; readonly input: object }
+  | { readonly type: 'usage'; readonly usage: OpenAIUsage };
 
 export class OmlxPromptTooLongError extends Error {
   constructor(
@@ -71,6 +73,7 @@ export class OmlxClient {
       model,
       messages,
       stream: true,
+      stream_options: { include_usage: true },
       temperature: options.temperature,
       max_tokens: options.maxTokens,
       tools: options.tools,
@@ -114,6 +117,12 @@ export class OmlxClient {
           const line = buffer.slice(0, newlineIndex);
           buffer = buffer.slice(newlineIndex + 1);
           const fragment = parseSseLine(line);
+          // oMLX emits the usage chunk (choices: [], usage: {...}) immediately
+          // before the [DONE] sentinel when stream_options.include_usage is set,
+          // so it is read here as a normal chunk ahead of the done check below.
+          if (fragment.usage) {
+            yield { type: 'usage', usage: fragment.usage };
+          }
           if (fragment.done) {
             yield* completeToolCalls(toolCallAccumulators);
             return;
@@ -199,6 +208,7 @@ interface ParsedSseLine {
   readonly done: boolean;
   readonly text?: string;
   readonly toolCalls: readonly OpenAIToolCallDelta[];
+  readonly usage?: OpenAIUsage;
 }
 
 function parseSseLine(line: string): ParsedSseLine {
@@ -217,7 +227,8 @@ function parseSseLine(line: string): ParsedSseLine {
   return {
     done: false,
     text: delta?.content ?? undefined,
-    toolCalls: delta?.tool_calls ?? []
+    toolCalls: delta?.tool_calls ?? [],
+    usage: chunk.usage ?? undefined
   };
 }
 

src/provider/OmlxLanguageModelProvider.ts

@@ -81,8 +81,13 @@ export class OmlxLanguageModelProvider implements vscode.LanguageModelChatProvid
       for await (const part of client.streamChatCompletion(model.id, convertedMessages, requestOptions, token)) {
         if (part.type === 'text') {
           progress.report(new vscode.LanguageModelTextPart(part.value));
-        } else {
+        } else if (part.type === 'toolCall') {
           progress.report(new vscode.LanguageModelToolCallPart(part.callId, part.name, part.input));
+        } else if (part.type === 'usage') {
+          this.output.appendLine(
+            `Usage ${model.id}: prompt_tokens=${part.usage.prompt_tokens}, completion_tokens=${part.usage.completion_tokens}, total_tokens=${part.usage.total_tokens}, cached_tokens=${part.usage.prompt_tokens_details?.cached_tokens ?? 0}`
+          );
+          progress.report(vscode.LanguageModelDataPart.json(part.usage, 'usage'));
         }
       }
     } catch (error) {

src/provider/openAICompatTypes.ts

@@ -55,6 +55,7 @@ export interface OpenAIChatCompletionRequest {
   readonly model: string;
   readonly messages: readonly OpenAIChatMessage[];
   readonly stream: boolean;
+  readonly stream_options?: { readonly include_usage: boolean };
   readonly temperature?: number;
   readonly max_tokens?: number;
   readonly tools?: readonly OpenAITool[];
@@ -88,13 +89,23 @@ export interface OpenAIChatCompletionResponse {
   }[];
 }
 
+export interface OpenAIUsage {
+  readonly prompt_tokens: number;
+  readonly completion_tokens: number;
+  readonly total_tokens: number;
+  readonly prompt_tokens_details?: {
+    readonly cached_tokens?: number;
+  };
+}
+
 export interface OpenAIChatCompletionChunk {
   readonly choices?: readonly {
     readonly delta?: {
       readonly content?: string | null;
       readonly tool_calls?: readonly OpenAIToolCallDelta[];
     };
   }[];
+  readonly usage?: OpenAIUsage | null;
 }
 
 export interface OpenAIToolCallDelta {