Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions packages/cli/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ import { hooksCommand } from '../commands/hooks.js';
import { gemmaCommand } from '../commands/gemma.js';
import {
setGeminiMdFilename as setServerGeminiMdFilename,
getCurrentGeminiMdFilename,
resetGeminiMdFilename,
Comment thread
devr0306 marked this conversation as resolved.
DEFAULT_CONTEXT_FILENAME,
ApprovalMode,
DEFAULT_GEMINI_EMBEDDING_MODEL,
DEFAULT_FILE_FILTERING_OPTIONS,
Expand Down Expand Up @@ -619,7 +620,7 @@ export async function loadCliConfig(
setServerGeminiMdFilename(settings.context.fileName);
} else {
// Reset to default if not provided in settings.
setServerGeminiMdFilename(getCurrentGeminiMdFilename());
resetGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
}

const fileService = new FileDiscoveryService(cwd);
Expand Down
74 changes: 57 additions & 17 deletions packages/core/src/tools/memoryTool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
import {
MemoryTool,
setGeminiMdFilename,
resetGeminiMdFilename,
getCurrentGeminiMdFilename,
getAllGeminiMdFilenames,
DEFAULT_CONTEXT_FILENAME,
Expand Down Expand Up @@ -45,14 +46,18 @@ vi.mock('node:fs/promises', async (importOriginal) => {
};
});

vi.mock('fs', () => ({
mkdirSync: vi.fn(),
createWriteStream: vi.fn(() => ({
on: vi.fn(),
write: vi.fn(),
end: vi.fn(),
})),
}));
vi.mock('fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('fs')>();
return {
...actual,
mkdirSync: vi.fn(),
createWriteStream: vi.fn(() => ({
on: vi.fn(),
write: vi.fn(),
end: vi.fn(),
})),
};
});

vi.mock('os');

