fix(providers): persist runtime-discovered Antigravity projectId to the connection (#8491)

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)

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) {

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<string, unknown> | null
+): Promise<void> {
+  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.
+  }
+}

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(),

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<string, unknown>,
+      }
+    );
+
+    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<string, unknown> | 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<string, unknown> | 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<string, unknown>,
+      }
+    );
+    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<string, unknown>,
+      }
+    );
+
+    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();
+  }
+});