From d7422c1142de03314eca6b0ca6539e62ae713d9c Mon Sep 17 00:00:00 2001 From: mkorwel Date: Thu, 19 Mar 2026 08:54:09 -0700 Subject: [PATCH] feat(workspaces): enhance wsr connect with wait polling, security fixes, and unit tests --- .../cli/src/commands/workspace/connect.ts | 76 ++++++++++++++----- packages/core/src/services/sshService.test.ts | 70 +++++++++++++++++ packages/core/src/services/sshService.ts | 11 ++- .../core/src/services/workspaceHubClient.ts | 10 +++ 4 files changed, 146 insertions(+), 21 deletions(-) create mode 100644 packages/core/src/services/sshService.test.ts diff --git a/packages/cli/src/commands/workspace/connect.ts b/packages/cli/src/commands/workspace/connect.ts index 8e8e9ddb0b..7ddc01ddc9 100644 --- a/packages/cli/src/commands/workspace/connect.ts +++ b/packages/cli/src/commands/workspace/connect.ts @@ -12,6 +12,27 @@ import chalk from 'chalk'; interface ConnectArgs { config?: Config; id: string; + forwardAgent?: boolean; + wait?: boolean; +} + +async function waitForReady(client: WorkspaceHubClient, id: string): Promise { + let attempts = 0; + const maxAttempts = 30; // 5 minutes with 10s interval + + while (attempts < maxAttempts) { + const ws = await client.getWorkspace(id); + if (!ws) throw new Error(`Workspace ${id} disappeared.`); + if (ws.status === 'READY') return ws; + if (ws.status === 'ERROR') throw new Error(`Workspace ${id} entered ERROR state.`); + + // eslint-disable-next-line no-console + console.log(chalk.blue(` Status: ${ws.status}. Waiting for READY... (${attempts + 1}/${maxAttempts})`)); + await new Promise(resolve => setTimeout(resolve, 10000)); + attempts++; + } + + throw new Error(`Timeout waiting for workspace ${id} to become READY.`); } export async function connectToWorkspace(args: ArgumentsCamelCase): Promise { @@ -21,16 +42,14 @@ export async function connectToWorkspace(args: ArgumentsCamelCase): return; } - const hubUrl = 'http://localhost:8080'; + const hubUrl = process.env['GEMINI_WORKSPACE_HUB_URL'] || 'http://localhost:8080'; const client = new WorkspaceHubClient(hubUrl); try { // eslint-disable-next-line no-console console.log(chalk.yellow(`Fetching workspace details for "${args.id}"...`)); - // We need to fetch the workspace info to get the instance name and zone - const workspaces = await client.listWorkspaces() as WorkspaceHubInfo[]; - const ws = workspaces.find(w => w.id === args.id || w.name === args.id); + let ws = await client.getWorkspace(args.id); if (!ws) { // eslint-disable-next-line no-console @@ -38,27 +57,35 @@ export async function connectToWorkspace(args: ArgumentsCamelCase): return; } - const { status, instance_name: instanceName, zone } = ws; - - if (status !== 'READY' && status !== 'PROVISIONING') { - // eslint-disable-next-line no-console - console.warn(chalk.yellow(`Warning: Workspace is in status ${status}. Connection might fail.`)); + if (ws.status !== 'READY') { + if (args.wait) { + ws = await waitForReady(client, args.id); + } else { + // eslint-disable-next-line no-console + console.warn(chalk.yellow(`Warning: Workspace is in status ${ws.status}.`)); + if (ws.status === 'PROVISIONING') { + // eslint-disable-next-line no-console + console.log(chalk.blue('Use --wait to automatically wait for provisioning to complete.')); + } + } } const ssh = new SSHService(); - const project = 'dev-project'; + // TODO: Get project from config once GCP settings are integrated into core Config + const project = process.env['GOOGLE_CLOUD_PROJECT'] || 'dev-project'; // eslint-disable-next-line no-console - console.log(chalk.green(`🚀 Teleporting to ${instanceName} (${zone})...`)); + console.log(chalk.green(`🚀 Teleporting to ${ws.instance_name} (${ws.zone})...`)); // Command to run on the remote VM: attach to the shpool session const remoteCommand = 'shpool attach main || shpool attach'; await ssh.connect({ - instanceName, - zone, + instanceName: ws.instance_name, + zone: ws.zone, project, command: remoteCommand, + forwardAgent: args.forwardAgent, }); } catch (error: unknown) { @@ -71,11 +98,24 @@ export async function connectToWorkspace(args: ArgumentsCamelCase): export const connectCommand: CommandModule = { command: 'connect ', describe: 'Connect to a remote workspace', - builder: (yargs) => yargs.positional('id', { - type: 'string', - describe: 'ID or Name of the workspace to connect to', - demandOption: true, - }), + builder: (yargs) => yargs + .positional('id', { + type: 'string', + describe: 'ID or Name of the workspace to connect to', + demandOption: true, + }) + .option('forward-agent', { + alias: 'A', + type: 'boolean', + describe: 'Forward SSH agent to the remote workspace', + default: false, + }) + .option('wait', { + alias: 'w', + type: 'boolean', + describe: 'Wait for the workspace to become READY if it is provisioning', + default: false, + }), handler: async (argv) => { await connectToWorkspace(argv); await exitCli(); diff --git a/packages/core/src/services/sshService.test.ts b/packages/core/src/services/sshService.test.ts new file mode 100644 index 0000000000..c13fcd1029 --- /dev/null +++ b/packages/core/src/services/sshService.test.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { SSHService } from './sshService.js'; +import { spawn } from 'node:child_process'; +import { EventEmitter } from 'node:events'; + +vi.mock('node:child_process', () => ({ + spawn: vi.fn(), +})); + +describe('SSHService', () => { + let service: SSHService; + + beforeEach(() => { + vi.clearAllMocks(); + service = new SSHService(); + }); + + it('should construct correct gcloud command and arguments', async () => { + const mockChild = new EventEmitter() as any; + vi.mocked(spawn).mockReturnValue(mockChild); + + const promise = service.connect({ + instanceName: 'test-inst', + zone: 'us-west1-a', + project: 'test-project', + forwardAgent: true, + }); + + // Simulate success exit + setTimeout(() => mockChild.emit('exit', 0), 10); + + const code = await promise; + expect(code).toBe(0); + + expect(spawn).toHaveBeenCalledWith( + 'gcloud', + [ + 'compute', + 'ssh', + 'test-inst', + '--zone=us-west1-a', + '--project=test-project', + '--tunnel-through-iap', + '--ssh-flag=-A', + ], + expect.objectContaining({ stdio: 'inherit' }) + ); + }); + + it('should handle failure exit code', async () => { + const mockChild = new EventEmitter() as any; + vi.mocked(spawn).mockReturnValue(mockChild); + + const promise = service.connect({ + instanceName: 'test-inst', + zone: 'us-west1-a', + project: 'test-project', + }); + + setTimeout(() => mockChild.emit('exit', 1), 10); + + await expect(promise).rejects.toThrow('gcloud ssh exited with code 1'); + }); +}); diff --git a/packages/core/src/services/sshService.ts b/packages/core/src/services/sshService.ts index 42a8c44825..bb05ac2ae0 100644 --- a/packages/core/src/services/sshService.ts +++ b/packages/core/src/services/sshService.ts @@ -21,7 +21,7 @@ export class SSHService { * This method spawns a child process and inherits stdio to allow interactive shell. */ async connect(options: SSHOptions): Promise { - const { instanceName, zone, project, command, forwardAgent = true } = options; + const { instanceName, zone, project, command, forwardAgent = false } = options; const args = [ 'compute', @@ -32,6 +32,7 @@ export class SSHService { '--tunnel-through-iap', ]; + // Security: Only forward agent if explicitly requested if (forwardAgent) { args.push('--ssh-flag=-A'); } @@ -43,9 +44,9 @@ export class SSHService { debugLogger.log(`[SSHService] Executing: gcloud ${args.join(' ')}`); return new Promise((resolve, reject) => { + // Security: Do NOT use shell: true to prevent command injection const child = spawn('gcloud', args, { stdio: 'inherit', - shell: true, }); child.on('exit', (code) => { @@ -57,7 +58,11 @@ export class SSHService { }); child.on('error', (err) => { - reject(err); + if ((err as any).code === 'ENOENT') { + reject(new Error('gcloud CLI not found. Please install the Google Cloud SDK.')); + } else { + reject(err); + } }); }); } diff --git a/packages/core/src/services/workspaceHubClient.ts b/packages/core/src/services/workspaceHubClient.ts index 73dc815ca4..081b9c8553 100644 --- a/packages/core/src/services/workspaceHubClient.ts +++ b/packages/core/src/services/workspaceHubClient.ts @@ -53,6 +53,16 @@ export class WorkspaceHubClient { } } + /** + * Fetch a specific workspace by ID or name + */ + async getWorkspace(idOrName: string): Promise { + const workspaces = await this.listWorkspaces(); + return ( + workspaces.find((w) => w.id === idOrName || w.name === idOrName) || null + ); + } + /** * Create a new workspace */