diff --git a/.changeset/fix-goal-turn-budget.md b/.changeset/fix-goal-turn-budget.md new file mode 100644 index 0000000000..b9189a8395 --- /dev/null +++ b/.changeset/fix-goal-turn-budget.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Allow goals to use every configured turn before the turn budget stops further work. diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index a6274ce5d7..8828c2cad5 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -204,6 +204,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { private readonly budgetGraceTurns = new Set(); private readonly pendingContinuationGoals = new Map(); private readonly goalTurnTargets = new Map(); + private readonly exhaustedTurnBudgetGoals = new Map(); private pendingContinuation?: PendingContinuation; constructor( @@ -504,12 +505,13 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { const next = this.requireState(); this.emitGoalUpdated(this.toSnapshot(next)); this.telemetry.track2('goal_continued', { turns_used: next.turnsUsed }); - return this.blockIfBudgetReached(next) ?? this.toSnapshot(next); + return this.toSnapshot(next); } private handleTurnLaunched(turnId: number, origin: TurnStartedEvent['origin']): void { this.liveTurnId = turnId; this.goalTurnTargets.delete(turnId); + this.exhaustedTurnBudgetGoals.delete(turnId); if (!this.goalDrivenTurns.has(turnId)) { const state = this.goalState; const continuationGoalId = isGoalContinuationOrigin(origin) @@ -533,6 +535,11 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { if (state === null || state.status !== 'active') return; const goalId = this.goalDrivenTurns.get(turnId); if (actor === 'model') this.goalTurnTargets.set(turnId, state.goalId); + if (this.toSnapshot(state).budget.turnBudgetReached) { + this.exhaustedTurnBudgetGoals.set(turnId, state.goalId); + } else { + this.exhaustedTurnBudgetGoals.delete(turnId); + } if (goalId !== undefined) return; this.goalDrivenTurns.set(turnId, state.goalId); this.countedGoalTurns.add(turnId); @@ -563,11 +570,20 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { private stopAfterBudgetReached(ctx: AfterStepContext): boolean { const goalId = this.goalTurnTarget(ctx.turnId); const state = this.goalState; + const budget = state === null ? null : this.toSnapshot(state).budget; + const turnBudgetBlocksCurrentTurn = + budget?.turnBudgetReached === true && + (this.exhaustedTurnBudgetGoals.get(ctx.turnId) === goalId || + (state?.status === 'blocked' && + state.terminalReason?.startsWith(GOAL_BUDGET_BLOCK_PREFIX) === true)); if ( goalId === undefined || state === null || state.goalId !== goalId || - !this.toSnapshot(state).budget.overBudget + budget === null || + (!budget.tokenBudgetReached && + !budget.wallClockBudgetReached && + !turnBudgetBlocksCurrentTurn) ) { return false; } @@ -643,6 +659,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.budgetGraceTurns.delete(turnId); this.pendingContinuationGoals.delete(turnId); this.goalTurnTargets.delete(turnId); + this.exhaustedTurnBudgetGoals.delete(turnId); return { goalId, lifecycleGoalId, starterTurn }; } diff --git a/packages/agent-core-v2/test/agent/goal/goal.test.ts b/packages/agent-core-v2/test/agent/goal/goal.test.ts index 572f590751..17bcdb260b 100644 --- a/packages/agent-core-v2/test/agent/goal/goal.test.ts +++ b/packages/agent-core-v2/test/agent/goal/goal.test.ts @@ -1042,13 +1042,36 @@ describe('AgentGoalService core workflow hooks', () => { expect(JSON.stringify(context.get().at(-1)?.content)).toContain('Continue working toward'); }); - it('blocks at the turn budget instead of launching a continuation', async () => { + it('blocks the next continuation only after the final allowed turn ends', async () => { await goals.createGoal({ objective: 'finish the task' }); await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); const turn = makeTurn(11); eventBus.publish({ type: 'turn.started', turnId: turn.id, origin: USER_PROMPT_ORIGIN }); - await runGoalStep(loopService, turn); + await loopService.hooks.onWillBeginStep.run({ + turnId: turn.id, + step: 1, + signal: turn.signal, + }); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'active', + turnsUsed: 1, + }); + + const afterStep: AfterStepContext = { + turnId: turn.id, + step: 1, + signal: turn.signal, + usage: zeroUsage, + finishReason: 'completed', + stopTurn: false, + }; + await loopService.hooks.onDidFinishStep.run(afterStep); + + expect(afterStep.stopTurn).toBe(false); + expect(goals.getGoal().goal?.status).toBe('active'); + endTurn(eventBus, turn); expect(goals.getGoal().goal).toMatchObject({ @@ -1059,6 +1082,64 @@ describe('AgentGoalService core workflow hooks', () => { expect(loopService.launches).toEqual([]); }); + it('completes on the final allowed continuation without applying the turn budget block', async () => { + await goals.createGoal({ objective: 'finish the task' }); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); + + const firstTurn = makeTurn(14); + eventBus.publish({ type: 'turn.started', turnId: firstTurn.id, origin: USER_PROMPT_ORIGIN }); + await runGoalStep(loopService, firstTurn); + endTurn(eventBus, firstTurn); + + await vi.waitFor(() => expect(loopService.launches).toHaveLength(1)); + const continuation = makeTurn(loopService.launches[0]!); + eventBus.publish({ + type: 'turn.started', + turnId: continuation.id, + origin: { kind: 'system_trigger', name: 'goal_continuation' }, + }); + await loopService.hooks.onWillBeginStep.run({ + turnId: continuation.id, + step: 1, + signal: continuation.signal, + }); + + const completed = await goals.markComplete({ reason: 'done' }, 'model'); + endTurn(eventBus, continuation); + + expect(completed).toMatchObject({ status: 'complete', turnsUsed: 2 }); + expect(goals.getGoal().goal).toBeNull(); + expect(loopService.launches).toHaveLength(1); + }); + + it('requests a blocked outcome step when the final allowed turn blocks the goal', async () => { + await goals.createGoal({ objective: 'finish the task' }); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); + + const turn = makeTurn(15); + eventBus.publish({ type: 'turn.started', turnId: turn.id, origin: USER_PROMPT_ORIGIN }); + await loopService.hooks.onWillBeginStep.run({ + turnId: turn.id, + step: 1, + signal: turn.signal, + }); + await goals.markBlocked({}, 'model'); + await runTerminalUpdateGoalResult(toolExecutor, turn, 'blocked', 'outcome prompt'); + + const afterStep: AfterStepContext = { + turnId: turn.id, + step: 1, + signal: turn.signal, + usage: zeroUsage, + finishReason: 'completed', + stopTurn: false, + }; + await loopService.hooks.onDidFinishStep.run(afterStep); + + expect(loopService.hasPendingRequests()).toBe(true); + expect(goals.getGoal().goal).toMatchObject({ status: 'blocked', turnsUsed: 1 }); + }); + it('accounts recorded turn usage for active goal turns', async () => { await goals.createGoal({ objective: 'finish the task' }); await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 7 } }, 'model'); @@ -1133,6 +1214,7 @@ describe('AgentGoalService core workflow hooks', () => { await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); endTurn(eventBus, turn); + await vi.waitFor(() => expect(goals.getGoal().goal?.status).toBe('blocked')); expect(goals.getGoal().goal).toMatchObject({ status: 'blocked', turnsUsed: 1, @@ -1557,7 +1639,47 @@ describe('AgentGoalService mid-turn budget stop', () => { } }); - it('blocks an over-budget goal at turn launch and runs the prompt as a normal turn', async () => { + it('rejects goal tool calls when an exhausted turn budget is resumed during a prompt', async () => { + const ctx = createTestAgent(); + try { + ctx.configure({ tools: ['UpdateGoal', 'SetGoalBudget'] }); + const goals = ctx.get(IAgentGoalService) as GoalServiceTestManager; + await goals.createGoal({ objective: 'work' }); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); + await goals.incrementTurn(); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); + + ctx.mockNextResponse({ + type: 'function', + id: 'resume', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'active' }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'raise-budget', + name: 'SetGoalBudget', + arguments: JSON.stringify({ value: 5, unit: 'turns' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'This step should never run.' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'resume the goal' }] }); + await ctx.untilTurnEnd(); + + expect(ctx.llmCalls).toHaveLength(2); + const history = ctx.get(IAgentContextMemoryService).get(); + expect(JSON.stringify(history)).toContain( + 'Goal budget exhausted; tool calls are rejected. Write your final message.', + ); + expect(JSON.stringify(history)).not.toContain('This step should never run.'); + await vi.waitFor(() => expect(goals.getGoal().goal?.status).toBe('blocked')); + expect(goals.getGoal().goal?.budget.turnBudget).toBe(1); + } finally { + await ctx.dispose(); + } + }); + + it("runs the prompt as a normal turn when the goal's turn budget was reached at launch", async () => { const telemetry: TelemetryRecord[] = []; const ctx = createTestAgent(telemetryServices(recordingTelemetry(telemetry))); try { @@ -1566,10 +1688,7 @@ describe('AgentGoalService mid-turn budget stop', () => { await goals.createGoal({ objective: 'work' }); await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); await goals.incrementTurn(); - expect(goals.getGoal().goal?.status).toBe('blocked'); - - const resumed = await goals.resumeGoal(); - expect(resumed.status).toBe('active'); + expect(goals.getGoal().goal?.status).toBe('active'); const telemetryAfterResume = telemetry.length; ctx.mockNextResponse({ type: 'text', text: 'Answering the prompt normally.' }); diff --git a/packages/agent-core-v2/test/agent/goal/injection/goalInjection.test.ts b/packages/agent-core-v2/test/agent/goal/injection/goalInjection.test.ts index 32cdb671c3..436362c928 100644 --- a/packages/agent-core-v2/test/agent/goal/injection/goalInjection.test.ts +++ b/packages/agent-core-v2/test/agent/goal/injection/goalInjection.test.ts @@ -178,6 +178,7 @@ describe('GoalInjection content', () => { await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); await goals.incrementTurn(); await goals.incrementTurn(); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); }))!; expect(text).toContain('currently blocked'); expect(text).toContain('Blocked after goal budget reached: turn budget 2');