fix(hooks): propagate stopHookActive in AfterAgent retry path (#20426) (#20439)

This commit is contained in:
Aarchi Kumari
2026-03-06 22:36:23 +05:30
committed by GitHub
parent d97eaf3420
commit 337e4bc8c6
4 changed files with 48 additions and 11 deletions
+11 -1
View File
@@ -182,6 +182,13 @@ describe('Hooks Agent Flow', () => {
); );
const afterAgentScript = ` const afterAgentScript = `
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({ console.log(JSON.stringify({
decision: 'block', decision: 'block',
reason: 'Security policy triggered', reason: 'Security policy triggered',
@@ -190,6 +197,7 @@ describe('Hooks Agent Flow', () => {
clearContext: true clearContext: true
} }
})); }));
}
`; `;
const afterAgentScriptPath = rig.createScript( const afterAgentScriptPath = rig.createScript(
'after_agent_clear.cjs', 'after_agent_clear.cjs',
@@ -198,8 +206,10 @@ describe('Hooks Agent Flow', () => {
rig.setup('should process clearContext in AfterAgent hook output', { rig.setup('should process clearContext in AfterAgent hook output', {
settings: { settings: {
hooks: { hooksConfig: {
enabled: true, enabled: true,
},
hooks: {
BeforeModel: [ BeforeModel: [
{ {
hooks: [ hooks: [
@@ -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":"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":"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}]}]}
+18
View File
@@ -3317,6 +3317,7 @@ ${JSON.stringify(
expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(
partToString(request), partToString(request),
'Hook Response', 'Hook Response',
false,
); );
// Map should be empty // Map should be empty
@@ -3358,6 +3359,7 @@ ${JSON.stringify(
expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(
partToString(request), partToString(request),
'Response 1\nResponse 2', 'Response 1\nResponse 2',
false,
); );
expect(client['hookStateMap'].size).toBe(0); expect(client['hookStateMap'].size).toBe(0);
@@ -3388,6 +3390,7 @@ ${JSON.stringify(
expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith( expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(
partToString(request), // Should be 'Do something' partToString(request), // Should be 'Do something'
expect.stringContaining('Ok'), expect.stringContaining('Ok'),
false,
); );
}); });
@@ -3558,6 +3561,21 @@ ${JSON.stringify(
expect.anything(), expect.anything(),
undefined, 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 () => { it('should call resetChat when AfterAgent hook returns shouldClearContext: true', async () => {
+10 -2
View File
@@ -191,10 +191,11 @@ export class GeminiClient {
currentRequest: PartListUnion, currentRequest: PartListUnion,
prompt_id: string, prompt_id: string,
turn?: Turn, turn?: Turn,
stopHookActive: boolean = false,
): Promise<DefaultHookOutput | undefined> { ): Promise<DefaultHookOutput | undefined> {
const hookState = this.hookStateMap.get(prompt_id); const hookState = this.hookStateMap.get(prompt_id);
// Only fire on the outermost call (when activeCalls is 1) // Only fire on the outermost call (when activeCalls is 1)
if (!hookState || hookState.activeCalls !== 1) { if (!hookState || (hookState.activeCalls !== 1 && !stopHookActive)) {
return undefined; return undefined;
} }
@@ -210,7 +211,11 @@ export class GeminiClient {
const hookOutput = await this.config const hookOutput = await this.config
.getHookSystem() .getHookSystem()
?.fireAfterAgentEvent(partToString(finalRequest), finalResponseText); ?.fireAfterAgentEvent(
partToString(finalRequest),
finalResponseText,
stopHookActive,
);
return hookOutput; return hookOutput;
} }
@@ -845,6 +850,7 @@ export class GeminiClient {
turns: number = MAX_TURNS, turns: number = MAX_TURNS,
isInvalidStreamRetry: boolean = false, isInvalidStreamRetry: boolean = false,
displayContent?: PartListUnion, displayContent?: PartListUnion,
stopHookActive: boolean = false,
): AsyncGenerator<ServerGeminiStreamEvent, Turn> { ): AsyncGenerator<ServerGeminiStreamEvent, Turn> {
if (!isInvalidStreamRetry) { if (!isInvalidStreamRetry) {
this.config.resetTurn(); this.config.resetTurn();
@@ -909,6 +915,7 @@ export class GeminiClient {
request, request,
prompt_id, prompt_id,
turn, turn,
stopHookActive,
); );
// Cast to AfterAgentHookOutput for access to shouldClearContext() // Cast to AfterAgentHookOutput for access to shouldClearContext()
@@ -954,6 +961,7 @@ export class GeminiClient {
boundedTurns - 1, boundedTurns - 1,
false, false,
displayContent, displayContent,
true, // stopHookActive: signal retry to AfterAgent hooks
); );
} }
} }