fix(core): Implement two-stage hook validation and improve error handling #15499
Conversation
Summary of ChangesHello @SandyTao520, I'm Gemini Code Assist1! 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 introduces a significant refinement to the system's hook validation process by establishing a clear two-stage pattern. This change aims to provide more precise control and feedback mechanisms for both security-related policy enforcement and general workflow management. By distinguishing between blocking an action due to policy and stopping an execution for logical reasons, the system can now communicate more effectively with the LLM and handle various intervention scenarios with greater clarity and robustness. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a two-stage validation pattern for hooks, separating security/policy decisions from workflow control. However, it critically introduces several high-severity prompt injection vulnerabilities: hook output is directly injected into LLM content, enabling malicious hooks to control LLM behavior and bypass security policies. Furthermore, the AfterTool hook implementation critically omits the policy-blocking validation stage, which could silently ignore security policies after tool execution.
| if (afterOutput?.shouldStopExecution()) { | ||
| const reason = afterOutput.getEffectiveReason(); | ||
| return { | ||
| llmContent: `Agent execution stopped by hook: ${reason}`, | ||
| returnDisplay: `Agent execution stopped by hook: ${reason}`, | ||
| llmContent: `Execution stopped: ${reason}`, | ||
| returnDisplay: `Execution stopped: ${reason}`, | ||
| error: { | ||
| type: ToolErrorType.EXECUTION_FAILED, | ||
| message: `Agent execution stopped: ${reason}`, | ||
| message: `Execution stopped: ${reason}`, | ||
| }, | ||
| }; |
There was a problem hiding this comment.
This AfterTool hook implementation has two critical issues. Firstly, the shouldStopExecution path is vulnerable to prompt injection because the reason from getEffectiveReason() is directly injected into llmContent, allowing malicious hooks to control the LLM. Secondly, it critically omits the Stage 1 validation for policy/security decisions (isBlockingDecision). This means a decision: 'block' from an AfterTool hook would be silently ignored, undermining the two-stage validation and potentially bypassing security policies.
| if (hookOutput?.isBlockingDecision()) { | ||
| const blockReason = `Policy Violation: ${hookOutput.getEffectiveReason()}`; | ||
| const continueRequest = [{ text: blockReason }]; | ||
| yield* this.sendMessageStream( | ||
| continueRequest, | ||
| signal, | ||
| prompt_id, | ||
| boundedTurns - 1, | ||
| ); |
There was a problem hiding this comment.
A prompt injection vulnerability exists in the AfterAgent hook processing. The reason provided by a hook, which can be controlled by a third-party extension, is directly concatenated into a new prompt sent to the LLM via sendMessageStream. A malicious hook could provide a crafted reason to manipulate the LLM, potentially leading to policy bypasses or unintended behavior. The getEffectiveReason() output is not sanitized or treated as untrusted data before being used as part of the next prompt.
| if (blockingError?.blocked) { | ||
| return { | ||
| llmContent: `Tool execution blocked: ${blockingError.reason}`, | ||
| returnDisplay: `Tool execution blocked: ${blockingError.reason}`, | ||
| llmContent: `Policy Block: ${blockingError.reason}`, | ||
| returnDisplay: `Policy Block: ${blockingError.reason}`, | ||
| error: { | ||
| type: ToolErrorType.EXECUTION_FAILED, | ||
| type: ToolErrorType.PERMISSION_DENIED, | ||
| message: blockingError.reason, | ||
| }, | ||
| }; |
There was a problem hiding this comment.
A prompt injection vulnerability exists in the BeforeTool hook processing. The blockingError.reason from a hook, which can be controlled by a third-party, is directly injected into the llmContent of the tool result. This content is then sent to the LLM, allowing a malicious hook to manipulate the LLM's behavior.
| if (beforeOutput?.shouldStopExecution()) { | ||
| const reason = beforeOutput.getEffectiveReason(); | ||
| return { | ||
| llmContent: `Agent execution stopped by hook: ${reason}`, | ||
| returnDisplay: `Agent execution stopped by hook: ${reason}`, | ||
| llmContent: `Execution stopped: ${reason}`, | ||
| returnDisplay: `Execution stopped: ${reason}`, | ||
| error: { | ||
| type: ToolErrorType.EXECUTION_FAILED, | ||
| message: `Agent execution stopped: ${reason}`, | ||
| message: `Execution stopped: ${reason}`, | ||
| }, | ||
| }; |
There was a problem hiding this comment.
Similar to the blocking case, the shouldStopExecution path in the BeforeTool hook is also vulnerable to prompt injection. The reason from getEffectiveReason() is injected into llmContent and sent to the LLM, creating a vector for malicious prompt injection.
|
Size Change: +1.04 kB (0%) Total Size: 22 MB
ℹ️ View Unchanged
|
1 participant
TLDR
Implemented a strict 'Two-Stage Validation Pattern' for hook responses (
BeforeAgent,AfterAgent,BeforeTool,AfterTool,BeforeModel). This distinguishes between Stage 1: Policy/Security Decisions (decision: 'block') and Stage 2: Workflow Lifecycle Decisions (continue: false).Dive Deeper
Previously, the distinction between blocking an action for security reasons and stopping a workflow for logic reasons was ambiguous. This PR unifies the behavior:
decision: 'block'. If blocked, it returns aPERMISSION_DENIEDerror (for tools) or aPolicy Blockerror (for agents). This signals to the LLM that the action was forbidden.continue: false. If false, it gracefully terminates the execution with anEXECUTION_FAILED(or 'Execution stopped') status, often with a 'stop reason'. This is for natural workflow endpoints.client.test.ts,geminiChat.test.ts, andcoreToolHookTriggers.test.ts.Reviewer Test Plan
BeforeToolhook that returnsdecision: 'block', reason: 'Unsafe'. Verify that the tool call fails withPERMISSION_DENIEDand the model receives 'Policy Block: Unsafe'.BeforeToolhook that returnscontinue: false, stopReason: 'Done'. Verify that the tool call fails withEXECUTION_FAILEDand the model receives 'Execution stopped: Done'.npm test packages/core/src/core/client.test.tsnpm test packages/core/src/core/geminiChat.test.tsnpm test packages/core/src/core/coreToolHookTriggers.test.tsTesting Matrix
Linked issues / bugs
#14718