Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions packages/cli/src/ui/hooks/useGeminiStream.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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),
Expand Down
23 changes: 12 additions & 11 deletions packages/cli/src/ui/hooks/useGeminiStream.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ import {
import { type Part, type PartListUnion, FinishReason } from '@google/genai';
import type {
HistoryItem,
HistoryItemGoalStatus,
HistoryItemWithoutId,
HistoryItemToolGroup,
SlashCommandProcessorResult,
Expand Down Expand Up @@ -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');
Expand Down Expand Up @@ -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(
Expand Down
33 changes: 30 additions & 3 deletions packages/cli/src/ui/utils/restoreGoal.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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', () => {
Expand Down Expand Up @@ -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);
Expand Down
33 changes: 22 additions & 11 deletions packages/cli/src/ui/utils/restoreGoal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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 };
}
Expand All @@ -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 };
}
47 changes: 47 additions & 0 deletions packages/core/src/goals/goalHook.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
9 changes: 8 additions & 1 deletion packages/core/src/goals/goalHook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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,
Expand Down
Loading