Improving memory tool instructions and eval testing

[alisa-alisa]

Summary

Improves instructions for memory tool and adds more evals to validate new behavior.

Details

Evals updates

• Prompt is refined.
• Eval tests are added.
• validateModelOutput is refactored to include forbiddenContent.
• validateModelOutput is split into assertModelHasOutput and checkModelOutputContent to follow the Single Responsibility Principle.

Security vulnerabilities in memory tool

Some security vulnerabilities are fixed as per /gemini review

Markdown Injection

The save_memory tool wasn't properly sanitizing the fact it was asked to remember. An attacker could provide a fact containing Markdown, like a new header (## New Instructions), which would be written directly into your memory file (GEMINI.md). When the agent reads this file for context in the future, it would treat these injected lines as trusted instructions, leading to a persistent prompt injection attack.

Bait and Switch

The process for saving a memory had a logic flaw.

• Step 1 (The "Bait"): The tool would first generate the change, create a diff, and show it to you for approval.
• Step 2 (The "Switch"): After you approved it, instead of writing the exact content you just approved, the tool would re-read the original fact and re-calculate the change from scratch before writing it to the file.

This creates a vulnerability. If the tool could be tricked into using a different fact during the final write than it used for the initial approval screen, it could save malicious content that you never approved. The fix is to ensure the exact content shown during approval is what gets written to the file.

Param injection

Fix:

1. Stricter Schema: update the tool's JSON schema to forbid any parameters other than fact. This serves as the first line of defense, instructing the model not to entertain these extra parameters.
2. Expose the Attack: As a vital second layer, modify the confirmation logic. If an attacker manages to bypass the schema and inject modified_content, the confirmation diff will now be generated from that malicious content instead of the harmless fact. This ensures you see exactly what the tool intends to write, exposing the attack and allowing you to cancel it.

Related Issues

Fixes #15678

How to Validate

I asked gemini cli to run tests 50 times and validate the tool call and then 20 times and validate the model output. It passed.

Also new tests were very flaky before the prompt refactoring.

Pre-Merge Checklist

• Updated relevant documentation and README (if needed)
• Added/updated tests (if needed)
• Noted breaking changes (if any)
• 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 @alisa-alisa, 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 significantly improves the robustness and clarity of the memory tool's functionality and its testing. It refines the tool's instructions to better guide its usage, while simultaneously expanding the suite of evaluation tests to cover a broader range of memory-related scenarios. A key enhancement is the update to the model output validation utility, which now allows for more precise control over expected and forbidden content in test assertions.

Highlights

• Enhanced Memory Tool Instructions: The description for the memoryTool has been updated to provide clearer guidance and examples, emphasizing its use for saving user preferences and important facts.
• Expanded Evaluation Tests for Memory Tool: A significant number of new evaluation tests have been added for the save_memory tool, covering various scenarios like remembering personal details, user preferences, command restrictions, and project-specific information.
• Refined Model Output Validation Utility: The validateModelOutput utility function has been refactored to accept an object as its second argument, allowing for more flexible validation, including the introduction of a forbiddenContent parameter.
• Introduction of Forbidden Content Validation: The validateModelOutput function now supports checking for content that should not be present in the model's output, improving the robustness of evaluation tests.
• Standardized Test Validation Calls: All existing integration tests have been updated to use the new object-based parameter structure for validateModelOutput.

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 significantly improves the save_memory tool by refining its instructional prompt and adding a comprehensive suite of evaluation tests. The refactoring of the validateModelOutput test helper to use a configuration object is a welcome change that improves readability and extensibility, as demonstrated by the addition of forbiddenContent validation. Overall, these changes enhance the tool's reliability and test coverage. I've identified one high-severity issue in a new test case where an assertion should be strengthened to prevent potential regressions.

[github-actions]

Size Change: -254 B (0%)

Total Size: 23.7 MB

ℹ️ View Unchanged

Filename	Size	Change
./bundle/gemini.js	23.7 MB	-254 B (0%)
./bundle/sandbox-macos-permissive-closed.sb	1.03 kB	0 B
./bundle/sandbox-macos-permissive-open.sb	890 B	0 B
./bundle/sandbox-macos-permissive-proxied.sb	1.31 kB	0 B
./bundle/sandbox-macos-restrictive-closed.sb	3.29 kB	0 B
./bundle/sandbox-macos-restrictive-open.sb	3.36 kB	0 B
./bundle/sandbox-macos-restrictive-proxied.sb	3.56 kB	0 B

_{compressed-size-action}

[alisa-alisa]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly improves the save_memory tool by refining its instructions and adding a comprehensive suite of evaluation tests. The refactoring of validateModelOutput into assertModelHasOutput and checkModelOutputContent further enhances test helpers by adhering to the Single Responsibility Principle and supporting forbidden content checks. However, a security audit identified two high-severity vulnerabilities in the memoryTool.ts implementation: a Markdown injection vulnerability in the save_memory tool that allows attackers to inject arbitrary content and instructions into the agent's long-term memory, and a logic flaw in the tool invocation process enabling a "bait and switch" attack where the content written to the memory file can differ from user approval. Both issues could lead to persistent prompt injection. Additionally, an existing negative test case could be made more robust. These critical issues should be addressed before merging.

[alisa-alisa]

/gemini review

[gemini-code-assist]

Code Review

This pull request significantly improves the security and reliability of the save_memory tool by addressing markdown injection and a 'Bait and Switch' logic flaw. While these fixes are well-implemented through input sanitization and content integrity checks, a critical security vulnerability remains: internal tool parameters (modified_by_user, modified_content) can still be injected by the LLM, potentially bypassing the new protections. Additionally, the tool's implementation has been refactored for better self-containment within MemoryToolInvocation, and the test suite has been substantially expanded and refactored. A minor issue with inconsistent sanitization for success messages was also noted.

[alisa-alisa]

/gemini review

[gemini-code-assist]

Code Review

This is an excellent pull request that significantly improves the security and reliability of the save_memory tool. The fixes for the Markdown Injection and "Bait and Switch" vulnerabilities are well-implemented and correctly address the root causes. Storing the proposed content and using it for the final write operation effectively closes the race condition, and sanitizing the input fact by removing newlines prevents prompt injection via new markdown blocks. The addition of additionalProperties: false to the tool's schema provides crucial defense-in-depth against parameter injection attacks.

The refactoring of the test helpers into assertModelHasOutput and checkModelOutputContent is a great improvement for code clarity and maintainability, and it's applied consistently across the test suite. The new evaluation tests for save_memory are comprehensive and greatly increase confidence in the tool's intended behavior.

Overall, these changes are high-quality and represent a major step forward for the tool's robustness. I have no further comments.