From ff012ff4204049f8156d80f68fef56a0a1cfc8b9 Mon Sep 17 00:00:00 2001 From: "Bob.Hou" Date: Wed, 5 Aug 2026 22:51:36 -0400 Subject: [PATCH] fix(antigravity): add onboardUser fallback when loadCodeAssist returns no project (#8886) Validated in local merge-train T4 (HouMinXi+Zartharas+Andrian+artickc) --- next.config.mjs | 4 + open-sse/config/antigravityUpstream.ts | 6 + .../services/antigravityProjectBootstrap.ts | 167 ++++++++++++++++-- .../antigravity-discovery-bootstrap.test.ts | 125 ++++++++++++- 4 files changed, 276 insertions(+), 26 deletions(-) diff --git a/next.config.mjs b/next.config.mjs index 8eccf141e3..e1c2e7cc1a 100644 --- a/next.config.mjs +++ b/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. diff --git a/open-sse/config/antigravityUpstream.ts b/open-sse/config/antigravityUpstream.ts index 1ed328d7f3..aa015ef3ae 100644 --- a/open-sse/config/antigravityUpstream.ts +++ b/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"; diff --git a/open-sse/services/antigravityProjectBootstrap.ts b/open-sse/services/antigravityProjectBootstrap.ts index 169066d42f..7d69216f1e 100644 --- a/open-sse/services/antigravityProjectBootstrap.ts +++ b/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): 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(); +/** Per-key lock to prevent concurrent onboard attempts for the same token. */ +const onboardLocks = new Map>(); + type FetchLike = (url: string, init?: RequestInit) => Promise; 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 { +): Promise { 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).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 { + 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(); + +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 { 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. */ diff --git a/tests/unit/antigravity-discovery-bootstrap.test.ts b/tests/unit/antigravity-discovery-bootstrap.test.ts index a02427796f..d2f5617dfe 100644 --- a/tests/unit/antigravity-discovery-bootstrap.test.ts +++ b/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 => { @@ -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 => { + 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 => { + 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 => { + 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 => { + 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"); + }); +});