From a6dca02344bae91108dd7297944203339af96ad0 Mon Sep 17 00:00:00 2001 From: Vedant Mahajan Date: Tue, 13 Jan 2026 22:20:54 +0530 Subject: [PATCH] Refactor beforeAgent and afterAgent hookEvents to follow desired output (#16495) --- packages/core/src/core/client.test.ts | 80 +++++++-------------------- packages/core/src/core/client.ts | 6 +- packages/core/src/hooks/hookSystem.ts | 11 ++-- 3 files changed, 28 insertions(+), 69 deletions(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 84418c9855..50ee7f765a 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -137,20 +137,8 @@ vi.mock('../telemetry/uiTelemetry.js', () => ({ })); vi.mock('../hooks/hookSystem.js'); const mockHookSystem = { - fireBeforeAgentEvent: vi.fn().mockResolvedValue({ - success: true, - finalOutput: undefined, - allOutputs: [], - errors: [], - totalDuration: 0, - }), - fireAfterAgentEvent: vi.fn().mockResolvedValue({ - success: true, - finalOutput: undefined, - allOutputs: [], - errors: [], - totalDuration: 0, - }), + fireBeforeAgentEvent: vi.fn().mockResolvedValue(undefined), + fireAfterAgentEvent: vi.fn().mockResolvedValue(undefined), }; /** @@ -2812,15 +2800,9 @@ ${JSON.stringify( it('should stop execution in BeforeAgent when hook returns continue: false', async () => { mockHookSystem.fireBeforeAgentEvent.mockResolvedValue({ - success: true, - finalOutput: { - shouldStopExecution: () => true, - getEffectiveReason: () => 'Stopped by hook', - systemMessage: undefined, - }, - allOutputs: [], - errors: [], - totalDuration: 0, + shouldStopExecution: () => true, + getEffectiveReason: () => 'Stopped by hook', + systemMessage: undefined, }); const mockChat: Partial = { @@ -2851,16 +2833,10 @@ ${JSON.stringify( it('should block execution in BeforeAgent when hook returns decision: block', async () => { mockHookSystem.fireBeforeAgentEvent.mockResolvedValue({ - success: true, - finalOutput: { - shouldStopExecution: () => false, - isBlockingDecision: () => true, - getEffectiveReason: () => 'Blocked by hook', - systemMessage: undefined, - }, - allOutputs: [], - errors: [], - totalDuration: 0, + shouldStopExecution: () => false, + isBlockingDecision: () => true, + getEffectiveReason: () => 'Blocked by hook', + systemMessage: undefined, }); const mockChat: Partial = { @@ -2890,15 +2866,9 @@ ${JSON.stringify( it('should stop execution in AfterAgent when hook returns continue: false', async () => { mockHookSystem.fireAfterAgentEvent.mockResolvedValue({ - success: true, - finalOutput: { - shouldStopExecution: () => true, - getEffectiveReason: () => 'Stopped after agent', - systemMessage: undefined, - }, - allOutputs: [], - errors: [], - totalDuration: 0, + shouldStopExecution: () => true, + getEffectiveReason: () => 'Stopped after agent', + systemMessage: undefined, }); mockTurnRunFn.mockImplementation(async function* () { @@ -2923,27 +2893,15 @@ ${JSON.stringify( it('should yield AgentExecutionBlocked and recurse in AfterAgent when hook returns decision: block', async () => { mockHookSystem.fireAfterAgentEvent .mockResolvedValueOnce({ - success: true, - finalOutput: { - shouldStopExecution: () => false, - isBlockingDecision: () => true, - getEffectiveReason: () => 'Please explain', - systemMessage: undefined, - }, - allOutputs: [], - errors: [], - totalDuration: 0, + shouldStopExecution: () => false, + isBlockingDecision: () => true, + getEffectiveReason: () => 'Please explain', + systemMessage: undefined, }) .mockResolvedValueOnce({ - success: true, - finalOutput: { - shouldStopExecution: () => false, - isBlockingDecision: () => false, - systemMessage: undefined, - }, - allOutputs: [], - errors: [], - totalDuration: 0, + shouldStopExecution: () => false, + isBlockingDecision: () => false, + systemMessage: undefined, }); mockTurnRunFn.mockImplementation(async function* () { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 535bf15ce7..dbf81aee2f 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -133,10 +133,9 @@ export class GeminiClient { return undefined; } - const hookResult = await this.config + const hookOutput = await this.config .getHookSystem() ?.fireBeforeAgentEvent(partToString(request)); - const hookOutput = hookResult?.finalOutput; hookState.hasFiredBeforeAgent = true; if (hookOutput?.shouldStopExecution()) { @@ -187,10 +186,9 @@ export class GeminiClient { '[no response text]'; const finalRequest = hookState.originalRequest || currentRequest; - const hookResult = await this.config + const hookOutput = await this.config .getHookSystem() ?.fireAfterAgentEvent(partToString(finalRequest), finalResponseText); - const hookOutput = hookResult?.finalOutput; return hookOutput; } diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index 547b44a923..7d5e7783f8 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -18,6 +18,7 @@ import type { SessionStartSource, SessionEndReason, PreCompressTrigger, + DefaultHookOutput, } from './types.js'; import type { AggregatedHookResult } from './hookAggregator.js'; /** @@ -120,25 +121,27 @@ export class HookSystem { async fireBeforeAgentEvent( prompt: string, - ): Promise { + ): Promise { if (!this.config.getEnableHooks()) { return undefined; } - return this.hookEventHandler.fireBeforeAgentEvent(prompt); + const result = await this.hookEventHandler.fireBeforeAgentEvent(prompt); + return result.finalOutput; } async fireAfterAgentEvent( prompt: string, response: string, stopHookActive: boolean = false, - ): Promise { + ): Promise { if (!this.config.getEnableHooks()) { return undefined; } - return this.hookEventHandler.fireAfterAgentEvent( + const result = await this.hookEventHandler.fireAfterAgentEvent( prompt, response, stopHookActive, ); + return result.finalOutput; } }