chore: clean up launched memory features (#26941)

Co-authored-by: Jenna Inouye <jinouye@google.com>
This commit is contained in:
Sandy Tao
2026-05-13 14:22:56 -07:00
committed by GitHub
parent 0750b01fe4
commit 7504259d72
83 changed files with 617 additions and 4183 deletions
+4 -14
View File
@@ -70,7 +70,6 @@ import {
debugLogger,
coreEvents,
CoreEvent,
refreshServerHierarchicalMemory,
flattenMemory,
type MemoryChangedPayload,
writeToStdout,
@@ -1075,19 +1074,10 @@ Logging in with Google... Restarting Gemini CLI to continue.
Date.now(),
);
try {
let flattenedMemory: string;
let fileCount: number;
if (config.isJitContextEnabled()) {
await config.getMemoryContextManager()?.refresh();
config.updateSystemInstructionIfInitialized();
flattenedMemory = flattenMemory(config.getUserMemory());
fileCount = config.getGeminiMdFileCount();
} else {
const result = await refreshServerHierarchicalMemory(config);
flattenedMemory = flattenMemory(result.memoryContent);
fileCount = result.fileCount;
}
await config.getMemoryContextManager()?.refresh();
config.updateSystemInstructionIfInitialized();
const flattenedMemory = flattenMemory(config.getUserMemory());
const fileCount = config.getGeminiMdFileCount();
historyManager.addItem(
{
@@ -80,6 +80,7 @@ describe('directoryCommand', () => {
}),
getWorkingDir: () => path.resolve('/test/dir'),
shouldLoadMemoryFromIncludeDirectories: () => false,
getMemoryContextManager: vi.fn(),
getDebugMode: () => false,
getFileService: () => ({}),
getFileFilteringOptions: () => ({ ignore: [], include: [] }),
@@ -15,10 +15,7 @@ import {
type CommandContext,
} from './types.js';
import { MessageType, type HistoryItem } from '../types.js';
import {
refreshServerHierarchicalMemory,
type Config,
} from '@google/gemini-cli-core';
import { type Config } from '@google/gemini-cli-core';
import {
expandHomeDir,
getDirectorySuggestions,
@@ -47,7 +44,7 @@ async function finishAddingDirectories(
if (added.length > 0) {
try {
if (config.shouldLoadMemoryFromIncludeDirectories()) {
await refreshServerHierarchicalMemory(config);
await config.getMemoryContextManager()?.refresh();
}
addItem({
type: MessageType.INFO,
@@ -11,13 +11,10 @@ import { createMockCommandContext } from '../../test-utils/mockCommandContext.js
import { MessageType } from '../types.js';
import type { LoadedSettings } from '../../config/settings.js';
import {
type Config,
refreshMemory,
refreshServerHierarchicalMemory,
SimpleExtensionLoader,
type FileDiscoveryService,
showMemory,
addMemory,
listMemoryFiles,
flattenMemory,
} from '@google/gemini-cli-core';
@@ -32,46 +29,28 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
return String(error);
}),
refreshMemory: vi.fn(async (config) => {
if (config.isJitContextEnabled()) {
await config.getContextManager()?.refresh();
const memoryContent = original.flattenMemory(config.getUserMemory());
const fileCount = config.getGeminiMdFileCount() || 0;
return {
type: 'message',
messageType: 'info',
content: `Memory reloaded successfully. Loaded ${memoryContent.length} characters from ${fileCount} file(s).`,
};
}
await config.getMemoryContextManager()?.refresh();
const memoryContent = original.flattenMemory(config.getUserMemory());
const fileCount = config.getGeminiMdFileCount() || 0;
return {
type: 'message',
messageType: 'info',
content: 'Memory reloaded successfully.',
content: `Memory reloaded successfully. Loaded ${memoryContent.length} characters from ${fileCount} file(s).`,
};
}),
showMemory: vi.fn(),
addMemory: vi.fn(),
listMemoryFiles: vi.fn(),
refreshServerHierarchicalMemory: vi.fn(),
};
});
const mockRefreshMemory = refreshMemory as Mock;
const mockRefreshServerHierarchicalMemory =
refreshServerHierarchicalMemory as Mock;
describe('memoryCommand', () => {
let mockContext: CommandContext;
const buildMemoryCommand = (isMemoryV2 = false): SlashCommand => {
const config: Pick<Config, 'isMemoryV2Enabled'> = {
isMemoryV2Enabled: () => isMemoryV2,
};
return memoryCommand(config as Config);
};
const buildMemoryCommand = (): SlashCommand => memoryCommand(null);
const getSubCommand = (
name: 'show' | 'add' | 'reload' | 'list',
): SlashCommand => {
const getSubCommand = (name: 'show' | 'reload' | 'list'): SlashCommand => {
const subCommand = buildMemoryCommand().subCommands?.find(
(cmd) => cmd.name === name,
);
@@ -81,23 +60,11 @@ describe('memoryCommand', () => {
return subCommand;
};
describe('Memory v2', () => {
it('omits the /memory add subcommand when memoryV2 is enabled', () => {
const command = buildMemoryCommand(true);
describe('subcommands', () => {
it('does not include the legacy add subcommand', () => {
const command = buildMemoryCommand();
const names = command.subCommands?.map((cmd) => cmd.name) ?? [];
expect(names).not.toContain('add');
});
it('includes the /memory add subcommand by default', () => {
const command = buildMemoryCommand(false);
const names = command.subCommands?.map((cmd) => cmd.name) ?? [];
expect(names).toContain('add');
});
it('includes the /memory add subcommand when no config is provided', () => {
const command = memoryCommand(null);
const names = command.subCommands?.map((cmd) => cmd.name) ?? [];
expect(names).toContain('add');
expect(names).toEqual(['show', 'reload', 'list', 'inbox']);
});
});
@@ -178,63 +145,6 @@ describe('memoryCommand', () => {
});
});
describe('/memory add', () => {
let addCommand: SlashCommand;
beforeEach(() => {
addCommand = getSubCommand('add');
vi.mocked(addMemory).mockImplementation((args) => {
if (!args || args.trim() === '') {
return {
type: 'message',
messageType: 'error',
content: 'Usage: /memory add <text to remember>',
};
}
return {
type: 'tool',
toolName: 'save_memory',
toolArgs: { fact: args.trim() },
};
});
mockContext = createMockCommandContext();
});
it('should return an error message if no arguments are provided', () => {
if (!addCommand.action) throw new Error('Command has no action');
const result = addCommand.action(mockContext, ' ');
expect(result).toEqual({
type: 'message',
messageType: 'error',
content: 'Usage: /memory add <text to remember>',
});
expect(mockContext.ui.addItem).not.toHaveBeenCalled();
});
it('should return a tool action and add an info message when arguments are provided', () => {
if (!addCommand.action) throw new Error('Command has no action');
const fact = 'remember this';
const result = addCommand.action(mockContext, ` ${fact} `);
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
{
type: MessageType.INFO,
text: `Attempting to save to memory: "${fact}"`,
},
expect.any(Number),
);
expect(result).toEqual({
type: 'tool',
toolName: 'save_memory',
toolArgs: { fact },
});
});
});
describe('/memory reload', () => {
let reloadCommand: SlashCommand;
let mockSetUserMemory: Mock;
@@ -270,8 +180,7 @@ describe('memoryCommand', () => {
updateSystemInstructionIfInitialized: vi
.fn()
.mockResolvedValue(undefined),
isJitContextEnabled: vi.fn().mockReturnValue(false),
getContextManager: vi.fn().mockReturnValue({
getMemoryContextManager: vi.fn().mockReturnValue({
refresh: mockContextManagerRefresh,
}),
getUserMemory: vi.fn().mockReturnValue(''),
@@ -294,21 +203,18 @@ describe('memoryCommand', () => {
mockRefreshMemory.mockClear();
});
it('should use ContextManager.refresh when JIT is enabled', async () => {
it('should use MemoryContextManager.refresh', async () => {
if (!reloadCommand.action) throw new Error('Command has no action');
// Enable JIT in mock config
const config = mockContext.services.agentContext?.config;
if (!config) throw new Error('Config is undefined');
vi.mocked(config.isJitContextEnabled).mockReturnValue(true);
vi.mocked(config.getUserMemory).mockReturnValue('JIT Memory Content');
vi.mocked(config.getGeminiMdFileCount).mockReturnValue(3);
await reloadCommand.action(mockContext, '');
expect(mockContextManagerRefresh).toHaveBeenCalledOnce();
expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
{
@@ -319,7 +225,7 @@ describe('memoryCommand', () => {
);
});
it('should display success message when memory is reloaded with content (Legacy)', async () => {
it('should display success message when memory is reloaded with content', async () => {
if (!reloadCommand.action) throw new Error('Command has no action');
const successMessage = {
+1 -31
View File
@@ -6,7 +6,6 @@
import React from 'react';
import {
addMemory,
type Config,
listMemoryFiles,
refreshMemory,
@@ -41,30 +40,6 @@ const showSubCommand: SlashCommand = {
},
};
const addSubCommand: SlashCommand = {
name: 'add',
description: 'Add content to the memory',
kind: CommandKind.BUILT_IN,
autoExecute: false,
action: (context, args): SlashCommandActionReturn | void => {
const result = addMemory(args);
if (result.type === 'message') {
return result;
}
context.ui.addItem(
{
type: MessageType.INFO,
text: `Attempting to save to memory: "${args.trim()}"`,
},
Date.now(),
);
return result;
},
};
const reloadSubCommand: SlashCommand = {
name: 'reload',
altNames: ['refresh'],
@@ -170,14 +145,9 @@ const inboxSubCommand: SlashCommand = {
},
};
export const memoryCommand = (config: Config | null): SlashCommand => {
// The `add` subcommand depends on the `save_memory` tool, which is not
// registered when Memory v2 is enabled. Omit it in that case.
const isMemoryV2 = config?.isMemoryV2Enabled() ?? false;
export const memoryCommand = (_config: Config | null): SlashCommand => {
const subCommands: SlashCommand[] = [
showSubCommand,
...(isMemoryV2 ? [] : [addSubCommand]),
reloadSubCommand,
listSubCommand,
inboxSubCommand,
@@ -1962,8 +1962,8 @@ describe('InputPrompt', () => {
},
{
name: 'should NOT trigger completion when cursor is after space following /',
text: '/memory add',
cursor: [0, 11],
text: '/memory list',
cursor: [0, 12],
showSuggestions: false,
},
{
@@ -174,27 +174,6 @@ exports[`InputPrompt > mouse interaction > should toggle paste expansion on doub
"
`;
exports[`InputPrompt > mouse interaction > should toggle paste expansion on double-click 4`] = `
"▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
> [Pasted Text: 10 lines]
▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
"
`;
exports[`InputPrompt > mouse interaction > should toggle paste expansion on double-click 5`] = `
"▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
> [Pasted Text: 10 lines]
▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
"
`;
exports[`InputPrompt > mouse interaction > should toggle paste expansion on double-click 6`] = `
"▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
> [Pasted Text: 10 lines]
▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
"
`;
exports[`InputPrompt > multiline rendering > should correctly render multiline input including blank lines 1`] = `
"────────────────────────────────────────────────────────────────────────────────────────────────────
> hello
-1
View File
@@ -144,7 +144,6 @@ export const INFORMATIVE_TIPS = [
'Authenticate with an OAuth-enabled MCP server with /mcp auth',
'Reload MCP servers with /mcp reload',
'See the current instructional context with /memory show',
'Add content to the instructional memory with /memory add',
'Reload instructional context from GEMINI.md files with /memory reload',
'List the paths of the GEMINI.md files in use with /memory list',
'Choose your Gemini model with /model',
@@ -94,6 +94,7 @@ describe('handleAtCommand', () => {
p.startsWith(testRootDir) || p.startsWith('/private' + testRootDir),
getDirectories: () => [testRootDir],
}),
getMemoryContextManager: () => undefined,
storage: {
getProjectTempDir: () => path.join(os.tmpdir(), 'gemini-cli-temp'),
},
@@ -352,7 +352,6 @@ describe('useGeminiStream', () => {
isInteractive: () => false,
getExperiments: () => {},
getMaxSessionTurns: vi.fn(() => 100),
isJitContextEnabled: vi.fn(() => false),
getGlobalMemory: vi.fn(() => ''),
getUserMemory: vi.fn(() => ''),
getMessageBus: vi.fn(() => mockMessageBus),
@@ -1951,23 +1950,23 @@ describe('useGeminiStream', () => {
it('should schedule a tool call when the command processor returns a schedule_tool action', async () => {
const clientToolRequest: SlashCommandProcessorResult = {
type: 'schedule_tool',
toolName: 'save_memory',
toolArgs: { fact: 'test fact' },
toolName: 'activate_skill',
toolArgs: { name: 'test-skill' },
};
mockHandleSlashCommand.mockResolvedValue(clientToolRequest);
const { result } = await renderTestHook();
await act(async () => {
await result.current.submitQuery('/memory add "test fact"');
await result.current.submitQuery('/memory show');
});
await waitFor(() => {
expect(mockScheduleToolCalls).toHaveBeenCalledWith(
[
expect.objectContaining({
name: 'save_memory',
args: { fact: 'test fact' },
name: 'activate_skill',
args: { name: 'test-skill' },
isClientInitiated: true,
}),
],
@@ -2195,25 +2194,25 @@ describe('useGeminiStream', () => {
});
});
it('should NOT record other client-initiated tool calls (like save_memory) in history', async () => {
it('should NOT record other client-initiated tool calls in history', async () => {
const { result, client: mockGeminiClient } = await renderTestHook();
mockHandleSlashCommand.mockResolvedValue({
type: 'schedule_tool',
toolName: 'save_memory',
toolArgs: { fact: 'test fact' },
toolName: 'write_todos',
toolArgs: { todos: [] },
});
await act(async () => {
await result.current.submitQuery('/memory add "test fact"');
await result.current.submitQuery('/todos');
});
// Simulate tool completion
const completedTool = {
request: {
callId: 'test-call-id',
name: 'save_memory',
args: { fact: 'test fact' },
name: 'write_todos',
args: { todos: [] },
isClientInitiated: true,
},
status: CoreToolCallStatus.Success,
@@ -2227,7 +2226,7 @@ describe('useGeminiStream', () => {
responseParts: [
{
functionResponse: {
name: 'save_memory',
name: 'write_todos',
response: { success: true },
},
},
@@ -2246,91 +2245,6 @@ describe('useGeminiStream', () => {
});
});
describe('Memory Refresh on save_memory', () => {
it('should call performMemoryRefresh when a save_memory tool call completes successfully', async () => {
const mockPerformMemoryRefresh = vi.fn();
const completedToolCall: TrackedCompletedToolCall = {
request: {
callId: 'save-mem-call-1',
name: 'save_memory',
args: { fact: 'test' },
isClientInitiated: true,
prompt_id: 'prompt-id-6',
},
status: CoreToolCallStatus.Success,
responseSubmittedToGemini: false,
response: {
callId: 'save-mem-call-1',
responseParts: [{ text: 'Memory saved' }],
resultDisplay: 'Success: Memory saved',
error: undefined,
errorType: undefined, // FIX: Added missing property
},
tool: {
name: 'save_memory',
displayName: 'save_memory',
description: 'Saves memory',
build: vi.fn(),
} as unknown as AnyDeclarativeTool,
invocation: {
getDescription: () => `Mock description`,
} as unknown as AnyToolInvocation,
};
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
await renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
mockAddItem,
mockConfig,
mockLoadedSettings,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
mockPerformMemoryRefresh,
false,
() => {},
() => {},
() => {},
80,
24,
),
);
// Trigger the onComplete callback with the completed save_memory tool
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete([completedToolCall]);
}
});
await waitFor(() => {
expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1);
});
});
});
describe('Error Handling', () => {
it('should call parseAndFormatApiError with the correct authType on stream initialization failure', async () => {
// 1. Setup
+3 -22
View File
@@ -226,7 +226,7 @@ export const useGeminiStream = (
shellModeActive: boolean,
getPreferredEditor: () => EditorType | undefined,
onAuthError: (error: string) => void,
performMemoryRefresh: () => Promise<void>,
_performMemoryRefresh: () => Promise<void>,
modelSwitchedFromQuotaError: boolean,
setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
onCancelSubmit: (
@@ -266,7 +266,6 @@ export const useGeminiStream = (
useStateAndRef<Set<string>>(new Set());
const [_isFirstToolInGroup, isFirstToolInGroupRef, setIsFirstToolInGroup] =
useStateAndRef<boolean>(true);
const processedMemoryToolsRef = useRef<Set<string>>(new Set());
const { startNewPrompt, getPromptCount } = useSessionStats();
const logger = useLogger(config);
const gitService = useMemo(() => {
@@ -1897,8 +1896,8 @@ export const useGeminiStream = (
if (geminiClient) {
for (const tool of clientTools) {
// Only manually record skill activations in the chat history.
// Other client-initiated tools (like save_memory) update the system
// prompt/context and don't strictly need to be in the history.
// Other client-initiated tools update context and don't strictly
// need to be in the history.
if (tool.request.name !== ACTIVATE_SKILL_TOOL_NAME) {
continue;
}
@@ -1925,14 +1924,6 @@ export const useGeminiStream = (
}
}
// Identify new, successful save_memory calls that we haven't processed yet.
const newSuccessfulMemorySaves = completedAndReadyToSubmitTools.filter(
(t) =>
t.request.name === 'save_memory' &&
t.status === 'success' &&
!processedMemoryToolsRef.current.has(t.request.callId),
);
for (const toolCall of completedAndReadyToSubmitTools) {
const backgroundedTool = getBackgroundedToolInfo(toolCall);
if (backgroundedTool) {
@@ -1944,15 +1935,6 @@ export const useGeminiStream = (
}
}
if (newSuccessfulMemorySaves.length > 0) {
// Perform the refresh only if there are new ones.
void performMemoryRefresh();
// Mark them as processed so we don't do this again on the next render.
newSuccessfulMemorySaves.forEach((t) =>
processedMemoryToolsRef.current.add(t.request.callId),
);
}
const geminiTools = completedAndReadyToSubmitTools.filter(
(t) => !t.request.isClientInitiated,
);
@@ -2076,7 +2058,6 @@ export const useGeminiStream = (
submitQuery,
markToolsAsSubmitted,
geminiClient,
performMemoryRefresh,
modelSwitchedFromQuotaError,
addItem,
registerBackgroundTask,
@@ -80,6 +80,8 @@ describe('useIncludeDirsTrust', () => {
clearPendingIncludeDirectories: vi.fn(),
getFolderTrust: vi.fn().mockReturnValue(true),
getWorkspaceContext: () => mockWorkspaceContext,
shouldLoadMemoryFromIncludeDirectories: vi.fn().mockReturnValue(false),
getMemoryContextManager: vi.fn(),
getGeminiClient: vi
.fn()
.mockReturnValue({ addDirectoryContext: vi.fn() }),
@@ -8,10 +8,7 @@ import { useEffect } from 'react';
import { type Config } from '@google/gemini-cli-core';
import { loadTrustedFolders } from '../../config/trustedFolders.js';
import { expandHomeDir, batchAddDirectories } from '../utils/directoryUtils.js';
import {
debugLogger,
refreshServerHierarchicalMemory,
} from '@google/gemini-cli-core';
import { debugLogger } from '@google/gemini-cli-core';
import { MultiFolderTrustDialog } from '../components/MultiFolderTrustDialog.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import { MessageType, type HistoryItem } from '../types.js';
@@ -35,7 +32,7 @@ async function finishAddingDirectories(
try {
if (config.shouldLoadMemoryFromIncludeDirectories()) {
await refreshServerHierarchicalMemory(config);
await config.getMemoryContextManager()?.refresh();
}
} catch (error) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
@@ -17,11 +17,6 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
return {
...original,
homedir: () => mockHomeDir,
loadServerHierarchicalMemory: vi.fn().mockResolvedValue({
memoryContent: 'mock memory',
fileCount: 10,
filePaths: ['/a/b/c.md'],
}),
};
});