fix(hooks): fix BeforeAgent/AfterAgent inconsistencies (#18514) (#21383)

Co-authored-by: Spencer <spencertang@google.com>
This commit is contained in:
krishdef7
2026-03-12 03:10:11 +05:30
committed by GitHub
parent 352bbc36c0
commit 926dddf0bf
+27 -16
View File
@@ -30,12 +30,6 @@ import { getCoreSystemPrompt } from './prompts.js';
import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
import { reportError } from '../utils/errorReporting.js'; import { reportError } from '../utils/errorReporting.js';
import { GeminiChat } from './geminiChat.js'; import { GeminiChat } from './geminiChat.js';
import { coreEvents, CoreEvent } from '../utils/events.js';
import {
getDisplayString,
resolveModel,
isGemini2Model,
} from '../config/models.js';
import { import {
retryWithBackoff, retryWithBackoff,
type RetryAvailabilityContext, type RetryAvailabilityContext,
@@ -76,7 +70,13 @@ import {
applyModelSelection, applyModelSelection,
createAvailabilityContextProvider, createAvailabilityContextProvider,
} from '../availability/policyHelpers.js'; } from '../availability/policyHelpers.js';
import {
getDisplayString,
resolveModel,
isGemini2Model,
} from '../config/models.js';
import { partToString } from '../utils/partUtils.js'; import { partToString } from '../utils/partUtils.js';
import { coreEvents, CoreEvent } from '../utils/events.js';
const MAX_TURNS = 100; const MAX_TURNS = 100;
@@ -907,6 +907,7 @@ export class GeminiClient {
const boundedTurns = Math.min(turns, MAX_TURNS); const boundedTurns = Math.min(turns, MAX_TURNS);
let turn = new Turn(this.getChat(), prompt_id); let turn = new Turn(this.getChat(), prompt_id);
let continuationHandled = false;
try { try {
turn = yield* this.processTurn( turn = yield* this.processTurn(
@@ -963,7 +964,15 @@ export class GeminiClient {
await this.resetChat(); await this.resetChat();
} }
const continueRequest = [{ text: continueReason }]; const continueRequest = [{ text: continueReason }];
yield* this.sendMessageStream( // Reset hook state so the continuation fires BeforeAgent fresh
// and fireAfterAgentHookSafe sees activeCalls=1, not 2.
const contHookState = this.hookStateMap.get(prompt_id);
if (contHookState) {
contHookState.hasFiredBeforeAgent = false;
contHookState.activeCalls--;
}
continuationHandled = true;
turn = yield* this.sendMessageStream(
continueRequest, continueRequest,
signal, signal,
prompt_id, prompt_id,
@@ -981,16 +990,18 @@ export class GeminiClient {
} }
throw error; throw error;
} finally { } finally {
const hookState = this.hookStateMap.get(prompt_id); if (!continuationHandled) {
if (hookState) { const hookState = this.hookStateMap.get(prompt_id);
hookState.activeCalls--; if (hookState) {
const isPendingTools = hookState.activeCalls--;
turn?.pendingToolCalls && turn.pendingToolCalls.length > 0; const isPendingTools =
const isAborted = signal?.aborted; turn?.pendingToolCalls && turn.pendingToolCalls.length > 0;
const isAborted = signal?.aborted;
if (hookState.activeCalls <= 0) { if (hookState.activeCalls <= 0) {
if (!isPendingTools || isAborted) { if (!isPendingTools || isAborted) {
this.hookStateMap.delete(prompt_id); this.hookStateMap.delete(prompt_id);
}
} }
} }
} }