Refac: Centralize storage file management (#4078)

Co-authored-by: Taylor Mullen <ntaylormullen@google.com>

packages/cli/src/ui/App.tsx

@@ -742,7 +742,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     }
   }, [config, config.getGeminiMdFileCount]);
 
-  const logger = useLogger();
+  const logger = useLogger(config.storage);
 
   useEffect(() => {
     const fetchUserMessages = async () => {

packages/cli/src/ui/commands/chatCommand.test.ts

@@ -67,11 +67,14 @@ describe('chatCommand', () => {
     mockContext = createMockCommandContext({
       services: {
         config: {
-          getProjectTempDir: () => '/tmp/gemini',
+          getProjectRoot: () => '/project/root',
           getGeminiClient: () =>
             ({
               getChat: mockGetChat,
             }) as unknown as GeminiClient,
+          storage: {
+            getProjectTempDir: () => '/project/root/.gemini/tmp/mockhash',
+          },
         },
         logger: {
           saveCheckpoint: mockSaveCheckpoint,

packages/cli/src/ui/commands/chatCommand.ts

@@ -28,7 +28,8 @@ const getSavedChatTags = async (
   context: CommandContext,
   mtSortDesc: boolean,
 ): Promise<ChatDetail[]> => {
-  const geminiDir = context.services.config?.getProjectTempDir();
+  const cfg = context.services.config;
+  const geminiDir = cfg?.storage?.getProjectTempDir();
   if (!geminiDir) {
     return [];
   }

packages/cli/src/ui/commands/ideCommand.test.ts

@@ -20,7 +20,14 @@ import * as core from '@google/gemini-cli-core';
 
 vi.mock('child_process');
 vi.mock('glob');
-vi.mock('@google/gemini-cli-core');
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const original = await importOriginal<typeof core>();
+  return {
+    ...original,
+    getOauthClient: vi.fn(original.getOauthClient),
+    getIdeInstaller: vi.fn(original.getIdeInstaller),
+  };
+});
 
 describe('ideCommand', () => {
   let mockContext: CommandContext;

packages/cli/src/ui/commands/restoreCommand.test.ts

@@ -39,7 +39,10 @@ describe('restoreCommand', () => {
 
     mockConfig = {
       getCheckpointingEnabled: vi.fn().mockReturnValue(true),
-      getProjectTempDir: vi.fn().mockReturnValue(geminiTempDir),
+      storage: {
+        getProjectTempCheckpointsDir: vi.fn().mockReturnValue(checkpointsDir),
+        getProjectTempDir: vi.fn().mockReturnValue(geminiTempDir),
+      },
       getGeminiClient: vi.fn().mockReturnValue({
         setHistory: mockSetHistory,
       }),
@@ -77,7 +80,9 @@ describe('restoreCommand', () => {
 
   describe('action', () => {
     it('should return an error if temp dir is not found', async () => {
-      vi.mocked(mockConfig.getProjectTempDir).mockReturnValue('');
+      vi.mocked(
+        mockConfig.storage.getProjectTempCheckpointsDir,
+      ).mockReturnValue('');
 
       expect(
         await restoreCommand(mockConfig)?.action?.(mockContext, ''),
@@ -219,7 +224,7 @@ describe('restoreCommand', () => {
 
   describe('completion', () => {
     it('should return an empty array if temp dir is not found', async () => {
-      vi.mocked(mockConfig.getProjectTempDir).mockReturnValue('');
+      vi.mocked(mockConfig.storage.getProjectTempDir).mockReturnValue('');
       const command = restoreCommand(mockConfig);
 
       expect(await command?.completion?.(mockContext, '')).toEqual([]);

packages/cli/src/ui/commands/restoreCommand.ts

@@ -22,9 +22,7 @@ async function restoreAction(
   const { config, git: gitService } = services;
   const { addItem, loadHistory } = ui;
 
-  const checkpointDir = config?.getProjectTempDir()
-    ? path.join(config.getProjectTempDir(), 'checkpoints')
-    : undefined;
+  const checkpointDir = config?.storage.getProjectTempCheckpointsDir();
 
   if (!checkpointDir) {
     return {
@@ -125,9 +123,7 @@ async function completion(
 ): Promise<string[]> {
   const { services } = context;
   const { config } = services;
-  const checkpointDir = config?.getProjectTempDir()
-    ? path.join(config.getProjectTempDir(), 'checkpoints')
-    : undefined;
+  const checkpointDir = config?.storage.getProjectTempCheckpointsDir();
   if (!checkpointDir) {
     return [];
   }

packages/cli/src/ui/components/InputPrompt.tsx

@@ -81,7 +81,7 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
   const [cursorPosition, setCursorPosition] = useState<[number, number]>([
     0, 0,
   ]);
-  const shellHistory = useShellHistory(config.getProjectRoot());
+  const shellHistory = useShellHistory(config.getProjectRoot(), config.storage);
   const historyData = shellHistory.history;
 
   const completion = useCommandCompletion(

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

@@ -17,15 +17,10 @@ import {
 
 const mockIsBinary = vi.hoisted(() => vi.fn());
 const mockShellExecutionService = vi.hoisted(() => vi.fn());
-vi.mock('@google/gemini-cli-core', async (importOriginal) => {
-  const original =
-    await importOriginal<typeof import('@google/gemini-cli-core')>();
-  return {
-    ...original,
-    ShellExecutionService: { execute: mockShellExecutionService },
-    isBinary: mockIsBinary,
-  };
-});
+vi.mock('@google/gemini-cli-core', () => ({
+  ShellExecutionService: { execute: mockShellExecutionService },
+  isBinary: mockIsBinary,
+}));
 vi.mock('fs');
 vi.mock('os');
 vi.mock('crypto');

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

@@ -16,6 +16,7 @@ import {
   makeSlashCommandEvent,
   SlashCommandStatus,
   ToolConfirmationOutcome,
+  Storage,
 } from '@google/gemini-cli-core';
 import { useSessionStats } from '../contexts/SessionContext.js';
 import { runExitCleanup } from '../../utils/cleanup.js';
@@ -82,11 +83,14 @@ export const useSlashCommandProcessor = (
     if (!config?.getProjectRoot()) {
       return;
     }
-    return new GitService(config.getProjectRoot());
+    return new GitService(config.getProjectRoot(), config.storage);
   }, [config]);
 
   const logger = useMemo(() => {
-    const l = new Logger(config?.getSessionId() || '');
+    const l = new Logger(
+      config?.getSessionId() || '',
+      config?.storage ?? new Storage(process.cwd()),
+    );
     // The logger's initialize is async, but we can create the instance
     // synchronously. Commands that use it will await its initialization.
     return l;

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

@@ -105,13 +105,14 @@ export const useGeminiStream = (
     useStateAndRef<HistoryItemWithoutId | null>(null);
   const processedMemoryToolsRef = useRef<Set<string>>(new Set());
   const { startNewPrompt, getPromptCount } = useSessionStats();
-  const logger = useLogger();
+  const storage = config.storage;
+  const logger = useLogger(storage);
   const gitService = useMemo(() => {
     if (!config.getProjectRoot()) {
       return;
     }
-    return new GitService(config.getProjectRoot());
-  }, [config]);
+    return new GitService(config.getProjectRoot(), storage);
+  }, [config, storage]);
 
   const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] =
     useReactToolScheduler(
@@ -877,9 +878,7 @@ export const useGeminiStream = (
       );
 
       if (restorableToolCalls.length > 0) {
-        const checkpointDir = config.getProjectTempDir()
-          ? path.join(config.getProjectTempDir(), 'checkpoints')
-          : undefined;
+        const checkpointDir = storage.getProjectTempCheckpointsDir();
 
         if (!checkpointDir) {
           return;
@@ -962,7 +961,15 @@ export const useGeminiStream = (
       }
     };
     saveRestorableToolCalls();
-  }, [toolCalls, config, onDebugMessage, gitService, history, geminiClient]);
+  }, [
+    toolCalls,
+    config,
+    onDebugMessage,
+    gitService,
+    history,
+    geminiClient,
+    storage,
+  ]);
 
   return {
     streamingState,

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

@@ -5,16 +5,16 @@
  */
 
 import { useState, useEffect } from 'react';
-import { sessionId, Logger } from '@google/gemini-cli-core';
+import { sessionId, Logger, Storage } from '@google/gemini-cli-core';
 
 /**
  * Hook to manage the logger instance.
  */
-export const useLogger = () => {
+export const useLogger = (storage: Storage) => {
   const [logger, setLogger] = useState<Logger | null>(null);
 
   useEffect(() => {
-    const newLogger = new Logger(sessionId);
+    const newLogger = new Logger(sessionId, storage);
     /**
      * Start async initialization, no need to await. Using await slows down the
      * time from launch to see the gemini-cli prompt and it's better to not save
@@ -26,7 +26,7 @@ export const useLogger = () => {
         setLogger(newLogger);
       })
       .catch(() => {});
-  }, []);
+  }, [storage]);
 
   return logger;
 };

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

@@ -11,9 +11,41 @@ import * as path from 'path';
 import * as os from 'os';
 import * as crypto from 'crypto';
 
-vi.mock('fs/promises');
+vi.mock('fs/promises', () => ({
+  readFile: vi.fn(),
+  writeFile: vi.fn(),
+  mkdir: vi.fn(),
+}));
 vi.mock('os');
 vi.mock('crypto');
+vi.mock('fs', async (importOriginal) => {
+  const actualFs = await importOriginal<typeof import('fs')>();
+  return {
+    ...actualFs,
+    mkdirSync: vi.fn(),
+  };
+});
+vi.mock('@google/gemini-cli-core', () => {
+  class Storage {
+    getProjectTempDir(): string {
+      return path.join('/test/home/', '.gemini', 'tmp', 'mocked_hash');
+    }
+    getHistoryFilePath(): string {
+      return path.join(
+        '/test/home/',
+        '.gemini',
+        'tmp',
+        'mocked_hash',
+        'shell_history',
+      );
+    }
+  }
+  return {
+    isNodeError: (err: unknown): err is NodeJS.ErrnoException =>
+      typeof err === 'object' && err !== null && 'code' in err,
+    Storage,
+  };
+});
 
 const MOCKED_PROJECT_ROOT = '/test/project';
 const MOCKED_HOME_DIR = '/test/home';

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

@@ -7,9 +7,8 @@
 import { useState, useEffect, useCallback } from 'react';
 import * as fs from 'fs/promises';
 import * as path from 'path';
-import { isNodeError, getProjectTempDir } from '@google/gemini-cli-core';
+import { isNodeError, Storage } from '@google/gemini-cli-core';
 
-const HISTORY_FILE = 'shell_history';
 const MAX_HISTORY_LENGTH = 100;
 
 export interface UseShellHistoryReturn {
@@ -20,9 +19,12 @@ export interface UseShellHistoryReturn {
   resetHistoryPosition: () => void;
 }
 
-async function getHistoryFilePath(projectRoot: string): Promise<string> {
-  const historyDir = getProjectTempDir(projectRoot);
-  return path.join(historyDir, HISTORY_FILE);
+async function getHistoryFilePath(
+  projectRoot: string,
+  configStorage?: Storage,
+): Promise<string> {
+  const storage = configStorage ?? new Storage(projectRoot);
+  return storage.getHistoryFilePath();
 }
 
 // Handle multiline commands
@@ -67,20 +69,23 @@ async function writeHistoryFile(
   }
 }
 
-export function useShellHistory(projectRoot: string): UseShellHistoryReturn {
+export function useShellHistory(
+  projectRoot: string,
+  storage?: Storage,
+): UseShellHistoryReturn {
   const [history, setHistory] = useState<string[]>([]);
   const [historyIndex, setHistoryIndex] = useState(-1);
   const [historyFilePath, setHistoryFilePath] = useState<string | null>(null);
 
   useEffect(() => {
     async function loadHistory() {
-      const filePath = await getHistoryFilePath(projectRoot);
+      const filePath = await getHistoryFilePath(projectRoot, storage);
       setHistoryFilePath(filePath);
       const loadedHistory = await readHistoryFile(filePath);
       setHistory(loadedHistory.reverse()); // Newest first
     }
     loadHistory();
-  }, [projectRoot]);
+  }, [projectRoot, storage]);
 
   const addCommandToHistory = useCallback(
     (command: string) => {