feat: Implement CLI and model memory management (#371)

Co-authored-by: N. Taylor Mullen <[email protected]>

1 files changed, 224 insertions, 0 deletions

diff --git a/packages/server/src/tools/memoryTool.test.ts b/packages/server/src/tools/memoryTool.test.ts
new file mode 100644
index 00000000..efbbb025
--- /dev/null
+++ b/packages/server/src/tools/memoryTool.test.ts
@@ -0,0 +1,224 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi, describe, it, expect, beforeEach, afterEach, Mock } from 'vitest';
+import { MemoryTool } from './memoryTool.js';
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import * as os from 'os';
+
+// Mock dependencies
+vi.mock('fs/promises');
+vi.mock('os');
+
+const MEMORY_SECTION_HEADER = '## Gemini Added Memories';
+
+// Define a type for our fsAdapter to ensure consistency
+interface FsAdapter {
+  readFile: (path: string, encoding: 'utf-8') => Promise<string>;
+  writeFile: (path: string, data: string, encoding: 'utf-8') => Promise<void>;
+  mkdir: (
+    path: string,
+    options: { recursive: boolean },
+  ) => Promise<string | undefined>;
+}
+
+describe('MemoryTool', () => {
+  const mockAbortSignal = new AbortController().signal;
+
+  const mockFsAdapter: {
+    readFile: Mock<FsAdapter['readFile']>;
+    writeFile: Mock<FsAdapter['writeFile']>;
+    mkdir: Mock<FsAdapter['mkdir']>;
+  } = {
+    readFile: vi.fn(),
+    writeFile: vi.fn(),
+    mkdir: vi.fn(),
+  };
+
+  beforeEach(() => {
+    vi.mocked(os.homedir).mockReturnValue('/mock/home');
+    mockFsAdapter.readFile.mockReset();
+    mockFsAdapter.writeFile.mockReset().mockResolvedValue(undefined);
+    mockFsAdapter.mkdir
+      .mockReset()
+      .mockResolvedValue(undefined as string | undefined);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe('performAddMemoryEntry (static method)', () => {
+    const testFilePath = path.join('/mock/home', '.gemini', 'GEMINI.md');
+
+    it('should create section and save a fact if file does not exist', async () => {
+      mockFsAdapter.readFile.mockRejectedValue({ code: 'ENOENT' }); // Simulate file not found
+      const fact = 'The sky is blue';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+
+      expect(mockFsAdapter.mkdir).toHaveBeenCalledWith(
+        path.dirname(testFilePath),
+        {
+          recursive: true,
+        },
+      );
+      expect(mockFsAdapter.writeFile).toHaveBeenCalledOnce();
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      expect(writeFileCall[0]).toBe(testFilePath);
+      const expectedContent = `${MEMORY_SECTION_HEADER}\n- ${fact}\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+      expect(writeFileCall[2]).toBe('utf-8');
+    });
+
+    it('should create section and save a fact if file is empty', async () => {
+      mockFsAdapter.readFile.mockResolvedValue(''); // Simulate empty file
+      const fact = 'The sky is blue';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      const expectedContent = `${MEMORY_SECTION_HEADER}\n- ${fact}\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+    });
+
+    it('should add a fact to an existing section', async () => {
+      const initialContent = `Some preamble.\n\n${MEMORY_SECTION_HEADER}\n- Existing fact 1\n`;
+      mockFsAdapter.readFile.mockResolvedValue(initialContent);
+      const fact = 'New fact 2';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+
+      expect(mockFsAdapter.writeFile).toHaveBeenCalledOnce();
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      const expectedContent = `Some preamble.\n\n${MEMORY_SECTION_HEADER}\n- Existing fact 1\n- ${fact}\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+    });
+
+    it('should add a fact to an existing empty section', async () => {
+      const initialContent = `Some preamble.\n\n${MEMORY_SECTION_HEADER}\n`; // Empty section
+      mockFsAdapter.readFile.mockResolvedValue(initialContent);
+      const fact = 'First fact in section';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+
+      expect(mockFsAdapter.writeFile).toHaveBeenCalledOnce();
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      const expectedContent = `Some preamble.\n\n${MEMORY_SECTION_HEADER}\n- ${fact}\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+    });
+
+    it('should add a fact when other ## sections exist and preserve spacing', async () => {
+      const initialContent = `${MEMORY_SECTION_HEADER}\n- Fact 1\n\n## Another Section\nSome other text.`;
+      mockFsAdapter.readFile.mockResolvedValue(initialContent);
+      const fact = 'Fact 2';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+
+      expect(mockFsAdapter.writeFile).toHaveBeenCalledOnce();
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      // Note: The implementation ensures a single newline at the end if content exists.
+      const expectedContent = `${MEMORY_SECTION_HEADER}\n- Fact 1\n- ${fact}\n\n## Another Section\nSome other text.\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+    });
+
+    it('should correctly trim and add a fact that starts with a dash', async () => {
+      mockFsAdapter.readFile.mockResolvedValue(`${MEMORY_SECTION_HEADER}\n`);
+      const fact = '- - My fact with dashes';
+      await MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter);
+      const writeFileCall = mockFsAdapter.writeFile.mock.calls[0];
+      const expectedContent = `${MEMORY_SECTION_HEADER}\n- My fact with dashes\n`;
+      expect(writeFileCall[1]).toBe(expectedContent);
+    });
+
+    it('should handle error from fsAdapter.writeFile', async () => {
+      mockFsAdapter.readFile.mockResolvedValue('');
+      mockFsAdapter.writeFile.mockRejectedValue(new Error('Disk full'));
+      const fact = 'This will fail';
+      await expect(
+        MemoryTool.performAddMemoryEntry(fact, testFilePath, mockFsAdapter),
+      ).rejects.toThrow('[MemoryTool] Failed to add memory entry: Disk full');
+    });
+  });
+
+  describe('execute (instance method)', () => {
+    let memoryTool: MemoryTool;
+    let performAddMemoryEntrySpy: Mock<typeof MemoryTool.performAddMemoryEntry>;
+
+    beforeEach(() => {
+      memoryTool = new MemoryTool();
+      // Spy on the static method for these tests
+      performAddMemoryEntrySpy = vi
+        .spyOn(MemoryTool, 'performAddMemoryEntry')
+        .mockResolvedValue(undefined) as Mock<
+        typeof MemoryTool.performAddMemoryEntry
+      >;
+      // Cast needed as spyOn returns MockInstance
+    });
+
+    it('should have correct name, displayName, description, and schema', () => {
+      expect(memoryTool.name).toBe('saveMemory');
+      expect(memoryTool.displayName).toBe('Save Memory');
+      expect(memoryTool.description).toContain(
+        'Saves a specific piece of information',
+      );
+      expect(memoryTool.schema).toBeDefined();
+      expect(memoryTool.schema.name).toBe('saveMemory');
+      expect(memoryTool.schema.parameters?.properties?.fact).toBeDefined();
+    });
+
+    it('should call performAddMemoryEntry with correct parameters and return success', async () => {
+      const params = { fact: 'The sky is blue' };
+      const result = await memoryTool.execute(params, mockAbortSignal);
+      const expectedFilePath = path.join('/mock/home', '.gemini', 'GEMINI.md');
+
+      // For this test, we expect the actual fs methods to be passed
+      const expectedFsArgument = {
+        readFile: fs.readFile,
+        writeFile: fs.writeFile,
+        mkdir: fs.mkdir,
+      };
+
+      expect(performAddMemoryEntrySpy).toHaveBeenCalledWith(
+        params.fact,
+        expectedFilePath,
+        expectedFsArgument,
+      );
+      const successMessage = `Okay, I've remembered that: "${params.fact}"`;
+      expect(result.llmContent).toBe(
+        JSON.stringify({ success: true, message: successMessage }),
+      );
+      expect(result.returnDisplay).toBe(successMessage);
+    });
+
+    it('should return an error if fact is empty', async () => {
+      const params = { fact: ' ' }; // Empty fact
+      const result = await memoryTool.execute(params, mockAbortSignal);
+      const errorMessage = 'Parameter "fact" must be a non-empty string.';
+
+      expect(performAddMemoryEntrySpy).not.toHaveBeenCalled();
+      expect(result.llmContent).toBe(
+        JSON.stringify({ success: false, error: errorMessage }),
+      );
+      expect(result.returnDisplay).toBe(`Error: ${errorMessage}`);
+    });
+
+    it('should handle errors from performAddMemoryEntry', async () => {
+      const params = { fact: 'This will fail' };
+      const underlyingError = new Error(
+        '[MemoryTool] Failed to add memory entry: Disk full',
+      );
+      performAddMemoryEntrySpy.mockRejectedValue(underlyingError);
+
+      const result = await memoryTool.execute(params, mockAbortSignal);
+
+      expect(result.llmContent).toBe(
+        JSON.stringify({
+          success: false,
+          error: `Failed to save memory. Detail: ${underlyingError.message}`,
+        }),
+      );
+      expect(result.returnDisplay).toBe(
+        `Error saving memory: ${underlyingError.message}`,
+      );
+    });
+  });
+});