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/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);