feat(core): add experimental memory manager agent to replace save_memory tool (#22726)

Co-authored-by: Christian Gunderman <gundermanc@gmail.com>

packages/core/src/agents/memory-manager-agent.test.ts

@@ -0,0 +1,153 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { MemoryManagerAgent } from './memory-manager-agent.js';
+import {
+  ASK_USER_TOOL_NAME,
+  EDIT_TOOL_NAME,
+  GLOB_TOOL_NAME,
+  GREP_TOOL_NAME,
+  LS_TOOL_NAME,
+  READ_FILE_TOOL_NAME,
+  WRITE_FILE_TOOL_NAME,
+} from '../tools/tool-names.js';
+import { Storage } from '../config/storage.js';
+import type { Config } from '../config/config.js';
+import type { HierarchicalMemory } from '../config/memory.js';
+
+function createMockConfig(memory: string | HierarchicalMemory = ''): Config {
+  return {
+    getUserMemory: vi.fn().mockReturnValue(memory),
+  } as unknown as Config;
+}
+
+describe('MemoryManagerAgent', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should have the correct name "save_memory"', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    expect(agent.name).toBe('save_memory');
+  });
+
+  it('should be a local agent', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    expect(agent.kind).toBe('local');
+  });
+
+  it('should have a description', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    expect(agent.description).toBeTruthy();
+    expect(agent.description).toContain('memory');
+  });
+
+  it('should have a system prompt with memory management instructions', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    const prompt = agent.promptConfig.systemPrompt;
+    const globalGeminiDir = Storage.getGlobalGeminiDir();
+    expect(prompt).toContain(`Global (${globalGeminiDir}`);
+    expect(prompt).toContain('Project (./');
+    expect(prompt).toContain('Memory Hierarchy');
+    expect(prompt).toContain('De-duplicating');
+    expect(prompt).toContain('Adding');
+    expect(prompt).toContain('Removing stale entries');
+    expect(prompt).toContain('Organizing');
+    expect(prompt).toContain('Routing');
+  });
+
+  it('should have efficiency guidelines in the system prompt', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    const prompt = agent.promptConfig.systemPrompt;
+    expect(prompt).toContain('Efficiency & Performance');
+    expect(prompt).toContain('Use as few turns as possible');
+    expect(prompt).toContain('Do not perform any exploration');
+    expect(prompt).toContain('Be strategic with your thinking');
+    expect(prompt).toContain('Context Awareness');
+  });
+
+  it('should inject hierarchical memory into initial context', () => {
+    const config = createMockConfig({
+      global:
+        '--- Context from: ../../.gemini/GEMINI.md ---\nglobal context\n--- End of Context from: ../../.gemini/GEMINI.md ---',
+      project:
+        '--- Context from: .gemini/GEMINI.md ---\nproject context\n--- End of Context from: .gemini/GEMINI.md ---',
+    });
+
+    const agent = MemoryManagerAgent(config);
+    const query = agent.promptConfig.query;
+
+    expect(query).toContain('# Initial Context');
+    expect(query).toContain('global context');
+    expect(query).toContain('project context');
+  });
+
+  it('should inject flat string memory into initial context', () => {
+    const config = createMockConfig('flat memory content');
+
+    const agent = MemoryManagerAgent(config);
+    const query = agent.promptConfig.query;
+
+    expect(query).toContain('# Initial Context');
+    expect(query).toContain('flat memory content');
+  });
+
+  it('should exclude extension memory from initial context', () => {
+    const config = createMockConfig({
+      global: 'global context',
+      extension: 'extension context that should be excluded',
+      project: 'project context',
+    });
+
+    const agent = MemoryManagerAgent(config);
+    const query = agent.promptConfig.query;
+
+    expect(query).toContain('global context');
+    expect(query).toContain('project context');
+    expect(query).not.toContain('extension context');
+  });
+
+  it('should not include initial context when memory is empty', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    const query = agent.promptConfig.query;
+
+    expect(query).not.toContain('# Initial Context');
+  });
+
+  it('should have file-management and search tools', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    expect(agent.toolConfig).toBeDefined();
+    expect(agent.toolConfig!.tools).toEqual(
+      expect.arrayContaining([
+        READ_FILE_TOOL_NAME,
+        EDIT_TOOL_NAME,
+        WRITE_FILE_TOOL_NAME,
+        LS_TOOL_NAME,
+        GLOB_TOOL_NAME,
+        GREP_TOOL_NAME,
+        ASK_USER_TOOL_NAME,
+      ]),
+    );
+  });
+
+  it('should require a "request" input parameter', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    const schema = agent.inputConfig.inputSchema as Record<string, unknown>;
+    expect(schema).toBeDefined();
+    expect(schema['properties']).toHaveProperty('request');
+    expect(schema['required']).toContain('request');
+  });
+
+  it('should use a fast model', () => {
+    const agent = MemoryManagerAgent(createMockConfig());
+    expect(agent.modelConfig.model).toBe('flash');
+  });
+});

