diff --git a/packages/cli/src/commands/workspace/connect.ts b/packages/cli/src/commands/workspace/connect.ts index b70315a748..5882073952 100644 --- a/packages/cli/src/commands/workspace/connect.ts +++ b/packages/cli/src/commands/workspace/connect.ts @@ -10,7 +10,8 @@ import { SSHService, SyncService, type Config, - type WorkspaceHubInfo + type WorkspaceHubInfo, + debugLogger } from '@google-gemini-cli-core'; import { exitCli } from '../utils.js'; diff --git a/packages/core/src/services/workspaceHubClient.ts b/packages/core/src/services/workspaceHubClient.ts index 4d98311b48..c5835a3533 100644 --- a/packages/core/src/services/workspaceHubClient.ts +++ b/packages/core/src/services/workspaceHubClient.ts @@ -97,7 +97,7 @@ export class WorkspaceHubClient { * Notify the hub that a user is connecting to a workspace */ async notifyConnect(id: string): Promise { - const url = `${this.hubUrl}/workspaces/${id}/connect`; + const url = new URL(`/workspaces/${id}/connect`, this.hubUrl).toString(); const response = await fetchWithTimeout(url, 5000, { method: 'POST', }); diff --git a/packages/workspace-manager/src/index.ts b/packages/workspace-manager/src/index.ts index 56c9f37744..f3eccd8fde 100644 --- a/packages/workspace-manager/src/index.ts +++ b/packages/workspace-manager/src/index.ts @@ -23,11 +23,12 @@ app.get('/health', (_req, res) => { * Endpoint to trigger cleanup of idle workspaces. * Typically called by a Cloud Scheduler job. */ -app.post('/cleanup', async (_req, res) => { +app.post('/cleanup', async (req, res) => { try { const cleanupService = new CleanupService(); - const count = await cleanupService.cleanupIdleWorkspaces(); - res.json({ status: 'ok', cleaned_count: count }); + const ttlMinutes = req.body.ttl_minutes ? Number(req.body.ttl_minutes) : 240; + const count = await cleanupService.cleanupIdleWorkspaces(ttlMinutes); + res.json({ status: 'ok', cleaned_count: count, ttl_minutes: ttlMinutes }); } catch (error) { const message = error instanceof Error ? error.message : String(error); res.status(500).json({ error: message }); diff --git a/packages/workspace-manager/src/routes/workspaceRoutes.test.ts b/packages/workspace-manager/src/routes/workspaceRoutes.test.ts index 6677c6d7df..d37e74e8f3 100644 --- a/packages/workspace-manager/src/routes/workspaceRoutes.test.ts +++ b/packages/workspace-manager/src/routes/workspaceRoutes.test.ts @@ -22,6 +22,7 @@ vi.mock('../services/computeService.js', () => ({ ComputeService: vi.fn().mockImplementation(() => ({ createWorkspaceInstance: vi.fn().mockResolvedValue(undefined), deleteWorkspaceInstance: vi.fn().mockResolvedValue(undefined), + getProjectId: vi.fn().mockReturnValue('dev-project'), })), })); diff --git a/packages/workspace-manager/src/services/cleanupService.test.ts b/packages/workspace-manager/src/services/cleanupService.test.ts new file mode 100644 index 0000000000..500850d3d6 --- /dev/null +++ b/packages/workspace-manager/src/services/cleanupService.test.ts @@ -0,0 +1,55 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { CleanupService } from './cleanupService.js'; +import { WorkspaceService } from './workspaceService.js'; +import { ComputeService } from './computeService.js'; + +vi.mock('./workspaceService.js'); +vi.mock('./computeService.js'); + +describe('CleanupService', () => { + let service: CleanupService; + let mockWorkspaceService: any; + let mockComputeService: any; + + beforeEach(() => { + vi.clearAllMocks(); + service = new CleanupService(); + mockWorkspaceService = vi.mocked(new WorkspaceService()); + mockComputeService = vi.mocked(new ComputeService()); + + // Wire up the internal instances + (service as any).workspaceService = mockWorkspaceService; + (service as any).computeService = mockComputeService; + }); + + it('should cleanup workspaces older than TTL', async () => { + const now = new Date(); + const oldDate = new Date(now.getTime() - 300 * 60 * 1000).toISOString(); // 5 hours ago + const newDate = new Date(now.getTime() - 60 * 60 * 1000).toISOString(); // 1 hour ago + + const workspaces = [ + { id: 'old-1', name: 'old', last_connected_at: oldDate, instance_name: 'inst-1', zone: 'z1', status: 'READY' }, + { id: 'new-1', name: 'new', last_connected_at: newDate, instance_name: 'inst-2', zone: 'z1', status: 'READY' }, + { id: 'stuck-1', name: 'stuck', last_connected_at: oldDate, instance_name: 'inst-3', zone: 'z1', status: 'PROVISIONING' }, + ]; + + mockWorkspaceService.listAllWorkspaces.mockResolvedValue(workspaces); + + const cleanedCount = await service.cleanupIdleWorkspaces(240); // 4 hour TTL + + expect(cleanedCount).toBe(2); + expect(mockComputeService.deleteWorkspaceInstance).toHaveBeenCalledWith('inst-1', 'z1'); + expect(mockComputeService.deleteWorkspaceInstance).toHaveBeenCalledWith('inst-3', 'z1'); + expect(mockWorkspaceService.deleteWorkspace).toHaveBeenCalledWith('old-1'); + expect(mockWorkspaceService.deleteWorkspace).toHaveBeenCalledWith('stuck-1'); + + // Should NOT have deleted the new one + expect(mockWorkspaceService.deleteWorkspace).not.toHaveBeenCalledWith('new-1'); + }); +}); diff --git a/packages/workspace-manager/src/services/cleanupService.ts b/packages/workspace-manager/src/services/cleanupService.ts index 8fb77d745a..e689101863 100644 --- a/packages/workspace-manager/src/services/cleanupService.ts +++ b/packages/workspace-manager/src/services/cleanupService.ts @@ -13,41 +13,37 @@ export class CleanupService { /** * Identifies and deletes workspaces that have been idle for too long. + * Targets READY, PROVISIONING, and ERROR statuses to prevent leaks. * @param ttlMinutes Threshold for idleness in minutes. */ async cleanupIdleWorkspaces(ttlMinutes: number = 240): Promise { - // 1. Get all workspaces (Note: in a huge system we'd need to paginate or use a query) - // For now, we list all, but in prod we'd query by last_connected_at const now = new Date(); const threshold = new Date(now.getTime() - ttlMinutes * 60 * 1000); - // Using simple approach: list all, filter in code. - // Real implementation should use Firestore .where('last_connected_at', '<', threshold.toISOString()) - // but Firestore where requires composite indexes for some queries. - - // We'll just list active ones for now - const snapshot = await (this.workspaceService as any).collection - .where('status', '==', 'READY') - .get(); + // Fetch all workspaces for evaluation + const workspaces = await this.workspaceService.listAllWorkspaces(); let cleanedCount = 0; - for (const doc of snapshot.docs) { - const data = doc.data(); - const lastConnected = new Date(data.last_connected_at); + for (const ws of workspaces) { + const lastConnected = new Date(ws.last_connected_at); + // Cleanup if idle past TTL if (lastConnected < threshold) { // eslint-disable-next-line no-console - console.log(`[Cleanup] Workspace ${doc.id} (${data.name}) is idle since ${data.last_connected_at}. Cleaning up...`); + console.log(`[Cleanup] Workspace ${ws.id} (${ws.name}) is idle since ${ws.last_connected_at} (Status: ${ws.status}). Cleaning up...`); try { - // Delete GCE - await this.computeService.deleteWorkspaceInstance(data.instance_name, data.zone); - // Update status or delete from DB - await this.workspaceService.deleteWorkspace(doc.id); + // 1. Delete GCE instance if it exists (or attempts to) + // ComputeService handles missing instances gracefully + await this.computeService.deleteWorkspaceInstance(ws.instance_name, ws.zone); + + // 2. Delete record from Firestore + await this.workspaceService.deleteWorkspace(ws.id); + cleanedCount++; } catch (err) { // eslint-disable-next-line no-console - console.error(`[Cleanup] Failed to clean up ${doc.id}:`, err); + console.error(`[Cleanup] Failed to clean up ${ws.id}:`, err); } } } diff --git a/packages/workspace-manager/src/services/workspaceService.ts b/packages/workspace-manager/src/services/workspaceService.ts index 7ae2230cfb..864d83108f 100644 --- a/packages/workspace-manager/src/services/workspaceService.ts +++ b/packages/workspace-manager/src/services/workspaceService.ts @@ -29,12 +29,25 @@ export class WorkspaceService { this.firestore = new Firestore(); } - private get collection() { + getCollection() { return this.firestore.collection('workspaces'); } + async listAllWorkspaces(): Promise { + const snapshot = await this.getCollection().get(); + + return snapshot.docs.map((doc) => { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + const data = doc.data() as WorkspaceData; + return { + id: doc.id, + ...data, + }; + }); + } + async listWorkspaces(ownerId: string): Promise { - const snapshot = await this.collection + const snapshot = await this.getCollection() .where('owner_id', '==', ownerId) .get(); @@ -56,14 +69,14 @@ export class WorkspaceService { } async createWorkspace(id: string, data: WorkspaceData): Promise { - await this.collection.doc(id).set(data); + await this.getCollection().doc(id).set(data); } async updateWorkspace(id: string, data: Partial): Promise { - await this.collection.doc(id).update(data); + await this.getCollection().doc(id).update(data); } async deleteWorkspace(id: string): Promise { - await this.collection.doc(id).delete(); + await this.getCollection().doc(id).delete(); } }