From 3c7a0c7334746e0a00cc0f386e5a6ab55acdb65c Mon Sep 17 00:00:00 2001 From: Marcel Stenkewitz Date: Thu, 15 Oct 2026 14:33:05 -0400 Subject: [PATCH] fix(models): report real token usage on streamed replies from any openai-compatible provider Upstream set includeUsage itself but scoped it to one provider (ollama-cloud), so a custom endpoint and other openai-compatible providers kept reporting output_tokens 0 on every streamed turn. The opt-in is now a supportsStreamUsage flag on the per-model compatibility hints, set in the Ollama Cloud catalog and read by provider-factory; a custom endpoint opts in through a confirm in providers add, stored on its model entries' compatibility. Unset keeps the option off, so every existing provider is unaffected. --- src/data/ollama-cloud-models.ts | 4 ++++ src/model-runtime-compatibility.ts | 7 +++++++ src/provider-factory.ts | 7 ++++--- src/providers-custom-add.ts | 13 +++++++++++++ src/registry/custom-endpoint.ts | 5 +++++ tests/ollama-cloud-stream-usage.test.ts | 26 +++++++++++++++++++++---- tests/ollama-cloud.test.ts | 1 + tests/providers-custom-add.test.ts | 18 +++++++++++++++-- 8 files changed, 72 insertions(+), 9 deletions(-) diff --git a/src/data/ollama-cloud-models.ts b/src/data/ollama-cloud-models.ts index 83e4ffa4..53a0f40c 100644 --- a/src/data/ollama-cloud-models.ts +++ b/src/data/ollama-cloud-models.ts @@ -51,6 +51,10 @@ const CATALOG: OllamaCloudModel[] = ([ maxTokensField: 'max_tokens' as const, // Effort mapping is unverified against ollama.com; send none. supportsReasoningEffort: false, + // ollama.com sends a streamed usage chunk only when the request carries + // stream_options.include_usage (observed 2026-10-13); without it every + // streamed turn reports output_tokens 0. + supportsStreamUsage: true, }, })); diff --git a/src/model-runtime-compatibility.ts b/src/model-runtime-compatibility.ts index ffa9bb4c..16b45557 100644 --- a/src/model-runtime-compatibility.ts +++ b/src/model-runtime-compatibility.ts @@ -33,6 +33,13 @@ export interface ModelRuntimeCompatibility { * endpoint that does implement it is unaffected. */ supportsCountTokens?: boolean; + /** + * Whether the upstream sends a streamed usage chunk only when the request + * carries `stream_options: {"include_usage": true}`. Unset keeps upstream's + * per-provider scoping, so existing providers are unaffected until their + * catalog opts in. + */ + supportsStreamUsage?: boolean; } export type OpenAiCompatibleRequestBody = Record; diff --git a/src/provider-factory.ts b/src/provider-factory.ts index d01bd80e..ad182305 100644 --- a/src/provider-factory.ts +++ b/src/provider-factory.ts @@ -15,7 +15,6 @@ import { } from './oauth/claude-identity.js'; import { isCredentialBearingHeader } from './credential-headers.js'; import { withAccessPrograms, type AccessPrograms } from './daybreak.js'; -import { OLLAMA_CLOUD_PROVIDER_ID } from './data/ollama-cloud-models.js'; import { transformOpenAiCompatibleRequestBody, type ModelRuntimeCompatibility, @@ -268,8 +267,10 @@ export async function createLanguageModel(spec: ProviderModelSpec): Promise) => diff --git a/src/providers-custom-add.ts b/src/providers-custom-add.ts index 49c05c9e..98235a32 100644 --- a/src/providers-custom-add.ts +++ b/src/providers-custom-add.ts @@ -58,6 +58,18 @@ export async function runCustomEndpointAddFlow( return 0; } + // Servers that only report token usage on streamed replies when the request + // asks for it need stream_options.include_usage; without it every streamed + // turn reports output_tokens 0. + const streamUsage = await p.confirm({ + message: 'Send stream_options.include_usage with streamed requests? (Needed when this server reports token usage only when the request asks for it.)', + initialValue: false, + }); + if (p.isCancel(streamUsage)) { + p.cancel('Cancelled.'); + return 0; + } + const spinner = p.spinner(); spinner.start('Testing connection...'); const result = await addCustomEndpointProvider({ @@ -66,6 +78,7 @@ export async function runCustomEndpointAddFlow( apiKey: String(apiKey ?? '').trim(), kind: 'openai', allowInsecureLocal, + supportsStreamUsage: streamUsage === true, }); spinner.stop(''); onResult(result); diff --git a/src/registry/custom-endpoint.ts b/src/registry/custom-endpoint.ts index 4abf4b72..6b28712c 100644 --- a/src/registry/custom-endpoint.ts +++ b/src/registry/custom-endpoint.ts @@ -38,6 +38,8 @@ export interface AddCustomEndpointInput { allowInsecureLocal?: boolean; /** Static headers this endpoint requires on every request (e.g. a plan/auth-tracking header). */ headers?: Record; + /** Opt this openai-compatible endpoint into stream_options.include_usage on streamed requests. */ + supportsStreamUsage?: boolean; } export interface AddCustomEndpointResult { @@ -249,6 +251,9 @@ export async function addCustomEndpointProvider(input: AddCustomEndpointInput): modelFormat: modelFormatForKind(input.kind), npm, apiUrl: fetched.baseUrl, + ...(input.supportsStreamUsage + ? { compatibility: { ...m.compatibility, supportsStreamUsage: true } } + : {}), })), }, }; diff --git a/tests/ollama-cloud-stream-usage.test.ts b/tests/ollama-cloud-stream-usage.test.ts index 4fc25d26..2e50bff2 100644 --- a/tests/ollama-cloud-stream-usage.test.ts +++ b/tests/ollama-cloud-stream-usage.test.ts @@ -3,6 +3,7 @@ import { createServer } from 'node:http'; import { createLanguageModel } from '../src/provider-factory.js'; import { generateAnthropicResponse, streamAnthropicResponse } from '../src/sdk-adapter.js'; import { buildOllamaCloudModels } from '../src/data/ollama-cloud-models.js'; +import type { ModelRuntimeCompatibility } from '../src/model-runtime-compatibility.js'; /** * ollama.com emits a streamed usage chunk (`choices: []` plus `usage`) only when the @@ -59,7 +60,10 @@ async function startListener(): Promise { }; } -async function streamOnce(providerId: string): Promise<{ +async function streamOnce( + providerId: string, + compatOverride?: ModelRuntimeCompatibility, +): Promise<{ body: Record; outputTokens: unknown; }> { @@ -73,7 +77,7 @@ async function streamOnce(providerId: string): Promise<{ baseURL: listener.baseUrl, providerId, authType: 'api', - compatibility: catalog.compatibility, + compatibility: compatOverride ?? catalog.compatibility, }); let sse = ''; await streamAnthropicResponse( @@ -101,10 +105,24 @@ describe('Ollama Cloud streamed usage', () => { expect(outputTokens).toBe(USAGE.completion_tokens); }); + it('asks for the usage chunk for any provider whose spec declares stream usage', async () => { + // The same ollama-cloud compatibility under a different provider id still + // sends stream_options: the flag on the catalog drives it, not the id. + const { body, outputTokens } = await streamOnce('openai'); + expect(body.stream_options).toEqual({ include_usage: true }); + expect(outputTokens).toBe(USAGE.completion_tokens); + }); + it.each(['opencode-go', 'custom-local', 'openai'])( - 'leaves the %s request without stream_options', + 'leaves the %s request without stream_options when its catalog does not opt in', async providerId => { - const { body } = await streamOnce(providerId); + // A catalog without supportsStreamUsage, like the real opencode-go one. + const { body } = await streamOnce(providerId, { + supportsStore: false, + supportsDeveloperRole: false, + maxTokensField: 'max_tokens' as const, + supportsReasoningEffort: false, + }); expect(body.stream).toBe(true); expect(body).not.toHaveProperty('stream_options'); }, diff --git a/tests/ollama-cloud.test.ts b/tests/ollama-cloud.test.ts index 596b55d8..a4947300 100644 --- a/tests/ollama-cloud.test.ts +++ b/tests/ollama-cloud.test.ts @@ -242,6 +242,7 @@ describe('Ollama Cloud per-model metadata', () => { supportsDeveloperRole: false, maxTokensField: 'max_tokens', supportsReasoningEffort: false, + supportsStreamUsage: true, }, }); } diff --git a/tests/providers-custom-add.test.ts b/tests/providers-custom-add.test.ts index d4eefc74..f9b6dddd 100644 --- a/tests/providers-custom-add.test.ts +++ b/tests/providers-custom-add.test.ts @@ -130,8 +130,11 @@ describe('providers add: custom OpenAI-compatible server', () => { apiKey: 'sk-or-test', kind: 'openai', allowInsecureLocal: false, + supportsStreamUsage: false, }); - expect(confirmMock).not.toHaveBeenCalled(); + // Only the stream-usage confirm is asked on the https path; the plain-http + // one is not. + expect(confirmMock).toHaveBeenCalledTimes(1); }); it('reports the provider name and how many models were found', async () => { @@ -205,7 +208,8 @@ describe('providers add: custom OpenAI-compatible server', () => { await runProvidersAdd(); - expect(confirmMock).toHaveBeenCalledTimes(1); + // Both confirms are asked now: plain http, then stream usage. + expect(confirmMock).toHaveBeenCalledTimes(2); }); it('does not allow plain http when the user declines', async () => { @@ -283,6 +287,16 @@ describe('providers add: custom OpenAI-compatible server', () => { confirmMock.mockResolvedValue(CANCEL); }, }, + { + at: 'the stream-usage confirmation', + answer: () => { + textMock.mockResolvedValueOnce('OpenRouter').mockResolvedValueOnce('https://openrouter.ai/api/v1'); + passwordMock.mockResolvedValue(''); + // On the https path the stream-usage confirm is the only confirm, + // so cancelling at it is the first confirm answer. + confirmMock.mockResolvedValueOnce(CANCEL); + }, + }, ])('cancelling at $at adds nothing and exits 0', async ({ answer }) => { selectMock.mockResolvedValue('custom'); answer();