From 7d7e66d5da5f764d6fc6f82c34d4c8b1343f3131 Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sun, 4 Oct 2026 01:01:43 -0700 Subject: [PATCH] Defer fresh executions during credential lock contention Co-Authored-By: Paperclip --- doc/connections/AI-CONNECTIONS.md | 7 ++-- .../agent-hire-ai-connections.test.ts | 32 +++++++++++++++++++ server/src/__tests__/ai-connections.test.ts | 31 +++++++++++++++++- server/src/services/ai-connections.ts | 13 ++++++++ server/src/services/heartbeat.ts | 5 ++- 5 files changed, 81 insertions(+), 7 deletions(-) diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index c74060853f9..29bd19a37eb 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -498,9 +498,10 @@ existing account name. Connecting installs access for that agent and resumes the work automatically. Explicit incompatible bindings and shared-account permission denials still fail; hiring never expands a restricted shared account's audience. -Concurrent runs of one subscription do not wait for each other. No credential -lease exists to hold them, so a fresh task execution cannot enter a contention -wait. A run that already entered this wait keeps its scheduled retries. It does +Concurrent runs of one subscription do not hold a credential lease. A fresh +task can briefly wait when a credential rotation holds the company file lock or +a grant/secret database row lock. Lock timeouts become `ai_connection_busy`, +keeping the task on its automatic pre-provider scheduled retry path. It does not request new credentials, and it does not consume the provider-failure retry allowance. Each retry revalidates the account, and existing run-dispatch rules still suppress cancelled, reassigned, or otherwise ineligible work. An assignee diff --git a/server/src/__tests__/agent-hire-ai-connections.test.ts b/server/src/__tests__/agent-hire-ai-connections.test.ts index d4f28de63d0..5c7d4ebda7b 100644 --- a/server/src/__tests__/agent-hire-ai-connections.test.ts +++ b/server/src/__tests__/agent-hire-ai-connections.test.ts @@ -1,3 +1,5 @@ +import * as codexAdapter from "@paperclipai/adapter-codex-local/server"; +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import { randomUUID } from "node:crypto"; import { mkdtemp, rm } from "node:fs/promises"; import os from "node:os"; @@ -323,6 +325,36 @@ describe("agent-created hires use managed AI connections", () => { }); describe("hired agents sharing a subscription", () => { + it("keeps a credential-lock timeout on automatic retry without blocking the task or starting a provider", async () => { + const f = await fixture("openai", "subscription"); + const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Waiting teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { cwd: home, engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); + const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Wait for credential rotation", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning(); + const execute = vi.fn(async () => ({ exitCode: 0, signal: null, timedOut: false, resultJson: {} })); + registerServerAdapter({ ...getServerAdapter(f.adapterType), execute }); + const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("quota refresh holds company lock"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })); + const heartbeat = heartbeatService(db); + try { + const run = await heartbeat.invoke(agent.id, "assignment", { issueId: issue.id, wakeReason: "issue_assigned", responsibleUserId: f.userId }, "system"); + expect(run).not.toBeNull(); + await expect.poll(async () => (await heartbeat.getRun(run!.id))?.status, { timeout: 20_000 }).toBe("cancelled"); + await heartbeat.drainActiveRunExecutions(); + expect(await heartbeat.getRun(run!.id)).toMatchObject({ errorCode: "ai_connection_busy", resultJson: { executionRecovery: { providerWorkStarted: false } } }); + const retries = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id)); + expect(retries).toHaveLength(1); + expect(retries[0]).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "ai_connection_busy" }); + const [savedIssue] = await db.select().from(issues).where(eq(issues.id, issue.id)); + expect(savedIssue.status).not.toBe("blocked"); + expect(savedIssue.executionRunId).toBe(retries[0].id); + expect(execute).not.toHaveBeenCalled(); + expect(await aiConnectionService(db).list(f.companyId, f.userId)).toEqual([expect.objectContaining({ status: "connected" })]); + } finally { + lock.mockRestore(); + await heartbeat.drainActiveRunExecutions(); + unregisterServerAdapter(f.adapterType); + } + }); + it.each(["openai", "anthropic"] as const)("runs the %s child alongside a live parent and inherits its connection", async (provider) => { const f = await fixture(provider, "subscription"); const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Concurrent teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { cwd: home, engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index f2004abac4b..ee125c81aea 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -16,7 +16,7 @@ import { aiConnectionService } from "../services/ai-connections.js"; import * as codexAdapter from "@paperclipai/adapter-codex-local/server"; import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import * as executionTarget from "@paperclipai/adapter-utils/execution-target"; -import { prepareManagedAiRuntime, assertManagedAiProjectAuth } from "../services/ai-connection-runtime.js"; +import { prepareManagedAiRuntime, assertManagedAiProjectAuth, isAiConnectionBusy } from "../services/ai-connection-runtime.js"; import { toolAccessService } from "../services/tool-access.js"; import { secretService } from "../services/secrets.js"; import { aiConnectionBindingSchema, connectionPurposeTransportSchema, isAiConnectionCompatible } from "@paperclipai/shared"; @@ -143,6 +143,35 @@ describe("managed AI connections", () => { } }); + it.each(["company_file", "grant_row", "secret_row"] as const)("defers runtime preparation on a contended %s lock and succeeds once it clears", async (kind) => { + const owner = `runtime-lock-${kind}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = JSON.stringify({ tokens: { access_token: "valid-access", refresh_token: "valid-refresh", account_id: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const held = deferredSignal(), release = deferredSignal(); + const lock = kind === "company_file" ? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("company lock held by another account"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) : undefined; + const holding = kind === "company_file" ? Promise.resolve() : db.transaction(async tx => { + if (kind === "grant_row") await tx.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)).for("update"); + else await tx.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)).for("update"); + held.resolve(); await release.promise; + }); + try { + if (kind !== "company_file") await held.promise; + const error = await prepare().then(() => { throw new Error("Runtime must defer"); }, error => error); + expect(isAiConnectionBusy(error)).toBe(true); + expect(error).toMatchObject({ status: 422, details: { code: "ai_connection_busy" } }); + expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active"); + } finally { release.resolve(); await holding; lock?.mockRestore(); } + const run = await prepare(); + try { expect(await readFile(path.join(String(run.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth); } + finally { await run.cleanup(); } + }); + it.each(["runtime", "failed_run"] as const)("recovers pending quota credentials before %s can reuse or invalidate the consumed token", async (consumer) => { const owner = `quota-consumer-${consumer}`; await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 323dd4652db..7e53fc78d31 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -1,3 +1,4 @@ +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js"; import { createHash, randomUUID } from "node:crypto"; @@ -861,6 +862,18 @@ export function aiConnectionService(db: Db) { if (recovered) await quotaCredentialRecovery(row.connection.companyId, row.grant.id).clear() .catch(error => logger.warn({ err: error, grantId: row.grant.id }, "Saved runtime credential recovery cleanup pending")); return value; + }).catch(error => { + // Drizzle wraps Postgres errors in `cause`. Contention is a temporary + // pre-provider wait, not evidence that this account needs reconnecting. + let current: unknown = error; + for (let depth = 0; depth < 4 && current && typeof current === "object"; depth += 1) { + const value = current as { code?: unknown; cause?: unknown }; + if (value.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE || value.code === "55P03") { + throw unprocessable("AI credentials are being updated. This execution will retry automatically.", { code: "ai_connection_busy" }); + } + current = value.cause; + } + throw error; }); } diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 4c7f3a281c6..1f165c3692c 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -16361,9 +16361,8 @@ export function heartbeatService( .then((rows) => rows[0] ?? null); } - // No code path raises the `ai_connection_busy` error any more. This - // function stays because a stored run row can still carry that error code - // from an earlier release. + // Credential rotation can briefly contend with a fresh runtime read. Keep + // the task on its automatic pre-provider retry path while the lock clears. async function finalizeAiConnectionBusyDeferral( run: typeof heartbeatRuns.$inferSelect, error: HttpError,