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

Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
This commit is contained in:
Allen Hutchison
2025-05-16 16:36:50 -07:00
committed by GitHub
parent d9bd2b0e14
commit 1bdec55fe1
11 changed files with 940 additions and 42 deletions

View File

@@ -0,0 +1,248 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
const { mockProcessExit } = vi.hoisted(() => ({
mockProcessExit: vi.fn((_code?: number): never => undefined as never),
}));
vi.mock('node:process', () => ({
exit: mockProcessExit,
cwd: vi.fn(() => '/mock/cwd'),
}));
vi.mock('node:fs/promises', () => ({
readFile: vi.fn(),
writeFile: vi.fn(),
mkdir: vi.fn(),
}));
import { act, renderHook } from '@testing-library/react';
import { vi, describe, it, expect, beforeEach, Mock } from 'vitest';
import { useSlashCommandProcessor } from './slashCommandProcessor.js';
import { MessageType } from '../types.js';
import * as memoryUtils from '../../config/memoryUtils.js';
import { type Config, MemoryTool } from '@gemini-code/server';
import * as fsPromises from 'node:fs/promises';
// Import the module for mocking its functions
import * as ShowMemoryCommandModule from './useShowMemoryCommand.js';
// Mock dependencies
vi.mock('./useShowMemoryCommand.js', () => ({
SHOW_MEMORY_COMMAND_NAME: '/memory show',
createShowMemoryAction: vi.fn(() => vi.fn()),
}));
// Spy on the static method we want to mock
const performAddMemoryEntrySpy = vi.spyOn(MemoryTool, 'performAddMemoryEntry');
describe('useSlashCommandProcessor', () => {
let mockAddItem: ReturnType<typeof vi.fn>;
let mockClearItems: ReturnType<typeof vi.fn>;
let mockRefreshStatic: ReturnType<typeof vi.fn>;
let mockSetShowHelp: ReturnType<typeof vi.fn>;
let mockOnDebugMessage: ReturnType<typeof vi.fn>;
let mockOpenThemeDialog: ReturnType<typeof vi.fn>;
let mockPerformMemoryRefresh: ReturnType<typeof vi.fn>;
let mockConfig: Config;
beforeEach(() => {
mockAddItem = vi.fn();
mockClearItems = vi.fn();
mockRefreshStatic = vi.fn();
mockSetShowHelp = vi.fn();
mockOnDebugMessage = vi.fn();
mockOpenThemeDialog = vi.fn();
mockPerformMemoryRefresh = vi.fn().mockResolvedValue(undefined);
mockConfig = { getDebugMode: vi.fn(() => false) } as unknown as Config;
// Clear mocks for fsPromises if they were used directly or indirectly
vi.mocked(fsPromises.readFile).mockClear();
vi.mocked(fsPromises.writeFile).mockClear();
vi.mocked(fsPromises.mkdir).mockClear();
performAddMemoryEntrySpy.mockReset(); // Reset the spy
vi.spyOn(memoryUtils, 'deleteLastMemoryEntry').mockImplementation(vi.fn());
vi.spyOn(memoryUtils, 'deleteAllAddedMemoryEntries').mockImplementation(
vi.fn(),
);
vi.mocked(memoryUtils.deleteLastMemoryEntry).mockClear();
vi.mocked(memoryUtils.deleteAllAddedMemoryEntries).mockClear();
mockProcessExit.mockClear();
(ShowMemoryCommandModule.createShowMemoryAction as Mock).mockClear();
mockPerformMemoryRefresh.mockClear();
});
const getProcessor = () => {
const { result } = renderHook(() =>
useSlashCommandProcessor(
mockConfig,
mockAddItem,
mockClearItems,
mockRefreshStatic,
mockSetShowHelp,
mockOnDebugMessage,
mockOpenThemeDialog,
mockPerformMemoryRefresh,
),
);
return result.current;
};
describe('/memory add', () => {
it('should call MemoryTool.performAddMemoryEntry and refresh on valid input', async () => {
performAddMemoryEntrySpy.mockResolvedValue(undefined);
const { handleSlashCommand } = getProcessor();
const fact = 'Remember this fact';
await act(async () => {
handleSlashCommand(`/memory add ${fact}`);
});
expect(mockAddItem).toHaveBeenNthCalledWith(
1,
expect.objectContaining({
type: MessageType.USER,
text: `/memory add ${fact}`,
}),
expect.any(Number),
);
expect(performAddMemoryEntrySpy).toHaveBeenCalledWith(
fact,
memoryUtils.getGlobalMemoryFilePath(), // Ensure this path is correct
{
readFile: fsPromises.readFile,
writeFile: fsPromises.writeFile,
mkdir: fsPromises.mkdir,
},
);
expect(mockPerformMemoryRefresh).toHaveBeenCalled();
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.INFO,
text: `Successfully added to memory: "${fact}"`,
}),
expect.any(Number),
);
});
it('should show usage error if no text is provided', async () => {
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/memory add ');
});
expect(performAddMemoryEntrySpy).not.toHaveBeenCalled();
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.ERROR,
text: 'Usage: /memory add <text to remember>',
}),
expect.any(Number),
);
});
it('should handle error from MemoryTool.performAddMemoryEntry', async () => {
const fact = 'Another fact';
performAddMemoryEntrySpy.mockRejectedValue(
new Error('[MemoryTool] Failed to add memory entry: Disk full'),
);
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand(`/memory add ${fact}`);
});
expect(performAddMemoryEntrySpy).toHaveBeenCalledWith(
fact,
memoryUtils.getGlobalMemoryFilePath(),
{
readFile: fsPromises.readFile,
writeFile: fsPromises.writeFile,
mkdir: fsPromises.mkdir,
},
);
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.ERROR,
text: 'Failed to add memory: [MemoryTool] Failed to add memory entry: Disk full',
}),
expect.any(Number),
);
});
});
describe('/memory show', () => {
it('should call the showMemoryAction', async () => {
const mockReturnedShowAction = vi.fn();
vi.mocked(ShowMemoryCommandModule.createShowMemoryAction).mockReturnValue(
mockReturnedShowAction,
);
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/memory show');
});
expect(
ShowMemoryCommandModule.createShowMemoryAction,
).toHaveBeenCalledWith(mockConfig, expect.any(Function));
expect(mockReturnedShowAction).toHaveBeenCalled();
});
});
describe('/memory refresh', () => {
it('should call performMemoryRefresh', async () => {
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/memory refresh');
});
expect(mockPerformMemoryRefresh).toHaveBeenCalled();
});
});
describe('Unknown /memory subcommand', () => {
it('should show an error for unknown /memory subcommand', async () => {
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/memory foobar');
});
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.ERROR,
text: 'Unknown /memory command: foobar. Available: show, refresh, add',
}),
expect.any(Number),
);
});
});
describe('Other commands', () => {
it('/help should open help', async () => {
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/help');
});
expect(mockSetShowHelp).toHaveBeenCalledWith(true);
});
});
describe('Unknown command', () => {
it('should show an error for a general unknown command', async () => {
const { handleSlashCommand } = getProcessor();
await act(async () => {
handleSlashCommand('/unknowncommand');
});
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.ERROR,
text: 'Unknown command: /unknowncommand',
}),
expect.any(Number),
);
});
});
});

