Refactor beforeAgent and afterAgent hookEvents to follow desired output (#16495)

This commit is contained in:
Vedant Mahajan
2026-01-13 22:20:54 +05:30
committed by GitHub
parent 304caa4e43
commit a6dca02344
3 changed files with 28 additions and 69 deletions

View File

@@ -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<GeminiChat> = {
@@ -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<GeminiChat> = {
@@ -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* () {

View File

@@ -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;
}

View File

@@ -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<AggregatedHookResult | undefined> {
): Promise<DefaultHookOutput | undefined> {
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<AggregatedHookResult | undefined> {
): Promise<DefaultHookOutput | undefined> {
if (!this.config.getEnableHooks()) {
return undefined;
}
return this.hookEventHandler.fireAfterAgentEvent(
const result = await this.hookEventHandler.fireAfterAgentEvent(
prompt,
response,
stopHookActive,
);
return result.finalOutput;
}
}