Skip to content
Merged
122 changes: 120 additions & 2 deletions packages/core/src/core/client-goal.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { beforeEach, describe, expect, it, vi } from 'vitest';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { Config } from '../config/config.js';
import type { GeminiChat } from './geminiChat.js';
import {
Expand All @@ -19,6 +19,10 @@ import type {
GoalStateRecordPayloadV2,
GoalTurnPermit,
} from '../goals/goal-protocol.js';
import {
__resetActiveGoalStoreForTests,
setActiveGoal,
} from '../goals/activeGoalStore.js';
import type { ChatRecord } from '../services/chatRecordingService.js';
import { ApprovalMode } from '../config/config.js';

Expand Down Expand Up @@ -1032,7 +1036,6 @@ describe('GeminiClient Goal admission', () => {
goalPermit: current,
goalTurnKey: `goal-runtime:${current.turnId}`,
},
0,
),
);
}
Expand All @@ -1041,4 +1044,119 @@ describe('GeminiClient Goal admission', () => {
expect(turnMocks.run).toHaveBeenCalledTimes(turns);
expect(client['sessionTurnCount']).toBe(0);
});

it('holds a runtime Goal turn to the caller recursion budget like any other send', async () => {
// Each Goal continuation is a fresh top-level send that starts from
// MAX_TURNS on its own, so a Goal has no reason to outlive one turn's
// recursion allowance. A caller that hands over an exhausted budget gets
// the same refusal every other message type gets -- and, because the
// turn was admitted, the interrupted-exit path pauses the Goal instead
// of leaving it running with a permit nobody will finish.
const { client, runtime } = setupGoalClient();
turnMocks.run.mockImplementation(async function* () {});

const events = await collect(
client.sendMessageStream(
[{ text: 'continue' }],
new AbortController().signal,
'goal-exhausted',
{
type: SendMessageType.Goal,
goalPermit: permit,
goalTurnKey: `goal-runtime:${permit.turnId}`,
},
0,
),
);

expect(turnMocks.run).not.toHaveBeenCalled();
expect(runtime.dispatch).toHaveBeenCalledWith(
expect.objectContaining({ action: 'pause' }),
);
expect(runtime.finishTurn).toHaveBeenCalledOnce();
expect(events).not.toContainEqual(
expect.objectContaining({ type: GeminiEventType.MaxSessionTurns }),
);
Comment thread
qqqys marked this conversation as resolved.
});

afterEach(() => __resetActiveGoalStoreForTests());

it('admits a runtime Goal turn to steer input at a hit session cap', async () => {
// Second half of the session-cap exclusion: runtime Goal turns skip the
// count increment (pinned by the 75-turn test above) and must also be
// admitted to steer input once the user's own turns hit the cap.
const { client } = setupGoalClient();
client['sessionTurnCount'] = 1;
const getSteerInput = vi.fn().mockResolvedValue(undefined);

await drain(
client.sendMessageStream(
[{ text: 'continue' }],
new AbortController().signal,
'goal-steer-cap',
{
type: SendMessageType.Goal,
goalPermit: permit,
goalTurnKey: `goal-runtime:${permit.turnId}`,
getSteerInput,
},
),
);

expect(getSteerInput).toHaveBeenCalled();
});

