fix(agent-core-v2): align goal turn budgets

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

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<number>();
   private readonly pendingContinuationGoals = new Map<number, string>();
   private readonly goalTurnTargets = new Map<number, string>();
+  private readonly exhaustedTurnBudgetGoals = new Map<number, string>();
   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 };
   }
 

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

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