5 files changed, 24 insertions, 12 deletions
diff --git a/packages/cli/src/ui/commands/chatCommand.test.ts b/packages/cli/src/ui/commands/chatCommand.test.ts
index c7299883..88808c92 100644
--- a/packages/cli/src/ui/commands/chatCommand.test.ts
+++ b/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,
diff --git a/packages/cli/src/ui/commands/chatCommand.ts b/packages/cli/src/ui/commands/chatCommand.ts
index 1c9029a9..fdd174ea 100644
--- a/packages/cli/src/ui/commands/chatCommand.ts
+++ b/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 [];
   }
diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts
index f61df287..104eafb3 100644
--- a/packages/cli/src/ui/commands/ideCommand.test.ts
+++ b/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;
diff --git a/packages/cli/src/ui/commands/restoreCommand.test.ts b/packages/cli/src/ui/commands/restoreCommand.test.ts
index 23f71499..b9ecc139 100644
--- a/packages/cli/src/ui/commands/restoreCommand.test.ts
+++ b/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([]);
diff --git a/packages/cli/src/ui/commands/restoreCommand.ts b/packages/cli/src/ui/commands/restoreCommand.ts
index 84259288..f34cdf25 100644
--- a/packages/cli/src/ui/commands/restoreCommand.ts
+++ b/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 [];
   }