fix: recheck credential versions after waits and defer conflicting refreshes

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/ai-connections.test.ts

@@ -103,9 +103,12 @@ describe("managed AI connections", () => {
       await exchangeStarted.promise;
       const edit = secretService(db).update(secretId, { description: "Changed during refresh" });
       const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
-      // Keep the provider exchange pending while both contenders reach their
-      // lock/read paths. Neither may hold a conflicting row lock or vend old auth.
-      await new Promise(resolve => setTimeout(resolve, 50));
+      // Wait for the runtime's old-token read to reach its bookkeeping update.
+      // After the refresh commits, resolution must detect the changed version.
+      await vi.waitFor(async () => {
+        const waiting = await db.execute(sql`select 1 from pg_stat_activity where datname = current_database() and wait_event_type = 'Lock' and query like '%last_resolved_at%'`);
+        expect(waiting.length).toBeGreaterThan(0);
+      });
       releaseExchange.resolve();
       const results = await Promise.all([refreshing, edit, preparing]);
       runtime = results[2];
@@ -124,6 +127,14 @@ describe("managed AI connections", () => {
     const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ access_token: "fresh-access", refresh_token: "rotated-once" })));
     const signal = new AbortController().signal;
     try {
+      // A pre-existing transaction may own a secret row before asking for the
+      // company file lock. Refresh must defer before exchanging its token.
+      await db.transaction(async tx => {
+        const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId;
+        await tx.select().from(companySecrets).where(eq(companySecrets.id, secretId)).for("update");
+        await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow();
+        expect(request).not.toHaveBeenCalled();
+      });
       const activeId = randomUUID();
       await db.insert(heartbeatRuns).values({ id: activeId, companyId, agentId, invocationSource: "on_demand", status: "running", contextSnapshot: { aiConnection: { grantId: account.grantId } } });
       await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow("provider_unavailable");

server/src/services/ai-connections.ts

@@ -402,10 +402,7 @@ export function aiConnectionService(db: Db) {
       } satisfies AiConnectionAttribution,
     };
   }
-  async function credential(row: Pick<Awaited<ReturnType<typeof select>>, "connection" | "grant">) {
-    return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => credentialUnlocked(row));
-  }
-  async function credentialUnlocked(row: Pick<Awaited<ReturnType<typeof select>>, "connection" | "grant">) {
+  async function credential(row: Pick<Awaited<ReturnType<typeof select>>, "connection" | "grant">, retry = 0): Promise<string> {
     const ref = row.grant.credentialSecretRefs.find(
       (r) => r.configPath === "ai.credential",
     );
@@ -435,6 +432,7 @@ export function aiConnectionService(db: Db) {
       responsibleUserId: row.grant.subjectUserId,
       actorType: "system" as const,
     };
+    let value: string;
     if (secret.scope === "user") {
       if (
         secret.ownerUserId !== row.grant.subjectUserId ||
@@ -452,14 +450,21 @@ export function aiConnectionService(db: Db) {
         context,
       );
       if (!result) throw unprocessable("Reconnect this AI account");
-      return result.value;
+      value = result.value;
+    } else {
+      value = await secrets.resolveSecretValue(row.connection.companyId, secret.id, "latest", context);
     }
-    return secrets.resolveSecretValue(
-      row.connection.companyId,
-      secret.id,
-      "latest",
-      context,
-    );
+    // Resolution records lastResolvedAt and can wait behind a rotating writer.
+    // Re-read after that wait instead of handing a run the pre-rotation token.
+    const [latest] = await db.select().from(companySecrets).where(and(
+      eq(companySecrets.id, secret.id), eq(companySecrets.companyId, row.connection.companyId),
+    ));
+    if (!latest || latest.status !== "active") throw unprocessable("Reconnect this AI account");
+    if (latest.latestVersion !== secret.latestVersion) {
+      if (retry >= 2) throw unprocessable("AI credentials are changing. Retry this execution.", { code: "ai_connection_busy" });
+      return credential(row, retry + 1);
+    }
+    return value;
   }
   async function save(
     companyId: string,
@@ -856,17 +861,19 @@ export function aiConnectionService(db: Db) {
   /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */
   async function refreshQuotaCredential(row: Parameters<typeof credential>[0], failedValue: string, signal: AbortSignal) {
     return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => db.transaction(async (tx) => {
+      // Do not wait on a transaction that might itself be waiting for the file
+      // lock. All database locks must be available before consuming OAuth tokens.
       const companyId = row.connection.companyId;
       const [grant] = await tx.select().from(connectionGrants).where(and(
         eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId),
         eq(connectionGrants.connectionId, row.connection.id),
-      )).for("update");
+      )).for("update", { noWait: true });
       if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
       const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
       if (!ref) throw new Error("credentials_unavailable");
       await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
         eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId),
-      )).for("update");
+      )).for("update", { noWait: true });
       const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant });
       signal.throwIfAborted();
       // Another poll, an agent, or reconnect already supplied a newer credential.