vibn-frontend/lib/cloud-run-workspace.ts

/**
 * Cloud Run Workspace Provisioning
 *
 * Provisions a dedicated Theia IDE instance per Vibn project using
 * Google Cloud Run. Each workspace:
 *   - Gets its own Cloud Run service: theia-{slug}
 *   - Scales to zero when idle (zero cost when not in use)
 *   - Starts in ~5-15s from cached image on demand
 *   - Is accessible at the Cloud Run URL stored on the project record
 *   - Auth is enforced by our Vibn session before the URL is revealed
 */

import { GoogleAuth, JWT } from 'google-auth-library';

const PROJECT_ID = 'master-ai-484822';
const REGION     = 'northamerica-northeast1';
const IMAGE      = `${REGION}-docker.pkg.dev/${PROJECT_ID}/vibn-ide/theia:latest`;
const VIBN_URL   = process.env.NEXTAUTH_URL ?? 'https://vibnai.com';

const CLOUD_RUN_API = `https://run.googleapis.com/v2/projects/${PROJECT_ID}/locations/${REGION}/services`;
const SCOPES = ['https://www.googleapis.com/auth/cloud-platform'];

async function getAccessToken(): Promise<string> {
  // Prefer an explicit service account key (avoids GCE metadata scope limitations).
  // Stored as base64 to survive Docker ARG/ENV special-character handling.
  const keyB64 = process.env.GOOGLE_SERVICE_ACCOUNT_KEY_B64;
  if (keyB64) {
    const keyJson = Buffer.from(keyB64, 'base64').toString('utf-8');
    const key = JSON.parse(keyJson) as {
      client_email: string;
      private_key: string;
    };
    const jwt = new JWT({
      email: key.client_email,
      key: key.private_key,
      scopes: SCOPES,
    });
    const token = await jwt.getAccessToken();
    if (!token.token) throw new Error('Failed to get GCP access token from service account key');
    return token.token as string;
  }

  // Fall back to ADC (works locally or on GCE with cloud-platform scope)
  const auth = new GoogleAuth({ scopes: SCOPES });
  const client = await auth.getClient();
  const token = await client.getAccessToken();
  if (!token.token) throw new Error('Failed to get GCP access token');
  return token.token;
}

export interface ProvisionResult {
  serviceUrl: string;   // https://theia-{slug}-xxx.run.app
  serviceName: string;  // theia-{slug}
}

/**
 * Creates a new Cloud Run service for a Vibn project workspace.
 * The service scales to zero when idle and starts on first request.
 */
export async function provisionTheiaWorkspace(
  slug: string,
  projectId: string,
  giteaRepo: string | null,
): Promise<ProvisionResult> {
  const token = await getAccessToken();
  const serviceName = `theia-${slug}`.slice(0, 49); // Cloud Run max 49 chars

  // Cloud Run v2: name must be empty in the body — it's passed via ?serviceId= in the URL
  const serviceBody = {
    template: {
      scaling: {
        minInstanceCount: 0,   // scale to zero when idle
        maxInstanceCount: 1,   // one instance per workspace
      },
      containers: [{
        image: IMAGE,
        ports: [{ containerPort: 3000 }],
        resources: {
          limits: { cpu: '1', memory: '2Gi' },
          cpuIdle: true,        // only allocate CPU during requests
        },
        env: [
          { name: 'VIBN_PROJECT_ID',   value: projectId },
          { name: 'VIBN_PROJECT_SLUG', value: slug },
          { name: 'VIBN_API_URL',      value: VIBN_URL },
          { name: 'GITEA_REPO',        value: giteaRepo ?? '' },
          { name: 'GITEA_API_URL',     value: process.env.GITEA_API_URL ?? 'https://git.vibnai.com' },
          // Token lets the startup script clone and push to the project's repo
          { name: 'GITEA_TOKEN',       value: process.env.GITEA_API_TOKEN ?? '' },
          // Gemini API key — needed by startup.sh to configure AI features in Theia
          { name: 'GOOGLE_API_KEY',    value: process.env.GOOGLE_API_KEY ?? '' },
        ],
        // 5 minute startup timeout — Theia needs time to initialise
        startupProbe: {
          httpGet: { path: '/', port: 3000 },
          failureThreshold: 30,
          periodSeconds: 10,
        },
      }],
      // Keep container alive for 15 minutes of idle before scaling to zero
      timeout: '900s',
    },
    ingress: 'INGRESS_TRAFFIC_ALL',
  };

  const createRes = await fetch(`${CLOUD_RUN_API}?serviceId=${serviceName}`, {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${token}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify(serviceBody),
  });

  if (!createRes.ok) {
    const body = await createRes.text();

    // 409 = service already exists — fetch its URL instead of failing
    if (createRes.status === 409) {
      console.log(`[workspace] Cloud Run service already exists: ${serviceName} — fetching existing URL`);
      const serviceUrl = await waitForServiceUrl(serviceName, token);
      await allowUnauthenticated(serviceName, token);
      console.log(`[workspace] Linked to existing service: ${serviceName} → ${serviceUrl}`);
      return { serviceUrl, serviceName };
    }

    throw new Error(`Cloud Run create service failed (${createRes.status}): ${body}`);
  }

  // Make service publicly accessible (auth handled by Vibn before URL is revealed)
  await allowUnauthenticated(serviceName, token);

  // Poll until the service URL is available (usually 10-30s)
  const serviceUrl = await waitForServiceUrl(serviceName, token);

  console.log(`[workspace] Cloud Run service ready: ${serviceName} → ${serviceUrl}`);
  return { serviceUrl, serviceName };
}

/**
 * Grants allUsers invoker access so the service URL works without GCP auth.
 * Vibn controls access by only sharing the URL with the project owner.
 */
async function allowUnauthenticated(serviceName: string, token: string): Promise<void> {
  await fetch(`${CLOUD_RUN_API}/${serviceName}:setIamPolicy`, {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${token}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      policy: {
        bindings: [{ role: 'roles/run.invoker', members: ['allUsers'] }],
      },
    }),
  });
}

/**
 * Polls until Cloud Run reports the service URL (service is ready).
 */
async function waitForServiceUrl(serviceName: string, token: string, maxWaitMs = 60_000): Promise<string> {
  const deadline = Date.now() + maxWaitMs;

  while (Date.now() < deadline) {
    await new Promise(r => setTimeout(r, 3000));

    const res = await fetch(`${CLOUD_RUN_API}/${serviceName}`, {
      headers: { Authorization: `Bearer ${token}` },
    });

    if (res.ok) {
      const svc = await res.json() as { urls?: string[] };
      if (svc.urls?.[0]) return svc.urls[0];
    }
  }

  // Return expected URL pattern even if polling timed out
  return `https://${serviceName}-${PROJECT_ID.slice(-6)}.${REGION}.run.app`;
}

/**
 * Triggers a warm-up request to a workspace so the container is ready
 * before the user clicks "Open IDE". Call this on user login.
 */
export async function prewarmWorkspace(serviceUrl: string): Promise<void> {
  try {
    await fetch(`${serviceUrl}/`, { signal: AbortSignal.timeout(5000) });
  } catch {
    // Ignore — fire and forget
  }
}

/**
 * Deletes a Cloud Run workspace service.
 */
export async function deleteTheiaWorkspace(serviceName: string): Promise<void> {
  const token = await getAccessToken();
  await fetch(`${CLOUD_RUN_API}/${serviceName}`, {
    method: 'DELETE',
    headers: { Authorization: `Bearer ${token}` },
  });
}