From 22ff923100426d02b97ae7f6bf35be90e7bc032a Mon Sep 17 00:00:00 2001 From: mkorwel Date: Thu, 19 Mar 2026 08:50:42 -0700 Subject: [PATCH] feat(workspaces): implement ssh tunneling and workspace connect command --- packages/cli/src/commands/workspace.ts | 2 + .../cli/src/commands/workspace/connect.ts | 83 +++++++++++++++++++ .../cli/src/ui/commands/workspaceCommand.ts | 26 +++++- packages/core/src/index.ts | 1 + packages/core/src/services/sshService.ts | 64 ++++++++++++++ plans/milestone-3-connectivity.md | 32 +++++++ plans/workspaces-implementation.md | 4 +- 7 files changed, 209 insertions(+), 3 deletions(-) create mode 100644 packages/cli/src/commands/workspace/connect.ts create mode 100644 packages/core/src/services/sshService.ts create mode 100644 plans/milestone-3-connectivity.md diff --git a/packages/cli/src/commands/workspace.ts b/packages/cli/src/commands/workspace.ts index 5d4e74d0c4..c914a67996 100644 --- a/packages/cli/src/commands/workspace.ts +++ b/packages/cli/src/commands/workspace.ts @@ -9,6 +9,7 @@ import type { CommandModule, Argv } from 'yargs'; import { listCommand } from './workspace/list.js'; import { createCommand } from './workspace/create.js'; import { deleteCommand } from './workspace/delete.js'; +import { connectCommand } from './workspace/connect.js'; import { defer } from '../deferred.js'; export const remoteWorkspaceCommand: CommandModule = { @@ -22,6 +23,7 @@ export const remoteWorkspaceCommand: CommandModule = { .command(defer(listCommand, 'wsr')) .command(defer(createCommand, 'wsr')) .command(defer(deleteCommand, 'wsr')) + .command(defer(connectCommand, 'wsr')) .demandCommand(1, 'You need at least one command before continuing.') .version(false), diff --git a/packages/cli/src/commands/workspace/connect.ts b/packages/cli/src/commands/workspace/connect.ts new file mode 100644 index 0000000000..8e8e9ddb0b --- /dev/null +++ b/packages/cli/src/commands/workspace/connect.ts @@ -0,0 +1,83 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { CommandModule, ArgumentsCamelCase } from 'yargs'; +import { WorkspaceHubClient, SSHService, type Config, type WorkspaceHubInfo } from '@google-gemini-cli-core'; +import { exitCli } from '../utils.js'; +import chalk from 'chalk'; + +interface ConnectArgs { + config?: Config; + id: string; +} + +export async function connectToWorkspace(args: ArgumentsCamelCase): Promise { + if (!args.config) { + // eslint-disable-next-line no-console + console.error(chalk.red('Internal error: Config not loaded.')); + return; + } + + const hubUrl = '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); + + if (!ws) { + // eslint-disable-next-line no-console + console.error(chalk.red(`Error: Workspace "${args.id}" not found.`)); + 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.`)); + } + + const ssh = new SSHService(); + const project = 'dev-project'; + + // eslint-disable-next-line no-console + console.log(chalk.green(`🚀 Teleporting to ${instanceName} (${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, + project, + command: remoteCommand, + }); + + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + // eslint-disable-next-line no-console + console.error(chalk.red('Connection failed:'), message); + } +} + +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, + }), + handler: async (argv) => { + await connectToWorkspace(argv); + await exitCli(); + }, +}; diff --git a/packages/cli/src/ui/commands/workspaceCommand.ts b/packages/cli/src/ui/commands/workspaceCommand.ts index 23a8234abd..a8f4e3b2a4 100644 --- a/packages/cli/src/ui/commands/workspaceCommand.ts +++ b/packages/cli/src/ui/commands/workspaceCommand.ts @@ -150,12 +150,36 @@ const deleteCommand: SlashCommand = { }, }; +const connectCommand: SlashCommand = { + name: 'connect', + description: 'Connect to a remote workspace', + kind: CommandKind.BUILT_IN, + autoExecute: true, + action: async ( + _context: CommandContext, + args: string, + ): Promise => { + const id = args.trim(); + if (!id) { + return { + type: 'message', + messageType: 'error', + content: 'Workspace ID is required. Usage: /workspace connect ', + }; + } + return { + type: 'submit_prompt', + content: `I want to connect to remote workspace "${id}". Please run the connect command.`, + }; + }, +}; + export const workspaceSlashCommand: SlashCommand = { name: 'workspace', altNames: ['wsr'], description: 'Manage remote workspaces', kind: CommandKind.BUILT_IN, autoExecute: false, - subCommands: [listCommand, createCommand, deleteCommand], + subCommands: [listCommand, createCommand, deleteCommand, connectCommand], action: async (context: CommandContext) => listAction(context), }; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index f0af5725b3..415ec0dec9 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -134,6 +134,7 @@ export * from './services/trackerTypes.js'; export * from './services/keychainService.js'; export * from './services/keychainTypes.js'; export * from './services/workspaceHubClient.js'; +export * from './services/sshService.js'; export * from './skills/skillManager.js'; export * from './skills/skillLoader.js'; diff --git a/packages/core/src/services/sshService.ts b/packages/core/src/services/sshService.ts new file mode 100644 index 0000000000..42a8c44825 --- /dev/null +++ b/packages/core/src/services/sshService.ts @@ -0,0 +1,64 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { spawn } from 'node:child_process'; +import { debugLogger } from '../utils/debugLogger.js'; + +export interface SSHOptions { + instanceName: string; + zone: string; + project: string; + command?: string; + forwardAgent?: boolean; +} + +export class SSHService { + /** + * Connect to a GCE instance using gcloud compute ssh with IAP tunneling. + * 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 args = [ + 'compute', + 'ssh', + instanceName, + `--zone=${zone}`, + `--project=${project}`, + '--tunnel-through-iap', + ]; + + if (forwardAgent) { + args.push('--ssh-flag=-A'); + } + + if (command) { + args.push('--command', command); + } + + debugLogger.log(`[SSHService] Executing: gcloud ${args.join(' ')}`); + + return new Promise((resolve, reject) => { + const child = spawn('gcloud', args, { + stdio: 'inherit', + shell: true, + }); + + child.on('exit', (code) => { + if (code === 0) { + resolve(0); + } else { + reject(new Error(`gcloud ssh exited with code ${code}`)); + } + }); + + child.on('error', (err) => { + reject(err); + }); + }); + } +} diff --git a/plans/milestone-3-connectivity.md b/plans/milestone-3-connectivity.md new file mode 100644 index 0000000000..f15e071ca5 --- /dev/null +++ b/plans/milestone-3-connectivity.md @@ -0,0 +1,32 @@ +# Milestone 3 Sub-plan: Connectivity & Persistence + +## 1. Objective +Enable the "Teleport" experience by allowing users to securely connect to their remote workspaces with persistent terminal sessions. + +## 2. Tasks + +### Task 3.1: SSH Tunneling Logic (IAP) +Implement the mechanism to securely tunnel SSH traffic through Google Identity-Aware Proxy. +- [ ] Implement `SSHService` in `packages/core/src/services/sshService.ts`. +- [ ] Logic to execute `gcloud compute ssh --tunnel-through-iap`. +- [ ] Handle SSH key generation and OS Login checks. + +### Task 3.2: Workspace Connect Command +Implement the primary `wsr connect ` command. +- [ ] Add `workspace connect` subcommand to `packages/cli/src/commands/workspace/connect.ts`. +- [ ] Logic to fetch instance details (name, zone) from the Hub before connecting. +- [ ] Pass necessary SSH flags (Agent Forwarding, Environment variables). + +### Task 3.3: Persistence Integration (shpool) +Ensure the remote session survives disconnects. +- [ ] Update `remoteCli.ts` to wrap the shell in `shpool attach`. +- [ ] Verify `shpool` daemon is running correctly via the container entrypoint. +- [ ] Logic to handle terminal resizing across local/remote. + +## 3. Verification & Success Criteria +- **Connect:** `gemini wsr connect [ID]` successfully drops the user into a remote shell. +- **Persistence:** User can disconnect (Ctrl+C or close terminal), reconnect, and find their previous state intact. +- **Security:** Connection only works for the workspace owner and requires a valid `gcloud` session. + +## 4. Next Steps +- Implement Task 3.1: Create the `SSHService` in the core package. diff --git a/plans/workspaces-implementation.md b/plans/workspaces-implementation.md index 07c90ab18b..e7bf0e3d0e 100644 --- a/plans/workspaces-implementation.md +++ b/plans/workspaces-implementation.md @@ -25,10 +25,10 @@ Enable developers to manage their remote fleet from the local CLI. See - [x] Add local configuration for Hub discovery (`settings.json`). ### Milestone 3: Connectivity & Persistence (Phase 3) - Enable the "Teleport" experience with session persistence. - +See [Milestone 3 Sub-plan](./milestone-3-connectivity.md) for details. - [ ] Implement `gemini workspace connect`. + - [ ] Setup `gcloud compute ssh --tunnel-through-iap` logic in the client. - [ ] Integrate `shpool` into the container entrypoint for session detachment.