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/a2a-server/src/agent/task.test.ts

@@ -18,6 +18,7 @@ import {
   GeminiEventType,
   type Config,
   type ToolCallRequestInfo,
+  type GitService,
   type CompletedToolCall,
 } from '@google/gemini-cli-core';
 import { createMockConfig } from '../utils/testing_utils.js';
@@ -25,6 +26,17 @@ import type { ExecutionEventBus, RequestContext } from '@a2a-js/sdk/server';
 import { CoderAgentEvent } from '../types.js';
 import type { ToolCall } from '@google/gemini-cli-core';
 
+const mockProcessRestorableToolCalls = vi.hoisted(() => vi.fn());
+
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...original,
+    processRestorableToolCalls: mockProcessRestorableToolCalls,
+  };
+});
+
 describe('Task', () => {
   it('scheduleToolCalls should not modify the input requests array', async () => {
     const mockConfig = createMockConfig();
@@ -72,6 +84,141 @@ describe('Task', () => {
     expect(requests).toEqual(originalRequests);
   });
 
+  describe('scheduleToolCalls', () => {
+    const mockConfig = createMockConfig();
+    const mockEventBus: ExecutionEventBus = {
+      publish: vi.fn(),
+      on: vi.fn(),
+      off: vi.fn(),
+      once: vi.fn(),
+      removeAllListeners: vi.fn(),
+      finished: vi.fn(),
+    };
+
+    beforeEach(() => {
+      vi.clearAllMocks();
+    });
+
+    it('should not create a checkpoint if no restorable tools are called', async () => {
+      // @ts-expect-error - Calling private constructor for test purposes.
+      const task = new Task(
+        'task-id',
+        'context-id',
+        mockConfig as Config,
+        mockEventBus,
+      );
+      const requests: ToolCallRequestInfo[] = [
+        {
+          callId: '1',
+          name: 'run_shell_command',
+          args: { command: 'ls' },
+          isClientInitiated: false,
+          prompt_id: 'prompt-id-1',
+        },
+      ];
+      const abortController = new AbortController();
+      await task.scheduleToolCalls(requests, abortController.signal);
+      expect(mockProcessRestorableToolCalls).not.toHaveBeenCalled();
+    });
+
+    it('should create a checkpoint if a restorable tool is called', async () => {
+      const mockConfig = createMockConfig({
+        getCheckpointingEnabled: () => true,
+        getGitService: () => Promise.resolve({} as GitService),
+      });
+      mockProcessRestorableToolCalls.mockResolvedValue({
+        checkpointsToWrite: new Map([['test.json', 'test content']]),
+        toolCallToCheckpointMap: new Map(),
+        errors: [],
+      });
+      // @ts-expect-error - Calling private constructor for test purposes.
+      const task = new Task(
+        'task-id',
+        'context-id',
+        mockConfig as Config,
+        mockEventBus,
+      );
+      const requests: ToolCallRequestInfo[] = [
+        {
+          callId: '1',
+          name: 'replace',
+          args: {
+            file_path: 'test.txt',
+            old_string: 'old',
+            new_string: 'new',
+          },
+          isClientInitiated: false,
+          prompt_id: 'prompt-id-1',
+        },
+      ];
+      const abortController = new AbortController();
+      await task.scheduleToolCalls(requests, abortController.signal);
+      expect(mockProcessRestorableToolCalls).toHaveBeenCalledOnce();
+    });
+
+    it('should process all restorable tools for checkpointing in a single batch', async () => {
+      const mockConfig = createMockConfig({
+        getCheckpointingEnabled: () => true,
+        getGitService: () => Promise.resolve({} as GitService),
+      });
+      mockProcessRestorableToolCalls.mockResolvedValue({
+        checkpointsToWrite: new Map([
+          ['test1.json', 'test content 1'],
+          ['test2.json', 'test content 2'],
+        ]),
+        toolCallToCheckpointMap: new Map([
+          ['1', 'test1'],
+          ['2', 'test2'],
+        ]),
+        errors: [],
+      });
+      // @ts-expect-error - Calling private constructor for test purposes.
+      const task = new Task(
+        'task-id',
+        'context-id',
+        mockConfig as Config,
+        mockEventBus,
+      );
+      const requests: ToolCallRequestInfo[] = [
+        {
+          callId: '1',
+          name: 'replace',
+          args: {
+            file_path: 'test.txt',
+            old_string: 'old',
+            new_string: 'new',
+          },
+          isClientInitiated: false,
+          prompt_id: 'prompt-id-1',
+        },
+        {
+          callId: '2',
+          name: 'write_file',
+          args: { file_path: 'test2.txt', content: 'new content' },
+          isClientInitiated: false,
+          prompt_id: 'prompt-id-2',
+        },
+        {
+          callId: '3',
+          name: 'not_restorable',
+          args: {},
+          isClientInitiated: false,
+          prompt_id: 'prompt-id-3',
+        },
+      ];
+      const abortController = new AbortController();
+      await task.scheduleToolCalls(requests, abortController.signal);
+      expect(mockProcessRestorableToolCalls).toHaveBeenCalledExactlyOnceWith(
+        [
+          expect.objectContaining({ callId: '1' }),
+          expect.objectContaining({ callId: '2' }),
+        ],
+        expect.anything(),
+        expect.anything(),
+      );
+    });
+  });
+
   describe('acceptAgentMessage', () => {
     it('should set currentTraceId when event has traceId', async () => {
       const mockConfig = createMockConfig();

packages/a2a-server/src/agent/task.ts

@@ -27,6 +27,8 @@ import {
   type Config,
   type UserTierId,
   type AnsiOutput,
+  EDIT_TOOL_NAMES,
+  processRestorableToolCalls,
 } from '@google/gemini-cli-core';
 import type { RequestContext } from '@a2a-js/sdk/server';
 import { type ExecutionEventBus } from '@a2a-js/sdk/server';
@@ -40,7 +42,8 @@ import type {
 } from '@a2a-js/sdk';
 import { v4 as uuidv4 } from 'uuid';
 import { logger } from '../utils/logger.js';
-import * as fs from 'node:fs';
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
 import { CoderAgentEvent } from '../types.js';
 import type {
   CoderAgentMessage,
@@ -511,7 +514,7 @@ export class Task {
     new_string: string,
   ): Promise<string> {
     try {
-      const currentContent = fs.readFileSync(file_path, 'utf8');
+      const currentContent = await fs.readFile(file_path, 'utf8');
       return this._applyReplacement(
         currentContent,
         old_string,
@@ -554,6 +557,44 @@ export class Task {
       return;
     }
 
+    // Set checkpoint file before any file modification tool executes
+    const restorableToolCalls = requests.filter((request) =>
+      EDIT_TOOL_NAMES.has(request.name),
+    );
+
+    if (restorableToolCalls.length > 0) {
+      const gitService = await this.config.getGitService();
+      if (gitService) {
+        const { checkpointsToWrite, toolCallToCheckpointMap, errors } =
+          await processRestorableToolCalls(
+            restorableToolCalls,
+            gitService,
+            this.geminiClient,
+          );
+
+        if (errors.length > 0) {
+          errors.forEach((error) => logger.error(error));
+        }
+
+        if (checkpointsToWrite.size > 0) {
+          const checkpointDir =
+            this.config.storage.getProjectTempCheckpointsDir();
+          await fs.mkdir(checkpointDir, { recursive: true });
+          for (const [fileName, content] of checkpointsToWrite) {
+            const filePath = path.join(checkpointDir, fileName);
+            await fs.writeFile(filePath, content);
+          }
+        }
+
+        for (const request of requests) {
+          const checkpoint = toolCallToCheckpointMap.get(request.callId);
+          if (checkpoint) {
+            request.checkpoint = checkpoint;
+          }
+        }
+      }
+    }
+
     const updatedRequests = await Promise.all(
       requests.map(async (request) => {
         if (

packages/a2a-server/src/commands/command-registry.ts

@@ -5,6 +5,7 @@
  */
 
 import { ExtensionsCommand } from './extensions.js';
+import { RestoreCommand } from './restore.js';
 import type { Command } from './types.js';
 
 class CommandRegistry {
@@ -12,6 +13,7 @@ class CommandRegistry {
 
   constructor() {
     this.register(new ExtensionsCommand());
+    this.register(new RestoreCommand());
   }
 
   register(command: Command) {

packages/a2a-server/src/commands/extensions.test.ts

@@ -6,7 +6,7 @@
 
 import { describe, it, expect, vi } from 'vitest';
 import { ExtensionsCommand, ListExtensionsCommand } from './extensions.js';
-import type { Config } from '@google/gemini-cli-core';
+import type { CommandContext } from './types.js';
 
 const mockListExtensions = vi.hoisted(() => vi.fn());
 vi.mock('@google/gemini-cli-core', async (importOriginal) => {
@@ -42,14 +42,14 @@ describe('ExtensionsCommand', () => {
 
   it('should default to listing extensions', async () => {
     const command = new ExtensionsCommand();
-    const mockConfig = {} as Config;
+    const mockConfig = { config: {} } as CommandContext;
     const mockExtensions = [{ name: 'ext1' }];
     mockListExtensions.mockReturnValue(mockExtensions);
 
     const result = await command.execute(mockConfig, []);
 
     expect(result).toEqual({ name: 'extensions list', data: mockExtensions });
-    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig);
+    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig.config);
   });
 });
 
@@ -61,19 +61,19 @@ describe('ListExtensionsCommand', () => {
 
   it('should call listExtensions with the provided config', async () => {
     const command = new ListExtensionsCommand();
-    const mockConfig = {} as Config;
+    const mockConfig = { config: {} } as CommandContext;
     const mockExtensions = [{ name: 'ext1' }];
     mockListExtensions.mockReturnValue(mockExtensions);
 
     const result = await command.execute(mockConfig, []);
 
     expect(result).toEqual({ name: 'extensions list', data: mockExtensions });
-    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig);
+    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig.config);
   });
 
   it('should return a message when no extensions are installed', async () => {
     const command = new ListExtensionsCommand();
-    const mockConfig = {} as Config;
+    const mockConfig = { config: {} } as CommandContext;
     mockListExtensions.mockReturnValue([]);
 
     const result = await command.execute(mockConfig, []);
@@ -82,6 +82,6 @@ describe('ListExtensionsCommand', () => {
       name: 'extensions list',
       data: 'No extensions installed.',
     });
-    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig);
+    expect(mockListExtensions).toHaveBeenCalledWith(mockConfig.config);
   });
 });

packages/a2a-server/src/commands/extensions.ts

@@ -4,8 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { listExtensions, type Config } from '@google/gemini-cli-core';
-import type { Command, CommandExecutionResponse } from './types.js';
+import { listExtensions } from '@google/gemini-cli-core';
+import type {
+  Command,
+  CommandContext,
+  CommandExecutionResponse,
+} from './types.js';
 
 export class ExtensionsCommand implements Command {
   readonly name = 'extensions';
@@ -14,10 +18,10 @@ export class ExtensionsCommand implements Command {
   readonly topLevel = true;
 
   async execute(
-    config: Config,
+    context: CommandContext,
     _: string[],
   ): Promise<CommandExecutionResponse> {
-    return new ListExtensionsCommand().execute(config, _);
+    return new ListExtensionsCommand().execute(context, _);
   }
 }
 
@@ -26,10 +30,10 @@ export class ListExtensionsCommand implements Command {
   readonly description = 'Lists all installed extensions.';
 
   async execute(
-    config: Config,
+    context: CommandContext,
     _: string[],
   ): Promise<CommandExecutionResponse> {
-    const extensions = listExtensions(config);
+    const extensions = listExtensions(context.config);
     const data = extensions.length ? extensions : 'No extensions installed.';
 
     return { name: this.name, data };

packages/a2a-server/src/commands/restore.test.ts

@@ -0,0 +1,137 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { RestoreCommand, ListCheckpointsCommand } from './restore.js';
+import type { CommandContext } from './types.js';
+import type { Config } from '@google/gemini-cli-core';
+import { createMockConfig } from '../utils/testing_utils.js';
+
+beforeEach(() => {
+  vi.clearAllMocks();
+});
+
+const mockPerformRestore = vi.hoisted(() => vi.fn());
+const mockLoggerInfo = vi.hoisted(() => vi.fn());
+const mockGetCheckpointInfoList = vi.hoisted(() => vi.fn());
+
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...original,
+    performRestore: mockPerformRestore,
+    getCheckpointInfoList: mockGetCheckpointInfoList,
+  };
+});
+
+const mockFs = vi.hoisted(() => ({
+  readFile: vi.fn(),
+  readdir: vi.fn(),
+  mkdir: vi.fn(),
+}));
+
+vi.mock('node:fs/promises', () => mockFs);
+
+vi.mock('../utils/logger.js', () => ({
+  logger: {
+    info: mockLoggerInfo,
+  },
+}));
+
+describe('RestoreCommand', () => {
+  const mockConfig = {
+    config: createMockConfig() as Config,
+    git: {},
+  } as CommandContext;
+
+  it('should return error if no checkpoint name is provided', async () => {
+    const command = new RestoreCommand();
+    const result = await command.execute(mockConfig, []);
+    expect(result.data).toEqual({
+      type: 'message',
+      messageType: 'error',
+      content: 'Please provide a checkpoint name to restore.',
+    });
+  });
+
+  it('should restore a checkpoint when a valid file is provided', async () => {
+    const command = new RestoreCommand();
+    const toolCallData = {
+      toolCall: {
+        name: 'test-tool',
+        args: {},
+      },
+      history: [],
+      clientHistory: [],
+      commitHash: '123',
+    };
+    mockFs.readFile.mockResolvedValue(JSON.stringify(toolCallData));
+    const restoreContent = {
+      type: 'message',
+      messageType: 'info',
+      content: 'Restored',
+    };
+    mockPerformRestore.mockReturnValue(
+      (async function* () {
+        yield restoreContent;
+      })(),
+    );
+    const result = await command.execute(mockConfig, ['checkpoint1.json']);
+    expect(result.data).toEqual([restoreContent]);
+  });
+
+  it('should show "file not found" error for a non-existent checkpoint', async () => {
+    const command = new RestoreCommand();
+    const error = new Error('File not found');
+    (error as NodeJS.ErrnoException).code = 'ENOENT';
+    mockFs.readFile.mockRejectedValue(error);
+    const result = await command.execute(mockConfig, ['checkpoint2.json']);
+    expect(result.data).toEqual({
+      type: 'message',
+      messageType: 'error',
+      content: 'File not found: checkpoint2.json',
+    });
+  });
+
+  it('should handle invalid JSON in checkpoint file', async () => {
+    const command = new RestoreCommand();
+    mockFs.readFile.mockResolvedValue('invalid json');
+    const result = await command.execute(mockConfig, ['checkpoint1.json']);
+    expect((result.data as { content: string }).content).toContain(
+      'An unexpected error occurred during restore.',
+    );
+  });
+});
+
+describe('ListCheckpointsCommand', () => {
+  const mockConfig = {
+    config: createMockConfig() as Config,
+  } as CommandContext;
+
+  it('should list all available checkpoints', async () => {
+    const command = new ListCheckpointsCommand();
+    const checkpointInfo = [{ file: 'checkpoint1.json', description: 'Test' }];
+    mockFs.readdir.mockResolvedValue(['checkpoint1.json']);
+    mockFs.readFile.mockResolvedValue(
+      JSON.stringify({ toolCall: { name: 'Test', args: {} } }),
+    );
+    mockGetCheckpointInfoList.mockReturnValue(checkpointInfo);
+    const result = await command.execute(mockConfig);
+    expect((result.data as { content: string }).content).toEqual(
+      JSON.stringify(checkpointInfo),
+    );
+  });
+
+  it('should handle errors when listing checkpoints', async () => {
+    const command = new ListCheckpointsCommand();
+    mockFs.readdir.mockRejectedValue(new Error('Read error'));
+    const result = await command.execute(mockConfig);
+    expect((result.data as { content: string }).content).toContain(
+      'An unexpected error occurred while listing checkpoints.',
+    );
+  });
+});

packages/a2a-server/src/commands/restore.ts

@@ -0,0 +1,155 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  getCheckpointInfoList,
+  getToolCallDataSchema,
+  isNodeError,
+  performRestore,
+} from '@google/gemini-cli-core';
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import type {
+  Command,
+  CommandContext,
+  CommandExecutionResponse,
+} from './types.js';
+
+export class RestoreCommand implements Command {
+  readonly name = 'restore';
+  readonly description =
+    'Restore to a previous checkpoint, or list available checkpoints to restore. This will reset the conversation and file history to the state it was in when the checkpoint was created';
+  readonly topLevel = true;
+  readonly requiresWorkspace = true;
+  readonly subCommands = [new ListCheckpointsCommand()];
+
+  async execute(
+    context: CommandContext,
+    args: string[],
+  ): Promise<CommandExecutionResponse> {
+    const { config, git: gitService } = context;
+    const argsStr = args.join(' ');
+
+    try {
+      if (!argsStr) {
+        return {
+          name: this.name,
+          data: {
+            type: 'message',
+            messageType: 'error',
+            content: 'Please provide a checkpoint name to restore.',
+          },
+        };
+      }
+
+      const selectedFile = argsStr.endsWith('.json')
+        ? argsStr
+        : `${argsStr}.json`;
+
+      const checkpointDir = config.storage.getProjectTempCheckpointsDir();
+      const filePath = path.join(checkpointDir, selectedFile);
+
+      let data: string;
+      try {
+        data = await fs.readFile(filePath, 'utf-8');
+      } catch (error) {
+        if (isNodeError(error) && error.code === 'ENOENT') {
+          return {
+            name: this.name,
+            data: {
+              type: 'message',
+              messageType: 'error',
+              content: `File not found: ${selectedFile}`,
+            },
+          };
+        }
+        throw error;
+      }
+
+      const toolCallData = JSON.parse(data);
+      const ToolCallDataSchema = getToolCallDataSchema();
+      const parseResult = ToolCallDataSchema.safeParse(toolCallData);
+
+      if (!parseResult.success) {
+        return {
+          name: this.name,
+          data: {
+            type: 'message',
+            messageType: 'error',
+            content: 'Checkpoint file is invalid or corrupted.',
+          },
+        };
+      }
+
+      const restoreResultGenerator = performRestore(
+        parseResult.data,
+        gitService,
+      );
+      const restoreResult = [];
+      for await (const result of restoreResultGenerator) {
+        restoreResult.push(result);
+      }
+
+      return {
+        name: this.name,
+        data: restoreResult,
+      };
+    } catch (_error) {
+      return {
+        name: this.name,
+        data: {
+          type: 'message',
+          messageType: 'error',
+          content: 'An unexpected error occurred during restore.',
+        },
+      };
+    }
+  }
+}
+
+export class ListCheckpointsCommand implements Command {
+  readonly name = 'restore list';
+  readonly description = 'Lists all available checkpoints.';
+  readonly topLevel = false;
+
+  async execute(context: CommandContext): Promise<CommandExecutionResponse> {
+    const { config } = context;
+
+    try {
+      const checkpointDir = config.storage.getProjectTempCheckpointsDir();
+      await fs.mkdir(checkpointDir, { recursive: true });
+      const files = await fs.readdir(checkpointDir);
+      const jsonFiles = files.filter((file) => file.endsWith('.json'));
+
+      const checkpointFiles = new Map<string, string>();
+      for (const file of jsonFiles) {
+        const filePath = path.join(checkpointDir, file);
+        const data = await fs.readFile(filePath, 'utf-8');
+        checkpointFiles.set(file, data);
+      }
+
+      const checkpointInfoList = getCheckpointInfoList(checkpointFiles);
+
+      return {
+        name: this.name,
+        data: {
+          type: 'message',
+          messageType: 'info',
+          content: JSON.stringify(checkpointInfoList),
+        },
+      };
+    } catch (_error) {
+      return {
+        name: this.name,
+        data: {
+          type: 'message',
+          messageType: 'error',
+          content: 'An unexpected error occurred while listing checkpoints.',
+        },
+      };
+    }
+  }
+}

packages/a2a-server/src/commands/types.ts

@@ -4,7 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { Config } from '@google/gemini-cli-core';
+import type { Config, GitService } from '@google/gemini-cli-core';
+
+export interface CommandContext {
+  config: Config;
+  git?: GitService;
+}
 
 export interface CommandArgument {
   readonly name: string;
@@ -18,8 +23,12 @@ export interface Command {
   readonly arguments?: CommandArgument[];
   readonly subCommands?: Command[];
   readonly topLevel?: boolean;
+  readonly requiresWorkspace?: boolean;
 
-  execute(config: Config, args: string[]): Promise<CommandExecutionResponse>;
+  execute(
+    config: CommandContext,
+    args: string[],
+  ): Promise<CommandExecutionResponse>;
 }
 
 export interface CommandExecutionResponse {

packages/a2a-server/src/config/config.ts

@@ -71,6 +71,9 @@ export async function loadConfig(
     ideMode: false,
     folderTrust: settings.folderTrust === true,
     extensionLoader,
+    checkpointing: process.env['CHECKPOINTING']
+      ? process.env['CHECKPOINTING'] === 'true'
+      : settings.checkpointing?.enabled,
     previewFeatures: settings.general?.previewFeatures,
   };
 

packages/a2a-server/src/http/app.test.ts

@@ -36,7 +36,7 @@ import {
   createMockConfig,
 } from '../utils/testing_utils.js';
 import { MockTool } from '@google/gemini-cli-core';
-import type { Command } from '../commands/types.js';
+import type { Command, CommandContext } from '../commands/types.js';
 
 const mockToolConfirmationFn = async () =>
   ({}) as unknown as ToolCallConfirmationDetails;
@@ -97,6 +97,7 @@ vi.mock('@google/gemini-cli-core', async () => {
       getUserTier: vi.fn().mockReturnValue('free'),
       initialize: vi.fn(),
     })),
+    performRestore: vi.fn(),
   };
 });
 
@@ -939,6 +940,17 @@ describe('E2E Tests', () => {
     });
 
     it('should return extensions for valid command', async () => {
+      const mockExtensionsCommand = {
+        name: 'extensions list',
+        description: 'a mock command',
+        execute: vi.fn(async (context: CommandContext) => {
+          // Simulate the actual command's behavior
+          const extensions = context.config.getExtensions();
+          return { name: 'extensions list', data: extensions };
+        }),
+      };
+      vi.spyOn(commandRegistry, 'get').mockReturnValue(mockExtensionsCommand);
+
       const agent = request.agent(app);
       const res = await agent
         .post('/executeCommand')
@@ -954,6 +966,8 @@ describe('E2E Tests', () => {
     });
 
     it('should return 404 for invalid command', async () => {
+      vi.spyOn(commandRegistry, 'get').mockReturnValue(undefined);
+
       const agent = request.agent(app);
       const res = await agent
         .post('/executeCommand')
@@ -986,5 +1000,66 @@ describe('E2E Tests', () => {
       expect(res.body.error).toBe('"args" field must be an array.');
       expect(getExtensionsSpy).not.toHaveBeenCalled();
     });
+
+    it('should execute a command that does not require a workspace when CODER_AGENT_WORKSPACE_PATH is not set', async () => {
+      const mockCommand = {
+        name: 'test-command',
+        description: 'a mock command',
+        execute: vi
+          .fn()
+          .mockResolvedValue({ name: 'test-command', data: 'success' }),
+      };
+      vi.spyOn(commandRegistry, 'get').mockReturnValue(mockCommand);
+
+      delete process.env['CODER_AGENT_WORKSPACE_PATH'];
+      const response = await request(app)
+        .post('/executeCommand')
+        .send({ command: 'test-command', args: [] });
+
+      expect(response.status).toBe(200);
+      expect(response.body.data).toBe('success');
+    });
+
+    it('should return 400 for a command that requires a workspace when CODER_AGENT_WORKSPACE_PATH is not set', async () => {
+      const mockWorkspaceCommand = {
+        name: 'workspace-command',
+        description: 'A command that requires a workspace',
+        requiresWorkspace: true,
+        execute: vi
+          .fn()
+          .mockResolvedValue({ name: 'workspace-command', data: 'success' }),
+      };
+      vi.spyOn(commandRegistry, 'get').mockReturnValue(mockWorkspaceCommand);
+
+      delete process.env['CODER_AGENT_WORKSPACE_PATH'];
+      const response = await request(app)
+        .post('/executeCommand')
+        .send({ command: 'workspace-command', args: [] });
+
+      expect(response.status).toBe(400);
+      expect(response.body.error).toBe(
+        'Command "workspace-command" requires a workspace, but CODER_AGENT_WORKSPACE_PATH is not set.',
+      );
+    });
+
+    it('should execute a command that requires a workspace when CODER_AGENT_WORKSPACE_PATH is set', async () => {
+      const mockWorkspaceCommand = {
+        name: 'workspace-command',
+        description: 'A command that requires a workspace',
+        requiresWorkspace: true,
+        execute: vi
+          .fn()
+          .mockResolvedValue({ name: 'workspace-command', data: 'success' }),
+      };
+      vi.spyOn(commandRegistry, 'get').mockReturnValue(mockWorkspaceCommand);
+
+      process.env['CODER_AGENT_WORKSPACE_PATH'] = '/tmp/test-workspace';
+      const response = await request(app)
+        .post('/executeCommand')
+        .send({ command: 'workspace-command', args: [] });
+
+      expect(response.status).toBe(200);
+      expect(response.body.data).toBe('success');
+    });
   });
 });

packages/a2a-server/src/http/app.ts

@@ -22,6 +22,7 @@ import { loadExtensions } from '../config/extension.js';
 import { commandRegistry } from '../commands/command-registry.js';
 import { SimpleExtensionLoader } from '@google/gemini-cli-core';
 import type { Command, CommandArgument } from '../commands/types.js';
+import { GitService } from '@google/gemini-cli-core';
 
 type CommandResponse = {
   name: string;
@@ -85,6 +86,14 @@ export async function createApp() {
       'a2a-server',
     );
 
+    let git: GitService | undefined;
+    if (config.getCheckpointingEnabled()) {
+      git = new GitService(config.getTargetDir(), config.storage);
+      await git.initialize();
+    }
+
+    const context = { config, git };
+
     // loadEnvironment() is called within getConfig now
     const bucketName = process.env['GCS_BUCKET_NAME'];
     let taskStoreForExecutor: TaskStore;
@@ -144,6 +153,7 @@ export async function createApp() {
     });
 
     expressApp.post('/executeCommand', async (req, res) => {
+      logger.info('[CoreAgent] Received /executeCommand request: ', req.body);
       try {
         const { command, args } = req.body;
 
@@ -159,13 +169,22 @@ export async function createApp() {
 
         const commandToExecute = commandRegistry.get(command);
 
+        if (commandToExecute?.requiresWorkspace) {
+          if (!process.env['CODER_AGENT_WORKSPACE_PATH']) {
+            return res.status(400).json({
+              error: `Command "${command}" requires a workspace, but CODER_AGENT_WORKSPACE_PATH is not set.`,
+            });
+          }
+        }
+
         if (!commandToExecute) {
           return res
             .status(404)
             .json({ error: `Command not found: ${command}` });
         }
 
-        const result = await commandToExecute.execute(config, args ?? []);
+        const result = await commandToExecute.execute(context, args ?? []);
+        logger.info('[CoreAgent] Sending /executeCommand response: ', result);
         return res.status(200).json(result);
       } catch (e) {
         logger.error('Error executing /executeCommand:', e);

packages/a2a-server/src/utils/testing_utils.ts

@@ -37,8 +37,10 @@ export function createMockConfig(
       isPathWithinWorkspace: () => true,
     }),
     getTargetDir: () => '/test',
+    getCheckpointingEnabled: vi.fn().mockReturnValue(false),
     storage: {
       getProjectTempDir: () => '/tmp',
+      getProjectTempCheckpointsDir: () => '/tmp/checkpoints',
     } as Storage,
     getTruncateToolOutputThreshold: () =>
       DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD,
@@ -62,6 +64,7 @@ export function createMockConfig(
     getMcpClientManager: vi.fn().mockReturnValue({
       getMcpServers: vi.fn().mockReturnValue({}),
     }),
+    getGitService: vi.fn(),
     ...overrides,
   } as unknown as Config;
   mockConfig.getMessageBus = vi.fn().mockReturnValue(createMockMessageBus());