From f2ebfeaecec0009b91cbfae672a8fc810fd4774b Mon Sep 17 00:00:00 2001 From: qqqys Date: Fri, 12 Jun 2026 13:10:06 +0800 Subject: [PATCH] fix(goal): persist iteration count across resume so MAX_GOAL_ITERATIONS bounds the whole session (#5000) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(goal): persist iteration count across resume so MAX_GOAL_ITERATIONS bounds the whole session On resume, restoreGoalFromHistory re-arms an unfinished /goal via registerGoalHook, which always primed the store with iterations: 0. Since findGoalToRestore only returned the condition, the running count recorded in the transcript was dropped, so the MAX_GOAL_ITERATIONS safety cap was re-granted in full on every resume — an unreachable goal could auto-loop another full budget after each /resume. The count is already persisted (checking goal_status items carry iterations), so the fix just reads it back: - findGoalToRestore returns { condition, iterations } from the latest non-terminal goal_status item (set items restore at 0). - registerGoalHook accepts an optional initialIterations (default 0, clamped at 0). - restoreGoalFromHistory threads the restored count through. Resume re-arm stays passive — continuation timing is unchanged. Closes #4999 Co-authored-by: Qwen-Coder * fix(goal): persist cumulative checking iterations * fix(goal): record checking status during continuations --------- Co-authored-by: Qwen-Coder --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 21 ++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 23 ++++----- packages/cli/src/ui/utils/restoreGoal.test.ts | 33 +++++++++++-- packages/cli/src/ui/utils/restoreGoal.ts | 33 ++++++++----- packages/core/src/goals/goalHook.test.ts | 47 +++++++++++++++++++ packages/core/src/goals/goalHook.ts | 9 +++- 6 files changed, 138 insertions(+), 28 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 051baf22a3..43a762bbd9 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -5611,9 +5611,13 @@ describe('useGeminiStream', () => { }); it('renders active goal StopHookLoop as a goal_status checking card', async () => { + const recordSlashCommand = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordSlashCommand, + }); mockGetActiveGoal.mockReturnValue({ condition: 'finish the refactor', - iterations: 2, + iterations: 7, setAt: 100, tokensAtStart: 0, hookId: 'goal-hook', @@ -5644,12 +5648,25 @@ describe('useGeminiStream', () => { type: 'goal_status', kind: 'checking', condition: 'finish the refactor', - iterations: 2, + iterations: 7, lastReason: 'not enough evidence yet', }), expect.any(Number), ); }); + expect(recordSlashCommand).toHaveBeenCalledWith({ + phase: 'result', + rawCommand: '/goal', + outputHistoryItems: [ + expect.objectContaining({ + type: 'goal_status', + kind: 'checking', + condition: 'finish the refactor', + iterations: 7, + lastReason: 'not enough evidence yet', + }), + ], + }); expect(mockAddItem).not.toHaveBeenCalledWith( expect.objectContaining({ type: 'stop_hook_loop' }), expect.any(Number), diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 2ba266382f..5a65db7f4e 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -60,6 +60,7 @@ import { import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { HistoryItem, + HistoryItemGoalStatus, HistoryItemWithoutId, HistoryItemToolGroup, SlashCommandProcessorResult, @@ -90,6 +91,7 @@ import { useSessionStats } from '../contexts/SessionContext.js'; import type { LoadedSettings } from '../../config/settings.js'; import { t } from '../../i18n/index.js'; import { useDualOutput } from '../../dualOutput/DualOutputContext.js'; +import { recordGoalStatusItem } from '../utils/restoreGoal.js'; import process from 'node:process'; const debugLogger = createDebugLogger('GEMINI_STREAM'); @@ -1365,17 +1367,16 @@ export const useGeminiStream = ( // continuation, not a hook failure. const activeGoal = getActiveGoal(config.getSessionId()); if (activeGoal && activeGoal.condition) { - addItem( - { - type: MessageType.GOAL_STATUS, - kind: 'checking', - condition: activeGoal.condition, - iterations: value.iterationCount, - lastReason: - activeGoal.lastReason ?? value.reasons[value.reasons.length - 1], - } as HistoryItemWithoutId, - userMessageTimestamp, - ); + const item: HistoryItemGoalStatus = { + type: MessageType.GOAL_STATUS, + kind: 'checking', + condition: activeGoal.condition, + iterations: activeGoal.iterations, + lastReason: + activeGoal.lastReason ?? value.reasons[value.reasons.length - 1], + }; + addItem(item, userMessageTimestamp); + recordGoalStatusItem(config, item); return; } addItem( diff --git a/packages/cli/src/ui/utils/restoreGoal.test.ts b/packages/cli/src/ui/utils/restoreGoal.test.ts index baf18f9b18..2bea549508 100644 --- a/packages/cli/src/ui/utils/restoreGoal.test.ts +++ b/packages/cli/src/ui/utils/restoreGoal.test.ts @@ -61,14 +61,14 @@ describe('findGoalToRestore', () => { ).toBeNull(); }); - it('returns the condition when last goal_status is set', () => { + it('returns the condition (iterations 0) when last goal_status is set', () => { expect( findGoalToRestore([ goalItem({ kind: 'achieved', condition: 'old goal' }), goalItem({ kind: 'set', condition: 'fresh goal' }), userItem(), ]), - ).toBe('fresh goal'); + ).toEqual({ condition: 'fresh goal', iterations: 0 }); }); it('returns the condition when last goal_status is checking', () => { @@ -78,7 +78,17 @@ describe('findGoalToRestore', () => { userItem(), goalItem({ kind: 'checking', condition: 'fresh goal' }), ]), - ).toBe('fresh goal'); + ).toEqual({ condition: 'fresh goal', iterations: 0 }); + }); + + it('carries the running iteration count from a checking item', () => { + expect( + findGoalToRestore([ + goalItem({ kind: 'set', condition: 'fresh goal' }), + userItem(), + goalItem({ kind: 'checking', condition: 'fresh goal', iterations: 7 }), + ]), + ).toEqual({ condition: 'fresh goal', iterations: 7 }); }); it('returns null when last goal_status is cleared', () => { @@ -123,6 +133,23 @@ describe('restoreGoalFromHistory', () => { expect(getActiveGoal('sess-1')).toMatchObject({ condition: 'write hello' }); }); + it('resumes the iteration count so the MAX cap is not reset on resume', () => { + const cfg = makeConfig(); + const result = restoreGoalFromHistory( + [ + goalItem({ kind: 'set', condition: 'write hello' }), + userItem(), + goalItem({ kind: 'checking', condition: 'write hello', iterations: 7 }), + ], + cfg, + ); + expect(result).toEqual({ restored: true, condition: 'write hello' }); + expect(getActiveGoal('sess-1')).toMatchObject({ + condition: 'write hello', + iterations: 7, + }); + }); + it('does nothing when no goal_status item exists', () => { const cfg = makeConfig(); const result = restoreGoalFromHistory([userItem()], cfg); diff --git a/packages/cli/src/ui/utils/restoreGoal.ts b/packages/cli/src/ui/utils/restoreGoal.ts index 09b5b5ad82..ab4afe3b86 100644 --- a/packages/cli/src/ui/utils/restoreGoal.ts +++ b/packages/cli/src/ui/utils/restoreGoal.ts @@ -22,18 +22,26 @@ import { /** * Finds the most recent `goal_status` history item. Returns the active - * condition when the latest goal event is non-terminal (`set` or `checking`), - * or `null` if the last goal_status was terminal/cancelled - * (achieved / failed / cleared / aborted) or none exists. + * condition plus the iteration count to resume from when the latest goal event + * is non-terminal (`set` or `checking`), or `null` if the last goal_status was + * terminal/cancelled (achieved / failed / cleared / aborted) or none exists. + * + * The iteration count is carried so the MAX_GOAL_ITERATIONS safety cap survives + * resume instead of resetting to zero. `checking` items persist the running + * count (see useGeminiStream's continuation handler); `set` items predate any + * iteration, so they restore at 0. */ -export function findGoalToRestore(history: HistoryItem[]): string | null { +export function findGoalToRestore( + history: HistoryItem[], +): { condition: string; iterations: number } | null { for (let i = history.length - 1; i >= 0; i--) { const item = history[i]; if (item?.type !== MessageType.GOAL_STATUS) continue; const goal = item as HistoryItemGoalStatus; - return goal.kind === 'set' || goal.kind === 'checking' - ? goal.condition - : null; + if (goal.kind === 'set' || goal.kind === 'checking') { + return { condition: goal.condition, iterations: goal.iterations ?? 0 }; + } + return null; } return null; } @@ -133,9 +141,9 @@ export function restoreGoalFromHistory( const lastTerminal = findLastTerminalGoal(history); setLastGoalTerminal(sessionId, lastTerminal ?? undefined); - const condition = findGoalToRestore(history); + const restorable = findGoalToRestore(history); - if (condition === null) { + if (restorable === null) { unregisterGoalHook(config, sessionId); return { restored: false }; } @@ -152,11 +160,14 @@ export function restoreGoalFromHistory( registerGoalHook({ config, sessionId, - condition, + condition: restorable.condition, tokensAtStart: 0, + // Resume the iteration count so MAX_GOAL_ITERATIONS is a cross-resume cap, + // not a per-resume one. + initialIterations: restorable.iterations, }); if (addItem) { installGoalTerminalObserver({ sessionId, config, addItem }); } - return { restored: true, condition }; + return { restored: true, condition: restorable.condition }; } diff --git a/packages/core/src/goals/goalHook.test.ts b/packages/core/src/goals/goalHook.test.ts index 43e908f409..d5c2319824 100644 --- a/packages/core/src/goals/goalHook.test.ts +++ b/packages/core/src/goals/goalHook.test.ts @@ -526,6 +526,53 @@ describe('registerGoalHook / unregisterGoalHook', () => { expect(getActiveGoal('sess-1')).toMatchObject({ condition: 'tests pass' }); }); + it('primes the store from initialIterations on resume', () => { + const goal = registerGoalHook({ + config, + sessionId: 'sess-1', + condition: 'tests pass', + tokensAtStart: 0, + initialIterations: 7, + }); + expect(goal.iterations).toBe(7); + expect(getActiveGoal('sess-1')?.iterations).toBe(7); + }); + + it('clamps a negative initialIterations to 0', () => { + const goal = registerGoalHook({ + config, + sessionId: 'sess-1', + condition: 'tests pass', + tokensAtStart: 0, + initialIterations: -3, + }); + expect(goal.iterations).toBe(0); + }); + + it('honors a resumed near-cap count so MAX survives resume (no fresh budget)', async () => { + // Simulate resume re-arming a goal that was already at the cap last session. + registerGoalHook({ + config, + sessionId: 'sess-1', + condition: 'do x', + tokensAtStart: 0, + initialIterations: MAX_GOAL_ITERATIONS, + }); + judgeMock.mockResolvedValue({ ok: false, reason: 'still not done' }); + const cb = createGoalStopHookCallback({ + config, + sessionId: 'sess-1', + condition: 'do x', + }); + const out = await cb(stopInput(), undefined); + // Without resumed iterations this would just block and continue; because the + // count survived resume, the very next not-met verdict hits the cap. + expect( + typeof out === 'object' && out !== null ? out.systemMessage : undefined, + ).toMatch(/max iterations/i); + expect(getActiveGoal('sess-1')).toBeUndefined(); + }); + it('replaces an existing goal cleanly', () => { registerGoalHook({ config, diff --git a/packages/core/src/goals/goalHook.ts b/packages/core/src/goals/goalHook.ts index ea6fd8657b..c53ede9139 100644 --- a/packages/core/src/goals/goalHook.ts +++ b/packages/core/src/goals/goalHook.ts @@ -276,6 +276,13 @@ export function registerGoalHook(args: { sessionId: string; condition: string; tokensAtStart: number; + /** + * Iteration count to resume from. Used on session resume so the + * MAX_GOAL_ITERATIONS safety cap survives a reload instead of resetting to + * zero (which would let an unreachable goal auto-loop another full budget + * every resume). Defaults to 0 for a freshly set goal. + */ + initialIterations?: number; }): ActiveGoal { const { config, sessionId, condition, tokensAtStart } = args; const system = config.getHookSystem(); @@ -310,7 +317,7 @@ export function registerGoalHook(args: { const goal: ActiveGoal = { condition, - iterations: 0, + iterations: Math.max(0, args.initialIterations ?? 0), setAt: Date.now(), tokensAtStart, hookId,