From decad0e17270c6748b47139b976488b4e60b82bb Mon Sep 17 00:00:00 2001 From: Bogdan Date: Mon, 21 Sep 2026 04:06:08 +0300 Subject: [PATCH] fix(openai-shim): request streaming usage from compatible local providers (#2144) * fix(openai-shim): request streaming usage from local providers * test(openai-shim): verify cached streaming usage accounting * fix(openai-shim): retry strict streams without usage options --- docs/integrations/overview.md | 10 + .../api/errors.openaiCompatibility.test.ts | 16 ++ src/services/api/errors.ts | 6 + .../api/openaiErrorClassification.test.ts | 41 ++++ src/services/api/openaiErrorClassification.ts | 104 ++++++++++ src/services/api/openaiShim.test.ts | 68 ++++++- .../requestExecutor.integration.test.ts | 192 +++++++++++++++++- .../api/openaiShim/requestExecutor.ts | 50 ++++- .../api/openaiShim/requestPreparation.test.ts | 131 ++++++++++++ .../api/openaiShim/requestPreparation.ts | 2 +- .../api/openaiShim/streamConversion.test.ts | 84 +++++++- 11 files changed, 693 insertions(+), 11 deletions(-) diff --git a/docs/integrations/overview.md b/docs/integrations/overview.md index b345fe90c..dcd759b63 100644 --- a/docs/integrations/overview.md +++ b/docs/integrations/overview.md @@ -118,6 +118,16 @@ routes. Fixed direct vendors usually set both to `false`; broad custom routes or gateways that intentionally accept user-supplied auth/header details set the relevant flag to `true`. +OpenAI-compatible streaming chat-completions routes request a terminal usage +chunk with `stream_options: { include_usage: true }` by default. Network +location does not change that protocol behavior. If a server explicitly rejects +the top-level `stream_options` field, the shim retries once without it; user +cancellations and unrelated client errors are not retried. Known routes that +reject the field should still declare `removeBodyFields: ['stream_options']` in +their `transportConfig.openaiShim` so the first request is compatible. Native +transports such as direct Ollama build their own request payloads and are +unaffected. + ### Reasoning support is per model and per route `capabilities.supportsReasoning` is descriptive. It says the model is known to diff --git a/src/services/api/errors.openaiCompatibility.test.ts b/src/services/api/errors.openaiCompatibility.test.ts index db366c640..0c71692cc 100644 --- a/src/services/api/errors.openaiCompatibility.test.ts +++ b/src/services/api/errors.openaiCompatibility.test.ts @@ -146,3 +146,19 @@ test('maps tool_stream_unsupported without promising a retry after failure', () expect(text).toMatch(/(\/model|--model)/) expect(text).not.toContain('Retrying') }) + +test('maps stream_options_unsupported after the bounded fallback is exhausted', () => { + const error = APIError.generate( + 400, + undefined, + 'OpenAI API error 400: stream_options is unsupported [openai_category=stream_options_unsupported]', + new Headers(), + ) + + const message = getAssistantMessageFromError(error, 'local-model') + const text = getFirstText(message) + + expect(text).toContain('rejected the `stream_options` parameter') + expect(text).toContain('without usage reporting') + expect(text).not.toContain('Retrying') +}) diff --git a/src/services/api/errors.ts b/src/services/api/errors.ts index 0dd49a8b3..0b542e69b 100644 --- a/src/services/api/errors.ts +++ b/src/services/api/errors.ts @@ -159,6 +159,12 @@ function mapOpenAICompatibilityFailureToAssistantMessage(options: { error: 'invalid_request', }) + case 'stream_options_unsupported': + return createAssistantAPIErrorMessage({ + content: 'The selected provider rejected the `stream_options` parameter and could not complete the compatibility fallback without usage reporting. Verify that the endpoint accepts ordinary streaming chat-completions requests.', + error: 'invalid_request', + }) + case 'malformed_provider_response': return createAssistantAPIErrorMessage({ content: `${API_ERROR_MESSAGE_PREFIX}: Provider returned a malformed response. Confirm endpoint compatibility and check local proxy/network middleware.`, diff --git a/src/services/api/openaiErrorClassification.test.ts b/src/services/api/openaiErrorClassification.test.ts index 0d6c13821..8b96c7497 100644 --- a/src/services/api/openaiErrorClassification.test.ts +++ b/src/services/api/openaiErrorClassification.test.ts @@ -158,6 +158,47 @@ test('classifies tool_stream rejection as tool_stream_unsupported (#1950)', () = expect(failure.retryable).toBe(false) }) +test.each([ + 'Unsupported parameter: stream_options', + 'Unrecognized request argument supplied: stream_options', + 'Parameter stream_options is not supported', + 'stream_options is not supported by this API', + '{"error":{"message":"Unknown parameter","param":"stream_options"}}', + '{"error":{"message":"Unknown parameter","param":"`stream_options`"}}', + '{"error":{"message":"Invalid parameter: stream_options","type":"invalid_request_error","param":"stream_options"}}', + '{"detail":[{"type":"extra_forbidden","loc":["body","stream_options"],"msg":"Extra inputs are not permitted"}]}', + '{"detail":[{"type":"extra_forbidden","loc":["stream_options"],"msg":"Extra inputs are not permitted"}]}', +])('classifies top-level stream_options rejections: %s', body => { + const failure = classifyOpenAIHttpFailure({ status: 400, body }) + + expect(failure.category).toBe('stream_options_unsupported') + expect(failure.retryable).toBe(false) +}) + +test('classifies a FastAPI stream_options rejection at status 422', () => { + const failure = classifyOpenAIHttpFailure({ + status: 422, + body: '{"detail":[{"type":"value_error.extra","loc":["body","stream_options"],"msg":"extra fields not permitted"}]}', + }) + + expect(failure.category).toBe('stream_options_unsupported') +}) + +test.each([ + 'Invalid request for model local-model; received stream_options in the request body', + 'Unsupported parameter: stream_options.include_usage', + 'Unknown parameter: stream_options[include_usage]', + 'Invalid parameter: stream_options/include_usage must be a boolean', + '{"error":{"message":"Invalid parameter: stream_options.include_usage must be a boolean","param":"stream_options.include_usage"}}', + '{"error":{"message":"stream_options.include_usage must be a boolean","param":"stream_options.include_usage"}}', + '{"detail":[{"type":"bool_type","loc":["body","stream_options","include_usage"],"msg":"Input should be a valid boolean"}]}', + '{"detail":[{"type":"dict_type","loc":["body","stream_options"],"msg":"Input should be an object"},{"type":"bool_type","loc":["body","stream_options","include_usage"],"msg":"Input should be a valid boolean"}]}', +])('does not classify unrelated or nested stream_options validation failures: %s', body => { + const failure = classifyOpenAIHttpFailure({ status: 400, body }) + + expect(failure.category).not.toBe('stream_options_unsupported') +}) + test('prioritizes tool_stream rejection over accompanying tool-call wording', () => { const failure = classifyOpenAIHttpFailure({ status: 400, diff --git a/src/services/api/openaiErrorClassification.ts b/src/services/api/openaiErrorClassification.ts index 1be20bfd1..a2cb8a456 100644 --- a/src/services/api/openaiErrorClassification.ts +++ b/src/services/api/openaiErrorClassification.ts @@ -12,6 +12,7 @@ export type OpenAICompatibilityFailureCategory = | 'context_overflow' | 'tool_call_incompatible' | 'tool_stream_unsupported' + | 'stream_options_unsupported' | 'malformed_provider_response' | 'provider_unavailable' | 'unknown' @@ -68,6 +69,7 @@ const OPENAI_COMPATIBILITY_FAILURE_CATEGORIES: ReadonlySet typeof value === 'string') + .join(' ') + if ( + /(?:unsupported|not supported|unknown|unrecognized|unexpected|not permitted|extra[_ -]?forbidden|invalid (?:request )?(?:argument|parameter|field|property))/i.test(signal) + ) { + return true + } + return undefined + } + if (parameter && /^stream_options(?:\.|\[|\/)/.test(parameter)) { + return false + } + } + return undefined + } catch { + return undefined + } +} + +function isStreamOptionsUnsupportedMessage(body: string): boolean { + const structuredValidation = getStructuredStreamOptionsValidationError(body) + if (structuredValidation !== undefined) return structuredValidation + + const normalized = body.toLowerCase().replace(/['"`]/g, '') + return ( + /(?:^|[\s:{,(])stream_options\b(?![.\[/])\s+(?:is\s+)?(?:an?\s+)?(?:unsupported|not\s+supported|unknown|unrecognized|not\s+permitted)\b/.test(normalized) || + /(?:unsupported|unknown|unrecognized|invalid)\s+(?:request\s+argument(?:\s+supplied)?|parameter(?:s|\(s\))?|field|property)(?:\s*[:=])?\s*(?:[\[(<]\s*)?stream_options\b(?![.\[/])/.test(normalized) || + /(?:request\s+argument(?:\s+supplied)?|parameter(?:s|\(s\))?|field|property)\s+(?:[\[(<]\s*)?stream_options\b(?![.\[/])(?:\s*[\])>])?\s+(?:is\s+)?(?:unsupported|not\s+supported|unknown|unrecognized|invalid|not\s+permitted)\b/.test(normalized) || + /\bstream_options\b(?![.\[/])\s+(?:is\s+)?(?:an?\s+)?(?:unsupported|not\s+supported|unknown|unrecognized|invalid)\s+(?:request\s+argument|parameter|field|property)\b/.test(normalized) || + /(?:additional properties? (?:are )?not allowed|extra inputs are not permitted|extra fields? (?:are )?not permitted|unexpected (?:field|property|parameter)).{0,80}\bstream_options\b(?![.\[/])/.test(normalized) || + /\bstream_options\b(?![.\[/]).{0,80}(?:additional properties? (?:are )?not allowed|extra inputs are not permitted|extra fields? (?:are )?not permitted|was unexpected)\b/.test(normalized) + ) +} + function isMalformedProviderResponse(body: string): boolean { const lower = body.toLowerCase() return ( @@ -563,6 +651,22 @@ export function classifyOpenAIHttpFailure(options: { } } + // Keep this classification narrower than a generic invalid-request match. + // The executor uses it to authorize one replay without `stream_options`. + if ( + (options.status === 400 || options.status === 422) && + isStreamOptionsUnsupportedMessage(body) + ) { + return { + source: 'http', + category: 'stream_options_unsupported', + retryable: false, + status: options.status, + message: body, + hint: 'Provider rejected the `stream_options` parameter; streaming usage is unavailable on this route.', + } + } + // `tool_stream` is a Z.AI-proprietary streaming extension. Some OpenAI- // compatible gateways (e.g. NVIDIA NIM) reject it with a 400 like // "Unsupported parameter(s): `tool_stream`". Classify it distinctly so the diff --git a/src/services/api/openaiShim.test.ts b/src/services/api/openaiShim.test.ts index f100f56ef..c23947c68 100644 --- a/src/services/api/openaiShim.test.ts +++ b/src/services/api/openaiShim.test.ts @@ -14,6 +14,8 @@ import { getAssistantMessageFromError, OPENCODE_GO_FREE_LIMIT_ERROR_MESSAGE, } from './errors.ts' +import { accumulateUsage, updateUsage } from './claude.js' +import { EMPTY_USAGE } from './emptyUsage.js' import { extractOpenAICategoryMarker, isOpenAIRequestNonReplayable, @@ -25,6 +27,8 @@ import { parseXmlToolCalls, } from './openaiShim.ts' import * as realGithubModelsCredentials from '../../utils/githubModelsCredentials.js' +import { tokenCountWithEstimation } from '../../utils/tokens.js' +import type { AssistantMessage } from '../../types/message.js' type FetchType = typeof globalThis.fetch @@ -1208,9 +1212,10 @@ test('uses correct empty input fallback schema for standard responses and respon // openaiShim test extraction seam 021 start: preserves usage from final OpenAI stream chunk with empty choices test('preserves usage from final OpenAI stream chunk with empty choices', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' globalThis.fetch = (async (_input, init) => { const url = typeof _input === 'string' ? _input : _input.url - expect(url).toBe('http://example.test/v1/chat/completions') + expect(url).toBe('http://127.0.0.1:8000/v1/chat/completions') const body = JSON.parse(String(init?.body)) expect(body.stream).toBe(true) @@ -1250,6 +1255,7 @@ test('preserves usage from final OpenAI stream chunk with empty choices', async prompt_tokens: 123, completion_tokens: 45, total_tokens: 168, + prompt_tokens_details: { cached_tokens: 23 }, }, }, ]) @@ -1274,13 +1280,63 @@ test('preserves usage from final OpenAI stream chunk with empty choices', async events.push(event) } - const usageEvent = events.find( + const usageEvents = events.filter( event => event.type === 'message_delta' && typeof event.usage === 'object' && event.usage !== null, - ) as { usage?: { input_tokens?: number; output_tokens?: number } } | undefined + ) as Array<{ + usage?: { + input_tokens?: number + output_tokens?: number + cache_read_input_tokens?: number + } + }> - expect(usageEvent).toBeDefined() - expect(usageEvent?.usage?.input_tokens).toBe(123) - expect(usageEvent?.usage?.output_tokens).toBe(45) + expect(usageEvents).toHaveLength(1) + expect(usageEvents[0]?.usage?.input_tokens).toBe(100) + expect(usageEvents[0]?.usage?.output_tokens).toBe(45) + expect(usageEvents[0]?.usage?.cache_read_input_tokens).toBe(23) + + let currentUsage = EMPTY_USAGE + let totalUsage = EMPTY_USAGE + for (const event of events) { + if (event.type === 'message_start') { + const message = event.message as { + usage?: Parameters[1] + } + currentUsage = updateUsage(EMPTY_USAGE, message.usage) + } else if (event.type === 'message_delta') { + currentUsage = updateUsage( + currentUsage, + event.usage as Parameters[1], + ) + } else if (event.type === 'message_stop') { + totalUsage = accumulateUsage(totalUsage, currentUsage) + } + } + + expect(events.filter(event => event.type === 'message_stop')).toHaveLength(1) + expect(totalUsage).toMatchObject({ + input_tokens: 100, + output_tokens: 45, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 23, + }) + + const accountedMessage = { + type: 'assistant', + uuid: '00000000-0000-0000-0000-000000000213', + timestamp: '2026-08-19T00:00:00.000Z', + message: { + id: 'msg_local_usage', + type: 'message', + role: 'assistant', + model: 'fake-model', + content: [{ type: 'text', text: 'hello world' }], + stop_reason: 'end_turn', + stop_sequence: null, + usage: currentUsage, + }, + } as unknown as AssistantMessage + expect(tokenCountWithEstimation([accountedMessage])).toBe(168) }) // openaiShim test extraction seam 021 end diff --git a/src/services/api/openaiShim/requestExecutor.integration.test.ts b/src/services/api/openaiShim/requestExecutor.integration.test.ts index a0099e6b4..4b21e0cfb 100644 --- a/src/services/api/openaiShim/requestExecutor.integration.test.ts +++ b/src/services/api/openaiShim/requestExecutor.integration.test.ts @@ -3087,7 +3087,7 @@ test('Local provider (vLLM/Ollama/etc.): strips unsupported store on chat_comple expect(requestBody?.store).toBeUndefined() }) -test('does not send stream_options to local OpenAI-compatible servers', async () => { +test('requests streaming usage from local OpenAI-compatible servers', async () => { process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' let requestBody: Record | undefined @@ -3107,7 +3107,195 @@ test('does not send stream_options to local OpenAI-compatible servers', async () }) expect(requestBody?.stream).toBe(true) - expect(requestBody).not.toHaveProperty('stream_options') + expect(requestBody?.stream_options).toEqual({ include_usage: true }) +}) + +test('retries a strict custom local stream once without rejected stream_options', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' + process.env.OPENAI_API_KEY = 'strict-local-key' + + const requestBodies: Array> = [] + const authorizations: Array = [] + globalThis.fetch = (async (_input, init) => { + requestBodies.push(JSON.parse(String(init?.body)) as Record) + const headers = init?.headers as Record | undefined + authorizations.push(headers?.Authorization ?? headers?.authorization ?? null) + if (requestBodies.length === 1) { + return new Response( + '{"error":{"message":"Unsupported parameter: stream_options","param":"stream_options"}}', + { status: 400, headers: { 'Content-Type': 'application/json' } }, + ) + } + return makeSseResponse(makeStreamChunks([ + { + id: 'chatcmpl-strict-local', + object: 'chat.completion.chunk', + model: 'strict-local-model', + choices: [{ index: 0, delta: { content: 'fallback works' }, finish_reason: null }], + }, + { + id: 'chatcmpl-strict-local', + object: 'chat.completion.chunk', + model: 'strict-local-model', + choices: [{ index: 0, delta: {}, finish_reason: 'stop' }], + }, + ])) + }) as unknown as FetchType + + const client = createOpenAIShimClient({}) as OpenAIShimClient + const result = await client.beta.messages + .create({ + model: 'strict-local-model', + messages: [{ role: 'user', content: 'hello' }], + tools: [{ + name: 'Read', + description: 'Read a file', + input_schema: { + type: 'object', + properties: { filePath: { type: 'string' } }, + required: ['filePath'], + }, + }], + max_tokens: 64, + stream: true, + }) + .withResponse() + + const textDeltas: string[] = [] + for await (const event of result.data) { + const delta = event.delta as { type?: string; text?: string } | undefined + if (delta?.type === 'text_delta' && typeof delta.text === 'string') { + textDeltas.push(delta.text) + } + } + + expect(requestBodies).toHaveLength(2) + expect(requestBodies[0]?.stream_options).toEqual({ include_usage: true }) + expect(requestBodies[1]).not.toHaveProperty('stream_options') + expect(requestBodies[1]?.tools).toEqual(requestBodies[0]?.tools) + expect(authorizations).toEqual(['Bearer strict-local-key', 'Bearer strict-local-key']) + expect(textDeltas.join('')).toBe('fallback works') +}) + +test('stops after one stream_options compatibility retry', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' + + const requestBodies: Array> = [] + globalThis.fetch = (async (_input, init) => { + requestBodies.push(JSON.parse(String(init?.body)) as Record) + return new Response('Unsupported parameter: stream_options', { + status: 400, + headers: { 'Content-Type': 'text/plain' }, + }) + }) as unknown as FetchType + + const client = createOpenAIShimClient({}) as OpenAIShimClient + await expect( + client.beta.messages.create({ + model: 'strict-local-model', + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }), + ).rejects.toThrow() + + expect(requestBodies).toHaveLength(2) + expect(requestBodies[0]?.stream_options).toEqual({ include_usage: true }) + expect(requestBodies[1]).not.toHaveProperty('stream_options') +}) + +test('does not remove stream_options or retry an unrelated custom-provider 400', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' + + const requestBodies: Array> = [] + globalThis.fetch = (async (_input, init) => { + requestBodies.push(JSON.parse(String(init?.body)) as Record) + return new Response( + '{"error":{"message":"Invalid request for model; received stream_options in the request body"}}', + { status: 400, headers: { 'Content-Type': 'application/json' } }, + ) + }) as unknown as FetchType + + const client = createOpenAIShimClient({}) as OpenAIShimClient + await expect( + client.beta.messages.create({ + model: 'strict-local-model', + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }), + ).rejects.toThrow() + + expect(requestBodies).toHaveLength(1) + expect(requestBodies[0]?.stream_options).toEqual({ include_usage: true }) +}) + +test('does not retry a nested stream_options validation failure', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' + + const requestBodies: Array> = [] + globalThis.fetch = (async (_input, init) => { + requestBodies.push(JSON.parse(String(init?.body)) as Record) + return new Response( + 'Invalid parameter: stream_options.include_usage must be a boolean', + { status: 400, headers: { 'Content-Type': 'text/plain' } }, + ) + }) as unknown as FetchType + + const client = createOpenAIShimClient({}) as OpenAIShimClient + await expect( + client.beta.messages.create({ + model: 'strict-local-model', + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }), + ).rejects.toThrow() + + expect(requestBodies).toHaveLength(1) + expect(requestBodies[0]?.stream_options).toEqual({ include_usage: true }) +}) + +test('does not retry a stream_options rejection after user cancellation', async () => { + process.env.OPENAI_BASE_URL = 'http://127.0.0.1:8000/v1' + + const controller = new AbortController() + const abortReason = new DOMException('The operation was aborted.', 'AbortError') + const requestBodies: Array> = [] + globalThis.fetch = (async (_input, init) => { + requestBodies.push(JSON.parse(String(init?.body)) as Record) + const response = new Response(null, { + status: 400, + headers: { 'Content-Type': 'text/plain' }, + }) + Object.defineProperty(response, 'text', { + value: async () => { + controller.abort(abortReason) + return 'Unsupported parameter: stream_options' + }, + configurable: true, + }) + return response + }) as unknown as FetchType + + const client = createOpenAIShimClient({}) as OpenAIShimClient + let caught: unknown + try { + await client.beta.messages.create( + { + model: 'strict-local-model', + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }, + { signal: controller.signal }, + ) + } catch (error) { + caught = error + } + + expect(requestBodies).toHaveLength(1) + expect(caught).toBe(abortReason) }) test('Mistral: strips unsupported store on chat_completions (#739)', async () => { diff --git a/src/services/api/openaiShim/requestExecutor.ts b/src/services/api/openaiShim/requestExecutor.ts index 373c70b54..b1b74c9dc 100644 --- a/src/services/api/openaiShim/requestExecutor.ts +++ b/src/services/api/openaiShim/requestExecutor.ts @@ -594,6 +594,7 @@ export async function executeOpenAIRequest( const attemptedLocalRequestUrls = new Set([requestUrl]) let didRetryWithoutTools = false let didRetryWithoutToolStream = false + let didRetryWithoutStreamOptions = false let retryCredentialLease: CredentialLease | null = null let didRefreshCopilotToken = false let refreshedCopilotToken: string | undefined @@ -637,6 +638,15 @@ export async function executeOpenAIRequest( } } + const serializedBodyHasTopLevelField = (field: string): boolean => { + try { + const payload = JSON.parse(serializedBody) as Record + return Object.prototype.hasOwnProperty.call(payload, field) + } catch { + return false + } + } + // WHY: byte-identity required for implicit prefix caching in // OpenAI/Kimi/DeepSeek. stableStringify sorts object keys at every // depth so spurious insertion-order differences across rebuilds of @@ -968,7 +978,21 @@ export async function executeOpenAIRequest( } // Read body exactly once here — Response body is a stream that can only // be consumed a single time. - const errorBody = await response.text().catch(() => 'unknown error') + let errorBody: string + try { + errorBody = await response.text() + } catch (error) { + if (options?.signal?.aborted) { + throw options.signal.reason ?? error + } + errorBody = 'unknown error' + } + if (options?.signal?.aborted) { + throw ( + options.signal.reason ?? + new DOMException('The operation was aborted.', 'AbortError') + ) + } const rateHint = isGithub && response.status === 429 ? formatRetryAfterHint(response) : '' @@ -1131,6 +1155,30 @@ export async function executeOpenAIRequest( continue } + // Some strict OpenAI-compatible servers reject the optional usage extension + // instead of ignoring it. Retry only an exact top-level field rejection, + // and only when the serialized request actually contained the field. This + // keeps usage enabled for capable custom routes without replaying unrelated + // client errors or a request the user has already cancelled. + if ( + !didRetryWithoutStreamOptions && + failure.category === 'stream_options_unsupported' && + options?.signal?.aborted !== true && + serializedBodyHasTopLevelField('stream_options') + ) { + didRetryWithoutStreamOptions = true + maxAttempts += 1 + delete body.stream_options + refreshSerializedBody() + retryCredentialLease = credentialLease + + logForDebugging( + `[OpenAIShim] self-heal retry reason=stream_options_unsupported method=POST url=${redactUrlForDiagnostics(requestUrl)} model=${request.resolvedModel}`, + { level: 'warn' }, + ) + continue + } + // `tool_stream` self-heal (#1950): some OpenAI-compatible gateways (e.g. // NVIDIA NIM) reject the Z.AI-proprietary `tool_stream` parameter with a // 400. Drop only that parameter and retry with tools intact — streaming diff --git a/src/services/api/openaiShim/requestPreparation.test.ts b/src/services/api/openaiShim/requestPreparation.test.ts index 1b89cd724..68ce7758a 100644 --- a/src/services/api/openaiShim/requestPreparation.test.ts +++ b/src/services/api/openaiShim/requestPreparation.test.ts @@ -126,3 +126,134 @@ test('prepares tools and streaming options for a remote chat route', async () => function: { name: 'Read' }, }) }) + +test('requests streaming usage from a custom loopback chat route', async () => { + await ensureIntegrationsLoaded() + const processEnv = { + OPENAI_BASE_URL: 'http://127.0.0.1:8000/v1', + OPENAI_API_KEY: 'test-key', + } + const request = resolveProviderRequest({ + model: 'local-openai-model', + processEnv, + }) + const prepared = prepareOpenAIRequest({ + request, + requestProcessEnv: processEnv, + params: { + model: 'local-openai-model', + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }, + dependencies, + }) + + expect(prepared.body.stream_options).toEqual({ include_usage: true }) +}) + +test.each([ + ['LM Studio', 'http://localhost:1234/v1', 'local-model'], + ['non-native Ollama', 'http://192.168.1.10:11434/v1', 'llama3.1:8b'], + ['remote Ollama', 'https://ollama.com/v1', 'qwen3-coder-next:cloud'], +])('requests streaming usage from a compatible %s route', async ( + _label, + baseUrl, + model, +) => { + await ensureIntegrationsLoaded() + const processEnv = { + OPENAI_BASE_URL: baseUrl, + OPENAI_API_KEY: 'test-key', + } + const request = resolveProviderRequest({ model, processEnv }) + const prepared = prepareOpenAIRequest({ + request, + requestProcessEnv: processEnv, + params: { + model, + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }, + dependencies, + }) + + expect(prepared.useNativeOllamaChat).toBe(false) + expect(JSON.parse(prepared.serializeBody()).stream_options).toEqual({ + include_usage: true, + }) +}) + +test('keeps native Ollama streaming options off the wire', async () => { + await ensureIntegrationsLoaded() + const processEnv = { + OPENAI_BASE_URL: 'http://localhost:11434/v1', + OPENAI_API_KEY: 'test-key', + } + const model = 'llama3.1:8b' + const request = resolveProviderRequest({ model, processEnv }) + const prepared = prepareOpenAIRequest({ + request, + requestProcessEnv: processEnv, + params: { + model, + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }, + dependencies, + }) + + expect(prepared.useNativeOllamaChat).toBe(true) + expect(JSON.parse(prepared.serializeBody())).not.toHaveProperty( + 'stream_options', + ) +}) + +test('lets explicit stream_options removal override streaming usage', async () => { + await ensureIntegrationsLoaded() + const processEnv = { + OPENAI_BASE_URL: 'https://api.xiaomimimo.com/v1', + OPENAI_API_KEY: 'test-key', + } + const model = 'mimo-v2.5-pro' + const request = resolveProviderRequest({ model, processEnv }) + const prepared = prepareOpenAIRequest({ + request, + requestProcessEnv: processEnv, + params: { + model, + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: true, + }, + dependencies, + }) + + expect(prepared.shimConfig.removeBodyFields).toContain('stream_options') + expect(prepared.body).not.toHaveProperty('stream_options') +}) + +test('omits streaming options from a non-streaming loopback request', async () => { + await ensureIntegrationsLoaded() + const processEnv = { + OPENAI_BASE_URL: 'http://127.0.0.1:8000/v1', + OPENAI_API_KEY: 'test-key', + } + const model = 'local-openai-model' + const request = resolveProviderRequest({ model, processEnv }) + const prepared = prepareOpenAIRequest({ + request, + requestProcessEnv: processEnv, + params: { + model, + messages: [{ role: 'user', content: 'hello' }], + max_tokens: 64, + stream: false, + }, + dependencies, + }) + + expect(prepared.body).not.toHaveProperty('stream_options') +}) diff --git a/src/services/api/openaiShim/requestPreparation.ts b/src/services/api/openaiShim/requestPreparation.ts index 383781c5a..eec877c47 100644 --- a/src/services/api/openaiShim/requestPreparation.ts +++ b/src/services/api/openaiShim/requestPreparation.ts @@ -250,7 +250,7 @@ export function prepareOpenAIRequest({ else if (maxCompletionTokensValue !== undefined) { body.max_completion_tokens = maxCompletionTokensValue } - if (params.stream && !isLocalProviderUrl(request.baseUrl)) { + if (params.stream) { body.stream_options = { include_usage: true } } diff --git a/src/services/api/openaiShim/streamConversion.test.ts b/src/services/api/openaiShim/streamConversion.test.ts index d341e060a..4140697e9 100644 --- a/src/services/api/openaiShim/streamConversion.test.ts +++ b/src/services/api/openaiShim/streamConversion.test.ts @@ -290,7 +290,10 @@ test('emits usage supplied by a final empty-choices chunk', async () => { createStreamDependencies(), ), ) - expect(events).toContainEqual({ + const usageEvents = events.filter( + event => event.type === 'message_delta' && event.usage !== undefined, + ) + expect(usageEvents).toEqual([{ type: 'message_delta', delta: { stop_reason: 'end_turn', stop_sequence: null }, usage: { @@ -299,7 +302,77 @@ test('emits usage supplied by a final empty-choices chunk', async () => { cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, + }]) + + const textIndex = events.findIndex( + event => + event.type === 'content_block_delta' && + (event.delta as { type?: string; text?: string } | undefined)?.type === + 'text_delta' && + (event.delta as { text?: string } | undefined)?.text === 'done', + ) + const firstStopIndex = events.findIndex( + event => + event.type === 'message_delta' && + (event.delta as { stop_reason?: string } | undefined)?.stop_reason === + 'end_turn', + ) + const usageIndex = events.indexOf(usageEvents[0]!) + const messageStopIndex = events.findIndex(event => event.type === 'message_stop') + expect(textIndex).toBeLessThan(firstStopIndex) + expect(firstStopIndex).toBeLessThan(usageIndex) + expect(usageIndex).toBeLessThan(messageStopIndex) +}) + +test('emits terminal usage once when both terminal chunks report it', async () => { + const usage = { prompt_tokens: 10, completion_tokens: 2 } + const events = await collect( + openaiStreamToAnthropic( + makeSseResponse([ + makeOpenAIChunk({ content: 'done' }), + makeOpenAIChunk({}, 'stop', usage), + { choices: [], usage }, + ]), + 'test-model', + undefined, + false, + undefined, + createStreamDependencies(), + ), + ) + + expect( + events.filter( + event => event.type === 'message_delta' && event.usage !== undefined, + ), + ).toHaveLength(1) +}) + +test('completes normally when a stream does not report usage', async () => { + const events = await collect( + openaiStreamToAnthropic( + makeSseResponse([ + makeOpenAIChunk({ content: 'done' }), + makeOpenAIChunk({}, 'stop'), + ]), + 'test-model', + undefined, + false, + undefined, + createStreamDependencies(), + ), + ) + + expect( + events.filter( + event => event.type === 'message_delta' && event.usage !== undefined, + ), + ).toHaveLength(0) + expect(events).toContainEqual({ + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, }) + expect(events.at(-1)).toEqual({ type: 'message_stop' }) }) test('strips think tags split across content chunks without phrase heuristics', async () => { @@ -385,6 +458,15 @@ test('routes provider JSON stream fallback through non-streaming conversion', as index: 0, delta: { type: 'text_delta', text: 'fallback' }, }) + expect( + events.filter( + event => event.type === 'message_delta' && event.usage !== undefined, + ), + ).toEqual([{ + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, + usage: { input_tokens: 2, output_tokens: 3 }, + }]) expect(events.at(-1)).toEqual({ type: 'message_stop' }) })