From 7cf8df553d8324efdd7d4a28ce0bf22a0d78f8ed Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:30:51 -0700 Subject: [PATCH 1/3] fix: serialize managed credential refresh with reads and secret edits Co-Authored-By: Paperclip --- .../src/server/codex-auth-cache.test.ts | 36 +++++++++++++++++- .../src/server/codex-auth-cache.ts | 19 +++++++++- server/src/__tests__/ai-connections.test.ts | 37 +++++++++++++++++++ server/src/services/ai-connection-runtime.ts | 6 +-- server/src/services/ai-connections.ts | 16 +++++--- ui/src/components/AccountQuotaPanels.test.tsx | 3 +- ui/src/components/AccountQuotaPanels.tsx | 9 +++-- 7 files changed, 109 insertions(+), 17 deletions(-) diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts index 0ca1bee28ca..2cd6ba7c173 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts @@ -463,13 +463,15 @@ describe("codex auth cache store", () => { releaseFirstCaller = resolve; }); + const entered = deferredSignal(); const firstCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("first-enter"); + entered.resolve(); await firstCallerGate; events.push("first-exit"); return "first"; }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const secondCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("second-enter"); events.push("second-exit"); @@ -497,12 +499,14 @@ describe("codex auth cache store", () => { releaseCompanyA = resolve; }); + const entered = deferredSignal(); const companyACall = withAccountHomeSecretMutationLock(env, "company-a", async () => { events.push("a-enter"); + entered.resolve(); await companyAGate; events.push("a-exit"); }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const companyBCall = withAccountHomeSecretMutationLock(env, "company-b", async () => { events.push("b-enter"); events.push("b-exit"); @@ -536,6 +540,28 @@ describe("codex auth cache store", () => { expect(events).toEqual(["promotion-enter", "mutation-enter", "mutation-exit", "promotion-exit"]); }); + it("allows awaited nested writes but expires inherited ownership after release", async () => { + const env = envFor(await makeInstanceRoot()); + const detachedGate = deferredSignal(); + let detached!: Promise; + let enteredDetached = false; + await withAccountHomeSecretMutationLock(env, "nested", async () => { + await expect(withAccountHomeSecretMutationLock(env, "nested", async () => "nested write")).resolves.toBe("nested write"); + detached = detachedGate.promise.then(() => withAccountHomeSecretMutationLock(env, "nested", async () => { enteredDetached = true; })); + }); + const holding = deferredSignal(); + const release = deferredSignal(); + const other = withAccountHomeSecretMutationLock(env, "nested", async () => { holding.resolve(); await release.promise; }); + await holding.promise; + detachedGate.resolve(); + try { + await new Promise(resolve => setTimeout(resolve, 30)); + expect(enteredDetached).toBe(false); + } finally { release.resolve(); } + await Promise.all([other, detached]); + expect(enteredDetached).toBe(true); + }); + }); describe("Phase 8: account-home directory validity before a secret write commits", () => { @@ -581,3 +607,9 @@ describe("codex auth cache store", () => { }); }); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.ts index c51921fff6b..d6dd757313a 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.ts @@ -1,4 +1,5 @@ -import { lstat, mkdir, open, readFile, rename, rm } from "node:fs/promises"; +import { AsyncLocalStorage } from "node:async_hooks"; +import { lstat, mkdir, open, readFile, realpath, rename, rm } from "node:fs/promises"; import path from "node:path"; import { randomUUID } from "node:crypto"; import { resolvePaperclipInstanceRootForAdapter } from "@paperclipai/adapter-utils/server-utils"; @@ -334,6 +335,11 @@ export async function withCodexAccountHomePromotionLock( * create) can still call into a secrets-service write that takes this lock * without deadlocking on itself. */ +// An outer credential transaction acquires this lock before database row locks. +// Nested secret-service writes must share that ownership, not acquire it again. +// Expire ownership on return so detached work cannot inherit a released lock. +const secretMutationOwnership = new AsyncLocalStorage>(); + export async function withAccountHomeSecretMutationLock( env: NodeJS.ProcessEnv = process.env, companyId: string, @@ -341,7 +347,16 @@ export async function withAccountHomeSecretMutationLock( ): Promise { const lockDir = resolveCodexAuthCacheNamedLockDir(env, companyId, "codex-account-home-secret-mutation-lock"); await ensurePrivateDir(lockDir); - return withDirectoryMergeLock(lockDir, fn, env); + const key = await realpath(lockDir); + const inherited = secretMutationOwnership.getStore(); + if (inherited?.get(key)?.active) return fn(); + return withDirectoryMergeLock(lockDir, async () => { + const ownership = { active: true }; + const scope = new Map(inherited); + scope.set(key, ownership); + try { return await secretMutationOwnership.run(scope, fn); } + finally { ownership.active = false; } + }, env); } /** diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 3d8632114a4..52d54a3bfb4 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -84,6 +84,37 @@ describe("managed AI connections", () => { } finally { fetchSpy.mockRestore(); } }); + it("serializes quota exchange with a concurrent secret edit and runtime credential read", async () => { + const owner = "quota-race-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "old-access", refresh_token: "single-use", id_token: "identity", account_id: "race-account" } }); + const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Racing refresh", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId; + const exchangeStarted = deferredSignal(); + const releaseExchange = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + exchangeStarted.resolve(); await releaseExchange.promise; + return Response.json({ access_token: "new-access", refresh_token: "new-refresh" }); + }); + let runtime: Awaited> | undefined; + const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal); + try { + 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)); + releaseExchange.resolve(); + const results = await Promise.all([refreshing, edit, preparing]); + runtime = results[2]; + expect(JSON.parse(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8"))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } }); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } }); + expect((await secretService(db).getById(secretId))?.description).toBe("Changed during refresh"); + } finally { releaseExchange.resolve(); await runtime?.cleanup(); request.mockRestore(); } + }, 10_000); + it("serializes quota credential refresh, persists rotation, and defers while the provider is active", async () => { const owner = "quota-refresh-owner"; await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); @@ -1069,3 +1100,9 @@ describe("AI connection recovery delivery", () => { }, 30000, ); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 534f5cf023a..b199dbc2bcf 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -11,7 +11,7 @@ import { } from "@paperclipai/shared"; import { aiConnectionService } from "./ai-connections.js"; import { secretService } from "./secrets.js"; -import { decideCodexAuthMerge } from "@paperclipai/adapter-codex-local/server"; +import { decideCodexAuthMerge, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; import type { AdapterExecutionTarget } from "@paperclipai/adapter-utils/execution-target"; import { runAdapterExecutionTargetProcess } from "@paperclipai/adapter-utils/execution-target"; import { decideGrokAuthMerge } from "@paperclipai/adapter-grok-local/server"; @@ -314,7 +314,7 @@ export async function prepareManagedAiRuntime( if (subscriptionFile) { const refreshed = await readFile(authFile, "utf8"); if (refreshed !== value) - await db.transaction(async (tx) => { + await withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => { const [grant] = await tx .select() .from(connectionGrants) @@ -374,7 +374,7 @@ export async function prepareManagedAiRuntime( .update(connectionGrants) .set({ updatedAt: new Date() }) .where(eq(connectionGrants.id, grant.id)); - }); + })); } } finally { if (home) await rm(home, { recursive: true, force: true }); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 767ae7ecb0c..770feba2419 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -1,3 +1,4 @@ +import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js"; import { createHash, randomUUID } from "node:crypto"; import { and, eq, inArray, or, sql } from "drizzle-orm"; @@ -402,6 +403,9 @@ export function aiConnectionService(db: Db) { }; } async function credential(row: Pick>, "connection" | "grant">) { + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => credentialUnlocked(row)); + } + async function credentialUnlocked(row: Pick>, "connection" | "grant">) { const ref = row.grant.credentialSecretRefs.find( (r) => r.configPath === "ai.credential", ); @@ -498,7 +502,7 @@ export function aiConnectionService(db: Db) { ); const id = reconnect?.connection.id ?? randomUUID(); const grantId = reconnect?.grant.id ?? randomUUID(); - return db.transaction(async (tx) => { + return withAccountHomeSecretMutationLock(undefined, companyId, () => db.transaction(async (tx) => { const secrets = secretService(tx); if (sessionId) { const [session] = await tx @@ -803,7 +807,7 @@ export function aiConnectionService(db: Db) { details: { provider: input.provider, method: input.method, grantId }, }); return { connectionId: id, grantId }; - }); + })); } /** A late failure must never invalidate credentials that were refreshed or reconnected meanwhile. */ async function markAuthenticationFailed(input: { @@ -813,7 +817,7 @@ export function aiConnectionService(db: Db) { runStartedAt: Date; attribution: AiConnectionAttribution & { identity: string }; }) { - return db.transaction(async (tx) => { + return withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => { const { attribution } = input; const [grant] = await tx.select().from(connectionGrants).where(and( eq(connectionGrants.companyId, input.companyId), @@ -847,11 +851,11 @@ export function aiConnectionService(db: Db) { entityType: "tool_connection", entityId: connection.id, details: { provider: attribution.provider, grantId: grant.id }, }); - }); + })); } /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */ async function refreshQuotaCredential(row: Parameters[0], failedValue: string, signal: AbortSignal) { - return db.transaction(async (tx) => { + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => db.transaction(async (tx) => { const companyId = row.connection.companyId; const [grant] = await tx.select().from(connectionGrants).where(and( eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId), @@ -904,7 +908,7 @@ export function aiConnectionService(db: Db) { action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, details: { provider: "openai", grantId: grant.id } }); return refreshed; - }); + })); } // A quota read uses the same credential audience as execution. Operator status // alone never grants access to another member's personal subscription. diff --git a/ui/src/components/AccountQuotaPanels.test.tsx b/ui/src/components/AccountQuotaPanels.test.tsx index c87a755b924..e5fcef037ae 100644 --- a/ui/src/components/AccountQuotaPanels.test.tsx +++ b/ui/src/components/AccountQuotaPanels.test.tsx @@ -15,9 +15,10 @@ describe("account quota panels", () => { }); it("displays unknown utilization without inventing zero percent", () => { const html = renderToStaticMarkup(); expect(html).toContain("$5 remaining"); + expect(html).toContain(`Resets ${new Date("2026-10-05T00:00:00Z").toLocaleString()}`); expect(html).not.toContain("0%"); expect(html).not.toContain('role="progressbar"'); }); diff --git a/ui/src/components/AccountQuotaPanels.tsx b/ui/src/components/AccountQuotaPanels.tsx index 31cd43d05d7..ea9b84c9230 100644 --- a/ui/src/components/AccountQuotaPanels.tsx +++ b/ui/src/components/AccountQuotaPanels.tsx @@ -27,9 +27,12 @@ export function AccountQuotaPanels({

)} {account.windows.map((window, index) => window.usedPercent == null ? ( -
- {window.label} - {window.valueLabel ?? "Usage not reported"} +
+
+ {window.label} + {window.valueLabel ?? "Usage not reported"} +
+ {window.resetsAt &&

Resets {new Date(window.resetsAt).toLocaleString()}

}
) : ( Date: Sat, 3 Oct 2026 21:33:54 -0700 Subject: [PATCH 2/3] fix: retain full-stream usage and scope biller inference to its provider Co-Authored-By: Paperclip --- .../opencode-local/src/server/execute.ts | 20 +++++++---- .../adapters/pi-local/src/server/execute.ts | 20 +++++++---- .../adapter-cost-failure-paths.test.ts | 33 +++++++++++++++++++ 3 files changed, 59 insertions(+), 14 deletions(-) diff --git a/packages/adapters/opencode-local/src/server/execute.ts b/packages/adapters/opencode-local/src/server/execute.ts index 72dd350f6b4..2fbbf9ea6f2 100644 --- a/packages/adapters/opencode-local/src/server/execute.ts +++ b/packages/adapters/opencode-local/src/server/execute.ts @@ -84,7 +84,7 @@ function parseModelProvider(model: string | null): string | null { } function resolveOpenCodeBiller(env: Record, provider: string | null): string { - return inferOpenAiCompatibleBiller(env, null) ?? provider ?? "unknown"; + return provider === "openai" ? inferOpenAiCompatibleBiller(env, "openai") ?? "unknown" : provider ?? "unknown"; } const REMOTE_OPENCODE_MODELS_PROBE_DEFAULT_TIMEOUT_SEC = 20; @@ -647,7 +647,9 @@ export async function execute(ctx: AdapterExecutionContext): Promise { + let hasAccounting = false; + const accountingLog = createUsageCheckpointLog(onLog, ctx.onUsage ?? (async () => {}), stdout => { + hasAccounting = true; const parsed = consumeAccounting(stdout); const provider = parseModelProvider(model || null); return { usage: parsed.usage, costUsd: parsed.costUsd, usageBasis: "per_run", provider, biller: resolveOpenCodeBiller(runtimeEnv, provider), billingType: "unknown", model, complete: false }; @@ -666,11 +668,15 @@ export async function execute(ctx: AdapterExecutionContext): Promise): Promise, provider: string | null): string { - return inferOpenAiCompatibleBiller(env, null) ?? provider ?? "unknown"; + return provider === "openai" ? inferOpenAiCompatibleBiller(env, "openai") ?? "unknown" : provider ?? "unknown"; } async function ensureSessionsDir(): Promise { @@ -738,7 +738,9 @@ export async function execute(ctx: AdapterExecutionContext): Promise { + let hasAccounting = false; + const accountingLog = createUsageCheckpointLog(bufferedOnLog, ctx.onUsage ?? (async () => {}), stdout => { + hasAccounting = true; const parsed = consumeAccounting(stdout); return { usage: parsed.usage, costUsd: parsed.usage.costUsd, usageBasis: "per_run", provider, biller: resolvePiBiller(runtimeEnv, provider), billingType: "unknown", model, complete: parsed.sawAgentEnd }; }); @@ -761,11 +763,15 @@ export async function execute(ctx: AdapterExecutionContext): Promise { onUsage.mockRejectedValue(new Error("Receipt persistence failed")); await expect(execute(context)).rejects.toThrow("Receipt persistence failed"); }); + it.each(cases.filter(f => ["pi", "opencode"].includes(f.name)))("$name preserves full-stream accounting after stdout truncation", async (fixture) => { + const dir = await mkdtemp(join(tmpdir(), "paperclip-accounting-tail-")); directories.push(dir); + const command = join(dir, "runtime"); await writeFile(command, "#!/bin/sh\nprintf 'anthropic test\\n'\n", { mode: 0o755 }); + const execute = await fixture.execute(); + for (const checkpoint of [true, false]) for (const timedOut of [false, true]) for (const missingPrice of [false, true]) { + const early = structuredClone(fixture.event) as any; + if (missingPrice) { + if (fixture.name === "pi") delete early.message.usage.cost; + else delete early.part.cost; + } + const end = fixture.name === "pi" && !timedOut ? '\n{"type":"agent_end","messages":[]}' : ""; + const first = JSON.stringify(early) + "\n"; + const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(5 * 1024 * 1024) } }) + "\n"; + const last = JSON.stringify(fixture.event) + end; + processResult.mockImplementation(async (_run, _target, _command, _args, options) => { + for (const chunk of [first, noise, last]) await options.onLog("stdout", chunk); + return { exitCode: timedOut ? null : 0, signal: timedOut ? "SIGTERM" : null, timedOut, stdout: (first + noise + last).slice(-4 * 1024 * 1024), stderr: "", pid: 123, startedAt: new Date().toISOString() }; + }); + const onUsage = checkpoint ? vi.fn() : undefined; + const result = await execute({ + runId: "tail-run", agent: { id: "test-agent", companyId: "test-company", name: "Accounting", adapterType: `${fixture.name}_local`, adapterConfig: {} }, + runtime: { sessionId: null, sessionParams: null, sessionDisplayId: null, taskKey: null }, + config: { engine: "cli", command, cwd: dir, model: "anthropic/test", paperclipRuntimeSkills: [], env: { ANTHROPIC_API_KEY: "fixture", OPENAI_BASE_URL: "https://unrelated-proxy.example/v1", OPENROUTER_API_KEY: "unrelated-key", OPENCODE_ALLOW_ALL_MODELS: "1" } }, + context: {}, onLog: async () => {}, onUsage, + }); + expect(result.usage).toMatchObject({ inputTokens: 40, outputTokens: 20, cachedInputTokens: 200 }); + expect(result.costUsd).toBe(missingPrice ? null : 0.008); + expect(result.biller).toBe("anthropic"); + expect(result.usageComplete).toBe(!timedOut); + if (onUsage) expect(onUsage.mock.calls.at(-1)![0]).toMatchObject({ usage: result.usage, costUsd: result.costUsd }); + } + }); + it.each(cases)("$name retains observed accounting when the process times out", async (fixture) => { const dir = await mkdtemp(join(tmpdir(), "paperclip-accounting-adapter-")); directories.push(dir); const command = join(dir, "runtime"); await writeFile(command, "#!/bin/sh\nprintf 'openai test\\n'\n", { mode: 0o755 }); From dfd62bd7e2f1189cb2a417f4b2c96b63a95e167f Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:38:46 -0700 Subject: [PATCH 3/3] 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.