View File

@@ -6,22 +6,17 @@
import { useCallback, useMemo } from 'react';
import { type PartListUnion } from '@google/genai';
import { getCommandFromQuery } from '../utils/commandUtils.js';
import { UseHistoryManagerReturn } from './useHistoryManager.js';
import { Config } from '@gemini-code/server'; // Import Config
import { Message, MessageType, HistoryItemWithoutId } from '../types.js'; // Import Message types
import {
createShowMemoryAction,
SHOW_MEMORY_COMMAND_NAME,
} from './useShowMemoryCommand.js';
import { REFRESH_MEMORY_COMMAND_NAME } from './useRefreshMemoryCommand.js'; // Only import name now
import process from 'node:process'; // For process.exit
import { Config } from '@gemini-code/server';
import { Message, MessageType, HistoryItemWithoutId } from '../types.js';
import { createShowMemoryAction } from './useShowMemoryCommand.js';
import { addMemoryEntry } from '../../config/memoryUtils.js';
export interface SlashCommand {
name: string;
altName?: string;
description: string;
action: (value: PartListUnion | string) => void; // Allow string for simpler actions
action: (mainCommand: string, subCommand?: string, args?: string) => void;
}
/**
@@ -54,13 +49,43 @@ export const useSlashCommandProcessor = (
await actionFn();
}, [config, addMessage]);
const addMemoryAction = useCallback(
async (_mainCommand: string, _subCommand?: string, args?: string) => {
if (!args || args.trim() === '') {
addMessage({
type: MessageType.ERROR,
content: 'Usage: /memory add <text to remember>',
timestamp: new Date(),
});
return;
}
try {
await addMemoryEntry(args);
addMessage({
type: MessageType.INFO,
content: `Successfully added to memory: "${args}"`,
timestamp: new Date(),
});
await performMemoryRefresh(); // Refresh memory to reflect changes
} catch (e) {
const errorMessage = e instanceof Error ? e.message : String(e);
addMessage({
type: MessageType.ERROR,
content: `Failed to add memory: ${errorMessage}`,
timestamp: new Date(),
});
}
},
[addMessage, performMemoryRefresh],
);
const slashCommands: SlashCommand[] = useMemo(
() => [
{
name: 'help',
altName: '?',
description: 'for help on gemini-code',
action: (_value: PartListUnion | string) => {
action: (_mainCommand, _subCommand, _args) => {
onDebugMessage('Opening help.');
setShowHelp(true);
},
@@ -68,7 +93,7 @@ export const useSlashCommandProcessor = (
{
name: 'clear',
description: 'clear the screen',
action: (_value: PartListUnion | string) => {
action: (_mainCommand, _subCommand, _args) => {
onDebugMessage('Clearing terminal.');
clearItems();
console.clear();
@@ -78,25 +103,39 @@ export const useSlashCommandProcessor = (
{
name: 'theme',
description: 'change the theme',
action: (_value) => {
action: (_mainCommand, _subCommand, _args) => {
openThemeDialog();
},
},
{
name: REFRESH_MEMORY_COMMAND_NAME.substring(1), // Remove leading '/'
description: 'Reloads instructions from all GEMINI.md files.',
action: performMemoryRefresh, // Use the passed in function
},
{
name: SHOW_MEMORY_COMMAND_NAME.substring(1), // Remove leading '/'
description: 'Displays the current hierarchical memory content.',
action: showMemoryAction,
name: 'memory',
description:
'Manage memory. Usage: /memory <show|refresh|add|delete_last|delete_all_added> [text for add]',
action: (mainCommand, subCommand, args) => {
switch (subCommand) {
case 'show':
showMemoryAction();
break;
case 'refresh':
performMemoryRefresh();
break;
case 'add':
addMemoryAction(mainCommand, subCommand, args);
break;
default:
addMessage({
type: MessageType.ERROR,
content: `Unknown /memory command: ${subCommand}. Available: show, refresh, add`,
timestamp: new Date(),
});
}
},
},
{
name: 'quit',
altName: 'exit',
description: 'exit the cli',
action: (_value: PartListUnion | string) => {
action: (_mainCommand, _subCommand, _args) => {
onDebugMessage('Quitting. Good-bye.');
process.exit(0);
},
@@ -108,8 +147,10 @@ export const useSlashCommandProcessor = (
refreshStatic,
openThemeDialog,
clearItems,
performMemoryRefresh, // Add to dependencies
performMemoryRefresh,
showMemoryAction,
addMemoryAction,
addMessage,
],
);
@@ -120,36 +161,51 @@ export const useSlashCommandProcessor = (
}
const trimmed = rawQuery.trim();
const [symbol, test] = getCommandFromQuery(trimmed);
if (symbol !== '/' && symbol !== '?') {
if (!trimmed.startsWith('/') && !trimmed.startsWith('?')) {
return false;
}
const userMessageTimestamp = Date.now();
// Add user message to history only if it's not a silent command or handled internally
// For now, adding all slash commands to history for transparency.
addItem({ type: MessageType.USER, text: trimmed }, userMessageTimestamp);
let subCommand: string | undefined;
let args: string | undefined;
const commandToMatch = (() => {
if (trimmed.startsWith('?')) {
return 'help'; // No subCommand or args for '?' acting as help
}
// For other slash commands like /memory add foo
const parts = trimmed.substring(1).trim().split(/\s+/);
if (parts.length > 1) {
subCommand = parts[1];
}
if (parts.length > 2) {
args = parts.slice(2).join(' ');
}
return parts[0]; // This is the main command name
})();
const mainCommand = commandToMatch;
for (const cmd of slashCommands) {
if (
test === cmd.name ||
test === cmd.altName ||
(symbol === '?' && cmd.altName === '?') // Special handling for ? as help
) {
cmd.action(trimmed); // Pass the full trimmed command for context if needed
if (mainCommand === cmd.name || mainCommand === cmd.altName) {
cmd.action(mainCommand, subCommand, args);
return true;
}
}
addItem(
{ type: MessageType.ERROR, text: `Unknown command: ${trimmed}` },
userMessageTimestamp,
);
addMessage({
type: MessageType.ERROR,
content: `Unknown command: ${trimmed}`,
timestamp: new Date(),
});
return true;
},
[addItem, slashCommands],
[addItem, slashCommands, addMessage],
);
return { handleSlashCommand, slashCommands };