Expand All @@ -77,30 +82,65 @@ describe('MemoryTool', () => {

afterEach(() => {
vi.restoreAllMocks();
setGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
resetGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
});

describe('setGeminiMdFilename', () => {
it('should update currentGeminiMdFilename when a valid new name is provided', () => {
it('should append to currentGeminiMdFilename when a valid new name is provided', () => {
const newName = 'CUSTOM_CONTEXT.md';
setGeminiMdFilename(newName);
expect(getCurrentGeminiMdFilename()).toBe(newName);
expect(getAllGeminiMdFilenames()).toEqual([
newName,
DEFAULT_CONTEXT_FILENAME,
]);
});

it('should not update currentGeminiMdFilename if the new name is empty or whitespace', () => {
const initialName = getCurrentGeminiMdFilename();
const initialNames = getAllGeminiMdFilenames();
setGeminiMdFilename(' ');
expect(getCurrentGeminiMdFilename()).toBe(initialName);
expect(getAllGeminiMdFilenames()).toEqual(initialNames);

setGeminiMdFilename('');
expect(getCurrentGeminiMdFilename()).toBe(initialName);
expect(getAllGeminiMdFilenames()).toEqual(initialNames);
});

it('should handle an array of filenames', () => {
it('should handle adding an array of filenames', () => {
const newNames = ['CUSTOM_CONTEXT.md', 'ANOTHER_CONTEXT.md'];
setGeminiMdFilename(newNames);
expect(getCurrentGeminiMdFilename()).toBe('CUSTOM_CONTEXT.md');
expect(getAllGeminiMdFilenames()).toEqual(newNames);
expect(getAllGeminiMdFilenames()).toEqual([
...newNames,
DEFAULT_CONTEXT_FILENAME,
]);
});

it('should ensure uniqueness when adding names', () => {
setGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
expect(getAllGeminiMdFilenames()).toEqual([DEFAULT_CONTEXT_FILENAME]);

setGeminiMdFilename(['NEW.md', 'NEW.md']);
expect(getAllGeminiMdFilenames()).toEqual([
'NEW.md',
DEFAULT_CONTEXT_FILENAME,
]);
});
});

describe('resetGeminiMdFilename', () => {
it('should replace all filenames with the provided one', () => {
setGeminiMdFilename('OTHER.md');
resetGeminiMdFilename('RESET.md');
expect(getAllGeminiMdFilenames()).toEqual(['RESET.md']);
});

it('should reset to default if no argument provided', () => {
resetGeminiMdFilename('OTHER.md');
resetGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
expect(getAllGeminiMdFilenames()).toEqual([DEFAULT_CONTEXT_FILENAME]);
});

it('should handle array reset', () => {
resetGeminiMdFilename(['A.md', 'B.md']);
expect(getAllGeminiMdFilenames()).toEqual(['A.md', 'B.md']);
});
});

Expand Down
64 changes: 56 additions & 8 deletions packages/core/src/tools/memoryTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ import * as path from 'node:path';
import { Storage } from '../config/storage.js';
import * as Diff from 'diff';
import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js';
import { tildeifyPath } from '../utils/paths.js';
import { resolveToRealPath, tildeifyPath } from '../utils/paths.js';
import type {
ModifiableDeclarativeTool,
ModifyContext,
Expand All @@ -33,17 +33,65 @@ export const DEFAULT_CONTEXT_FILENAME = 'GEMINI.md';
export const MEMORY_SECTION_HEADER = '## Gemini Added Memories';
export const PROJECT_MEMORY_INDEX_FILENAME = 'MEMORY.md';

// This variable will hold the currently configured filename for GEMINI.md context files.
// It defaults to DEFAULT_CONTEXT_FILENAME but can be overridden by setGeminiMdFilename.
// This variable will hold the currently configured filenames for GEMINI.md context files.
// It defaults to DEFAULT_CONTEXT_FILENAME but can be extended by setGeminiMdFilename.
let currentGeminiMdFilename: string | string[] = DEFAULT_CONTEXT_FILENAME;
Comment thread
devr0306 marked this conversation as resolved.
Comment thread
devr0306 marked this conversation as resolved.

/**
* Adds one or more filenames to the current context filenames.
* Ensures uniqueness and maintains order.
*/
export function setGeminiMdFilename(newFilename: string | string[]): void {
if (Array.isArray(newFilename)) {
if (newFilename.length > 0) {
currentGeminiMdFilename = newFilename.map((name) => name.trim());
const filenames = Array.isArray(newFilename) ? newFilename : [newFilename];
const current = getAllGeminiMdFilenames();
const next = new Set<string>();

for (const filename of filenames) {
const trimmed = filename.trim();
if (trimmed !== '') {
const normalized = path.normalize(trimmed);
// Sanitize to prevent path traversal while allowing subdirectories
const validatedPath = resolveToRealPath(normalized);
if (validatedPath) {
next.add(normalized);
}
}
} else if (newFilename && newFilename.trim() !== '') {
currentGeminiMdFilename = newFilename.trim();
}

for (const filename of current) {
next.add(filename);
}

const result = Array.from(next);
if (result.length > 1) {
currentGeminiMdFilename = result;
} else if (result.length === 1) {
currentGeminiMdFilename = result[0];
}
}
Comment thread
devr0306 marked this conversation as resolved.
Comment thread
devr0306 marked this conversation as resolved.
Comment thread
devr0306 marked this conversation as resolved.

/**
* Resets the context filenames to the provided value, or the default if none provided.
* This replaces all current filenames.
*/
export function resetGeminiMdFilename(
filename: string | string[] = DEFAULT_CONTEXT_FILENAME,
): void {
const filenames = Array.isArray(filename) ? filename : [filename];
const cleaned = Array.from(
new Set(
filenames
.map((f) => path.normalize(f.trim()))
.filter((f) => !!resolveToRealPath(f)),
),
);

if (cleaned.length === 0) {
currentGeminiMdFilename = DEFAULT_CONTEXT_FILENAME;
} else if (cleaned.length === 1) {
currentGeminiMdFilename = cleaned[0];
} else {
currentGeminiMdFilename = cleaned;
}
}
Comment thread
devr0306 marked this conversation as resolved.
Comment thread
devr0306 marked this conversation as resolved.
Comment thread
devr0306 marked this conversation as resolved.

Expand Down
Loading