chore: clean up launched memory features (#26941)

Co-authored-by: Jenna Inouye <jinouye@google.com>

packages/cli/src/ui/hooks/atCommandProcessor.test.ts

@@ -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'),
       },

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -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

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -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,

packages/cli/src/ui/hooks/useIncludeDirsTrust.test.tsx

@@ -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() }),

packages/cli/src/ui/hooks/useIncludeDirsTrust.tsx

@@ -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