fix(antigravity): add onboardUser fallback when loadCodeAssist returns no project (#8886)

Validated in local merge-train T4 (HouMinXi+Zartharas+Andrian+artickc)

next.config.mjs

@@ -173,6 +173,10 @@ const nextConfig = {
     serverActions: {
       bodySizeLimit: process.env.OMNIROUTE_SERVER_ACTIONS_BODY_LIMIT || "50mb",
     },
+    // Reduce peak heap during production builds (Next.js 15+).
+    webpackMemoryOptimizations: true,
+    // Run webpack in a separate Node worker, lowering main-process memory.
+    webpackBuildWorker: true,
     // Next.js proxy (middleware) has a default 10MB body clone limit. File
     // uploads (OpenAI-compatible /v1/files) routinely exceed this. Match the
     // 512 MB server-side cap; tune via env if needed.

open-sse/config/antigravityUpstream.ts

@@ -12,6 +12,12 @@ export const ANTIGRAVITY_BOOTSTRAP_BASE_URLS = Object.freeze([
   "https://cloudcode-pa.googleapis.com",
 ]);
 
+export const ANTIGRAVITY_ONBOARD_PATH = "/v1internal:onboardUser";
+
+export function getAntigravityOnboardUrls(): string[] {
+  return ANTIGRAVITY_BOOTSTRAP_BASE_URLS.map((base) => `${base}${ANTIGRAVITY_ONBOARD_PATH}`);
+}
+
 const ANTIGRAVITY_MODELS_PATH = "/v1internal:models";
 const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH = "/v1internal:fetchAvailableModels";
 

open-sse/services/antigravityProjectBootstrap.ts

@@ -1,5 +1,5 @@
 /**
- * Antigravity project bootstrap — loadCodeAssist.
+ * Antigravity project bootstrap — loadCodeAssist + onboardUser.
  *
  * The Google Cloud Code Assist API (/v1internal:models) requires a prior
  * /v1internal:loadCodeAssist call to assign a project context to the
@@ -10,52 +10,70 @@
  * attempt. Results are memoized per-token for the process lifetime to
  * avoid redundant round-trips.
  *
- * Based on the Antigravity loadCodeAssist flow and the CLIProxyAPI reference
- * implementation in internal/runtime/executor/antigravity_executor.go.
+ * When loadCodeAssist returns no project (account never onboarded),
+ * the fallback calls onboardUser to create the project, then retries.
  */
 
 import {
   getAntigravityContentHeaders,
   getAntigravityLoadCodeAssistMetadata,
 } from "./antigravityHeaders.ts";
+import { extractCodeAssistOnboardTierId } from "./codeAssistSubscription.ts";
 import type { AntigravityClientProfile } from "./antigravityClientProfile.ts";
-import { ANTIGRAVITY_BOOTSTRAP_BASE_URLS } from "../config/antigravityUpstream.ts";
+import { ANTIGRAVITY_BOOTSTRAP_BASE_URLS, getAntigravityOnboardUrls } from "../config/antigravityUpstream.ts";
 
 const LOAD_CODE_ASSIST_PATH = "/v1internal:loadCodeAssist";
 const BOOTSTRAP_TIMEOUT_MS = 8_000;
+const ONBOARD_TIMEOUT_MS = 15_000;
+const DEFAULT_TIER_ID = "legacy-tier";
 
-/** Ordered list of loadCodeAssist endpoint URLs (mirrors the models discovery order). */
+/** Ordered list of loadCodeAssist endpoint URLs. */
 export function getAntigravityLoadCodeAssistUrls(): string[] {
   return ANTIGRAVITY_BOOTSTRAP_BASE_URLS.map((base) => `${base}${LOAD_CODE_ASSIST_PATH}`);
 }
 
+/** Max entries in the per-token caches (prevents unbounded growth). */
+const MAX_CACHE_SIZE = 256;
+
+/** LRU-style Map: deleting and re-inserting moves the key to the end. */
+function evictOldest(cache: Map<string, unknown>): void {
+  if (cache.size >= MAX_CACHE_SIZE) {
+    const oldest = cache.keys().next().value;
+    if (oldest !== undefined) cache.delete(oldest);
+  }
+}
+
 /** Per-token memoization cache (lives for the process lifetime). */
 const projectCache = new Map<string, string>();
 
+/** Per-key lock to prevent concurrent onboard attempts for the same token. */
+const onboardLocks = new Map<string, Promise<boolean>>();
+
 type FetchLike = (url: string, init?: RequestInit) => Promise<Response>;
 
 function getProjectCacheKey(accessToken: string, clientProfile: AntigravityClientProfile): string {
   return `${clientProfile}:${accessToken}`;
 }
 
+type LoadCodeAssistResult = { projectId: string | null; tierId: string };
+
 /**
  * Attempt loadCodeAssist against each known base URL in order.
- * Returns the discovered project id, or null if all endpoints fail.
+ * Returns the discovered project id and tier id, or null projectId if all endpoints fail.
  */
 async function tryLoadCodeAssist(
   accessToken: string,
   fetchImpl: FetchLike,
   clientProfile: AntigravityClientProfile,
   signal?: AbortSignal
-): Promise<string | null> {
+): Promise<LoadCodeAssistResult> {
   const urls = getAntigravityLoadCodeAssistUrls();
   const headers = getAntigravityContentHeaders(clientProfile, accessToken);
 
-  for (const url of urls) {
+  for (let i = 0; i < urls.length; i++) {
+    const url = urls[i];
     if (signal?.aborted) throw signal.reason;
     try {
-      // Combine the caller's cancellation signal (#8098) with the per-attempt
-      // bootstrap timeout so an aborted request tears down immediately.
       const timeoutSignal = AbortSignal.timeout(BOOTSTRAP_TIMEOUT_MS);
       const response = await fetchImpl(url, {
         method: "POST",
@@ -75,7 +93,7 @@ async function tryLoadCodeAssist(
 
       // cloudaicompanionProject may be a plain string or an object with an id field.
       const raw = data.cloudaicompanionProject;
-      let projectId =
+      const projectId =
         typeof raw === "string"
           ? raw.trim()
           : raw &&
@@ -84,16 +102,21 @@ async function tryLoadCodeAssist(
             ? ((raw as Record<string, unknown>).id as string).trim()
             : "";
 
+      const tierId = extractCodeAssistOnboardTierId(data) || DEFAULT_TIER_ID;
+
       if (projectId) {
-        return projectId;
+        return { projectId, tierId };
       }
 
+      // Continue to next URL if available — a different endpoint might
+      // have the project. Only return empty when this is the last URL.
+      if (i === urls.length - 1) {
+        return { projectId: null, tierId };
+      }
       console.warn(
         `[models] antigravity loadCodeAssist at ${url} returned no project id — trying next`
       );
     } catch (error) {
-      // A caller-initiated abort (#8098) must propagate, not be swallowed as a
-      // "try next URL" transient — otherwise a cancelled request silently proceeds.
       if (signal?.aborted || (error instanceof Error && error.name === "AbortError")) {
         throw signal?.reason ?? error;
       }
@@ -101,7 +124,65 @@ async function tryLoadCodeAssist(
       console.warn(`[models] antigravity loadCodeAssist threw for ${url}: ${msg} — trying next`);
     }
   }
-  return null;
+  return { projectId: null, tierId: DEFAULT_TIER_ID };
+}
+
+/**
+ * Attempt onboardUser to create a Cloud Code project for the account.
+ * Called when loadCodeAssist returns no project — the account has never
+ * been onboarded. Returns true if any endpoint reports success.
+ */
+async function tryOnboardUser(
+  accessToken: string,
+  fetchImpl: FetchLike,
+  clientProfile: AntigravityClientProfile,
+  tierId: string,
+  signal?: AbortSignal
+): Promise<boolean> {
+  const urls = getAntigravityOnboardUrls();
+  const headers = getAntigravityContentHeaders(clientProfile, accessToken);
+
+  for (const url of urls) {
+    if (signal?.aborted) throw signal.reason;
+    try {
+      const timeoutSignal = AbortSignal.timeout(ONBOARD_TIMEOUT_MS);
+      const response = await fetchImpl(url, {
+        method: "POST",
+        headers,
+        body: JSON.stringify({
+          tier_id: tierId,
+          metadata: getAntigravityLoadCodeAssistMetadata(),
+        }),
+        signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal,
+      });
+
+      if (response.ok) {
+        return true;
+      }
+
+      console.warn(
+        `[models] antigravity onboardUser failed at ${url} (${response.status}) — trying next`
+      );
+    } catch (error) {
+      if (signal?.aborted || (error instanceof Error && error.name === "AbortError")) {
+        throw signal?.reason ?? error;
+      }
+      const msg = error instanceof Error ? error.message : String(error);
+      console.warn(`[models] antigravity onboardUser threw for ${url}: ${msg} — trying next`);
+    }
+  }
+  return false;
+}
+
+/** Per-token memoization for accounts we already tried onboarding (avoid repeated calls). */
+const onboardAttemptedCache = new Set<string>();
+
+function addToOnboardAttemptedCache(key: string): void {
+  if (onboardAttemptedCache.size >= MAX_CACHE_SIZE) {
+    const oldest = onboardAttemptedCache.values().next().value;
+    if (oldest !== undefined) onboardAttemptedCache.delete(oldest);
+  }
+  onboardAttemptedCache.add(key);
 }
 
 /**
@@ -123,22 +204,72 @@ export async function ensureAntigravityProjectAssigned(
 ): Promise<string | undefined> {
   const cacheKey = getProjectCacheKey(accessToken, clientProfile);
   if (projectCache.has(cacheKey)) {
-    return projectCache.get(cacheKey); // already bootstrapped for this token
+    const cached = projectCache.get(cacheKey)!;
+    // Touch on read: delete+reinsert moves this entry to the end (LRU).
+    projectCache.delete(cacheKey);
+    projectCache.set(cacheKey, cached);
+    return cached;
   }
 
-  const projectId = await tryLoadCodeAssist(accessToken, fetchImpl, clientProfile, signal);
+  const { projectId: initialProjectId, tierId } = await tryLoadCodeAssist(
+    accessToken, fetchImpl, clientProfile, signal
+  );
+
+  let projectId = initialProjectId;
+
+  // loadCodeAssist is read-only — if the account was never onboarded, it returns
+  // empty. Call onboardUser to create the project, then retry discovery.
+  if (!projectId && !onboardAttemptedCache.has(cacheKey)) {
+    // Per-key lock: concurrent calls for the same token share one onboard attempt.
+    let lock = onboardLocks.get(cacheKey);
+    if (!lock) {
+      lock = (async () => {
+        let aborted = false;
+        try {
+          const onboarded = await tryOnboardUser(
+            accessToken, fetchImpl, clientProfile, tierId, signal
+          );
+          if (onboarded) {
+            const retry = await tryLoadCodeAssist(
+              accessToken, fetchImpl, clientProfile, signal
+            );
+            if (retry.projectId) {
+              evictOldest(projectCache);
+              projectCache.set(cacheKey, retry.projectId);
+              return true;
+            }
+          }
+          return false;
+        } catch (e) {
+          aborted = signal?.aborted === true;
+          return false;
+        } finally {
+          onboardLocks.delete(cacheKey);
+          if (!aborted) addToOnboardAttemptedCache(cacheKey);
+        }
+      })();
+      onboardLocks.set(cacheKey, lock);
+    }
+    const success = await lock;
+    if (success) {
+      const cached = projectCache.get(cacheKey);
+      if (cached) return cached;
+    }
+  }
 
   if (projectId) {
+    evictOldest(projectCache);
     projectCache.set(cacheKey, projectId);
     return projectId;
   }
-  // Non-fatal: if all endpoints failed, we proceed without caching.
   return undefined;
 }
 
 /** Exported for tests. */
 export function clearAntigravityProjectCache(): void {
   projectCache.clear();
+  onboardAttemptedCache.clear();
+  onboardLocks.clear();
 }
 
 /** Exported for tests — inspect cache state. */

tests/unit/antigravity-discovery-bootstrap.test.ts

@@ -161,7 +161,7 @@ describe("ensureAntigravityProjectAssigned", () => {
     assert.equal(capturedHeaders?.get("Client-Metadata"), null);
   });
 
-  test("bootstrap uses the single stable production loadCodeAssist endpoint and stays non-fatal on 404", async () => {
+  test("bootstrap tries loadCodeAssist then onboardUser on 404, non-fatal", async () => {
     const hitUrls: string[] = [];
 
     const mockFetch = async (url: string, _init?: RequestInit): Promise<Response> => {
@@ -171,13 +171,12 @@ describe("ensureAntigravityProjectAssigned", () => {
 
     const projectId = await ensureAntigravityProjectAssigned("bootstrap-404-token", mockFetch);
 
-    // #8098 narrowed the bootstrap to the single stable production endpoint (no
-    // daily/sandbox fallback), so a 404 has no next URL to try — the call fails closed
-    // (undefined) and the caller proceeds with any DB-stored project id.
-    assert.equal(hitUrls.length, 1, "bootstrap tries exactly the one dedicated production URL");
-    // Exact hostname match (not substring .includes) so the check can't be fooled by a
-    // look-alike host (CodeQL js/incomplete-url-substring-sanitization).
-    assert.equal(new URL(hitUrls[0]).hostname, "cloudcode-pa.googleapis.com");
+    // loadCodeAssist returns no project on 404, so the fallback calls
+    // onboardUser (also 404). Total: 2 URLs (loadCodeAssist + onboardUser).
+    assert.equal(hitUrls.length, 2, "must try loadCodeAssist then onboardUser");
+    for (const url of hitUrls) {
+      assert.equal(new URL(url).hostname, "cloudcode-pa.googleapis.com");
+    }
     assert.equal(projectId, undefined, "a 404 bootstrap is non-fatal and returns undefined");
   });
 
@@ -237,3 +236,113 @@ describe("ordering guarantee: loadCodeAssist before :models", () => {
     assert.ok(loadIdx < modelsIdx, ":loadCodeAssist must be called BEFORE :models");
   });
 });
+
+// ── onboardUser fallback when loadCodeAssist returns no project ──────────
+
+describe("onboardUser fallback", () => {
+  test("calls onboardUser when loadCodeAssist returns empty, then retries loadCodeAssist", async () => {
+    let loadCalls = 0;
+    let onboardCalls = 0;
+
+    const mockFetch = async (url: string, _init?: RequestInit): Promise<Response> => {
+      if (url.endsWith(":loadCodeAssist")) {
+        loadCalls++;
+        // First call returns empty, second returns project after onboarding.
+        if (loadCalls >= 2) {
+          return new Response(JSON.stringify({ cloudaicompanionProject: "proj-after-onboard" }), {
+            status: 200,
+            headers: { "Content-Type": "application/json" },
+          });
+        }
+        return new Response(JSON.stringify({}), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      if (url.endsWith(":onboardUser")) {
+        onboardCalls++;
+        return new Response(JSON.stringify({ done: true }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      return new Response("Not Found", { status: 404 });
+    };
+
+    const projectId = await ensureAntigravityProjectAssigned("onboard-test-token", mockFetch);
+
+    assert.equal(projectId, "proj-after-onboard");
+    assert.equal(onboardCalls, 1, "onboardUser must be called exactly once");
+    assert.ok(loadCalls >= 2, "loadCodeAssist must be called twice (before and after onboard)");
+  });
+
+  test("returns undefined when both loadCodeAssist and onboardUser fail", async () => {
+    const mockFetch = async (url: string, _init?: RequestInit): Promise<Response> => {
+      if (url.endsWith(":loadCodeAssist")) {
+        return new Response(JSON.stringify({}), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      if (url.endsWith(":onboardUser")) {
+        return new Response("Forbidden", { status: 403 });
+      }
+      return new Response("Not Found", { status: 404 });
+    };
+
+    const projectId = await ensureAntigravityProjectAssigned("both-fail-token", mockFetch);
+    assert.equal(projectId, undefined, "must return undefined when both fail");
+  });
+
+  test("does not retry onboardUser for the same token", async () => {
+    let onboardCalls = 0;
+
+    const mockFetch = async (url: string, _init?: RequestInit): Promise<Response> => {
+      if (url.endsWith(":loadCodeAssist")) {
+        return new Response(JSON.stringify({}), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      if (url.endsWith(":onboardUser")) {
+        onboardCalls++;
+        return new Response(JSON.stringify({ done: true }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      return new Response("Not Found", { status: 404 });
+    };
+
+    await ensureAntigravityProjectAssigned("dedup-token", mockFetch);
+    await ensureAntigravityProjectAssigned("dedup-token", mockFetch);
+
+    assert.equal(onboardCalls, 1, "onboardUser must be called only once per token");
+  });
+
+  test("skips onboardUser when loadCodeAssist succeeds on first try", async () => {
+    let onboardCalls = 0;
+
+    const mockFetch = async (url: string, _init?: RequestInit): Promise<Response> => {
+      if (url.endsWith(":loadCodeAssist")) {
+        return new Response(JSON.stringify({ cloudaicompanionProject: "proj-exists" }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      if (url.endsWith(":onboardUser")) {
+        onboardCalls++;
+        return new Response(JSON.stringify({ done: true }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+      return new Response("Not Found", { status: 404 });
+    };
+
+    const projectId = await ensureAntigravityProjectAssigned("already-ok-token", mockFetch);
+
+    assert.equal(projectId, "proj-exists");
+    assert.equal(onboardCalls, 0, "onboardUser must NOT be called when loadCodeAssist succeeds");
+  });
+});