packages/core/src/agents/memory-manager-agent.ts

@@ -0,0 +1,156 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { z } from 'zod';
+import type { LocalAgentDefinition } from './types.js';
+import {
+  ASK_USER_TOOL_NAME,
+  EDIT_TOOL_NAME,
+  GLOB_TOOL_NAME,
+  GREP_TOOL_NAME,
+  LS_TOOL_NAME,
+  READ_FILE_TOOL_NAME,
+  WRITE_FILE_TOOL_NAME,
+} from '../tools/tool-names.js';
+import { Storage } from '../config/storage.js';
+import { flattenMemory } from '../config/memory.js';
+import { GEMINI_MODEL_ALIAS_FLASH } from '../config/models.js';
+import type { Config } from '../config/config.js';
+
+const MemoryManagerSchema = z.object({
+  response: z
+    .string()
+    .describe('A summary of the memory operations performed.'),
+});
+
+/**
+ * A memory management agent that replaces the built-in save_memory tool.
+ * It provides richer memory operations: adding, removing, de-duplicating,
+ * and organizing memories in the global GEMINI.md file.
+ *
+ * Users can override this agent by placing a custom save_memory.md
+ * in ~/.gemini/agents/ or .gemini/agents/.
+ */
+export const MemoryManagerAgent = (
+  config: Config,
+): LocalAgentDefinition<typeof MemoryManagerSchema> => {
+  const globalGeminiDir = Storage.getGlobalGeminiDir();
+
+  const getInitialContext = (): string => {
+    const memory = config.getUserMemory();
+    // Only include global and project memory — extension memory is read-only
+    // and not relevant to the memory manager.
+    const content =
+      typeof memory === 'string'
+        ? memory
+        : flattenMemory({ global: memory.global, project: memory.project });
+    if (!content.trim()) return '';
+    return `\n# Initial Context\n\n${content}\n`;
+  };
+
+  const buildSystemPrompt = (): string =>
+    `
+You are a memory management agent maintaining user memories in GEMINI.md files.
+
+# Memory Hierarchy
+
+## Global (${globalGeminiDir})
+- \`${globalGeminiDir}/GEMINI.md\` — Cross-project user preferences, key personal info,
+  and habits that apply everywhere.
+
+## Project (./)
+- \`./GEMINI.md\` — **Table of Contents** for project-specific context:
+  architecture decisions, conventions, key contacts, and references to
+  subdirectory GEMINI.md files for detailed context.
+- Subdirectory GEMINI.md files (e.g. \`src/GEMINI.md\`, \`docs/GEMINI.md\`) —
+  detailed, domain-specific context for that part of the project. Reference
+  these from the root \`./GEMINI.md\`.
+
+## Routing
+
+When adding a memory, route it to the right store:
+- **Global**: User preferences, personal info, tool aliases, cross-project habits → **global**
+- **Project Root**: Project architecture, conventions, workflows, team info → **project root**
+- **Subdirectory**: Detailed context about a specific module or directory → **subdirectory
+  GEMINI.md**, with a reference added to the project root
+
+- **Ambiguity**: If a memory (like a coding preference or workflow) could be interpreted as either a global habit or a project-specific convention, you **MUST** use \`${ASK_USER_TOOL_NAME}\` to clarify the user's intent. Do NOT make a unilateral decision when ambiguity exists between Global and Project stores.
+
+# Operations
+
+1. **Adding** — Route to the correct store and file. Check for duplicates in your provided context first.
+2. **Removing stale entries** — Delete outdated or unwanted entries. Clean up
+   dangling references.
+3. **De-duplicating** — Semantically equivalent entries should be combined. Keep the most informative version.
+4. **Organizing** — Restructure for clarity. Update references between files.
+
+# Restrictions
+- Keep GEMINI.md files lean — they are loaded into context every session.
+- Keep entries concise.
+- Edit surgically — preserve existing structure and user-authored content.
+- NEVER write or read any files other than GEMINI.md files.
+
+# Efficiency & Performance
+- **Use as few turns as possible.** Execute independent reads and writes to different files in parallel by calling multiple tools in a single turn.
+- **Do not perform any exploration of the codebase.** Try to use the provided file context and only search additional GEMINI.md files as needed to accomplish your task.
+- **Be strategic with your thinking.** carefully decide where to route memories and how to de-duplicate memories, but be decisive with simple memory writes.
+- **Minimize file system operations.** You should typically only modify the GEMINI.md files that are already provided in your context. Only read or write to other files if explicitly directed or if you are following a specific reference from an existing memory file.
+- **Context Awareness.** If a file's content is already provided in the "Initial Context" section, you do not need to call \`read_file\` for it.
+
+# Insufficient context
+If you find that you have insufficient context to read or modify the memories as described,
+reply with what you need, and exit. Do not search the codebase for the missing context.
+`.trim();
+
+  return {
+    kind: 'local',
+    name: 'save_memory',
+    displayName: 'Memory Manager',
+    description: `Writes and reads memory, preferences or facts across ALL future sessions. Use this for recurring instructions like coding styles or tool aliases.`,
+    inputConfig: {
+      inputSchema: {
+        type: 'object',
+        properties: {
+          request: {
+            type: 'string',
+            description:
+              'The memory operation to perform. Examples: "Remember that I prefer tabs over spaces", "Clean up stale memories", "De-duplicate my memories", "Organize my memories".',
+          },
+        },
+        required: ['request'],
+      },
+    },
+    outputConfig: {
+      outputName: 'result',
+      description: 'A summary of the memory operations performed.',
+      schema: MemoryManagerSchema,
+    },
+    modelConfig: {
+      model: GEMINI_MODEL_ALIAS_FLASH,
+    },
+    toolConfig: {
+      tools: [
+        READ_FILE_TOOL_NAME,
+        EDIT_TOOL_NAME,
+        WRITE_FILE_TOOL_NAME,
+        LS_TOOL_NAME,
+        GLOB_TOOL_NAME,
+        GREP_TOOL_NAME,
+        ASK_USER_TOOL_NAME,
+      ],
+    },
+    get promptConfig() {
+      return {
+        systemPrompt: buildSystemPrompt(),
+        query: `${getInitialContext()}\${request}`,
+      };
+    },
+    runConfig: {
+      maxTimeMinutes: 5,
+      maxTurns: 10,
+    },
+  };
+};

