Refac: Centralize storage file management (#4078)

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

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) => {