From dfd62bd7e2f1189cb2a417f4b2c96b63a95e167f Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:38:46 -0700 Subject: [PATCH] fix: recheck credential versions after waits and defer conflicting refreshes Co-Authored-By: Paperclip --- doc/connections/AI-CONNECTIONS.md | 11 +++++++ server/src/__tests__/ai-connections.test.ts | 17 +++++++++-- server/src/services/ai-connections.ts | 33 +++++++++++++-------- 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index d31ce5fea97..c0ca25fa2c6 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -186,6 +186,17 @@ and contain neither credentials nor raw provider errors. Expired credentials report `authentication_required`; the probe does not exchange refresh tokens or change connection health. +The Costs dashboard uses the company quota endpoint. It reads each authorized +managed subscription separately, retaining the last successful observation on +transient errors. For managed Codex accounts, a 401 can trigger one persisted +OAuth refresh; it defers while potentially competing OpenAI work is active. +Refresh takes the company secret-mutation lock before database row locks and +requires those row locks immediately, before exchanging a single-use token. +Reconnect and runtime credential write-back follow the same lock order. +Credential resolution checks its version after any database wait and reloads a +rotated value before giving it to a new run. Unknown utilization has no progress +bar, but any provider-reported reset time remains visible. + Verification: ```sh diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 52d54a3bfb4..dbd467740bf 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/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"); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 770feba2419..3b452933784 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -402,10 +402,7 @@ export function aiConnectionService(db: Db) { } satisfies AiConnectionAttribution, }; } - async function credential(row: Pick>, "connection" | "grant">) { - return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => credentialUnlocked(row)); - } - async function credentialUnlocked(row: Pick>, "connection" | "grant">) { + async function credential(row: Pick>, "connection" | "grant">, retry = 0): Promise { 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[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.