From 011d7df6f14b0cfc58d21b6158e48da3677cd916 Mon Sep 17 00:00:00 2001 From: tt-a1i <53142663+tt-a1i@users.noreply.github.com> Date: Sat, 20 Jun 2026 22:41:53 +0800 Subject: [PATCH] fix(core): parse max output token env strictly --- .../anthropicContentGenerator.test.ts | 78 +++++++++++++++++++ .../anthropicContentGenerator.ts | 8 +- .../provider/default.test.ts | 40 ++++++++++ .../provider/default.ts | 8 +- packages/core/src/core/tokenLimits.ts | 14 ++++ 5 files changed, 142 insertions(+), 6 deletions(-) diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts index b75b1414cc4..fc75a7e80a6 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts @@ -70,16 +70,20 @@ const importConverter = async (): Promise<{ }> => import('./converter.js'); describe('AnthropicContentGenerator', () => { + const MAX_OUTPUT_TOKENS_ENV = 'QWEN_CODE_MAX_OUTPUT_TOKENS'; let mockConfig: Config; let anthropicState: { constructorOptions?: Record; lastCreateArgs?: AnthropicCreateArgs; createImpl: ReturnType; }; + let savedMaxOutputTokensEnv: string | undefined; beforeEach(async () => { vi.clearAllMocks(); vi.resetModules(); + savedMaxOutputTokensEnv = process.env[MAX_OUTPUT_TOKENS_ENV]; + delete process.env[MAX_OUTPUT_TOKENS_ENV]; mockTokenizer.calculateTokens.mockResolvedValue({ totalTokens: 50, @@ -106,6 +110,11 @@ describe('AnthropicContentGenerator', () => { }); afterEach(() => { + if (savedMaxOutputTokensEnv === undefined) { + delete process.env[MAX_OUTPUT_TOKENS_ENV]; + } else { + process.env[MAX_OUTPUT_TOKENS_ENV] = savedMaxOutputTokensEnv; + } vi.restoreAllMocks(); }); @@ -1620,6 +1629,75 @@ describe('AnthropicContentGenerator', () => { ); }); + it('ignores malformed QWEN_CODE_MAX_OUTPUT_TOKENS values', async () => { + const { AnthropicContentGenerator } = await importGenerator(); + + for (const envValue of ['1.5', '2k', 'abc']) { + process.env[MAX_OUTPUT_TOKENS_ENV] = envValue; + anthropicState.createImpl.mockResolvedValueOnce({ + id: `anthropic-${envValue}`, + model: 'claude-sonnet-4', + content: [{ type: 'text', text: 'hi' }], + }); + + const generator = new AnthropicContentGenerator( + { + model: 'claude-sonnet-4', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: {}, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + await generator.generateContent({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + + const [anthropicRequest] = + anthropicState.lastCreateArgs as AnthropicCreateArgs; + expect(anthropicRequest).toEqual( + expect.objectContaining({ max_tokens: 8000 }), + ); + } + }); + + it('respects a valid QWEN_CODE_MAX_OUTPUT_TOKENS value', async () => { + const { AnthropicContentGenerator } = await importGenerator(); + process.env[MAX_OUTPUT_TOKENS_ENV] = '9000'; + anthropicState.createImpl.mockResolvedValue({ + id: 'anthropic-1', + model: 'claude-sonnet-4', + content: [{ type: 'text', text: 'hi' }], + }); + + const generator = new AnthropicContentGenerator( + { + model: 'claude-sonnet-4', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: {}, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + await generator.generateContent({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + + const [anthropicRequest] = + anthropicState.lastCreateArgs as AnthropicCreateArgs; + expect(anthropicRequest).toEqual( + expect.objectContaining({ max_tokens: 9000 }), + ); + }); + it('respects configured max_tokens for unknown models', async () => { const { AnthropicContentGenerator } = await importGenerator(); anthropicState.createImpl.mockResolvedValue({ diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts index a736385eb20..f4edde4b333 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts @@ -40,6 +40,7 @@ import { tokenLimit, CAPPED_DEFAULT_MAX_TOKENS, hasExplicitOutputLimit, + parsePositiveIntegerEnvValue, } from '../tokenLimits.js'; const debugLogger = createDebugLogger('ANTHROPIC'); @@ -594,9 +595,10 @@ export class AnthropicContentGenerator implements ContentGenerator { : userMaxTokens; } else { // No explicit user config — check env var, then use capped default. - const envVal = process.env['QWEN_CODE_MAX_OUTPUT_TOKENS']; - const envMaxTokens = envVal ? parseInt(envVal, 10) : NaN; - if (!isNaN(envMaxTokens) && envMaxTokens > 0) { + const envMaxTokens = parsePositiveIntegerEnvValue( + process.env['QWEN_CODE_MAX_OUTPUT_TOKENS'], + ); + if (envMaxTokens !== undefined) { maxTokens = isKnownModel ? Math.min(envMaxTokens, modelLimit) : envMaxTokens; diff --git a/packages/core/src/core/openaiContentGenerator/provider/default.test.ts b/packages/core/src/core/openaiContentGenerator/provider/default.test.ts index ce57b5965f2..58b160b175d 100644 --- a/packages/core/src/core/openaiContentGenerator/provider/default.test.ts +++ b/packages/core/src/core/openaiContentGenerator/provider/default.test.ts @@ -10,6 +10,7 @@ import { expect, vi, beforeEach, + afterEach, type MockedFunction, } from 'vitest'; import OpenAI from 'openai'; @@ -37,12 +38,16 @@ vi.mock('../../../utils/runtimeFetchOptions.js', () => ({ })); describe('DefaultOpenAICompatibleProvider', () => { + const MAX_OUTPUT_TOKENS_ENV = 'QWEN_CODE_MAX_OUTPUT_TOKENS'; let provider: DefaultOpenAICompatibleProvider; let mockContentGeneratorConfig: ContentGeneratorConfig; let mockCliConfig: Config; + let savedMaxOutputTokensEnv: string | undefined; beforeEach(() => { vi.clearAllMocks(); + savedMaxOutputTokensEnv = process.env[MAX_OUTPUT_TOKENS_ENV]; + delete process.env[MAX_OUTPUT_TOKENS_ENV]; const mockedBuildRuntimeFetchOptions = buildRuntimeFetchOptions as unknown as MockedFunction< (sdkType: 'openai', proxyUrl?: string) => OpenAIRuntimeFetchOptions @@ -70,6 +75,14 @@ describe('DefaultOpenAICompatibleProvider', () => { ); }); + afterEach(() => { + if (savedMaxOutputTokensEnv === undefined) { + delete process.env[MAX_OUTPUT_TOKENS_ENV]; + } else { + process.env[MAX_OUTPUT_TOKENS_ENV] = savedMaxOutputTokensEnv; + } + }); + describe('constructor', () => { it('should initialize with provided configs', () => { expect(provider).toBeInstanceOf(DefaultOpenAICompatibleProvider); @@ -209,6 +222,33 @@ describe('DefaultOpenAICompatibleProvider', () => { expect(result.max_tokens).toBe(8000); }); + it('should ignore malformed QWEN_CODE_MAX_OUTPUT_TOKENS values', () => { + const request: OpenAI.Chat.ChatCompletionCreateParams = { + model: 'gpt-4', + messages: [{ role: 'user', content: 'Hello' }], + }; + + for (const envValue of ['1.5', '2k', 'abc']) { + process.env[MAX_OUTPUT_TOKENS_ENV] = envValue; + + const result = provider.buildRequest(request, 'prompt-id'); + + expect(result.max_tokens).toBe(8000); + } + }); + + it('should respect a valid QWEN_CODE_MAX_OUTPUT_TOKENS value', () => { + process.env[MAX_OUTPUT_TOKENS_ENV] = '9000'; + const request: OpenAI.Chat.ChatCompletionCreateParams = { + model: 'gpt-4', + messages: [{ role: 'user', content: 'Hello' }], + }; + + const result = provider.buildRequest(request, 'prompt-id'); + + expect(result.max_tokens).toBe(9000); + }); + it('should respect user max_tokens for unknown models (deployment aliases, self-hosted)', () => { // Unknown models: user config is respected entirely (backend may support larger limits) const request: OpenAI.Chat.ChatCompletionCreateParams = { diff --git a/packages/core/src/core/openaiContentGenerator/provider/default.ts b/packages/core/src/core/openaiContentGenerator/provider/default.ts index 9ddd8d6b01f..d659eadb6e5 100644 --- a/packages/core/src/core/openaiContentGenerator/provider/default.ts +++ b/packages/core/src/core/openaiContentGenerator/provider/default.ts @@ -9,6 +9,7 @@ import { tokenLimit, CAPPED_DEFAULT_MAX_TOKENS, hasExplicitOutputLimit, + parsePositiveIntegerEnvValue, } from '../../tokenLimits.js'; type AssistantMessageWithReasoningFields = @@ -187,9 +188,10 @@ export class DefaultOpenAICompatibleProvider // No explicit user config — check env var, then use capped default. // Capped default (8K) reduces GPU slot over-reservation by ~4×. // Requests hitting the cap get one clean retry at 64K (geminiChat.ts). - const envVal = process.env['QWEN_CODE_MAX_OUTPUT_TOKENS']; - const envMaxTokens = envVal ? parseInt(envVal, 10) : NaN; - if (!isNaN(envMaxTokens) && envMaxTokens > 0) { + const envMaxTokens = parsePositiveIntegerEnvValue( + process.env['QWEN_CODE_MAX_OUTPUT_TOKENS'], + ); + if (envMaxTokens !== undefined) { effectiveMaxTokens = isKnownModel ? Math.min(envMaxTokens, modelLimit) : envMaxTokens; diff --git a/packages/core/src/core/tokenLimits.ts b/packages/core/src/core/tokenLimits.ts index 824de8560a5..0365e3feb57 100644 --- a/packages/core/src/core/tokenLimits.ts +++ b/packages/core/src/core/tokenLimits.ts @@ -18,6 +18,20 @@ export const DEFAULT_OUTPUT_TOKEN_LIMIT: TokenCount = 32_000; // 32K tokens export const CAPPED_DEFAULT_MAX_TOKENS: TokenCount = 8_000; export const ESCALATED_MAX_TOKENS: TokenCount = 64_000; +export function parsePositiveIntegerEnvValue( + raw: string | undefined, +): number | undefined { + if (raw === undefined) return undefined; + + const trimmed = raw.trim(); + if (!/^\d+$/.test(trimmed)) return undefined; + + const parsed = Number(trimmed); + if (!Number.isSafeInteger(parsed) || parsed <= 0) return undefined; + + return parsed; +} + /** * Accurate numeric limits: * - power-of-two approximations (128K -> 131072, 256K -> 262144, etc.)