From 5647fb09aeb7eb3d9d8b090f87941b78e05890e7 Mon Sep 17 00:00:00 2001 From: mkorwel Date: Thu, 19 Mar 2026 09:50:43 -0700 Subject: [PATCH] feat(workspaces): implement Hub auto-cleanup and connection heartbeat --- .../cli/src/commands/workspace/connect.ts | 9 ++- .../core/src/services/workspaceHubClient.ts | 15 +++++ packages/workspace-manager/src/index.ts | 16 ++++++ .../src/routes/workspaceRoutes.ts | 31 +++++++++- .../src/services/cleanupService.ts | 57 +++++++++++++++++++ .../src/services/workspaceService.test.ts | 1 + .../src/services/workspaceService.ts | 18 ++++-- 7 files changed, 140 insertions(+), 7 deletions(-) create mode 100644 packages/workspace-manager/src/services/cleanupService.ts diff --git a/packages/cli/src/commands/workspace/connect.ts b/packages/cli/src/commands/workspace/connect.ts index d297e6a640..b70315a748 100644 --- a/packages/cli/src/commands/workspace/connect.ts +++ b/packages/cli/src/commands/workspace/connect.ts @@ -130,7 +130,14 @@ export async function connectToWorkspace(args: ArgumentsCamelCase): } } - // 3. Connect via SSH + // 3. Notify Hub of connection (refresh TTL) + try { + await client.notifyConnect(readyWs.id); + } catch (err) { + debugLogger.warn(`[Connect] Failed to notify Hub of connection:`, err); + } + + // 4. Connect via SSH // eslint-disable-next-line no-console console.log(chalk.green(`🚀 Teleporting to ${instanceName} (${zone})...`)); diff --git a/packages/core/src/services/workspaceHubClient.ts b/packages/core/src/services/workspaceHubClient.ts index 7ce280f888..4d98311b48 100644 --- a/packages/core/src/services/workspaceHubClient.ts +++ b/packages/core/src/services/workspaceHubClient.ts @@ -93,6 +93,21 @@ export class WorkspaceHubClient { return (await response.json()) as WorkspaceHubInfo; } + /** + * Notify the hub that a user is connecting to a workspace + */ + async notifyConnect(id: string): Promise { + const url = `${this.hubUrl}/workspaces/${id}/connect`; + const response = await fetchWithTimeout(url, 5000, { + method: 'POST', + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`Hub API error (${response.status}): ${errorText}`); + } + } + /** * Delete a workspace */ diff --git a/packages/workspace-manager/src/index.ts b/packages/workspace-manager/src/index.ts index 2e60f3df0b..56c9f37744 100644 --- a/packages/workspace-manager/src/index.ts +++ b/packages/workspace-manager/src/index.ts @@ -7,6 +7,7 @@ import express from 'express'; import { workspaceRouter } from './routes/workspaceRoutes.js'; import { iapMiddleware } from './middleware/iap.js'; +import { CleanupService } from './services/cleanupService.js'; export const app = express(); app.use(express.json()); @@ -18,6 +19,21 @@ app.get('/health', (_req, res) => { res.send({ status: 'ok' }); }); +/** + * Endpoint to trigger cleanup of idle workspaces. + * Typically called by a Cloud Scheduler job. + */ +app.post('/cleanup', async (_req, res) => { + try { + const cleanupService = new CleanupService(); + const count = await cleanupService.cleanupIdleWorkspaces(); + res.json({ status: 'ok', cleaned_count: count }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + res.status(500).json({ error: message }); + } +}); + // Register Workspace Routes app.use('/workspaces', workspaceRouter); diff --git a/packages/workspace-manager/src/routes/workspaceRoutes.ts b/packages/workspace-manager/src/routes/workspaceRoutes.ts index 05fb24b347..147ce14911 100644 --- a/packages/workspace-manager/src/routes/workspaceRoutes.ts +++ b/packages/workspace-manager/src/routes/workspaceRoutes.ts @@ -46,6 +46,7 @@ router.post('/', async (req, res) => { const workspaceId = uuidv4(); const instanceName = `workspace-${workspaceId.slice(0, 8)}`; + const now = new Date().toISOString(); const workspaceData = { owner_id: authReq.user.id, name, @@ -54,7 +55,8 @@ router.post('/', async (req, res) => { machine_type: machineType, zone, project_id: computeService.getProjectId(), - created_at: new Date().toISOString(), + created_at: now, + last_connected_at: now, }; // 1. Save to state store @@ -78,6 +80,33 @@ router.post('/', async (req, res) => { } }); +router.post('/:id/connect', async (req, res) => { + try { + const authReq = req as unknown as AuthenticatedRequest; + const { id } = req.params; + const workspace = await workspaceService.getWorkspace(id); + + if (!workspace) { + res.status(404).json({ error: 'Workspace not found' }); + return; + } + + if (workspace.owner_id !== authReq.user.id) { + res.status(403).json({ error: 'Unauthorized' }); + return; + } + + // Update last_connected_at timestamp + const now = new Date().toISOString(); + await workspaceService.updateWorkspace(id, { last_connected_at: now }); + + res.json({ status: 'ok', last_connected_at: now }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + res.status(500).json({ error: message }); + } +}); + router.delete('/:id', async (req, res) => { try { const authReq = req as unknown as AuthenticatedRequest; diff --git a/packages/workspace-manager/src/services/cleanupService.ts b/packages/workspace-manager/src/services/cleanupService.ts new file mode 100644 index 0000000000..8fb77d745a --- /dev/null +++ b/packages/workspace-manager/src/services/cleanupService.ts @@ -0,0 +1,57 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { WorkspaceService } from './workspaceService.js'; +import { ComputeService } from './computeService.js'; + +export class CleanupService { + private workspaceService = new WorkspaceService(); + private computeService = new ComputeService(); + + /** + * Identifies and deletes workspaces that have been idle for too long. + * @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(); + + let cleanedCount = 0; + for (const doc of snapshot.docs) { + const data = doc.data(); + const lastConnected = new Date(data.last_connected_at); + + 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...`); + + try { + // Delete GCE + await this.computeService.deleteWorkspaceInstance(data.instance_name, data.zone); + // Update status or delete from DB + await this.workspaceService.deleteWorkspace(doc.id); + cleanedCount++; + } catch (err) { + // eslint-disable-next-line no-console + console.error(`[Cleanup] Failed to clean up ${doc.id}:`, err); + } + } + } + + return cleanedCount; + } +} diff --git a/packages/workspace-manager/src/services/workspaceService.test.ts b/packages/workspace-manager/src/services/workspaceService.test.ts index b9cd4fdd12..a3c98b2e09 100644 --- a/packages/workspace-manager/src/services/workspaceService.test.ts +++ b/packages/workspace-manager/src/services/workspaceService.test.ts @@ -57,6 +57,7 @@ describe('WorkspaceService', () => { zone: 'us1', project_id: 'p1', created_at: 'now', + last_connected_at: 'now', }; await service.createWorkspace('id1', data); expect(mockCollection.doc).toHaveBeenCalledWith('id1'); diff --git a/packages/workspace-manager/src/services/workspaceService.ts b/packages/workspace-manager/src/services/workspaceService.ts index a440813da0..7ae2230cfb 100644 --- a/packages/workspace-manager/src/services/workspaceService.ts +++ b/packages/workspace-manager/src/services/workspaceService.ts @@ -15,6 +15,7 @@ export interface WorkspaceData { zone: string; project_id: string; created_at: string; + last_connected_at: string; } export interface WorkspaceRecord extends WorkspaceData { @@ -28,9 +29,12 @@ export class WorkspaceService { this.firestore = new Firestore(); } + private get collection() { + return this.firestore.collection('workspaces'); + } + async listWorkspaces(ownerId: string): Promise { - const snapshot = await this.firestore - .collection('workspaces') + const snapshot = await this.collection .where('owner_id', '==', ownerId) .get(); @@ -45,17 +49,21 @@ export class WorkspaceService { } async getWorkspace(id: string): Promise { - const doc = await this.firestore.collection('workspaces').doc(id).get(); + const doc = await this.collection.doc(id).get(); if (!doc.exists) return null; // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion return { id: doc.id, ...(doc.data() as WorkspaceData) }; } async createWorkspace(id: string, data: WorkspaceData): Promise { - await this.firestore.collection('workspaces').doc(id).set(data); + await this.collection.doc(id).set(data); + } + + async updateWorkspace(id: string, data: Partial): Promise { + await this.collection.doc(id).update(data); } async deleteWorkspace(id: string): Promise { - await this.firestore.collection('workspaces').doc(id).delete(); + await this.collection.doc(id).delete(); } }