diff --git a/integration-tests/hooks-agent-flow.test.ts b/integration-tests/hooks-agent-flow.test.ts index 949770308b..b602737a39 100644 --- a/integration-tests/hooks-agent-flow.test.ts +++ b/integration-tests/hooks-agent-flow.test.ts @@ -182,14 +182,22 @@ describe('Hooks Agent Flow', () => { ); const afterAgentScript = ` - console.log(JSON.stringify({ - decision: 'block', - reason: 'Security policy triggered', - hookSpecificOutput: { - hookEventName: 'AfterAgent', - clearContext: true - } - })); + const fs = require('fs'); + const input = JSON.parse(fs.readFileSync(0, 'utf-8')); + if (input.stop_hook_active) { + // Retry turn: allow execution to proceed (breaks the loop) + console.log(JSON.stringify({ decision: 'allow' })); + } else { + // First call: block and clear context to trigger the retry + console.log(JSON.stringify({ + decision: 'block', + reason: 'Security policy triggered', + hookSpecificOutput: { + hookEventName: 'AfterAgent', + clearContext: true + } + })); + } `; const afterAgentScriptPath = rig.createScript( 'after_agent_clear.cjs', @@ -198,8 +206,10 @@ describe('Hooks Agent Flow', () => { rig.setup('should process clearContext in AfterAgent hook output', { settings: { - hooks: { + hooksConfig: { enabled: true, + }, + hooks: { BeforeModel: [ { hooks: [ diff --git a/integration-tests/hooks-system.after-agent.responses b/integration-tests/hooks-system.after-agent.responses index 1475070c3d..526c59362d 100644 --- a/integration-tests/hooks-system.after-agent.responses +++ b/integration-tests/hooks-system.after-agent.responses @@ -1,2 +1,3 @@ {"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"Hi there!"}],"role":"model"},"finishReason":"STOP","index":0}]}]} {"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"Clarification: I am a bot."}],"role":"model"},"finishReason":"STOP","index":0}]}]} +{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"Security policy triggered"}],"role":"model"},"finishReason":"STOP","index":0}]}]} diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 4dee75b93b..58e9645b28 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -3317,6 +3317,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), 'Hook Response', + false, ); // Map should be empty @@ -3358,6 +3359,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), 'Response 1\nResponse 2', + false, ); expect(client['hookStateMap'].size).toBe(0); @@ -3388,6 +3390,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), // Should be 'Do something' expect.stringContaining('Ok'), + false, ); }); @@ -3558,6 +3561,21 @@ ${JSON.stringify( expect.anything(), undefined, ); + + // First call should have stopHookActive=false, retry should have stopHookActive=true + expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledTimes(2); + expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenNthCalledWith( + 1, + expect.any(String), + expect.any(String), + false, + ); + expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenNthCalledWith( + 2, + expect.any(String), + expect.any(String), + true, + ); }); it('should call resetChat when AfterAgent hook returns shouldClearContext: true', async () => { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 2ab75cf365..db6c5fb574 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -191,10 +191,11 @@ export class GeminiClient { currentRequest: PartListUnion, prompt_id: string, turn?: Turn, + stopHookActive: boolean = false, ): Promise { const hookState = this.hookStateMap.get(prompt_id); // Only fire on the outermost call (when activeCalls is 1) - if (!hookState || hookState.activeCalls !== 1) { + if (!hookState || (hookState.activeCalls !== 1 && !stopHookActive)) { return undefined; } @@ -210,7 +211,11 @@ export class GeminiClient { const hookOutput = await this.config .getHookSystem() - ?.fireAfterAgentEvent(partToString(finalRequest), finalResponseText); + ?.fireAfterAgentEvent( + partToString(finalRequest), + finalResponseText, + stopHookActive, + ); return hookOutput; } @@ -845,6 +850,7 @@ export class GeminiClient { turns: number = MAX_TURNS, isInvalidStreamRetry: boolean = false, displayContent?: PartListUnion, + stopHookActive: boolean = false, ): AsyncGenerator { if (!isInvalidStreamRetry) { this.config.resetTurn(); @@ -909,6 +915,7 @@ export class GeminiClient { request, prompt_id, turn, + stopHookActive, ); // Cast to AfterAgentHookOutput for access to shouldClearContext() @@ -954,6 +961,7 @@ export class GeminiClient { boundedTurns - 1, false, displayContent, + true, // stopHookActive: signal retry to AfterAgent hooks ); } }