From 6e276a3c06f83ff9f8792cf49c73257b0f8768c7 Mon Sep 17 00:00:00 2001 From: Luyu Cheng <2239547+chengluyu@users.noreply.github.com> Date: Wed, 3 Jun 2026 13:44:29 +0800 Subject: [PATCH 1/3] fix: simplify goal budget tool schema --- .changeset/fix-goal-tool-object-schema.md | 6 +++ .../src/tools/builtin/goal/set-goal-budget.ts | 48 +++++++++++-------- packages/agent-core/test/tools/goal.test.ts | 42 +++++++++++++++- 3 files changed, 76 insertions(+), 20 deletions(-) create mode 100644 .changeset/fix-goal-tool-object-schema.md diff --git a/.changeset/fix-goal-tool-object-schema.md b/.changeset/fix-goal-tool-object-schema.md new file mode 100644 index 0000000000..6caf9baa2c --- /dev/null +++ b/.changeset/fix-goal-tool-object-schema.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code": patch +--- + +Fix goal budget tool schemas for OpenAI-compatible providers. diff --git a/packages/agent-core/src/tools/builtin/goal/set-goal-budget.ts b/packages/agent-core/src/tools/builtin/goal/set-goal-budget.ts index 4a90e55a1e..acf09c410f 100644 --- a/packages/agent-core/src/tools/builtin/goal/set-goal-budget.ts +++ b/packages/agent-core/src/tools/builtin/goal/set-goal-budget.ts @@ -16,22 +16,16 @@ import DESCRIPTION from './set-goal-budget.md'; const MIN_REASONABLE_TIME_BUDGET_MS = 1_000; const MAX_REASONABLE_TIME_BUDGET_MS = 24 * 60 * 60 * 1000; +const BUDGET_UNITS = ['turns', 'tokens', 'milliseconds', 'seconds', 'minutes', 'hours'] as const; -const WholeNumberBudgetValueSchema = z - .number() - .int() - .positive() - .describe('The positive whole-number budget value.'); -const TimeBudgetValueSchema = z.number().positive().describe('The positive numeric time budget value.'); - -export const SetGoalBudgetToolInputSchema = z.discriminatedUnion('unit', [ - z.object({ value: WholeNumberBudgetValueSchema, unit: z.literal('turns') }).strict(), - z.object({ value: WholeNumberBudgetValueSchema, unit: z.literal('tokens') }).strict(), - z.object({ value: TimeBudgetValueSchema, unit: z.literal('milliseconds') }).strict(), - z.object({ value: TimeBudgetValueSchema, unit: z.literal('seconds') }).strict(), - z.object({ value: TimeBudgetValueSchema, unit: z.literal('minutes') }).strict(), - z.object({ value: TimeBudgetValueSchema, unit: z.literal('hours') }).strict(), -]); +export const SetGoalBudgetToolInputSchema = z + .object({ + // Keep the provider-facing schema simple. Fractional turn/token budgets + // are normalized during execution instead of rejected at schema validation. + value: z.number().positive().describe('The positive numeric budget value.'), + unit: z.enum(BUDGET_UNITS), + }) + .strict(); export type SetGoalBudgetToolInput = z.infer; @@ -46,21 +40,27 @@ export class SetGoalBudgetTool implements BuiltinTool { const store = requireGoalStore(this.agent, this.name); if (isGoalToolError(store)) return store; + const normalizedArgs = normalizeBudgetInput(args); return { - description: `Setting goal budget: ${formatBudget(args.value, args.unit)}`, + description: `Setting goal budget: ${formatBudget( + normalizedArgs.value, + normalizedArgs.unit, + )}`, approvalRule: this.name, execute: async () => { try { - const budget = budgetLimitsFromInput(args); + const budget = budgetLimitsFromInput(normalizedArgs); if (budget === null) { return { output: - `Goal budget not set: ${formatBudget(args.value, args.unit)} is not a ` + + `Goal budget not set: ${formatBudget(normalizedArgs.value, normalizedArgs.unit)} is not a ` + 'reasonable goal budget.', }; } await store.setBudgetLimits({ budgetLimits: budget, actor: 'model' }); - return { output: `Goal budget set: ${formatBudget(args.value, args.unit)}.` }; + return { + output: `Goal budget set: ${formatBudget(normalizedArgs.value, normalizedArgs.unit)}.`, + }; } catch (error) { return goalErrorResult(error); } @@ -69,6 +69,16 @@ export class SetGoalBudgetTool implements BuiltinTool { } } +function normalizeBudgetInput(input: SetGoalBudgetToolInput): SetGoalBudgetToolInput { + switch (input.unit) { + case 'turns': + case 'tokens': + return { ...input, value: Math.max(1, Math.round(input.value)) }; + default: + return input; + } +} + function budgetLimitsFromInput(input: SetGoalBudgetToolInput): GoalBudgetLimits | null { switch (input.unit) { case 'turns': diff --git a/packages/agent-core/test/tools/goal.test.ts b/packages/agent-core/test/tools/goal.test.ts index 60cdf2396d..e05dbf8dfb 100644 --- a/packages/agent-core/test/tools/goal.test.ts +++ b/packages/agent-core/test/tools/goal.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, it } from 'vitest'; import type { Agent } from '../../src/agent'; import { ErrorCodes } from '../../src/errors'; +import { compileToolArgsValidator, validateToolArgs } from '../../src/tools/args-validator'; import { CreateGoalTool, CreateGoalToolInputSchema, @@ -129,13 +130,36 @@ describe('GetGoalTool', () => { }); describe('SetGoalBudgetTool', () => { + it('advertises an object parameter schema for OpenAI-compatible providers', () => { + const parameters = new SetGoalBudgetTool(fakeAgent()).parameters; + + expect(parameters).toMatchObject({ + type: 'object', + required: ['value', 'unit'], + additionalProperties: false, + properties: { + value: expect.objectContaining({ type: 'number', exclusiveMinimum: 0 }), + unit: expect.objectContaining({ + type: 'string', + enum: ['turns', 'tokens', 'milliseconds', 'seconds', 'minutes', 'hours'], + }), + }, + }); + expect(parameters).not.toHaveProperty('oneOf'); + expect(parameters).not.toHaveProperty('anyOf'); + + const validator = compileToolArgsValidator(parameters); + expect(validateToolArgs(validator, { value: 1.5, unit: 'turns' })).toBeNull(); + expect(validateToolArgs(validator, { value: 1.5, unit: 'hours' })).toBeNull(); + }); + it('accepts a value with a supported budget unit', () => { for (const unit of ['turns', 'tokens', 'milliseconds', 'seconds', 'minutes', 'hours']) { expect(SetGoalBudgetToolInputSchema.safeParse({ value: 20, unit }).success).toBe(true); } expect(SetGoalBudgetToolInputSchema.safeParse({ value: 0, unit: 'turns' }).success).toBe(false); expect(SetGoalBudgetToolInputSchema.safeParse({ value: 1, unit: 'years' }).success).toBe(false); - expect(SetGoalBudgetToolInputSchema.safeParse({ value: 1.5, unit: 'turns' }).success).toBe(false); + expect(SetGoalBudgetToolInputSchema.safeParse({ value: 1.5, unit: 'turns' }).success).toBe(true); expect(SetGoalBudgetToolInputSchema.safeParse({ value: 1.5, unit: 'hours' }).success).toBe(true); }); @@ -160,6 +184,22 @@ describe('SetGoalBudgetTool', () => { expect(store.getGoal().goal?.budget.wallClockBudgetMs).toBe(30 * 60 * 1000); }); + it('rounds fractional turn and token budgets before setting them', async () => { + const store = makeStore(); + await store.createGoal({ objective: 'work' }); + const tool = new SetGoalBudgetTool(fakeAgent({ goals: store })); + + expect((await executeTool(tool, ctx({ value: 1.5, unit: 'turns' }))).output).toBe( + 'Goal budget set: 2 turns.', + ); + expect(store.getGoal().goal?.budget.turnBudget).toBe(2); + + expect((await executeTool(tool, ctx({ value: 0.4, unit: 'tokens' }))).output).toBe( + 'Goal budget set: 1 token.', + ); + expect(store.getGoal().goal?.budget.tokenBudget).toBe(1); + }); + it('ignores unreasonable time budgets and tells the model why', async () => { const store = makeStore(); await store.createGoal({ objective: 'work' }); From 018b1a1a04532e68215c91212ec3c320f21c7ffd Mon Sep 17 00:00:00 2001 From: Luyu Cheng <2239547+chengluyu@users.noreply.github.com> Date: Wed, 3 Jun 2026 13:56:00 +0800 Subject: [PATCH 2/3] fix: use model output cap for completions --- .changeset/use-model-output-cap.md | 6 +++ packages/agent-core/src/agent/config/index.ts | 4 ++ packages/agent-core/src/agent/index.ts | 1 + .../src/session/provider-manager.ts | 4 +- .../agent-core/src/utils/completion-budget.ts | 4 ++ .../test/agent/config-state.test.ts | 53 ++++++++++++++++++- .../test/harness/runtime-provider.test.ts | 28 ++++++++++ .../test/utils/completion-budget.test.ts | 10 ++++ 8 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 .changeset/use-model-output-cap.md diff --git a/.changeset/use-model-output-cap.md b/.changeset/use-model-output-cap.md new file mode 100644 index 0000000000..0fb49069fb --- /dev/null +++ b/.changeset/use-model-output-cap.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code": patch +--- + +Use configured model output limits for completion token caps. diff --git a/packages/agent-core/src/agent/config/index.ts b/packages/agent-core/src/agent/config/index.ts index 45aca21367..3f1aafafab 100644 --- a/packages/agent-core/src/agent/config/index.ts +++ b/packages/agent-core/src/agent/config/index.ts @@ -127,6 +127,10 @@ export class ConfigState { return this.tryResolvedProviderConfig()?.modelCapabilities ?? UNKNOWN_CAPABILITY; } + get maxOutputSize(): number | undefined { + return this.tryResolvedProviderConfig()?.maxOutputSize; + } + private get resolvedProviderConfig(): ResolvedRuntimeProvider | undefined { if (this._modelAlias === undefined) return undefined; return this.agent.modelProvider?.resolveProviderConfig(this._modelAlias); diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index 8ee06f260f..9f9323053b 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/packages/agent-core/src/agent/index.ts @@ -209,6 +209,7 @@ export class Agent { const provider = this.config.provider.withThinking(this.config.thinkingLevel); const loopControl = this.kimiConfig?.loopControl; const completionBudgetConfig = resolveCompletionBudget({ + maxOutputSize: this.config.maxOutputSize, reservedContextSize: loopControl?.reservedContextSize, }); return new KosongLLM({ diff --git a/packages/agent-core/src/session/provider-manager.ts b/packages/agent-core/src/session/provider-manager.ts index f675c12669..38d3e1cd1f 100644 --- a/packages/agent-core/src/session/provider-manager.ts +++ b/packages/agent-core/src/session/provider-manager.ts @@ -17,6 +17,7 @@ export interface ResolvedRuntimeProvider { readonly providerName: string; readonly provider: KosongProviderConfig; readonly modelCapabilities: ModelCapability; + readonly maxOutputSize?: number; } interface ProviderManagerOptions { @@ -57,7 +58,7 @@ export class SingleModelProvider implements ModelProvider { modelCapabilities: this.modelCapabilities, providerName: 'single-model-provider', provider: this.providerConfig, - } + }; } } @@ -115,6 +116,7 @@ export class ProviderManager implements ModelProvider { providerName, provider, modelCapabilities: resolveModelCapabilities(alias, provider), + maxOutputSize: alias.maxOutputSize, }; } diff --git a/packages/agent-core/src/utils/completion-budget.ts b/packages/agent-core/src/utils/completion-budget.ts index 5136ec60e4..ceb086ef29 100644 --- a/packages/agent-core/src/utils/completion-budget.ts +++ b/packages/agent-core/src/utils/completion-budget.ts @@ -16,6 +16,7 @@ const DEFAULT_UNKNOWN_CONTEXT_FALLBACK = 32000; * non-positive env values disable clamping. */ export function resolveCompletionBudget(args: { + readonly maxOutputSize?: number; readonly reservedContextSize?: number; readonly env?: NodeJS.ProcessEnv; }): CompletionBudgetConfig | undefined { @@ -28,6 +29,9 @@ export function resolveCompletionBudget(args: { if (fromLegacy !== 'absent') { return fromLegacy === 'disabled' ? undefined : { hardCap: fromLegacy }; } + if (args.maxOutputSize !== undefined && args.maxOutputSize > 0) { + return { hardCap: args.maxOutputSize }; + } if (args.reservedContextSize !== undefined && args.reservedContextSize > 0) { return { fallback: args.reservedContextSize }; } diff --git a/packages/agent-core/test/agent/config-state.test.ts b/packages/agent-core/test/agent/config-state.test.ts index f130f57cfb..4aa3c2c45e 100644 --- a/packages/agent-core/test/agent/config-state.test.ts +++ b/packages/agent-core/test/agent/config-state.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from 'vitest'; +import { emptyUsage } from '@moonshot-ai/kosong'; import { ProviderManager } from '../../src/session/provider-manager'; import { testAgent } from './harness'; @@ -73,7 +74,57 @@ describe('ConfigState model capabilities', () => { }); }); -it('uses session id as a provider prompt cache hint without storing it on Agent', () => { + it('uses model max output size as the LLM completion cap', async () => { + let requestMaxTokens: unknown; + const ctx = testAgent({ + generate: async (provider) => { + requestMaxTokens = ( + provider as unknown as { readonly modelParameters: Record } + ).modelParameters['max_tokens']; + return { + id: 'response-1', + message: { role: 'assistant', content: [], toolCalls: [] }, + usage: emptyUsage(), + finishReason: 'completed', + rawFinishReason: 'stop', + }; + }, + providerManager: new ProviderManager({ + config: { + providers: { + deepseek: { + type: 'openai', + apiKey: 'test-key', + baseUrl: 'https://api.deepseek.example/v1', + }, + }, + models: { + 'deepseek/deepseek-v4-flash': { + provider: 'deepseek', + model: 'deepseek-v4-flash', + maxContextSize: 1_000_000, + maxOutputSize: 384000, + }, + }, + }, + }), + }); + + ctx.agent.config.update({ + modelAlias: 'deepseek/deepseek-v4-flash', + systemPrompt: 'system', + thinkingLevel: 'off', + }); + await ctx.agent.llm.chat({ + messages: [], + tools: [], + signal: new AbortController().signal, + }); + + expect(requestMaxTokens).toBe(384000); + }); + + it('uses session id as a provider prompt cache hint without storing it on Agent', () => { const ctx = testAgent({ providerManager: new ProviderManager({ promptCacheKey: 'session-test', diff --git a/packages/agent-core/test/harness/runtime-provider.test.ts b/packages/agent-core/test/harness/runtime-provider.test.ts index 6ef2c3c78a..986357bf87 100644 --- a/packages/agent-core/test/harness/runtime-provider.test.ts +++ b/packages/agent-core/test/harness/runtime-provider.test.ts @@ -227,6 +227,34 @@ describe('resolveRuntimeProvider model metadata', () => { }); describe('resolveRuntimeProvider maxOutputSize forwarding', () => { + it('returns alias.maxOutputSize for request completion budgeting', () => { + const resolved = resolveRuntimeProvider({ + config: { + ...BASE_CONFIG, + providers: { + ...BASE_CONFIG.providers, + openai: { + type: 'openai', + apiKey: 'sk-openai', + baseUrl: 'https://openai.example/v1', + }, + }, + models: { + ...BASE_CONFIG.models!, + 'deepseek-alias': { + provider: 'openai', + model: 'deepseek-v4-flash', + maxContextSize: 1_000_000, + maxOutputSize: 384000, + }, + }, + }, + model: 'deepseek-alias', + }); + + expect(resolved.maxOutputSize).toBe(384000); + }); + it('forwards alias.maxOutputSize to the anthropic provider config as defaultMaxTokens', () => { const resolved = resolveRuntimeProvider({ config: { diff --git a/packages/agent-core/test/utils/completion-budget.test.ts b/packages/agent-core/test/utils/completion-budget.test.ts index 9b75f5f1bb..7df91f5d01 100644 --- a/packages/agent-core/test/utils/completion-budget.test.ts +++ b/packages/agent-core/test/utils/completion-budget.test.ts @@ -181,6 +181,16 @@ describe('resolveCompletionBudget', () => { expect(budget?.fallback).toBe(12345); }); + it('uses model max output size as the default hard cap when no env var is set', () => { + const budget = resolveCompletionBudget({ + maxOutputSize: 384000, + reservedContextSize: 12345, + env: {}, + }); + expect(budget?.hardCap).toBe(384000); + expect(budget?.fallback).toBeUndefined(); + }); + it('falls back to 32000 only for unknown context when nothing is configured', () => { const budget = resolveCompletionBudget({ env: {} }); expect(budget?.hardCap).toBeUndefined(); From 241bdcd784f61e4938c0565c4c8fa913c5195f9c Mon Sep 17 00:00:00 2001 From: Luyu Cheng <2239547+chengluyu@users.noreply.github.com> Date: Wed, 3 Jun 2026 14:18:32 +0800 Subject: [PATCH 3/3] fix: use max completion tokens for OpenAI models --- .../fix-openai-completion-token-field.md | 6 +++ .../kosong/src/providers/openai-legacy.ts | 46 +++++++++++++++---- packages/kosong/test/openai-legacy.test.ts | 27 +++++++++++ 3 files changed, 70 insertions(+), 9 deletions(-) create mode 100644 .changeset/fix-openai-completion-token-field.md diff --git a/.changeset/fix-openai-completion-token-field.md b/.changeset/fix-openai-completion-token-field.md new file mode 100644 index 0000000000..f19491a50b --- /dev/null +++ b/.changeset/fix-openai-completion-token-field.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/kosong": patch +"@moonshot-ai/kimi-code": patch +--- + +Use the OpenAI completion token field required by newer Chat Completions models. diff --git a/packages/kosong/src/providers/openai-legacy.ts b/packages/kosong/src/providers/openai-legacy.ts index e050e9ede5..23820c8dbc 100644 --- a/packages/kosong/src/providers/openai-legacy.ts +++ b/packages/kosong/src/providers/openai-legacy.ts @@ -80,6 +80,7 @@ export interface OpenAILegacyOptions { export interface OpenAILegacyGenerationKwargs { max_tokens?: number | undefined; + max_completion_tokens?: number | undefined; temperature?: number | undefined; top_p?: number | undefined; n?: number | undefined; @@ -103,6 +104,34 @@ interface OpenAIToolCallOut { function: { name: string; arguments: string | null }; } +function usesMaxCompletionTokens(model: string): boolean { + const normalized = model.toLowerCase(); + return /^o\d(?:$|[-.])/.test(normalized) || /^gpt-5(?:$|[-.])/.test(normalized); +} + +function completionTokenKwargs( + model: string, + maxCompletionTokens: number, +): OpenAILegacyGenerationKwargs { + return usesMaxCompletionTokens(model) + ? { max_completion_tokens: maxCompletionTokens } + : { max_tokens: maxCompletionTokens }; +} + +function normalizeGenerationKwargs( + model: string, + source: OpenAILegacyGenerationKwargs, +): OpenAILegacyGenerationKwargs { + const kwargs = { ...source }; + if (usesMaxCompletionTokens(model)) { + if (kwargs.max_completion_tokens === undefined && kwargs.max_tokens !== undefined) { + kwargs.max_completion_tokens = kwargs.max_tokens; + } + delete kwargs.max_tokens; + } + return kwargs; +} + function convertMessage( message: Message, reasoningKey: string | undefined, @@ -365,10 +394,8 @@ export class OpenAILegacyChatProvider implements ChatProvider { ? normalizedReasoningKey : undefined; this._reasoningEffort = undefined; - this._generationKwargs = {}; - if (options.maxTokens !== undefined) { - this._generationKwargs.max_tokens = options.maxTokens; - } + this._generationKwargs = + options.maxTokens !== undefined ? completionTokenKwargs(this._model, options.maxTokens) : {}; this._toolMessageConversion = options.toolMessageConversion ?? null; this._httpClient = options.httpClient; this._clientFactory = options.clientFactory; @@ -388,7 +415,7 @@ export class OpenAILegacyChatProvider implements ChatProvider { return { model: this._model, baseUrl: this._baseUrl, - ...this._generationKwargs, + ...normalizeGenerationKwargs(this._model, this._generationKwargs), }; } @@ -414,9 +441,10 @@ export class OpenAILegacyChatProvider implements ChatProvider { messages.push(convertMessage(msg, this._reasoningKey, this._toolMessageConversion)); } - const kwargs: Record = { - ...this._generationKwargs, - }; + const kwargs: Record = normalizeGenerationKwargs( + this._model, + this._generationKwargs, + ); // Determine reasoning_effort let reasoningEffort: string | undefined = this._reasoningEffort; @@ -490,7 +518,7 @@ export class OpenAILegacyChatProvider implements ChatProvider { } withMaxCompletionTokens(maxCompletionTokens: number): OpenAILegacyChatProvider { - return this.withGenerationKwargs({ max_tokens: maxCompletionTokens }); + return this.withGenerationKwargs(completionTokenKwargs(this._model, maxCompletionTokens)); } private _clone(): OpenAILegacyChatProvider { diff --git a/packages/kosong/test/openai-legacy.test.ts b/packages/kosong/test/openai-legacy.test.ts index d51f8d8a6c..cd04a57585 100644 --- a/packages/kosong/test/openai-legacy.test.ts +++ b/packages/kosong/test/openai-legacy.test.ts @@ -478,6 +478,33 @@ describe('OpenAILegacyChatProvider', () => { expect(provider).not.toBe(original); expect(body['max_tokens']).toBe(1024); }); + + it.each(['gpt-5', 'gpt-5-codex', 'o3'])( + 'withMaxCompletionTokens sets max_completion_tokens for %s', + async (model) => { + const provider = createProvider({ model }).withMaxCompletionTokens(1024); + const history: Message[] = [ + { role: 'user', content: [{ type: 'text', text: 'Hi' }], toolCalls: [] }, + ]; + const body = await captureRequestBody(provider, '', [], history); + + expect(body['max_completion_tokens']).toBe(1024); + expect(body['max_tokens']).toBeUndefined(); + }, + ); + + it('keeps max_tokens for OpenAI-compatible non-OpenAI reasoning models', async () => { + const provider = createProvider({ model: 'deepseek-reasoner' }).withMaxCompletionTokens( + 1024, + ); + const history: Message[] = [ + { role: 'user', content: [{ type: 'text', text: 'Hi' }], toolCalls: [] }, + ]; + const body = await captureRequestBody(provider, '', [], history); + + expect(body['max_tokens']).toBe(1024); + expect(body['max_completion_tokens']).toBeUndefined(); + }); }); describe('maxTokens option', () => {