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

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 {