feat(a2a): Introduce restore command for a2a server (#13015)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Shreya Keshive <shreyakeshive@google.com>

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

@@ -9,6 +9,9 @@ import path from 'node:path';
 import { z } from 'zod';
 import {
   type Config,
+  formatCheckpointDisplayList,
+  getToolCallDataSchema,
+  getTruncatedCheckpointNames,
   performRestore,
   type ToolCallData,
 } from '@google/gemini-cli-core';
@@ -28,23 +31,7 @@ const HistoryItemSchema = z
   })
   .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(),
-});
+const ToolCallDataSchema = getToolCallDataSchema(HistoryItemSchema);
 
 async function restoreAction(
   context: CommandContext,
@@ -78,15 +65,7 @@ async function restoreAction(
           content: 'No restorable tool calls found.',
         };
       }
-      const truncatedFiles = jsonFiles.map((file) => {
-        const components = file.split('.');
-        if (components.length <= 1) {
-          return file;
-        }
-        components.pop();
-        return components.join('.');
-      });
-      const fileList = truncatedFiles.join('\n');
+      const fileList = formatCheckpointDisplayList(jsonFiles);
       return {
         type: 'message',
         messageType: 'info',
@@ -171,9 +150,8 @@ async function completion(
   }
   try {
     const files = await fs.readdir(checkpointDir);
-    return files
-      .filter((file) => file.endsWith('.json'))
-      .map((file) => file.replace('.json', ''));
+    const jsonFiles = files.filter((file) => file.endsWith('.json'));
+    return getTruncatedCheckpointNames(jsonFiles);
   } catch (_err) {
     return [];
   }

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

@@ -16,7 +16,6 @@ import type {
   ThoughtSummary,
   ToolCallRequestInfo,
   GeminiErrorEventValue,
-  ToolCallData,
 } from '@google/gemini-cli-core';
 import {
   GeminiEventType as ServerGeminiEventType,
@@ -34,10 +33,11 @@ import {
   parseAndFormatApiError,
   ToolConfirmationOutcome,
   promptIdContext,
-  WRITE_FILE_TOOL_NAME,
   tokenLimit,
   debugLogger,
   runInDevTraceSpan,
+  EDIT_TOOL_NAMES,
+  processRestorableToolCalls,
 } from '@google/gemini-cli-core';
 import { type Part, type PartListUnion, FinishReason } from '@google/genai';
 import type {
@@ -76,8 +76,6 @@ enum StreamProcessingStatus {
   Error,
 }
 
-const EDIT_TOOL_NAMES = new Set(['replace', WRITE_FILE_TOOL_NAME]);
-
 function showCitations(settings: LoadedSettings): boolean {
   const enabled = settings?.merged?.ui?.showCitations;
   if (enabled !== undefined) {
@@ -1248,98 +1246,37 @@ export const useGeminiStream = (
       );
 
       if (restorableToolCalls.length > 0) {
-        const checkpointDir = storage.getProjectTempCheckpointsDir();
-
-        if (!checkpointDir) {
+        if (!gitService) {
+          onDebugMessage(
+            'Checkpointing is enabled but Git service is not available. Failed to create snapshot. Ensure Git is installed and working properly.',
+          );
           return;
         }
 
-        try {
-          await fs.mkdir(checkpointDir, { recursive: true });
-        } catch (error) {
-          if (!isNodeError(error) || error.code !== 'EEXIST') {
-            onDebugMessage(
-              `Failed to create checkpoint directory: ${getErrorMessage(error)}`,
-            );
-            return;
-          }
+        const { checkpointsToWrite, errors } = await processRestorableToolCalls<
+          HistoryItem[]
+        >(
+          restorableToolCalls.map((call) => call.request),
+          gitService,
+          geminiClient,
+          history,
+        );
+
+        if (errors.length > 0) {
+          errors.forEach(onDebugMessage);
         }
 
-        for (const toolCall of restorableToolCalls) {
-          const filePath = toolCall.request.args['file_path'] as string;
-          if (!filePath) {
-            onDebugMessage(
-              `Skipping restorable tool call due to missing file_path: ${toolCall.request.name}`,
-            );
-            continue;
-          }
-
+        if (checkpointsToWrite.size > 0) {
+          const checkpointDir = storage.getProjectTempCheckpointsDir();
           try {
-            if (!gitService) {
-              onDebugMessage(
-                `Checkpointing is enabled but Git service is not available. Failed to create snapshot for ${filePath}. Ensure Git is installed and working properly.`,
-              );
-              continue;
+            await fs.mkdir(checkpointDir, { recursive: true });
+            for (const [fileName, content] of checkpointsToWrite) {
+              const filePath = path.join(checkpointDir, fileName);
+              await fs.writeFile(filePath, content);
             }
-
-            let commitHash: string | undefined;
-            try {
-              commitHash = await gitService.createFileSnapshot(
-                `Snapshot for ${toolCall.request.name}`,
-              );
-            } catch (error) {
-              onDebugMessage(
-                `Failed to create new snapshot: ${getErrorMessage(error)}. Attempting to use current commit.`,
-              );
-            }
-
-            if (!commitHash) {
-              commitHash = await gitService.getCurrentCommitHash();
-            }
-
-            if (!commitHash) {
-              onDebugMessage(
-                `Failed to create snapshot for ${filePath}. Checkpointing may not be working properly. Ensure Git is installed and the project directory is accessible.`,
-              );
-              continue;
-            }
-
-            const timestamp = new Date()
-              .toISOString()
-              .replace(/:/g, '-')
-              .replace(/\./g, '_');
-            const toolName = toolCall.request.name;
-            const fileName = path.basename(filePath);
-            const toolCallWithSnapshotFileName = `${timestamp}-${fileName}-${toolName}.json`;
-            const clientHistory = await geminiClient?.getHistory();
-            const toolCallWithSnapshotFilePath = path.join(
-              checkpointDir,
-              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(checkpointData, null, 2),
-            );
           } catch (error) {
             onDebugMessage(
-              `Failed to create checkpoint for ${filePath}: ${getErrorMessage(
-                error,
-              )}. This may indicate a problem with Git or file system permissions.`,
+              `Failed to write checkpoint file: ${getErrorMessage(error)}`,
             );
           }
         }