it('does not decrement the caller recursion budget inside a legacy hook Goal chain', async () => {
// A legacy /goal chain recurses inside one sendMessageStream call. With
// the caller budget of 2 preserved at every hop, two blocked stops still
// run three model turns; a decremented budget would refuse the third.
// Unsupported Goal runtime: the legacy hook chain owns the send, so the
// branch under test is the one that keys the budget on the active goal.
const { client, config } = setupGoalClient();
vi.mocked(config.getGoalRuntimeReady).mockRejectedValue(
new GoalPersistenceUnavailableError('legacy-only session'),
);
vi.mocked(config.getSkipNextSpeakerCheck).mockReturnValue(true);
vi.mocked(config.getDisableAllHooks).mockReturnValue(false);
vi.mocked(config.getMaxSessionTurns).mockReturnValue(0);
vi.mocked(config.hasHooksForEvent).mockImplementation(
(event) => event === 'Stop',
);
let stopRequestCount = 0;
const messageBus = {
request: vi.fn(async () => {
stopRequestCount += 1;
if (stopRequestCount <= 2) {
return {
output: { decision: 'block', reason: 'keep going' },
stopHookCount: 1,
};
}
return { output: undefined, stopHookCount: 1 };
}),
};
vi.mocked(config.getMessageBus).mockReturnValue(
messageBus as unknown as ReturnType<Config['getMessageBus']>,
);
setActiveGoal('goal-test-session', {
condition: 'ship',
iterations: 0,
setAt: 1,
tokensAtStart: 0,
hookId: 'goal-hook:test',
});

await drain(
client.sendMessageStream(
[{ text: 'start the chain' }],
new AbortController().signal,
'legacy-goal-chain',
undefined,
2,
),
);

expect(messageBus.request).toHaveBeenCalledTimes(3);
expect(turnMocks.run).toHaveBeenCalledTimes(3);
});
});
31 changes: 25 additions & 6 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3257,6 +3257,15 @@ export class GeminiClient {
}
}

// A runtime-scheduled Goal turn is not a session turn. `maxSessionTurns`
// counts every model call the user's own prompts drive, tool
// continuations included; a Goal that reads a few files per
// continuation would spend a user-set cap of N in N/4 continuations
// and die mid-run with no resume path in headless. Autonomous Goal
// spend is bounded by the Goal's own token budget instead (armed at
// creation, re-armed only by an explicit resume or edit), and the
// headless host excludes runtime Goal turns from the same cap for the
// same reason, so counting them here would split the two ceilings.
if (messageType !== SendMessageType.Retry && !isGoalRuntimeTurn) {
// Attribution snapshots are recorded on every non-retry turn. File
// history snapshots are created only at UserQuery boundaries; later
Expand Down Expand Up @@ -3305,11 +3314,12 @@ export class GeminiClient {
}
}

// Ensure turns never exceeds MAX_TURNS to prevent infinite loops
const boundedTurns =
messageType === SendMessageType.Goal
? MAX_TURNS
: Math.min(turns, MAX_TURNS);
// Ensure turns never exceeds MAX_TURNS to prevent infinite loops. A
// runtime Goal turn honours the caller's budget like every other
// message type: each continuation is a fresh top-level send that
// starts from MAX_TURNS on its own, so nothing about a Goal needs to
// outlive one turn's recursion allowance.
const boundedTurns = Math.min(turns, MAX_TURNS);
if (!boundedTurns) {
this.cancelPendingMemoryPrefetch('no_safe_delivery_point');
endCurrentInteraction('error', 'max turns exhausted', 'max_turns');
Expand All @@ -3327,6 +3337,9 @@ export class GeminiClient {
) {
return undefined;
}
// Same ceiling as the session-turn check above, same exclusion: a
// runtime Goal turn does not count toward `maxSessionTurns`, so it
// must not be refused steer input on that count either.
const maxSessionTurns = this.config.getMaxSessionTurns();
if (
!isGoalRuntimeTurn &&
Expand Down Expand Up @@ -4155,7 +4168,13 @@ export class GeminiClient {
// these semantics (fresh DaemonToolLoopState per continuation).
// Runaway protection is preserved: the cap still bounds each
// iteration, and the chain itself is bounded by
// stopHookBlockingCap / MAX_GOAL_ITERATIONS.
// stopHookBlockingCap / MAX_GOAL_ITERATIONS. Those are the only
// bounds on this path: the legacy hook Goal recurses inside one
Comment thread
qqqys marked this conversation as resolved.
Outdated
// sendMessageStream call, so the runtime's token budget (which
// meters continuations the Goal runtime schedules) never sees it,
// and the recursion budget is not decremented below because a
// 50-iteration chain with steer and next-speaker continues would
// otherwise exhaust MAX_TURNS before its own iteration cap.
this.loopDetector.reset(prompt_id);

const activeGoal = getActiveGoal(this.config.getSessionId());
Expand Down
Loading