mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-12 12:54:07 -07:00
feat(workspaces): transform offload into repository-agnostic Gemini Workspaces
This commit is contained in:
mkorwel
@@ -0,0 +1,81 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* WorkspaceProvider interface defines the contract for different remote
|
||||
* execution environments (GCE, Workstations, etc.).
|
||||
*/
|
||||
export interface WorkspaceProvider {
|
||||
/**
|
||||
* Provisions the underlying infrastructure.
|
||||
*/
|
||||
provision(): Promise<number>;
|
||||
|
||||
/**
|
||||
* Ensures the workspace is running and accessible.
|
||||
*/
|
||||
ensureReady(): Promise<number>;
|
||||
|
||||
/**
|
||||
* Performs the initial setup of the workspace (SSH, scripts, auth).
|
||||
*/
|
||||
setup(options: SetupOptions): Promise<number>;
|
||||
|
||||
/**
|
||||
* Returns the raw command string that would be used to execute a command.
|
||||
*/
|
||||
getRunCommand(command: string, options?: ExecOptions): string;
|
||||
|
||||
/**
|
||||
* Executes a command on the workspace.
|
||||
*/
|
||||
exec(command: string, options?: ExecOptions): Promise<number>;
|
||||
|
||||
/**
|
||||
* Executes a command on the workspace and returns the output.
|
||||
*/
|
||||
getExecOutput(command: string, options?: ExecOptions): Promise<{ status: number; stdout: string; stderr: string }>;
|
||||
|
||||
/**
|
||||
* Synchronizes local files to the workspace.
|
||||
*/
|
||||
sync(localPath: string, remotePath: string, options?: SyncOptions): Promise<number>;
|
||||
|
||||
/**
|
||||
* Returns the status of the workspace.
|
||||
*/
|
||||
getStatus(): Promise<WorkspaceStatus>;
|
||||
|
||||
/**
|
||||
* Stops the workspace to save costs.
|
||||
*/
|
||||
stop(): Promise<number>;
|
||||
}
|
||||
|
||||
export interface SetupOptions {
|
||||
projectId: string;
|
||||
zone: string;
|
||||
dnsSuffix?: string;
|
||||
syncAuth?: boolean;
|
||||
}
|
||||
|
||||
export interface ExecOptions {
|
||||
interactive?: boolean;
|
||||
cwd?: string;
|
||||
wrapContainer?: string;
|
||||
}
|
||||
|
||||
export interface SyncOptions {
|
||||
delete?: boolean;
|
||||
exclude?: string[];
|
||||
}
|
||||
|
||||
export interface WorkspaceStatus {
|
||||
name: string;
|
||||
status: string;
|
||||
internalIp?: string;
|
||||
externalIp?: string;
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { spawnSync } from 'child_process';
|
||||
import os from 'os';
|
||||
|
||||
/**
|
||||
* Centralized SSH/RSYNC management for GCE Workers.
|
||||
* Handles Magic Hostname routing with Zero-Knowledge security.
|
||||
* STRICTLY uses Direct Internal connection (Corporate Magic).
|
||||
*/
|
||||
export class GceConnectionManager {
|
||||
private projectId: string;
|
||||
private zone: string;
|
||||
private instanceName: string;
|
||||
|
||||
constructor(projectId: string, zone: string, instanceName: string) {
|
||||
this.projectId = projectId;
|
||||
this.zone = zone;
|
||||
this.instanceName = instanceName;
|
||||
}
|
||||
|
||||
getMagicRemote(): string {
|
||||
const user = `${process.env.USER || 'node'}_google_com`;
|
||||
const dnsSuffix = '.internal.gcpnode.com';
|
||||
return `${user}@nic0.${this.instanceName}.${this.zone}.c.${this.projectId}${dnsSuffix}`;
|
||||
}
|
||||
|
||||
getCommonArgs(): string[] {
|
||||
return [
|
||||
'-o', 'StrictHostKeyChecking=no',
|
||||
'-o', 'UserKnownHostsFile=/dev/null',
|
||||
'-o', 'LogLevel=ERROR',
|
||||
'-o', 'ConnectTimeout=15',
|
||||
'-i', `${os.homedir()}/.ssh/google_compute_engine`
|
||||
];
|
||||
}
|
||||
|
||||
getRunCommand(command: string, options: { interactive?: boolean } = {}): string {
|
||||
const fullRemote = this.getMagicRemote();
|
||||
return `ssh ${this.getCommonArgs().join(' ')} ${options.interactive ? '-t' : ''} ${fullRemote} ${this.quote(command)}`;
|
||||
}
|
||||
|
||||
run(command: string, options: { interactive?: boolean; stdio?: 'pipe' | 'inherit' } = {}): { status: number; stdout: string; stderr: string } {
|
||||
const sshCmd = this.getRunCommand(command, options);
|
||||
const res = spawnSync(sshCmd, { stdio: options.stdio || 'pipe', shell: true });
|
||||
return {
|
||||
status: res.status ?? 1,
|
||||
stdout: res.stdout?.toString() || '',
|
||||
stderr: res.stderr?.toString() || ''
|
||||
};
|
||||
}
|
||||
|
||||
sync(localPath: string, remotePath: string, options: { delete?: boolean; exclude?: string[] } = {}): number {
|
||||
const fullRemote = this.getMagicRemote();
|
||||
// We use --no-t and --no-perms to avoid "Operation not permitted" errors
|
||||
// when syncing to volumes that might have UID mismatches with the container.
|
||||
const rsyncArgs = ['-rvz', '--quiet', '--no-t', '--no-perms', '--no-owner', '--no-group'];
|
||||
if (options.delete) rsyncArgs.push('--delete');
|
||||
if (options.exclude) options.exclude.forEach(ex => rsyncArgs.push(`--exclude="${ex}"`));
|
||||
|
||||
const sshCmd = `ssh ${this.getCommonArgs().join(' ')}`;
|
||||
const directRsync = `rsync ${rsyncArgs.join(' ')} -e ${this.quote(sshCmd)} ${localPath} ${fullRemote}:${remotePath}`;
|
||||
|
||||
const res = spawnSync(directRsync, { stdio: 'inherit', shell: true });
|
||||
return res.status ?? 1;
|
||||
}
|
||||
|
||||
private quote(str: string) {
|
||||
return `'${str.replace(/'/g, "'\\''")}'`;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { spawnSync } from 'child_process';
|
||||
import path from 'path';
|
||||
import fs from 'fs';
|
||||
import os from 'os';
|
||||
import { WorkspaceProvider, SetupOptions, ExecOptions, SyncOptions, WorkspaceStatus } from './BaseProvider.ts';
|
||||
import { GceConnectionManager } from './GceConnectionManager.ts';
|
||||
|
||||
export class GceCosProvider implements WorkspaceProvider {
|
||||
private projectId: string;
|
||||
private zone: string;
|
||||
private instanceName: string;
|
||||
private sshConfigPath: string;
|
||||
private knownHostsPath: string;
|
||||
private sshAlias = 'gcli-worker';
|
||||
private conn: GceConnectionManager;
|
||||
|
||||
constructor(projectId: string, zone: string, instanceName: string, repoRoot: string) {
|
||||
this.projectId = projectId;
|
||||
this.zone = zone;
|
||||
this.instanceName = instanceName;
|
||||
const workspacesDir = path.join(repoRoot, '.gemini/workspaces');
|
||||
if (!fs.existsSync(workspacesDir)) fs.mkdirSync(workspacesDir, { recursive: true });
|
||||
this.sshConfigPath = path.join(workspacesDir, 'ssh_config');
|
||||
this.knownHostsPath = path.join(workspacesDir, 'known_hosts');
|
||||
this.conn = new GceConnectionManager(projectId, zone, instanceName);
|
||||
}
|
||||
|
||||
async provision(): Promise<number> {
|
||||
const imageUri = 'us-docker.pkg.dev/gemini-code-dev/gemini-cli/maintainer:latest';
|
||||
const region = this.zone.split('-').slice(0, 2).join('-');
|
||||
const vpcName = 'iap-vpc';
|
||||
const subnetName = 'iap-subnet';
|
||||
|
||||
console.log(`🏗️ Ensuring "Magic" Network Infrastructure in ${this.projectId}...`);
|
||||
|
||||
const vpcCheck = spawnSync('gcloud', ['compute', 'networks', 'describe', vpcName, '--project', this.projectId], { stdio: 'pipe' });
|
||||
if (vpcCheck.status !== 0) {
|
||||
spawnSync('gcloud', ['compute', 'networks', 'create', vpcName, '--project', this.projectId, '--subnet-mode=custom'], { stdio: 'inherit' });
|
||||
}
|
||||
|
||||
const subnetCheck = spawnSync('gcloud', ['compute', 'networks', 'subnets', 'describe', subnetName, '--project', this.projectId, '--region', region], { stdio: 'pipe' });
|
||||
if (subnetCheck.status !== 0) {
|
||||
spawnSync('gcloud', ['compute', 'networks', 'subnets', 'create', subnetName,
|
||||
'--project', this.projectId, '--network', vpcName, '--region', region,
|
||||
'--range=10.0.0.0/24', '--enable-private-ip-google-access'], { stdio: 'inherit' });
|
||||
} else {
|
||||
spawnSync('gcloud', ['compute', 'networks', 'subnets', 'update', subnetName, '--project', this.projectId, '--region', region, '--enable-private-ip-google-access'], { stdio: 'pipe' });
|
||||
}
|
||||
|
||||
const fwCheck = spawnSync('gcloud', ['compute', 'firewall-rules', 'describe', 'allow-corporate-ssh', '--project', this.projectId], { stdio: 'pipe' });
|
||||
if (fwCheck.status !== 0) {
|
||||
spawnSync('gcloud', ['compute', 'firewall-rules', 'create', 'allow-corporate-ssh',
|
||||
'--project', this.projectId, '--network', vpcName, '--allow=tcp:22', '--source-ranges=0.0.0.0/0'], { stdio: 'inherit' });
|
||||
}
|
||||
|
||||
console.log(`🚀 Provisioning GCE COS worker: ${this.instanceName}...`);
|
||||
|
||||
const startupScriptContent = `#!/bin/bash
|
||||
set -e
|
||||
echo "🚀 Starting Maintainer Worker Resilience Loop..."
|
||||
|
||||
# Wait for Docker to be ready
|
||||
until docker info >/dev/null 2>&1; do echo "Waiting for docker..."; sleep 2; done
|
||||
|
||||
# Pull with retries
|
||||
for i in {1..5}; do
|
||||
docker pull ${imageUri} && break || (echo "Pull failed, retry $i..." && sleep 5)
|
||||
done
|
||||
|
||||
# Run if not already exists
|
||||
if ! docker ps -a | grep -q "maintainer-worker"; then
|
||||
docker run -d --name maintainer-worker --restart always \\
|
||||
-v ~/.workspace:/home/node/.workspace:rw \\
|
||||
-v ~/dev:/home/node/dev:rw \\
|
||||
-v ~/.gemini:/home/node/.gemini:rw \\
|
||||
${imageUri} /bin/bash -c "while true; do sleep 1000; done"
|
||||
fi
|
||||
echo "✅ Maintainer Worker is active."
|
||||
`;
|
||||
|
||||
const tmpScriptPath = path.join(os.tmpdir(), `gcli-startup-${Date.now()}.sh`);
|
||||
fs.writeFileSync(tmpScriptPath, startupScriptContent);
|
||||
|
||||
const result = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'create', this.instanceName,
|
||||
'--project', this.projectId,
|
||||
'--zone', this.zone,
|
||||
'--machine-type', 'n2-standard-8',
|
||||
'--create-disk', `auto-delete=yes,boot=yes,device-name=${this.instanceName},image=projects/cos-cloud/global/images/family/cos-stable,mode=rw,size=200,type=projects/${this.projectId}/zones/${this.zone}/diskTypes/pd-balanced`,
|
||||
'--metadata-from-file', `startup-script=${tmpScriptPath}`,
|
||||
'--metadata', 'enable-oslogin=TRUE',
|
||||
'--network-interface', `network=${vpcName},subnet=${subnetName},no-address`,
|
||||
'--scopes', 'https://www.googleapis.com/auth/cloud-platform',
|
||||
'--quiet'
|
||||
], { stdio: 'inherit' });
|
||||
|
||||
fs.unlinkSync(tmpScriptPath);
|
||||
|
||||
|
||||
if (result.status === 0) {
|
||||
console.log('⏳ Waiting for OS Login and SSH to initialize (this takes ~45s)...');
|
||||
await new Promise(r => setTimeout(r, 45000));
|
||||
}
|
||||
|
||||
return result.status ?? 1;
|
||||
}
|
||||
|
||||
async ensureReady(): Promise<number> {
|
||||
const status = await this.getStatus();
|
||||
if (status.status !== 'RUNNING') {
|
||||
console.log(`⚠️ Worker ${this.instanceName} is ${status.status}. Waking it up...`);
|
||||
const res = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'start', this.instanceName,
|
||||
'--project', this.projectId,
|
||||
'--zone', this.zone
|
||||
], { stdio: 'inherit' });
|
||||
if (res.status !== 0) return res.status ?? 1;
|
||||
|
||||
console.log('⏳ Waiting for boot...');
|
||||
await new Promise(r => setTimeout(r, 20000));
|
||||
}
|
||||
|
||||
// NEW: Verify the container is actually running AND up to date
|
||||
console.log(' - Verifying remote container health and image version...');
|
||||
const containerCheck = await this.getExecOutput('sudo docker ps -q --filter "name=maintainer-worker"');
|
||||
|
||||
let needsUpdate = false;
|
||||
if (containerCheck.status === 0 && containerCheck.stdout.trim()) {
|
||||
// Check if the running image is stale
|
||||
const imageUri = 'us-docker.pkg.dev/gemini-code-dev/gemini-cli/maintainer:latest';
|
||||
const remoteDigest = await this.getExecOutput(`sudo docker inspect --format='{{index .Config.Labels "org.opencontainers.image.revision"}}' maintainer-worker || sudo docker inspect --format='{{.Image}}' maintainer-worker`);
|
||||
// We'll pull the latest tag to see if it's different (or just force pull if it's been a while)
|
||||
// For simplicity in this environment, we'll just check if tmux is missing as a proxy for "stale image"
|
||||
const tmuxCheck = await this.getExecOutput('sudo docker exec maintainer-worker which tmux');
|
||||
if (tmuxCheck.status !== 0) {
|
||||
console.log(' ⚠️ Remote container is stale (missing tmux). Triggering update...');
|
||||
needsUpdate = true;
|
||||
}
|
||||
} else {
|
||||
needsUpdate = true;
|
||||
}
|
||||
|
||||
if (needsUpdate) {
|
||||
console.log(' ⚠️ Container missing or stale. Attempting refresh...');
|
||||
const imageUri = 'us-docker.pkg.dev/gemini-code-dev/gemini-cli/maintainer:latest';
|
||||
const recoverCmd = `sudo docker pull ${imageUri} && (sudo docker rm -f maintainer-worker || true) && sudo docker run -d --name maintainer-worker --restart always -v ~/.workspace:/home/node/.workspace:rw -v ~/dev:/home/node/dev:rw -v ~/.gemini:/home/node/.gemini:rw ${imageUri} /bin/bash -c "while true; do sleep 1000; done"`;
|
||||
const recoverRes = await this.exec(recoverCmd);
|
||||
if (recoverRes !== 0) {
|
||||
console.error(' ❌ Critical: Failed to refresh maintainer container.');
|
||||
return 1;
|
||||
}
|
||||
console.log(' ✅ Container refreshed.');
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
async setup(options: SetupOptions): Promise<number> {
|
||||
const dnsSuffix = options.dnsSuffix || '.internal.gcpnode.com';
|
||||
const internalHostname = `nic0.${this.instanceName}.${this.zone}.c.${this.projectId}${dnsSuffix.startsWith('.') ? dnsSuffix : '.' + dnsSuffix}`;
|
||||
const user = `${process.env.USER || 'node'}_google_com`;
|
||||
|
||||
const sshEntry = `
|
||||
Host ${this.sshAlias}
|
||||
HostName ${internalHostname}
|
||||
IdentityFile ~/.ssh/google_compute_engine
|
||||
User ${user}
|
||||
UserKnownHostsFile /dev/null
|
||||
CheckHostIP no
|
||||
StrictHostKeyChecking no
|
||||
ConnectTimeout 15
|
||||
`;
|
||||
|
||||
fs.writeFileSync(this.sshConfigPath, sshEntry);
|
||||
console.log(` ✅ Created project SSH config: ${this.sshConfigPath}`);
|
||||
|
||||
console.log(' - Verifying direct connection (may trigger corporate SSO prompt)...');
|
||||
const res = this.conn.run('echo 1');
|
||||
if (res.status !== 0) {
|
||||
console.error('\n❌ All connection attempts failed. Please ensure you have "gcert" and IAP permissions.');
|
||||
return 1;
|
||||
}
|
||||
console.log(' ✅ Connection verified.');
|
||||
return 0;
|
||||
}
|
||||
|
||||
getRunCommand(command: string, options: ExecOptions = {}): string {
|
||||
let finalCmd = command;
|
||||
if (options.wrapContainer) {
|
||||
finalCmd = `sudo docker exec ${options.interactive ? '-it' : ''} ${options.cwd ? `-w ${options.cwd}` : ''} ${options.wrapContainer} sh -c ${this.quote(command)}`;
|
||||
}
|
||||
return this.conn.getRunCommand(finalCmd, { interactive: options.interactive });
|
||||
}
|
||||
|
||||
async exec(command: string, options: ExecOptions = {}): Promise<number> {
|
||||
const res = await this.getExecOutput(command, options);
|
||||
return res.status;
|
||||
}
|
||||
|
||||
async getExecOutput(command: string, options: ExecOptions = {}): Promise<{ status: number; stdout: string; stderr: string }> {
|
||||
let finalCmd = command;
|
||||
if (options.wrapContainer) {
|
||||
finalCmd = `docker exec ${options.interactive ? '-it' : ''} ${options.cwd ? `-w ${options.cwd}` : ''} ${options.wrapContainer} sh -c ${this.quote(command)}`;
|
||||
}
|
||||
|
||||
return this.conn.run(finalCmd, { interactive: options.interactive, stdio: options.interactive ? 'inherit' : 'pipe' });
|
||||
}
|
||||
|
||||
async sync(localPath: string, remotePath: string, options: SyncOptions = {}): Promise<number> {
|
||||
console.log(`📦 Syncing ${localPath} to remote:${remotePath}...`);
|
||||
return this.conn.sync(localPath, remotePath, options);
|
||||
}
|
||||
|
||||
async getStatus(): Promise<WorkerStatus> {
|
||||
const res = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'describe', this.instanceName,
|
||||
'--project', this.projectId,
|
||||
'--zone', this.zone,
|
||||
'--format', 'json(name,status,networkInterfaces[0].networkIP)'
|
||||
], { stdio: 'pipe' });
|
||||
|
||||
if (res.status !== 0) {
|
||||
return { name: this.instanceName, status: 'UNKNOWN' };
|
||||
}
|
||||
|
||||
try {
|
||||
const data = JSON.parse(res.stdout.toString());
|
||||
return {
|
||||
name: data.name,
|
||||
status: data.status,
|
||||
internalIp: data.networkInterfaces?.[0]?.networkIP
|
||||
};
|
||||
} catch (e) {
|
||||
return { name: this.instanceName, status: 'ERROR' };
|
||||
}
|
||||
}
|
||||
|
||||
async stop(): Promise<number> {
|
||||
const res = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'stop', this.instanceName,
|
||||
'--project', this.projectId,
|
||||
'--zone', this.zone
|
||||
], { stdio: 'inherit' });
|
||||
return res.status ?? 1;
|
||||
}
|
||||
|
||||
private quote(str: string) {
|
||||
return `'${str.replace(/'/g, "'\\''")}'`;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { GceCosProvider } from './GceCosProvider.ts';
|
||||
import { WorkspaceProvider } from './BaseProvider.ts';
|
||||
import path from 'path';
|
||||
import { fileURLToPath } from 'url';
|
||||
|
||||
const __dirname = path.dirname(fileURLToPath(import.meta.url));
|
||||
const REPO_ROOT = path.resolve(__dirname, '../../../../..');
|
||||
|
||||
export class ProviderFactory {
|
||||
static getProvider(config: { projectId: string; zone: string; instanceName: string }): WorkspaceProvider {
|
||||
// Currently we only have GceCosProvider, but this is where we'd branch
|
||||
return new GceCosProvider(config.projectId, config.zone, config.instanceName, REPO_ROOT);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user