diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index c910556ca82..e75a7d2c547 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -2988,6 +2988,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), 'Hook Response', + false, ); // Map should be empty @@ -3029,6 +3030,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), 'Response 1\nResponse 2', + false, ); expect(client['hookStateMap'].size).toBe(0); @@ -3059,6 +3061,7 @@ ${JSON.stringify( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( partToString(request), // Should be 'Do something' expect.stringContaining('Ok'), + false, ); }); @@ -3229,6 +3232,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 c94dd5c04d0..a460c72070b 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -183,10 +183,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) { + // Fire on the outermost call (when activeCalls is 1) OR if it's a retry (stopHookActive) + if (!hookState || (hookState.activeCalls !== 1 && !stopHookActive)) { return undefined; } @@ -202,7 +203,11 @@ export class GeminiClient { const hookOutput = await this.config .getHookSystem() - ?.fireAfterAgentEvent(partToString(finalRequest), finalResponseText); + ?.fireAfterAgentEvent( + partToString(finalRequest), + finalResponseText, + stopHookActive, + ); return hookOutput; } @@ -793,6 +798,7 @@ export class GeminiClient { turns: number = MAX_TURNS, isInvalidStreamRetry: boolean = false, displayContent?: PartListUnion, + stopHookActive: boolean = false, ): AsyncGenerator { if (!isInvalidStreamRetry) { this.config.resetTurn(); @@ -857,6 +863,7 @@ export class GeminiClient { request, prompt_id, turn, + stopHookActive, ); // Cast to AfterAgentHookOutput for access to shouldClearContext() @@ -902,6 +909,7 @@ export class GeminiClient { boundedTurns - 1, false, displayContent, + true, // stopHookActive: signal retry to AfterAgent hooks ); } }