packages/core/src/agents/registry.ts

@@ -13,6 +13,7 @@ import { CodebaseInvestigatorAgent } from './codebase-investigator.js';
 import { CliHelpAgent } from './cli-help-agent.js';
 import { GeneralistAgent } from './generalist-agent.js';
 import { BrowserAgentDefinition } from './browser/browserAgentDefinition.js';
+import { MemoryManagerAgent } from './memory-manager-agent.js';
 import { A2AAuthProviderFactory } from './auth-provider/factory.js';
 import type { AuthenticationHandler } from '@a2a-js/sdk/client';
 import { type z } from 'zod';
@@ -249,6 +250,24 @@ export class AgentRegistry {
     if (browserConfig.enabled) {
       this.registerLocalAgent(BrowserAgentDefinition(this.config));
     }
+
+    // Register the memory manager agent as a replacement for the save_memory tool.
+    if (this.config.isMemoryManagerEnabled()) {
+      this.registerLocalAgent(MemoryManagerAgent(this.config));
+
+      // Ensure the global .gemini directory is accessible to tools.
+      // This allows the save_memory agent to read and write to it.
+      // Access control is enforced by the Policy Engine (memory-manager.toml).
+      try {
+        const globalDir = Storage.getGlobalGeminiDir();
+        this.config.getWorkspaceContext().addDirectory(globalDir);
+      } catch (e) {
+        debugLogger.warn(
+          `[AgentRegistry] Could not add global .gemini directory to workspace:`,
+          e,
+        );
+      }
+    }
   }
 
   private async refreshAgents(): Promise<void> {