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

[Aarchi-07]

Summary

The AfterAgent hook's stop_hook_active field was never set to true on retries, causing hooks that rely on it to create infinite deny loops.

Root cause: fireAfterAgentHookSafe called fireAfterAgentEvent without passing stopHookActive, and the activeCalls guard prevented the hook from firing on recursive retry calls.

Fix:

• Add stopHookActive parameter to fireAfterAgentHookSafe and sendMessageStream
• Decrement activeCalls before retry recursion so the inner sendMessageStream fires AfterAgent again
• Pass stopHookActive=true on the retry path so hooks receive stop_hook_active: true and can break the loop

Related Issues

Fixes #20426

Pre-Merge Checklist

• Updated relevant documentation and README (if needed)
• Added/updated tests (if needed)
• Noted breaking changes (if any): none
• Validated on required platforms/methods:
  • MacOS
    • npm run
    • npx
    • Docker
    • Podman
    • Seatbelt
  • Windows
    • npm run
    • npx
    • Docker
  • Linux
    • npm run
    • npx
    • Docker

[gemini-code-assist]

Summary of Changes

Hello @Aarchi-07, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical bug where AfterAgent hooks could enter infinite deny loops during retry operations because the stop_hook_active field was not correctly propagated. The changes introduce a new parameter to track the hook's active state, adjust the internal call counter to allow hooks to re-fire on retries, and explicitly signal when a retry is occurring, thereby enabling hooks to properly terminate the loop.

Highlights

• Propagated stopHookActive: Introduced a stopHookActive parameter to fireAfterAgentHookSafe and sendMessageStream functions to correctly propagate the hook's active state during retries.
• Adjusted Retry Logic: Modified the retry logic within sendMessageStream to decrement the activeCalls counter, ensuring that AfterAgent hooks are re-fired during recursive retry calls.
• Signaled Retry Path: Ensured that stopHookActive is explicitly set to true when sendMessageStream is called on the retry path, allowing hooks to properly break out of potential infinite loops.

Changelog

• packages/core/src/core/client.test.ts
  • Updated existing fireAfterAgentEvent mock calls to include the new false argument for stopHookActive.
  • Added new test assertions to verify that fireAfterAgentEvent is called twice during a retry scenario, once with stopHookActive: false and once with stopHookActive: true.
• packages/core/src/core/client.ts
  • Added a stopHookActive boolean parameter with a default of false to the fireAfterAgentHookSafe method signature.
  • Modified the call to fireAfterAgentEvent within fireAfterAgentHookSafe to pass the new stopHookActive parameter.
  • Added a stopHookActive boolean parameter with a default of false to the sendMessageStream method signature.
  • Introduced logic to decrement activeCalls for the relevant hookState before recursively calling sendMessageStream during a retry.
  • Passed true for the stopHookActive parameter when sendMessageStream is recursively called for a retry.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to fix an infinite loop in the AfterAgent hook retry path by propagating a stopHookActive flag. However, the implementation introduces a critical logic flaw by manually decrementing the activeCalls counter. This leads to a double-decrement, corrupting the counter (potentially making it negative or zero), and consequently preventing AfterAgent hooks from firing in subsequent turns of the same session, effectively disabling security guardrails. The recommended fix is to remove the manual decrement and adjust the firing condition in fireAfterAgentHookSafe to account for the stopHookActive flag.

[scidomino]

This PR addresses a crucial issue with the AfterAgent hook retry logic. However, as noted by the automated code assist, manually decrementing activeCalls right before the recursive sendMessageStream call introduces a bug.

Because sendMessageStream increments activeCalls at the beginning (via fireBeforeAgentHookSafe) and decrements it in a finally block at the end, decrementing it manually before the recursive call will result in a double decrement once the call stacks unwind. This will corrupt the hookStateMap and break hook behavior for subsequent turns.

To fix this, we should remove the manual decrement:

// Remove this block
const retryHookState = this.hookStateMap.get(prompt_id);
if (retryHookState) {
  retryHookState.activeCalls--;
}

Instead, we should adjust the condition within fireAfterAgentHookSafe to allow the hook to fire even when activeCalls > 1 if it is a retry (indicated by stopHookActive).

For example, you could update fireAfterAgentHookSafe to:

  private async fireAfterAgentHookSafe(
    currentRequest: PartListUnion,
    prompt_id: string,
    turn?: Turn,
    stopHookActive: boolean = false,
  ): Promise<DefaultHookOutput | undefined> {
    const hookState = this.hookStateMap.get(prompt_id);
    // Fire on the outermost call (when activeCalls is 1) OR if it's a retry (stopHookActive)
    if (!hookState || (hookState.activeCalls !== 1 && !stopHookActive)) {
      return undefined;
    }
// ...

Could you update the PR with this approach?