mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-13 05:12:55 -07:00
feat(workspaces): complete task 1.2 with actual gce provisioning and unit tests
This commit is contained in:
@@ -0,0 +1,61 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||||
|
import { ComputeService } from './computeService.js';
|
||||||
|
|
||||||
|
const mockInstancesClient = {
|
||||||
|
insert: vi.fn().mockResolvedValue([{ latestResponse: { name: 'op-1' } }]),
|
||||||
|
delete: vi.fn().mockResolvedValue([{ latestResponse: { name: 'op-2' } }]),
|
||||||
|
};
|
||||||
|
|
||||||
|
vi.mock('@google-cloud/compute', () => ({
|
||||||
|
InstancesClient: vi.fn().mockImplementation(() => mockInstancesClient),
|
||||||
|
}));
|
||||||
|
|
||||||
|
describe('ComputeService', () => {
|
||||||
|
let service: ComputeService;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
vi.clearAllMocks();
|
||||||
|
service = new ComputeService();
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('createWorkspaceInstance', () => {
|
||||||
|
it('should call instancesClient.insert with correct parameters', async () => {
|
||||||
|
const options = {
|
||||||
|
instanceName: 'test-inst',
|
||||||
|
machineType: 'e2-medium',
|
||||||
|
imageTag: 'latest',
|
||||||
|
zone: 'us-west1-a',
|
||||||
|
};
|
||||||
|
|
||||||
|
await service.createWorkspaceInstance(options);
|
||||||
|
|
||||||
|
expect(mockInstancesClient.insert).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({
|
||||||
|
project: expect.any(String),
|
||||||
|
zone: 'us-west1-a',
|
||||||
|
instanceResource: expect.objectContaining({
|
||||||
|
name: 'test-inst',
|
||||||
|
}),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('deleteWorkspaceInstance', () => {
|
||||||
|
it('should call instancesClient.delete', async () => {
|
||||||
|
await service.deleteWorkspaceInstance('inst1', 'zone1');
|
||||||
|
expect(mockInstancesClient.delete).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({
|
||||||
|
instance: 'inst1',
|
||||||
|
zone: 'zone1',
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -15,20 +15,80 @@ export interface ProvisionOptions {
|
|||||||
|
|
||||||
export class ComputeService {
|
export class ComputeService {
|
||||||
private client: InstancesClient;
|
private client: InstancesClient;
|
||||||
|
private projectId: string;
|
||||||
|
|
||||||
constructor() {
|
constructor() {
|
||||||
this.client = new InstancesClient();
|
this.client = new InstancesClient();
|
||||||
|
// In a real GCP environment, this is usually available via metadata server or env
|
||||||
|
this.projectId = process.env['GOOGLE_CLOUD_PROJECT'] || 'dev-project';
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Provision a new GCE VM with the Workspace Container
|
* Provision a new GCE VM with the Workspace Container
|
||||||
*/
|
*/
|
||||||
async createWorkspaceInstance(options: ProvisionOptions): Promise<void> {
|
async createWorkspaceInstance(options: ProvisionOptions): Promise<void> {
|
||||||
// TODO: Implement actual instancesClient.insert call
|
const { instanceName, machineType, imageTag, zone } = options;
|
||||||
// For now, we just log and return
|
|
||||||
|
// TODO: Get the actual base image URL from config
|
||||||
|
const containerImage = `us-west1-docker.pkg.dev/${this.projectId}/workspaces/gemini-workspace:${imageTag}`;
|
||||||
|
|
||||||
|
// The container declaration for Container-Optimized OS
|
||||||
|
const containerDeclaration = `
|
||||||
|
spec:
|
||||||
|
containers:
|
||||||
|
- name: workspace
|
||||||
|
image: ${containerImage}
|
||||||
|
securityContext:
|
||||||
|
privileged: false
|
||||||
|
stdin: true
|
||||||
|
tty: true
|
||||||
|
restartPolicy: Always
|
||||||
|
`;
|
||||||
|
|
||||||
|
const [operation] = await this.client.insert({
|
||||||
|
project: this.projectId,
|
||||||
|
zone,
|
||||||
|
instanceResource: {
|
||||||
|
name: instanceName,
|
||||||
|
machineType: `zones/${zone}/machineTypes/${machineType}`,
|
||||||
|
disks: [
|
||||||
|
{
|
||||||
|
boot: true,
|
||||||
|
autoDelete: true,
|
||||||
|
initializeParams: {
|
||||||
|
sourceImage: 'projects/cos-cloud/global/images/family/cos-stable',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
],
|
||||||
|
networkInterfaces: [
|
||||||
|
{
|
||||||
|
network: 'global/networks/default',
|
||||||
|
// We use IAP for access, but a NAT might be needed for outbound internet
|
||||||
|
accessConfigs: [{ name: 'External NAT', type: 'ONE_TO_ONE_NAT' }],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
metadata: {
|
||||||
|
items: [
|
||||||
|
{
|
||||||
|
key: 'gce-container-declaration',
|
||||||
|
value: containerDeclaration,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
key: 'google-logging-enabled',
|
||||||
|
value: 'true',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
// Security: Tag for IAP access
|
||||||
|
tags: {
|
||||||
|
items: ['allow-ssh-iap'],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
// eslint-disable-next-line no-console
|
// eslint-disable-next-line no-console
|
||||||
console.log(
|
console.log(
|
||||||
`[ComputeService] Mocking creation of ${options.instanceName} in ${options.zone}`,
|
`[ComputeService] Creation started for ${instanceName}. Op ID: ${operation.latestResponse.name}`,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -39,10 +99,15 @@ export class ComputeService {
|
|||||||
instanceName: string,
|
instanceName: string,
|
||||||
zone: string,
|
zone: string,
|
||||||
): Promise<void> {
|
): Promise<void> {
|
||||||
// TODO: Implement actual instancesClient.delete call
|
const [operation] = await this.client.delete({
|
||||||
|
project: this.projectId,
|
||||||
|
zone,
|
||||||
|
instance: instanceName,
|
||||||
|
});
|
||||||
|
|
||||||
// eslint-disable-next-line no-console
|
// eslint-disable-next-line no-console
|
||||||
console.log(
|
console.log(
|
||||||
`[ComputeService] Mocking deletion of ${instanceName} in ${zone}`,
|
`[ComputeService] Deletion started for ${instanceName}. Op ID: ${operation.latestResponse.name}`,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,65 @@
|
|||||||
|
/**
|
||||||
|
* @license
|
||||||
|
* Copyright 2026 Google LLC
|
||||||
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||||
|
import { WorkspaceService } from './workspaceService.js';
|
||||||
|
import { Firestore } from '@google-cloud/firestore';
|
||||||
|
|
||||||
|
const mockCollection = {
|
||||||
|
where: vi.fn().mockReturnThis(),
|
||||||
|
get: vi.fn(),
|
||||||
|
doc: vi.fn().mockReturnThis(),
|
||||||
|
set: vi.fn(),
|
||||||
|
delete: vi.fn(),
|
||||||
|
};
|
||||||
|
|
||||||
|
vi.mock('@google-cloud/firestore', () => ({
|
||||||
|
Firestore: vi.fn().mockImplementation(() => ({
|
||||||
|
collection: vi.fn().mockReturnValue(mockCollection),
|
||||||
|
})),
|
||||||
|
}));
|
||||||
|
|
||||||
|
describe('WorkspaceService', () => {
|
||||||
|
let service: WorkspaceService;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
vi.clearAllMocks();
|
||||||
|
service = new WorkspaceService();
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('listWorkspaces', () => {
|
||||||
|
it('should query Firestore for workspaces by ownerId', async () => {
|
||||||
|
const mockDocs = [{ id: '1', data: () => ({ name: 'ws1' }) }];
|
||||||
|
mockCollection.get.mockResolvedValue({ docs: mockDocs });
|
||||||
|
|
||||||
|
const result = await service.listWorkspaces('user1');
|
||||||
|
|
||||||
|
expect(result).toHaveLength(1);
|
||||||
|
expect(result[0].id).toBe('1');
|
||||||
|
expect(result[0].name).toBe('ws1');
|
||||||
|
// Verify the collection name was 'workspaces'
|
||||||
|
const firestoreInstance = vi.mocked(Firestore).mock.results[0].value;
|
||||||
|
expect(firestoreInstance.collection).toHaveBeenCalledWith('workspaces');
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('createWorkspace', () => {
|
||||||
|
it('should save workspace data to Firestore', async () => {
|
||||||
|
const data = {
|
||||||
|
owner_id: 'u1',
|
||||||
|
name: 'test',
|
||||||
|
instance_name: 'inst',
|
||||||
|
status: 'READY',
|
||||||
|
machine_type: 'e2',
|
||||||
|
zone: 'us1',
|
||||||
|
created_at: 'now',
|
||||||
|
};
|
||||||
|
await service.createWorkspace('id1', data);
|
||||||
|
expect(mockCollection.doc).toHaveBeenCalledWith('id1');
|
||||||
|
expect(mockCollection.set).toHaveBeenCalledWith(data);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -25,8 +25,8 @@ Implement the core API to manage GCE-based workspaces.
|
|||||||
- [x] Initialize `packages/workspace-manager/`.
|
- [x] Initialize `packages/workspace-manager/`.
|
||||||
- [x] Implement Express server for `/workspaces` (List, Create, Delete).
|
- [x] Implement Express server for `/workspaces` (List, Create, Delete).
|
||||||
- [x] Integrate Firestore to track workspace state (owner, instance_id, status).
|
- [x] Integrate Firestore to track workspace state (owner, instance_id, status).
|
||||||
- [ ] Integrate `@google-cloud/compute` for GCE instance lifecycle.
|
- [x] Integrate `@google-cloud/compute` for GCE instance lifecycle.
|
||||||
- [ ] Provision a VM with `Container-on-VM` settings pointing to the
|
- [x] Provision a VM with `Container-on-VM` settings pointing to the
|
||||||
`gemini-workspace` image.
|
`gemini-workspace` image.
|
||||||
|
|
||||||
### Task 1.3: Cloud Run Deployment (v1)
|
### Task 1.3: Cloud Run Deployment (v1)
|
||||||
|
|||||||
Reference in New Issue
Block a user