From da7eb17ba782787dada1ae62d5fbff98dbce2ac2 Mon Sep 17 00:00:00 2001 From: ikelvingo Date: Sat, 25 Jul 2026 08:35:25 -0300 Subject: [PATCH] fix(providers): persist runtime-discovered Antigravity projectId to the connection (#8491) --- .../8491-antigravity-projectid-persist.md | 1 + open-sse/executors/antigravity.ts | 12 +- .../services/antigravityProjectPersist.ts | 42 ++++ .../[id]/models/discovery/normalizers.ts | 12 +- ...y-projectid-discovery-persist-8491.test.ts | 191 ++++++++++++++++++ 5 files changed, 256 insertions(+), 2 deletions(-) create mode 100644 changelog.d/fixes/8491-antigravity-projectid-persist.md create mode 100644 open-sse/services/antigravityProjectPersist.ts create mode 100644 tests/integration/antigravity-projectid-discovery-persist-8491.test.ts diff --git a/changelog.d/fixes/8491-antigravity-projectid-persist.md b/changelog.d/fixes/8491-antigravity-projectid-persist.md new file mode 100644 index 0000000000..71ef22b43f --- /dev/null +++ b/changelog.d/fixes/8491-antigravity-projectid-persist.md @@ -0,0 +1 @@ +- fix(providers): persist a runtime-discovered Antigravity projectId back onto the connection so it survives token refreshes and restarts instead of being rediscovered or lost (#8491) diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index 26e2c69edf..ada9cd9c29 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -23,6 +23,7 @@ import { persistCreditBalance, getAllPersistedCreditBalances } from "@/lib/db/cr import { setConnectionRateLimitUntil } from "@/lib/db/providers"; import { getMitmAlias } from "@/lib/db/models"; import { ensureAntigravityProjectAssigned } from "../services/antigravityProjectBootstrap.ts"; +import { persistDiscoveredAntigravityProjectId } from "../services/antigravityProjectPersist.ts"; import { resolveAntigravityModelId, getAntigravityModelFallbacks, @@ -541,7 +542,16 @@ export class AntigravityExecutor extends BaseExecutor { getAntigravityClientProfile(credentials), signal ); - if (discovered) projectId = discovered; + if (discovered) { + projectId = discovered; + // #8491: persist the recovered id so it survives the next token refresh + // or process restart instead of being silently rediscovered every time. + await persistDiscoveredAntigravityProjectId( + credentials.connectionId, + discovered, + credentials.providerSpecificData + ); + } } if (!projectId) { diff --git a/open-sse/services/antigravityProjectPersist.ts b/open-sse/services/antigravityProjectPersist.ts new file mode 100644 index 0000000000..8a0d1070a5 --- /dev/null +++ b/open-sse/services/antigravityProjectPersist.ts @@ -0,0 +1,42 @@ +/** + * Persist a runtime-discovered Antigravity projectId back onto its connection row + * (#8491). + * + * `ensureAntigravityProjectAssigned()` recovers a missing projectId via a + * `loadCodeAssist` round-trip and hands it back to the caller for the in-flight + * request only — nothing wrote it back to the connection record, so every + * subsequent token refresh (or process restart) lost the discovery and forced a + * fresh round-trip. This module is the single best-effort write path both call + * sites (`open-sse/executors/antigravity.ts` and the models-discovery + * normalizer) funnel through, mirroring the shape `mapAntigravityTokens()` + * already persists at OAuth-exchange time (`src/lib/oauth/providers/antigravity.ts`). + */ + +import { updateProviderConnection } from "@/lib/db/providers"; + +/** + * Write `discoveredProjectId` onto both the `projectId` column and + * `providerSpecificData.projectId` for `connectionId`, preserving any other + * `providerSpecificData` fields already on the connection. + * + * Best-effort / non-fatal by design: a persistence failure must never block + * the in-flight request, which already has the discovered id in hand. + */ +export async function persistDiscoveredAntigravityProjectId( + connectionId: string | undefined | null, + discoveredProjectId: string | undefined | null, + existingProviderSpecificData?: Record | null +): Promise { + if (!connectionId || !discoveredProjectId) return; + try { + await updateProviderConnection(connectionId, { + projectId: discoveredProjectId, + providerSpecificData: { + ...(existingProviderSpecificData || {}), + projectId: discoveredProjectId, + }, + }); + } catch { + // Non-fatal: persistence failure must never block the in-flight request. + } +} diff --git a/src/app/api/providers/[id]/models/discovery/normalizers.ts b/src/app/api/providers/[id]/models/discovery/normalizers.ts index ab6c06e477..5b553fb3fa 100644 --- a/src/app/api/providers/[id]/models/discovery/normalizers.ts +++ b/src/app/api/providers/[id]/models/discovery/normalizers.ts @@ -17,6 +17,7 @@ import { } from "@omniroute/open-sse/config/agyModels.ts"; import { normalizeAntigravityClientProfile } from "@/shared/constants/antigravityClientProfile"; import { ensureAntigravityProjectAssigned } from "@omniroute/open-sse/services/antigravityProjectBootstrap.ts"; +import { persistDiscoveredAntigravityProjectId } from "@omniroute/open-sse/services/antigravityProjectPersist.ts"; import { asRecord, toNonEmptyString } from "./helpers"; const antigravityDiscoveryInflight = new Map< @@ -115,7 +116,16 @@ export async function fetchAntigravityDiscoveryModelsCached( const promise = (async () => { await resolveAntigravityClientVersion(profile); - await ensureAntigravityProjectAssigned(accessToken, fetch, profile); + const discovered = await ensureAntigravityProjectAssigned(accessToken, fetch, profile); + if (discovered) { + // #8491: persist the recovered id so it survives the next token refresh + // or process restart instead of being silently rediscovered every time. + await persistDiscoveredAntigravityProjectId( + connectionId, + discovered, + asRecord(providerSpecificData) + ); + } for (const discoveryUrl of [ ...getAntigravityFetchAvailableModelsUrls(), diff --git a/tests/integration/antigravity-projectid-discovery-persist-8491.test.ts b/tests/integration/antigravity-projectid-discovery-persist-8491.test.ts new file mode 100644 index 0000000000..e2d5b402b5 --- /dev/null +++ b/tests/integration/antigravity-projectid-discovery-persist-8491.test.ts @@ -0,0 +1,191 @@ +// #8491 — a runtime-discovered Antigravity projectId must be persisted onto the +// connection so it survives the next token refresh / process restart, instead of +// being silently rediscovered (or lost) on every subsequent request. +// +// PART A: a fresh connection with an empty projectId, a mocked loadCodeAssist that +// returns a project id — after transformRequest() resolves, the connection row must +// have the discovered id written back (both the projectId column and +// providerSpecificData.projectId). +// +// PART B: once persisted, a SECOND request that re-reads credentials from the DB +// (simulating a token refresh / new request cycle, exactly what +// getProviderCredentials does in src/sse/services/auth.ts) must find the persisted +// projectId and must NOT re-invoke loadCodeAssist — the discovery branch never +// triggers because credentials.projectId is now populated from the DB. +import { test } from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-8491-antigravity-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "test-8491-antigravity-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const { AntigravityExecutor } = await import("../../open-sse/executors/antigravity.ts"); +const { clearAntigravityProjectCache } = await import( + "../../open-sse/services/antigravityProjectBootstrap.ts" +); + +test.after(() => { + core.resetDbInstance(); + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } +}); + +const BOOTSTRAP_URL = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; +const DISCOVERED_PROJECT_ID = "discovered-project-8491"; + +async function seedConnection() { + const connection = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + name: "antigravity-8491", + email: "antigravity-8491@example.test", + accessToken: "fake-antigravity-8491-token", + refreshToken: "fake-antigravity-8491-refresh", + expiresAt: new Date(Date.now() + 60 * 60 * 1000).toISOString(), + providerSpecificData: { clientProfile: "ide" }, + isActive: true, + testStatus: "active", + }); + assert(connection && typeof connection.id === "string"); + return connection; +} + +test("#8491 PART A: runtime-discovered projectId must be persisted to the connection", async () => { + clearAntigravityProjectCache(); + const connection = await seedConnection(); + const executor = new AntigravityExecutor(); + + const originalFetch = globalThis.fetch; + let loadCodeAssistCalls = 0; + globalThis.fetch = (async (input: string | URL | Request) => { + const url = String(input instanceof Request ? input.url : input); + if (url === BOOTSTRAP_URL) { + loadCodeAssistCalls += 1; + return new Response(JSON.stringify({ cloudaicompanionProject: DISCOVERED_PROJECT_ID }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + throw new Error(`Unexpected fetch in PART A: ${url}`); + }) as typeof fetch; + + try { + const result = await executor.transformRequest( + "antigravity/gemini-3.1-pro", + { request: { contents: [] } }, + true, + { + accessToken: connection.accessToken as string, + connectionId: connection.id, + providerSpecificData: connection.providerSpecificData as Record, + } + ); + + if (result instanceof Response) { + throw new Error(`Expected an envelope but got a ${result.status} Response`); + } + assert.equal(loadCodeAssistCalls, 1, "loadCodeAssist must be called to recover the project"); + assert.equal(result.project, DISCOVERED_PROJECT_ID, "the in-flight request uses the discovered id"); + + const persisted = await providersDb.getProviderConnectionById(connection.id); + assert.equal( + persisted?.projectId, + DISCOVERED_PROJECT_ID, + "discovered projectId must be persisted onto the connection" + ); + assert.equal( + (persisted?.providerSpecificData as Record | undefined)?.projectId, + DISCOVERED_PROJECT_ID, + "discovered projectId must also be persisted onto providerSpecificData.projectId" + ); + // The pre-existing providerSpecificData field must survive the persistence write. + assert.equal( + (persisted?.providerSpecificData as Record | undefined)?.clientProfile, + "ide" + ); + } finally { + globalThis.fetch = originalFetch; + clearAntigravityProjectCache(); + } +}); + +test("#8491 PART B: a second request re-reading credentials from the DB must not need to re-discover", async () => { + clearAntigravityProjectCache(); + const connection = await seedConnection(); + const executor = new AntigravityExecutor(); + + const originalFetch = globalThis.fetch; + let loadCodeAssistCalls = 0; + globalThis.fetch = (async (input: string | URL | Request) => { + const url = String(input instanceof Request ? input.url : input); + if (url === BOOTSTRAP_URL) { + loadCodeAssistCalls += 1; + return new Response(JSON.stringify({ cloudaicompanionProject: DISCOVERED_PROJECT_ID }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + throw new Error(`Unexpected fetch in PART B: ${url}`); + }) as typeof fetch; + + try { + // Request 1: empty projectId, discovers + persists. + await executor.transformRequest( + "antigravity/gemini-3.1-pro", + { request: { contents: [] } }, + true, + { + accessToken: connection.accessToken as string, + connectionId: connection.id, + providerSpecificData: connection.providerSpecificData as Record, + } + ); + assert.equal(loadCodeAssistCalls, 1, "first request must discover via loadCodeAssist"); + + // Simulate a token refresh between requests (rotates the access token, exactly + // what AntigravityExecutor.refreshCredentials()'s DB write does). + await providersDb.updateProviderConnection(connection.id, { + accessToken: "fake-antigravity-8491-token-ROTATED", + }); + + // Request 2: credentials re-read from the DB (what getProviderCredentials does + // for every request) — the persisted projectId must already be populated. + const refreshed = await providersDb.getProviderConnectionById(connection.id); + assert(refreshed); + const result2 = await executor.transformRequest( + "antigravity/gemini-3.1-pro", + { request: { contents: [] } }, + true, + { + accessToken: refreshed.accessToken as string, + connectionId: refreshed.id, + projectId: refreshed.projectId as string | undefined, + providerSpecificData: refreshed.providerSpecificData as Record, + } + ); + + if (result2 instanceof Response) { + throw new Error(`Expected an envelope but got a ${result2.status} Response`); + } + assert.equal( + result2.project, + DISCOVERED_PROJECT_ID, + "the second request must reuse the persisted projectId" + ); + assert.equal( + loadCodeAssistCalls, + 1, + "a second request with a rotated access token should not need to re-discover " + + "(#8491 cache-key half) once the id was already persisted" + ); + } finally { + globalThis.fetch = originalFetch; + clearAntigravityProjectCache(); + } +});