fix: simplify goal budget schema and fix output caps

.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.

.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.

.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.

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);

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({

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,
     };
   }
 

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<typeof SetGoalBudgetToolInputSchema>;
 
@@ -46,21 +40,27 @@ export class SetGoalBudgetTool implements BuiltinTool<SetGoalBudgetToolInput> {
     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<SetGoalBudgetToolInput> {
   }
 }
 
+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':

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 };
   }

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<string, unknown> }
+        ).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',

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: {

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' });

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();