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

[krishdef7]

Fixes #18514

Root Cause

All three bugs trace to hook state management in sendMessageStream in packages/core/src/core/client.ts.

When isBlockingDecision() fires and triggers a recursive sendMessageStream call with the same prompt_id:

1. BeforeAgent not firing on injected prompt, hasFiredBeforeAgent is still true from the prior turn. The guard at fireBeforeAgentHookSafe returns early, silently skipping BeforeAgent for the injected prompt.

2. AfterAgent not firing after text-only continuation, the recursive call increments activeCalls to 2. fireAfterAgentHookSafe guards on activeCalls !== 1 and returns undefined, so AfterAgent never fires for that turn.

3. BeforeAgent firing spuriously after first tool call — the finally block checks pendingToolCalls on the outer turn (not the continuation turn), which can prematurely delete the hookState mid-sequence. The next call creates a fresh hookState with hasFiredBeforeAgent=false, causing an unexpected BeforeAgent fire.

Fix

Three surgical changes, all in sendMessageStream:

• Before the recursive sendMessageStream call in the isBlockingDecision branch: reset hasFiredBeforeAgent = false and decrement activeCalls on the hookState, so the continuation is treated as a fresh outermost turn.
• Set a continuationHandled flag to prevent the finally block from double-decrementing activeCalls when the continuation has already managed its own state.

Testing

• All 71 existing tests pass (npx vitest run packages/core/src/core/client.test.ts)
• Full preflight passes (npm run preflight) - 0 errors
• The fix specifically covers the exact event sequences described in the issue: hook injection via exit code 2, text-only responses, and tool call sequences

[gemini-code-assist]

Summary of Changes

Hello, 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 resolves several inconsistencies in the BeforeAgent and AfterAgent hooks within the sendMessageStream process, particularly concerning recursive calls and state management. It ensures that these hooks fire reliably and as expected during various interaction flows, including injected prompts, text-only continuations, and tool calls. Additionally, it improves the handling of '@' symbols in the CLI to prevent unintended command parsing and enhances transcript flushing logic in GeminiChat for better tool hook integration.

Highlights

• Corrected BeforeAgent hook firing: Ensured the BeforeAgent hook fires correctly for injected prompts by resetting its state before recursive calls within sendMessageStream.
• Resolved AfterAgent hook inconsistencies: Fixed an issue where the AfterAgent hook would not fire after text-only continuations by correctly managing activeCalls during recursive sendMessageStream calls.
• Prevented spurious BeforeAgent calls: Addressed the problem of BeforeAgent firing unexpectedly after tool calls by preventing premature deletion of hook state in the finally block of sendMessageStream.
• Improved '@' command handling in CLI: Implemented escaping and unescaping of '@' symbols in the CLI input prompt and command processor to prevent misinterpretation of literal '@' characters.
• Enhanced transcript flushing for tool calls: Modified GeminiChat to ensure the transcript is flushed before tool dispatch, even for pure tool calls without explicit text or thoughts, to maintain up-to-date state for BeforeTool hooks.

Changelog

• packages/cli/src/ui/components/InputPrompt.tsx
  • Imported unescapeLiteralAt utility.
  • Modified paste handling to escape '@' symbols in pasted text.
  • Updated input handling to escape '@' symbols in key sequences.
  • Applied unescapeLiteralAt to display text segments.
• packages/cli/src/ui/hooks/atCommandProcessor.ts
  • Imported unescapeLiteralAt utility.
  • Applied unescapeLiteralAt to text content when parsing '@' commands.
• packages/cli/src/utils/unescapeLiteralAt.ts
  • Added a new utility function unescapeLiteralAt to correctly unescape '@' symbols.
• packages/core/src/core/client.ts
  • Introduced continuationHandled flag to manage hook state.
  • Reset hasFiredBeforeAgent and decremented activeCalls for continuation calls.
  • Conditionally executed the finally block's hook state cleanup based on continuationHandled.
• packages/core/src/core/geminiChat.test.ts
  • Added a new test case to verify transcript flushing before tool dispatch for pure tool calls.
• packages/core/src/core/geminiChat.ts
  • Introduced hasThoughts flag to track thought presence.
  • Updated the condition for recording model response to include hasThoughts or hasToolCall for flushing.

Activity

• All 71 existing tests for packages/core/src/core/client.test.ts passed.
• Full preflight checks passed with 0 errors.
• The fix specifically addresses event sequences described in the issue, including hook injection via exit code 2, text-only responses, and tool call sequences.

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 addresses several issues. The main fix resolves inconsistencies with BeforeAgent/AfterAgent hooks during recursive calls in sendMessageStream by improving state management with a continuationHandled flag and resetting hook state appropriately. This is a good, targeted fix for the described problem. The PR also includes two other unrelated changes: a feature to allow escaping @ characters in the input prompt to treat them as literals, and a fix to ensure the chat transcript is flushed even for model responses that only contain tool calls or thoughts.

However, the changes related to handling @ characters introduce a critical regression. By automatically escaping @ on paste, it's no longer possible to paste and execute at-commands, which is a core feature. I've left specific comments on how to address this, referencing the rule about maintaining UI consistency.

While all these changes are valuable, bundling unrelated fixes and features into a single pull request goes against the repository's guideline to keep PRs small and focused (line 67). This makes the changes harder to review and understand. In the future, please submit distinct changes in separate pull requests.

[krishdef7]

Cleaned up - the PR now touches only packages/core/src/core/client.ts. The paste-related files (InputPrompt.tsx, atCommandProcessor.ts, unescapeLiteralAt.ts) were accidentally bundled from a prior branch and have been removed. The previous bot feedback on those files is now outdated.

[spencer426]

Thanks for the PR! The core logic fix for the hook states looks good, but it appears your branch is out of sync with main and is actively reverting some very recent changes. Could you please rebase or restore these missing lines?

[spencer426]

It looks like this reverts the recent introduction of AgentLoopContext. This constructor should still accept private readonly context: AgentLoopContext instead of falling back to just Config.

[spencer426]

This reverts the method signature by removing the readonly modifier. Please keep readonly Content[] to preserve type safety.

[spencer426]

Same here, please keep the readonly modifier in the parameter signature: setHistory(history: readonly Content[]).

[spencer426]

This unintentionally reverts PR #22027 by removing retryFetchErrors and onRetry from the generateContentWithRetry call, which are needed for network retry telemetry. Please restore these properties.

[krishdef7]

@spencer426 all four points addressed, AgentLoopContext constructor, readonly Content[] on getHistory/setHistory/tryMaskToolOutputs, retryFetchErrors + onRetry from #22027, and this.context.toolRegistry at all 3 call sites. Single commit on top of upstream/main, 79 tests passing, preflight green.

[spencer426]

LGTM