From bbdcb97a06c729a708ed29e599d54454aa40b422 Mon Sep 17 00:00:00 2001 From: Anton <39598727+NomenAK@users.noreply.github.com> Date: Thu, 14 May 2026 05:09:26 +0200 Subject: [PATCH] fix(antigravity): bootstrap project via loadCodeAssist + fetchAvailableModels fallback (#2219) Integrated into release/v3.8.0 --- .../services/antigravityProjectBootstrap.ts | 125 +++++++++++ src/app/api/providers/[id]/models/route.ts | 12 +- .../antigravity-discovery-bootstrap.test.ts | 204 ++++++++++++++++++ 3 files changed, 339 insertions(+), 2 deletions(-) create mode 100644 open-sse/services/antigravityProjectBootstrap.ts create mode 100644 tests/unit/antigravity-discovery-bootstrap.test.ts diff --git a/open-sse/services/antigravityProjectBootstrap.ts b/open-sse/services/antigravityProjectBootstrap.ts new file mode 100644 index 0000000000..ce045461a0 --- /dev/null +++ b/open-sse/services/antigravityProjectBootstrap.ts @@ -0,0 +1,125 @@ +/** + * Antigravity project bootstrap — loadCodeAssist. + * + * The Google Cloud Code Assist API (/v1internal:models) requires a prior + * /v1internal:loadCodeAssist call to assign a project context to the + * OAuth token. Without this bootstrap, :models returns 404. + * + * This module provides an idempotent ensureAntigravityProjectAssigned() + * helper that is called once per access-token before every discovery + * attempt. Results are memoized per-token for the process lifetime to + * avoid redundant round-trips. + * + * Based on AntigravityService.loadCodeAssist() in + * src/lib/oauth/services/antigravity.ts and the CLIProxyAPI reference + * implementation in internal/runtime/executor/antigravity_executor.go. + */ + +import { + getAntigravityHeaders, + getAntigravityLoadCodeAssistMetadata, +} from "./antigravityHeaders.ts"; +import { ANTIGRAVITY_BASE_URLS } from "../config/antigravityUpstream.ts"; + +const LOAD_CODE_ASSIST_PATH = "/v1internal:loadCodeAssist"; +const BOOTSTRAP_TIMEOUT_MS = 8_000; + +/** Ordered list of loadCodeAssist endpoint URLs (mirrors the models discovery order). */ +export function getAntigravityLoadCodeAssistUrls(): string[] { + return ANTIGRAVITY_BASE_URLS.map((base) => `${base}${LOAD_CODE_ASSIST_PATH}`); +} + +/** Per-token memoization cache (lives for the process lifetime). */ +const projectCache = new Map(); + +type FetchLike = (url: string, init?: RequestInit) => Promise; + +/** + * Attempt loadCodeAssist against each known base URL in order. + * Returns the discovered project id, or null if all endpoints fail. + */ +async function tryLoadCodeAssist( + accessToken: string, + fetchImpl: FetchLike +): Promise { + const urls = getAntigravityLoadCodeAssistUrls(); + for (const url of urls) { + try { + const response = await fetchImpl(url, { + method: "POST", + headers: getAntigravityHeaders("loadCodeAssist", accessToken), + body: JSON.stringify({ metadata: getAntigravityLoadCodeAssistMetadata() }), + signal: AbortSignal.timeout(BOOTSTRAP_TIMEOUT_MS), + }); + + if (!response.ok) { + console.warn( + `[models] antigravity loadCodeAssist failed at ${url} (${response.status}) — trying next` + ); + continue; + } + + const data = (await response.json()) as Record; + + // cloudaicompanionProject may be a plain string or an object with an id field. + const raw = data.cloudaicompanionProject; + let projectId = + typeof raw === "string" + ? raw.trim() + : raw && + typeof raw === "object" && + typeof (raw as Record).id === "string" + ? ((raw as Record).id as string).trim() + : ""; + + if (projectId) { + return projectId; + } + + console.warn( + `[models] antigravity loadCodeAssist at ${url} returned no project id — trying next` + ); + } catch (error) { + const msg = error instanceof Error ? error.message : String(error); + console.warn(`[models] antigravity loadCodeAssist threw for ${url}: ${msg} — trying next`); + } + } + return null; +} + +/** + * Ensure a project is assigned to the given access token by calling + * loadCodeAssist if not already cached. Idempotent — repeated calls + * for the same token return the cached result without a network round-trip. + * + * Failures are non-fatal: the caller should proceed with the :models + * request regardless (the stored project_id in the DB may still be valid). + * + * @param accessToken The OAuth bearer token for the current connection. + * @param fetchImpl Injected fetch implementation (defaults to globalThis.fetch). + */ +export async function ensureAntigravityProjectAssigned( + accessToken: string, + fetchImpl: FetchLike = fetch +): Promise { + if (projectCache.has(accessToken)) { + return; // already bootstrapped for this token + } + + const projectId = await tryLoadCodeAssist(accessToken, fetchImpl); + + if (projectId) { + projectCache.set(accessToken, projectId); + } + // Non-fatal: if all endpoints failed, we proceed without caching. +} + +/** Exported for tests. */ +export function clearAntigravityProjectCache(): void { + projectCache.clear(); +} + +/** Exported for tests — inspect cache state. */ +export function getAntigravityProjectFromCache(accessToken: string): string | undefined { + return projectCache.get(accessToken); +} diff --git a/src/app/api/providers/[id]/models/route.ts b/src/app/api/providers/[id]/models/route.ts index 71dd9583ef..d86c6daaca 100755 --- a/src/app/api/providers/[id]/models/route.ts +++ b/src/app/api/providers/[id]/models/route.ts @@ -21,7 +21,11 @@ import { import { getProviderOutboundGuard } from "@/shared/network/outboundUrlGuard"; import { getStaticQoderModels } from "@omniroute/open-sse/services/qoderCli.ts"; import { getAntigravityHeaders } from "@omniroute/open-sse/services/antigravityHeaders.ts"; -import { getAntigravityModelsDiscoveryUrls } from "@omniroute/open-sse/config/antigravityUpstream.ts"; +import { ensureAntigravityProjectAssigned } from "@omniroute/open-sse/services/antigravityProjectBootstrap.ts"; +import { + getAntigravityModelsDiscoveryUrls, + getAntigravityFetchAvailableModelsUrls, +} from "@omniroute/open-sse/config/antigravityUpstream.ts"; import { buildGlmCodingHeaders, buildGlmModelsUrl, @@ -229,8 +233,12 @@ async function fetchAntigravityDiscoveryModelsCached( const promise = (async () => { await resolveAntigravityVersion(); + await ensureAntigravityProjectAssigned(accessToken); - for (const discoveryUrl of getAntigravityModelsDiscoveryUrls()) { + for (const discoveryUrl of [ + ...getAntigravityFetchAvailableModelsUrls(), + ...getAntigravityModelsDiscoveryUrls(), + ]) { try { const response = await safeOutboundFetch(discoveryUrl, { ...SAFE_OUTBOUND_FETCH_PRESETS.modelsDiscovery, diff --git a/tests/unit/antigravity-discovery-bootstrap.test.ts b/tests/unit/antigravity-discovery-bootstrap.test.ts new file mode 100644 index 0000000000..6128eec2e2 --- /dev/null +++ b/tests/unit/antigravity-discovery-bootstrap.test.ts @@ -0,0 +1,204 @@ +/** + * Tests: antigravity loadCodeAssist bootstrap before :models discovery. + * + * The Google Cloud Code Assist /v1internal:models endpoint requires a prior + * /v1internal:loadCodeAssist call to assign a project context to the OAuth + * token. Without this bootstrap, :models returns 404 for all three base URLs. + * + * These tests verify: + * 1. ensureAntigravityProjectAssigned calls loadCodeAssist before returning. + * 2. The call is memoized — repeated calls for the same token do not re-hit + * the network. + * 3. Non-fatal: if loadCodeAssist fails, the function resolves without throwing. + * 4. The loadCodeAssist request uses the correct headers (Authorization, User-Agent). + * 5. Ordering guarantee — in a full discovery flow, loadCodeAssist is called + * BEFORE any :models request. + */ + +import { test, describe, beforeEach } from "node:test"; +import assert from "node:assert/strict"; + +import { + ensureAntigravityProjectAssigned, + clearAntigravityProjectCache, + getAntigravityProjectFromCache, + getAntigravityLoadCodeAssistUrls, +} from "../../open-sse/services/antigravityProjectBootstrap.ts"; + +// Reset the module-level memoization cache between tests. +beforeEach(() => { + clearAntigravityProjectCache(); +}); + +describe("ensureAntigravityProjectAssigned", () => { + test("calls loadCodeAssist and caches the returned project id", async () => { + const calls: string[] = []; + + const mockFetch = async (url: string, _init?: RequestInit): Promise => { + calls.push(url); + if (url.endsWith(":loadCodeAssist")) { + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-from-bootstrap" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("Not Found", { status: 404 }); + }; + + await ensureAntigravityProjectAssigned("fake-token-1", mockFetch); + + const loadCalls = calls.filter((u) => u.endsWith(":loadCodeAssist")); + assert.ok(loadCalls.length >= 1, ":loadCodeAssist must be called at least once"); + + const cached = getAntigravityProjectFromCache("fake-token-1"); + assert.equal(cached, "proj-from-bootstrap", "project id must be memoized after first call"); + }); + + test("subsequent calls for the same token skip the network", async () => { + let networkCalls = 0; + + const mockFetch = async (url: string, _init?: RequestInit): Promise => { + networkCalls += 1; + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-cached" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + await ensureAntigravityProjectAssigned("fake-token-2", mockFetch); + await ensureAntigravityProjectAssigned("fake-token-2", mockFetch); + await ensureAntigravityProjectAssigned("fake-token-2", mockFetch); + + assert.equal(networkCalls, 1, "network must be called exactly once for the same token"); + }); + + test("different tokens each trigger their own loadCodeAssist call", async () => { + const calledFor: string[] = []; + + const mockFetch = async (url: string, init?: RequestInit): Promise => { + const auth = new Headers(init?.headers).get("Authorization") ?? ""; + calledFor.push(auth); + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-x" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + await ensureAntigravityProjectAssigned("token-A", mockFetch); + await ensureAntigravityProjectAssigned("token-B", mockFetch); + + assert.equal(calledFor.length, 2, "each unique token should trigger one network call"); + }); + + test("does not throw when loadCodeAssist returns non-200", async () => { + const mockFetch = async (_url: string, _init?: RequestInit): Promise => { + return new Response("Service Unavailable", { status: 503 }); + }; + + // Must resolve without throwing even if all endpoints fail. + await assert.doesNotReject(ensureAntigravityProjectAssigned("fail-token", mockFetch)); + }); + + test("does not throw when fetch rejects (network error)", async () => { + const mockFetch = async (_url: string, _init?: RequestInit): Promise => { + throw new Error("ECONNREFUSED"); + }; + + await assert.doesNotReject(ensureAntigravityProjectAssigned("throw-token", mockFetch)); + }); + + test("sets Authorization header with Bearer token", async () => { + let capturedAuth: string | null = null; + + const mockFetch = async (_url: string, init?: RequestInit): Promise => { + capturedAuth = new Headers(init?.headers).get("Authorization") ?? null; + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-auth-check" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + await ensureAntigravityProjectAssigned("my-secret-token", mockFetch); + + assert.equal(capturedAuth, "Bearer my-secret-token", "Authorization header must be set"); + }); + + test("falls through to next URL when first loadCodeAssist returns 404", async () => { + const hitUrls: string[] = []; + + const mockFetch = async (url: string, _init?: RequestInit): Promise => { + hitUrls.push(url); + if (url.includes("sandbox")) { + // First URL fails + return new Response("not found", { status: 404 }); + } + // Second URL succeeds + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-fallback" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + await ensureAntigravityProjectAssigned("fallback-token", mockFetch); + + assert.ok(hitUrls.length >= 2, "should try at least two URLs on the first failure"); + const cached = getAntigravityProjectFromCache("fallback-token"); + assert.equal(cached, "proj-fallback", "should cache the project from the successful URL"); + }); + + test("getAntigravityLoadCodeAssistUrls returns URLs matching ANTIGRAVITY_BASE_URLS", () => { + const urls = getAntigravityLoadCodeAssistUrls(); + assert.ok(urls.length >= 1, "must return at least one URL"); + for (const url of urls) { + assert.ok(url.endsWith(":loadCodeAssist"), `URL must end with :loadCodeAssist, got: ${url}`); + assert.ok(url.startsWith("https://"), `URL must be HTTPS, got: ${url}`); + } + }); +}); + +// ── Ordering guarantee: loadCodeAssist BEFORE :models ───────────────────────── +// +// This test simulates the full discovery flow: a test-controlled fetch +// that records call order, and verifies that :loadCodeAssist precedes +// any :models request. The integration is verified by calling +// ensureAntigravityProjectAssigned then simulating a :models request. + +describe("ordering guarantee: loadCodeAssist before :models", () => { + test("loadCodeAssist is called before :models in a simulated discovery flow", async () => { + const callOrder: string[] = []; + + const mockFetch = async (url: string, _init?: RequestInit): Promise => { + if (url.endsWith(":loadCodeAssist")) { + callOrder.push("loadCodeAssist"); + return new Response(JSON.stringify({ cloudaicompanionProject: "proj-order-test" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url.endsWith(":models")) { + callOrder.push("models"); + return new Response( + JSON.stringify({ + models: [{ id: "gemini-3-pro-antigravity", displayName: "Gemini 3 Pro" }], + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + } + return new Response("not found", { status: 404 }); + }; + + // Step 1: bootstrap project (what route.ts now does before the models loop). + await ensureAntigravityProjectAssigned("order-token", mockFetch); + + // Step 2: simulate a :models discovery request (what the loop does). + const modelsUrl = "https://cloudcode-pa.googleapis.com/v1internal:models"; + await mockFetch(modelsUrl); + + const loadIdx = callOrder.indexOf("loadCodeAssist"); + const modelsIdx = callOrder.indexOf("models"); + + assert.ok(loadIdx >= 0, ":loadCodeAssist must be called"); + assert.ok(modelsIdx >= 0, ":models must be called"); + assert.ok(loadIdx < modelsIdx, ":loadCodeAssist must be called BEFORE :models"); + }); +});