feat(cli): Move key restore logic to core (#13013)

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

@@ -11,11 +11,13 @@ import { theme } from '../semantic-colors.js';
 import type {
   CommandContext,
   SlashCommand,
-  MessageActionReturn,
   SlashCommandActionReturn,
 } from './types.js';
 import { CommandKind } from './types.js';
-import { decodeTagName } from '@google/gemini-cli-core';
+import {
+  decodeTagName,
+  type MessageActionReturn,
+} from '@google/gemini-cli-core';
 import path from 'node:path';
 import type {
   HistoryItemWithoutId,

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

@@ -4,14 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {
-  SlashCommand,
-  CommandContext,
-  MessageActionReturn,
-} from './types.js';
+import type { SlashCommand, CommandContext } from './types.js';
 import { CommandKind } from './types.js';
 import { MessageType, type HistoryItemHooksList } from '../types.js';
-import type { HookRegistryEntry } from '@google/gemini-cli-core';
+import type {
+  HookRegistryEntry,
+  MessageActionReturn,
+} from '@google/gemini-cli-core';
 import { getErrorMessage } from '@google/gemini-cli-core';
 import { SettingScope } from '../../config/settings.js';
 

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

@@ -9,7 +9,8 @@ import * as fs from 'node:fs';
 import * as path from 'node:path';
 import { initCommand } from './initCommand.js';
 import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
-import type { SubmitPromptActionReturn, CommandContext } from './types.js';
+import type { CommandContext } from './types.js';
+import type { SubmitPromptActionReturn } from '@google/gemini-cli-core';
 
 // Mock the 'fs' module
 vi.mock('fs', () => ({

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

@@ -8,10 +8,12 @@ import type {
   SlashCommand,
   SlashCommandActionReturn,
   CommandContext,
-  MessageActionReturn,
 } from './types.js';
 import { CommandKind } from './types.js';
-import type { DiscoveredMCPPrompt } from '@google/gemini-cli-core';
+import type {
+  DiscoveredMCPPrompt,
+  MessageActionReturn,
+} from '@google/gemini-cli-core';
 import {
   DiscoveredMCPTool,
   getMCPDiscoveryState,

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

@@ -155,10 +155,10 @@ describe('restoreCommand', () => {
 
     it('should restore a tool call and project state', async () => {
       const toolCallData = {
-        history: [{ type: 'user', text: 'do a thing' }],
+        history: [{ type: 'user', text: 'do a thing', id: 123 }],
         clientHistory: [{ role: 'user', parts: [{ text: 'do a thing' }] }],
         commitHash: 'abcdef123',
-        toolCall: { name: 'run_shell_command', args: 'ls' },
+        toolCall: { name: 'run_shell_command', args: { command: 'ls' } },
       };
       await fs.writeFile(
         path.join(checkpointsDir, 'my-checkpoint.json'),
@@ -169,7 +169,7 @@ describe('restoreCommand', () => {
       expect(await command?.action?.(mockContext, 'my-checkpoint')).toEqual({
         type: 'tool',
         toolName: 'run_shell_command',
-        toolArgs: 'ls',
+        toolArgs: { command: 'ls' },
       });
       expect(mockContext.ui.loadHistory).toHaveBeenCalledWith(
         toolCallData.history,
@@ -189,7 +189,7 @@ describe('restoreCommand', () => {
 
     it('should restore even if only toolCall is present', async () => {
       const toolCallData = {
-        toolCall: { name: 'run_shell_command', args: 'ls' },
+        toolCall: { name: 'run_shell_command', args: { command: 'ls' } },
       };
       await fs.writeFile(
         path.join(checkpointsDir, 'my-checkpoint.json'),
@@ -201,7 +201,7 @@ describe('restoreCommand', () => {
       expect(await command?.action?.(mockContext, 'my-checkpoint')).toEqual({
         type: 'tool',
         toolName: 'run_shell_command',
-        toolArgs: 'ls',
+        toolArgs: { command: 'ls' },
       });
 
       expect(mockContext.ui.loadHistory).not.toHaveBeenCalled();
@@ -222,7 +222,7 @@ describe('restoreCommand', () => {
       type: 'message',
       messageType: 'error',
       // A more specific error message would be ideal, but for now, we can assert the current behavior.
-      content: expect.stringContaining('Could not read restorable tool calls.'),
+      content: expect.stringContaining('Checkpoint file is invalid'),
     });
   });
 

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

@@ -6,13 +6,45 @@
 
 import * as fs from 'node:fs/promises';
 import path from 'node:path';
+import { z } from 'zod';
+import {
+  type Config,
+  performRestore,
+  type ToolCallData,
+} from '@google/gemini-cli-core';
 import {
   type CommandContext,
   type SlashCommand,
   type SlashCommandActionReturn,
   CommandKind,
 } from './types.js';
-import type { Config } from '@google/gemini-cli-core';
+import type { HistoryItem } from '../types.js';
+import type { Content } from '@google/genai';
+
+const HistoryItemSchema = z
+  .object({
+    type: z.string(),
+    id: z.number(),
+  })
+  .passthrough();
+
+const ContentSchema = z
+  .object({
+    role: z.string().optional(),
+    parts: z.array(z.record(z.unknown())),
+  })
+  .passthrough();
+
+const ToolCallDataSchema = z.object({
+  history: z.array(HistoryItemSchema).optional(),
+  clientHistory: z.array(ContentSchema).optional(),
+  commitHash: z.string().optional(),
+  toolCall: z.object({
+    name: z.string(),
+    args: z.record(z.unknown()),
+  }),
+  messageId: z.string().optional(),
+});
 
 async function restoreAction(
   context: CommandContext,
@@ -74,33 +106,43 @@ async function restoreAction(
 
     const filePath = path.join(checkpointDir, selectedFile);
     const data = await fs.readFile(filePath, 'utf-8');
-    const toolCallData = JSON.parse(data);
+    const parseResult = ToolCallDataSchema.safeParse(JSON.parse(data));
 
-    if (toolCallData.history) {
-      if (!loadHistory) {
-        // This should not happen
-        return {
-          type: 'message',
-          messageType: 'error',
-          content: 'loadHistory function is not available.',
-        };
+    if (!parseResult.success) {
+      return {
+        type: 'message',
+        messageType: 'error',
+        content: `Checkpoint file is invalid: ${parseResult.error.message}`,
+      };
+    }
+
+    // We safely cast here because:
+    // 1. ToolCallDataSchema strictly validates the existence of 'history' as an array and 'id'/'type' on each item.
+    // 2. We trust that files valid according to this schema (written by useGeminiStream) contain the full HistoryItem structure.
+    const toolCallData = parseResult.data as ToolCallData<
+      HistoryItem[],
+      Record<string, unknown>
+    >;
+
+    const actionStream = performRestore(toolCallData, gitService);
+
+    for await (const action of actionStream) {
+      if (action.type === 'message') {
+        addItem(
+          {
+            type: action.messageType,
+            text: action.content,
+          },
+          Date.now(),
+        );
+      } else if (action.type === 'load_history' && loadHistory) {
+        loadHistory(action.history);
+        if (action.clientHistory) {
+          await config
+            ?.getGeminiClient()
+            ?.setHistory(action.clientHistory as Content[]);
+        }
       }
-      loadHistory(toolCallData.history);
-    }
-
-    if (toolCallData.clientHistory) {
-      await config?.getGeminiClient()?.setHistory(toolCallData.clientHistory);
-    }
-
-    if (toolCallData.commitHash) {
-      await gitService?.restoreProjectFromSnapshot(toolCallData.commitHash);
-      addItem(
-        {
-          type: 'info',
-          text: 'Restored project to the state before the tool call.',
-        },
-        Date.now(),
-      );
     }
 
     return {

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

@@ -15,8 +15,9 @@ import {
   updateGitignore,
   GITHUB_WORKFLOW_PATHS,
 } from './setupGithubCommand.js';
-import type { CommandContext, ToolActionReturn } from './types.js';
+import type { CommandContext } from './types.js';
 import * as commandUtils from '../utils/commandUtils.js';
+import type { ToolActionReturn } from '@google/gemini-cli-core';
 import { debugLogger } from '@google/gemini-cli-core';
 
 vi.mock('child_process');

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

@@ -4,9 +4,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { MessageActionReturn, SlashCommand } from './types.js';
+import type { SlashCommand } from './types.js';
 import { CommandKind } from './types.js';
 import { terminalSetup } from '../utils/terminalSetup.js';
+import { type MessageActionReturn } from '@google/gemini-cli-core';
 
 /**
  * Command to configure terminal keybindings for multiline input support.

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

@@ -5,13 +5,17 @@
  */
 
 import type { ReactNode } from 'react';
-import type { Content, PartListUnion } from '@google/genai';
 import type {
   HistoryItemWithoutId,
   HistoryItem,
   ConfirmationRequest,
 } from '../types.js';
-import type { Config, GitService, Logger } from '@google/gemini-cli-core';
+import type {
+  Config,
+  GitService,
+  Logger,
+  CommandActionReturn,
+} from '@google/gemini-cli-core';
 import type { LoadedSettings } from '../../config/settings.js';
 import type { UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
 import type { SessionStatsState } from '../contexts/SessionContext.js';
@@ -84,31 +88,12 @@ export interface CommandContext {
   overwriteConfirmed?: boolean;
 }
 
-/**
- * The return type for a command action that results in scheduling a tool call.
- */
-export interface ToolActionReturn {
-  type: 'tool';
-  toolName: string;
-  toolArgs: Record<string, unknown>;
-}
-
 /** The return type for a command action that results in the app quitting. */
 export interface QuitActionReturn {
   type: 'quit';
   messages: HistoryItem[];
 }
 
-/**
- * The return type for a command action that results in a simple message
- * being displayed to the user.
- */
-export interface MessageActionReturn {
-  type: 'message';
-  messageType: 'info' | 'error';
-  content: string;
-}
-
 /**
  * The return type for a command action that needs to open a dialog.
  */
@@ -128,25 +113,6 @@ export interface OpenDialogActionReturn {
     | 'permissions';
 }
 
-/**
- * The return type for a command action that results in replacing
- * the entire conversation history.
- */
-export interface LoadHistoryActionReturn {
-  type: 'load_history';
-  history: HistoryItemWithoutId[];
-  clientHistory: Content[]; // The history for the generative client
-}
-
-/**
- * The return type for a command action that should immediately submit
- * content as a prompt to the Gemini model.
- */
-export interface SubmitPromptActionReturn {
-  type: 'submit_prompt';
-  content: PartListUnion;
-}
-
 /**
  * The return type for a command action that needs to pause and request
  * confirmation for a set of shell commands before proceeding.
@@ -177,12 +143,9 @@ export interface OpenCustomDialogActionReturn {
 }
 
 export type SlashCommandActionReturn =
-  | ToolActionReturn
-  | MessageActionReturn
+  | CommandActionReturn<HistoryItemWithoutId[]>
   | QuitActionReturn
   | OpenDialogActionReturn
-  | LoadHistoryActionReturn
-  | SubmitPromptActionReturn
   | ConfirmShellCommandsActionReturn
   | ConfirmActionReturn
   | OpenCustomDialogActionReturn;

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

@@ -16,6 +16,7 @@ import type {
   ThoughtSummary,
   ToolCallRequestInfo,
   GeminiErrorEventValue,
+  ToolCallData,
 } from '@google/gemini-cli-core';
 import {
   GeminiEventType as ServerGeminiEventType,
@@ -1313,22 +1314,23 @@ export const useGeminiStream = (
               toolCallWithSnapshotFileName,
             );
 
+            const checkpointData: ToolCallData<
+              HistoryItem[],
+              Record<string, unknown>
+            > & { filePath: string } = {
+              history,
+              clientHistory,
+              toolCall: {
+                name: toolCall.request.name,
+                args: toolCall.request.args,
+              },
+              commitHash,
+              filePath,
+            };
+
             await fs.writeFile(
               toolCallWithSnapshotFilePath,
-              JSON.stringify(
-                {
-                  history,
-                  clientHistory,
-                  toolCall: {
-                    name: toolCall.request.name,
-                    args: toolCall.request.args,
-                  },
-                  commitHash,
-                  filePath,
-                },
-                null,
-                2,
-              ),
+              JSON.stringify(checkpointData, null, 2),
             );
           } catch (error) {
             onDebugMessage(