/** * @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; writeFile: (path: string, data: string, encoding: 'utf-8') => Promise; mkdir: ( path: string, options: { recursive: boolean }, ) => Promise; } describe('MemoryTool', () => { const mockAbortSignal = new AbortController().signal; const mockFsAdapter: { readFile: Mock; writeFile: Mock; mkdir: Mock; } = { 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; 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('save_memory'); 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('save_memory'); 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}`, ); }); }); });