feat(workspaces): enhance Hub cleanup with better encapsulation and robustness

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';

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<void> {
-    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',
     });

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 });

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'),
     })),
   }));
 

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');
+  });
+});

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<number> {
-    // 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);
             }
         }
     }

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<WorkspaceRecord[]> {
+    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<WorkspaceRecord[]> {
-    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<void> {
-    await this.collection.doc(id).set(data);
+    await this.getCollection().doc(id).set(data);
   }
 
   async updateWorkspace(id: string, data: Partial<WorkspaceData>): Promise<void> {
-    await this.collection.doc(id).update(data);
+    await this.getCollection().doc(id).update(data);
   }
 
   async deleteWorkspace(id: string): Promise<void> {
-    await this.collection.doc(id).delete();
+    await this.getCollection().doc(id).delete();
   }
 }