fix(agent-core-v2): align goal budget handling

packages/agent-core-v2/src/agent/goal/goalService.ts

@@ -83,6 +83,7 @@ declare module '#/agent/wireRecord/wireRecord' {
 }
 
 const MAX_GOAL_OBJECTIVE_LENGTH = 4000;
+const MAX_GOAL_COMPLETION_CRITERION_LENGTH = MAX_GOAL_OBJECTIVE_LENGTH;
 
 const GOAL_CANCELLED_REMINDER = [
   'The user cancelled the current goal.',
@@ -614,12 +615,15 @@ function budgetTelemetryProperties(limits: GoalBudgetLimits): TelemetryPropertie
 }
 
 function tokenUsageTotal(usage: TokenUsage): number {
-  return usage.inputCacheRead + usage.inputCacheCreation + usage.inputOther + usage.output;
+  return usage.output;
 }
 
 function normalizeCompletionCriterion(value: string | undefined): string | undefined {
   const trimmed = value?.trim();
-  return trimmed?.length ? trimmed : undefined;
+  if (!trimmed?.length) return undefined;
+  return trimmed.length > MAX_GOAL_COMPLETION_CRITERION_LENGTH
+    ? trimmed.slice(0, MAX_GOAL_COMPLETION_CRITERION_LENGTH)
+    : trimmed;
 }
 
 function isGoalOutcomeReminder(message: ContextMessage | undefined): boolean {

packages/agent-core-v2/src/agent/goal/tools/set-goal-budget.ts

@@ -38,14 +38,19 @@ export class SetGoalBudgetTool implements BuiltinTool<SetGoalBudgetToolInput> {
 
   resolveExecution(args: SetGoalBudgetToolInput): ToolExecution {
     const normalizedArgs = normalizeBudgetInput(args);
+    const budget = budgetLimitsFromInput(normalizedArgs);
+    const overBudgetAfterSet = budget !== null && this.wouldExceedBudget(budget);
     return {
       description: `Setting goal budget: ${formatBudget(
         normalizedArgs.value,
         normalizedArgs.unit,
       )}`,
+      stopBatchAfterThis: overBudgetAfterSet,
       approvalRule: this.name,
       execute: async () => {
-        const budget = budgetLimitsFromInput(normalizedArgs);
+        if (this.goal.getGoal().goal === null) {
+          return { output: 'Goal budget not set: no current goal.' };
+        }
         if (budget === null) {
           return {
             output:
@@ -53,13 +58,35 @@ export class SetGoalBudgetTool implements BuiltinTool<SetGoalBudgetToolInput> {
               'reasonable goal budget.',
           };
         }
-        await this.goal.setBudgetLimits({ budgetLimits: budget }, 'model');
+        const snapshot = await this.goal.setBudgetLimits({ budgetLimits: budget }, 'model');
+        if (snapshot.budget.overBudget) {
+          return {
+            output:
+              `Goal budget set: ${formatBudget(normalizedArgs.value, normalizedArgs.unit)}. ` +
+              'The goal has already reached this budget and will stop now.',
+            stopTurn: true,
+          };
+        }
         return {
           output: `Goal budget set: ${formatBudget(normalizedArgs.value, normalizedArgs.unit)}.`,
         };
       },
     };
   }
+
+  private wouldExceedBudget(newLimits: GoalBudgetLimits): boolean {
+    const goal = this.goal.getGoal().goal;
+    if (goal === null) return false;
+    const current = goal.budget;
+    const turnBudget = newLimits.turnBudget ?? current.turnBudget;
+    const tokenBudget = newLimits.tokenBudget ?? current.tokenBudget;
+    const wallClockBudgetMs = newLimits.wallClockBudgetMs ?? current.wallClockBudgetMs;
+    return (
+      (turnBudget !== null && goal.turnsUsed >= turnBudget) ||
+      (tokenBudget !== null && goal.tokensUsed >= tokenBudget) ||
+      (wallClockBudgetMs !== null && goal.wallClockMs >= wallClockBudgetMs)
+    );
+  }
 }
 
 registerTool(SetGoalBudgetTool);

packages/agent-core-v2/test/goal/goal-tools.test.ts

@@ -0,0 +1,93 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { USER_PROMPT_ORIGIN } from '#/agent/contextMemory/types';
+import { IAgentGoalService } from '#/agent/goal/goal';
+import { SetGoalBudgetTool } from '#/agent/goal/tools/set-goal-budget';
+import { IAgentLoopService } from '#/agent/loop/loop';
+import { IAgentTurnService } from '#/agent/turn/turn';
+import { IEventBus } from '#/app/event/eventBus';
+
+import { agentService, createTestAgent, type TestAgentContext } from '../harness';
+import { stubLoopWithHooks, stubTurn } from '../turn/stubs';
+
+const signal = new AbortController().signal;
+
+describe('goal tools', () => {
+  let ctx: TestAgentContext;
+  let goals: IAgentGoalService;
+  let loopService: IAgentLoopService;
+  let eventBus: IEventBus;
+  let tool: SetGoalBudgetTool;
+
+  beforeEach(() => {
+    loopService = stubLoopWithHooks();
+    ctx = createTestAgent(
+      agentService(IAgentTurnService, stubTurn({ hasActiveTurn: true })),
+      agentService(IAgentLoopService, loopService),
+    );
+    goals = ctx.get(IAgentGoalService);
+    eventBus = ctx.get(IEventBus);
+    tool = new SetGoalBudgetTool(goals);
+  });
+
+  afterEach(async () => {
+    await ctx.dispose();
+  });
+
+  it('SetGoalBudget reports no current goal without failing', async () => {
+    const execution = tool.resolveExecution({ value: 20, unit: 'turns' });
+    if (execution.isError === true) throw new Error('execution should not be an error');
+
+    const result = await execution.execute({ turnId: 0, toolCallId: 'call_1', signal });
+
+    expect(result.isError).toBeFalsy();
+    expect(result.stopTurn).toBeFalsy();
+    expect(result.output).toBe('Goal budget not set: no current goal.');
+  });
+
+  it('SetGoalBudget returns stop signals when the requested limit is already exhausted', async () => {
+    await goals.createGoal({ objective: 'work' });
+    await countGoalTurn(1);
+
+    const execution = tool.resolveExecution({ value: 1, unit: 'turns' });
+    if (execution.isError === true) throw new Error('execution should not be an error');
+
+    expect(execution.stopBatchAfterThis).toBe(true);
+    const result = await execution.execute({ turnId: 0, toolCallId: 'call_1', signal });
+
+    expect(result.stopTurn).toBe(true);
+    expect(result.output).toContain('will stop now');
+    expect(goals.getGoal().goal).toMatchObject({
+      status: 'blocked',
+      budget: { overBudget: true },
+    });
+  });
+
+  it('SetGoalBudget leaves the turn running when the requested limit has room', async () => {
+    await goals.createGoal({ objective: 'work' });
+    await countGoalTurn(2);
+
+    const execution = tool.resolveExecution({ value: 5, unit: 'turns' });
+    if (execution.isError === true) throw new Error('execution should not be an error');
+
+    expect(execution.stopBatchAfterThis).toBeFalsy();
+    const result = await execution.execute({ turnId: 0, toolCallId: 'call_1', signal });
+
+    expect(result.stopTurn).toBeFalsy();
+    expect(result.output).toBe('Goal budget set: 5 turns.');
+    expect(goals.getGoal().goal).toMatchObject({
+      status: 'active',
+      budget: { turnBudget: 5, overBudget: false },
+    });
+  });
+
+  async function countGoalTurn(turnId: number): Promise<void> {
+    const abortController = new AbortController();
+    eventBus.publish({ type: 'turn.started', turnId, origin: USER_PROMPT_ORIGIN });
+    await loopService.hooks.beforeStep.run({
+      turnId,
+      step: 1,
+      signal: abortController.signal,
+    });
+  }
+});

packages/agent-core-v2/test/goal/goal.test.ts

@@ -160,6 +160,15 @@ describe('AgentGoalService', () => {
       expect(goals.getGoal().goal?.completionCriterion).toBe('tests pass');
     });
 
+    it('truncates an over-long completion criterion instead of failing', async () => {
+      const snapshot = await goals.createGoal({
+        objective: 'Ship feature X',
+        completionCriterion: 'c'.repeat(4001),
+      });
+
+      expect(snapshot.completionCriterion).toBe('c'.repeat(4000));
+    });
+
     it('sets no default work caps when none is provided', async () => {
       const snapshot = await goals.createGoal({ objective: 'Do work' });
 
@@ -609,12 +618,16 @@ describe('AgentGoalService core workflow hooks', () => {
         output: 0,
       }),
     ).toBe(false);
+    expect(goals.getGoal().goal).toMatchObject({
+      status: 'active',
+      tokensUsed: 0,
+    });
     expect(
       await runStepUsageHooks(loopService, goals, turn, {
         inputCacheRead: 0,
         inputCacheCreation: 0,
         inputOther: 0,
-        output: 3,
+        output: 7,
       }),
     ).toBe(true);