diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index 47ebfb7fc..5697f92f4 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/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 { diff --git a/packages/agent-core-v2/src/agent/goal/tools/set-goal-budget.ts b/packages/agent-core-v2/src/agent/goal/tools/set-goal-budget.ts index fbd16cfc8..a22c62c4e 100644 --- a/packages/agent-core-v2/src/agent/goal/tools/set-goal-budget.ts +++ b/packages/agent-core-v2/src/agent/goal/tools/set-goal-budget.ts @@ -38,14 +38,19 @@ export class SetGoalBudgetTool implements BuiltinTool { 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 { '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); diff --git a/packages/agent-core-v2/test/goal/goal-tools.test.ts b/packages/agent-core-v2/test/goal/goal-tools.test.ts new file mode 100644 index 000000000..fb704afe6 --- /dev/null +++ b/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 { + 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, + }); + } +}); diff --git a/packages/agent-core-v2/test/goal/goal.test.ts b/packages/agent-core-v2/test/goal/goal.test.ts index 269c95dae..79b7630ad 100644 --- a/packages/agent-core-v2/test/goal/goal.test.ts +++ b/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);