fix: retain and retry contended credential write-back

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

doc/connections/AI-CONNECTIONS.md

@@ -532,3 +532,11 @@ A reconnect during credential resolution defers the poll instead of associating
 one account's quota with another revision. Both Costs surfaces retain matching
 successful observations on transient failures and clear them on authentication
 failure or credential rotation.
+
+Managed runtime token write-back retries a company-lock timeout twice (three
+30-second acquisition attempts), so a slow quota exchange does not discard a
+different account's replacement tokens. Any remaining write-back error preserves
+the private runtime home for retry; cleanup deletes it only after a successful
+transaction or an intentional freshness/authorization discard. These retained
+homes are recovery evidence, not a background replay queue; persistent database
+or lock failures still require operator intervention.

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

@@ -12,6 +12,8 @@ import { createDb, companies, agents, heartbeatRuns, companyMemberships, connect
 import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres";
 import { fetchCompanyQuotaWindows } from "../services/quota-windows.js";
 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 { toolAccessService } from "../services/tool-access.js";
@@ -131,6 +133,66 @@ describe("managed AI connections", () => {
     } finally { releaseBody.resolve(); await saving.catch(() => {}); request.mockRestore(); }
   });
 
+  it("retries another account's token write-back while a slow quota exchange holds the company lock", async () => {
+    const auth = (account: string, marker: string) => JSON.stringify({ tokens: { account_id: account, id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: marker === "old" ? "2026-09-10T10:00:00Z" : "2026-09-10T11:00:00Z" });
+    const quotaOwner = "slow-quota-owner", runOwner = "slow-quota-other-run";
+    await db.insert(companyMemberships).values([quotaOwner, runOwner].map(principalId => ({ companyId, principalId, principalType: "user" as const, status: "active" as const, membershipRole: "member" as const })));
+    for (const owner of [quotaOwner, runOwner]) await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth(owner, "old"));
+    const [quotaRow] = await service.quotaAccounts(companyId, quotaOwner);
+    const started = deferredSignal(), release = deferredSignal();
+    const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
+      started.resolve(); await release.promise;
+      return Response.json({ access_token: "quota-fresh", refresh_token: "quota-replacement" });
+    });
+    const refreshing = service.refreshQuotaCredential(quotaRow, auth(quotaOwner, "old"), new AbortController().signal);
+    let run: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
+    let cleaning: Promise<void> | undefined;
+    let lock: ReturnType<typeof vi.spyOn> | undefined;
+    try {
+      await started.promise;
+      // This different account starts after the refresh's active-run check.
+      run = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: runOwner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
+      const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json");
+      await writeFile(authFile, auth(runOwner, "replacement"));
+      // Inject the two 30-second acquisition expirations, then use the real
+      // still-held lock and real Postgres write-back for the final attempt.
+      lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock")
+        .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }))
+        .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }));
+      cleaning = run.cleanup();
+      await vi.waitFor(() => expect(lock).toHaveBeenCalledTimes(3));
+      expect(await readFile(authFile, "utf8")).toBe(auth(runOwner, "replacement"));
+      release.resolve();
+      await Promise.all([refreshing, cleaning]);
+      const [runRow] = await service.quotaAccounts(companyId, runOwner);
+      expect(await service.credential(runRow)).toBe(auth(runOwner, "replacement"));
+      expect(JSON.parse(await service.credential(quotaRow))).toMatchObject({ tokens: { access_token: "quota-fresh", refresh_token: "quota-replacement" } });
+      await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" });
+    } finally { release.resolve(); await Promise.allSettled([refreshing, cleaning]); lock?.mockRestore(); request.mockRestore(); }
+  });
+
+  it.each(["lock", "database"])("preserves replacement tokens after exhausted %s write-back and permits a later retry", async failure => {
+    const owner = `retained-writeback-${failure}`;
+    await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
+    const auth = (hour: string) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` });
+    await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("10"));
+    const run = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
+    const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json");
+    await writeFile(authFile, auth("11"));
+    const injected = failure === "lock"
+      ? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValue(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }))
+      : vi.spyOn(db, "transaction").mockRejectedValue(new Error("database unavailable"));
+    try {
+      await expect(run.cleanup()).rejects.toThrow(failure === "lock" ? "lock timeout" : "database unavailable");
+      expect(injected).toHaveBeenCalledTimes(failure === "lock" ? 3 : 1);
+      expect(await readFile(authFile, "utf8")).toBe(auth("11"));
+    } finally { injected.mockRestore(); }
+    await run.cleanup();
+    const [row] = await service.quotaAccounts(companyId, owner);
+    expect(await service.credential(row)).toBe(auth("11"));
+    await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" });
+  });
+
   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" });

server/src/services/ai-connection-runtime.ts

@@ -12,6 +12,7 @@ import {
 import { aiConnectionService } from "./ai-connections.js";
 import { secretService } from "./secrets.js";
 import { decideCodexAuthMerge, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server";
+import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
 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";
@@ -310,75 +311,86 @@ export async function prepareManagedAiRuntime(
       identity,
       home,
       cleanup: async () => {
-        try {
-          if (subscriptionFile) {
-            const refreshed = await readFile(authFile, "utf8");
-            if (refreshed !== value)
-              await withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => {
-                const [grant] = await tx
-                  .select()
-                  .from(connectionGrants)
-                  .where(
-                    and(
-                      eq(connectionGrants.id, selection.grant.id),
-                      eq(connectionGrants.companyId, input.companyId),
-                    ),
-                  )
-                  .for("update");
-                // A missing or revoked grant blocks the write-back. Among
-                // active copies, the merge decision below keeps the
-                // credential with the newest provider freshness field.
-                if (!grant || grant.status !== "active") return;
-                const ref = grant.credentialSecretRefs.find(
-                  (r) => r.configPath === "ai.credential",
-                );
-                if (!ref) return;
-                // Lock the referenced secret row for the rest of this
-                // transaction. The grant-row lock above does not cover it,
-                // so an authorized rotation of this secret could otherwise
-                // land between the read and the write below and be
-                // overwritten by this stale write-back.
-                await tx
-                  .select({ id: companySecrets.id })
-                  .from(companySecrets)
-                  .where(
-                    and(
-                      eq(companySecrets.id, ref.secretId),
-                      eq(companySecrets.companyId, input.companyId),
-                    ),
-                  )
-                  .for("update");
-                const current = await aiConnectionService(
-                  tx as unknown as Db,
-                ).credential({ ...selection, grant });
-                const destination = path.join(
-                  providerHome,
-                  "current-auth.json",
-                );
-                await writeFile(destination, current, { mode: 0o600 });
-                const decision =
-                  input.binding.provider === "openai"
-                    ? await decideCodexAuthMerge(authFile, destination, {
-                        errorLabel: "AI account refresh",
-                      })
-                    : await decideGrokAuthMerge(authFile, destination, {
-                        errorLabel: "AI account refresh",
-                      });
-                if (decision !== 10) return;
-                await secretService(tx).rotate(
-                  ref.secretId,
-                  { value: refreshed },
-                  { userId: grant.subjectUserId },
-                );
-                await tx
-                  .update(connectionGrants)
-                  .set({ updatedAt: new Date() })
-                  .where(eq(connectionGrants.id, grant.id));
-              }));
+        if (subscriptionFile) {
+          const refreshed = await readFile(authFile, "utf8");
+          if (refreshed !== value) {
+            // A quota exchange can hold this company lock for 60 seconds.
+            // Retry its 30-second acquisition timeout without discarding the
+            // provider's only copy of a rotated, single-use refresh token.
+            const writeBack = () => withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => {
+              const [grant] = await tx
+                .select()
+                .from(connectionGrants)
+                .where(
+                  and(
+                    eq(connectionGrants.id, selection.grant.id),
+                    eq(connectionGrants.companyId, input.companyId),
+                  ),
+                )
+                .for("update");
+              // A missing or revoked grant blocks the write-back. Among
+              // active copies, the merge decision below keeps the
+              // credential with the newest provider freshness field.
+              if (!grant || grant.status !== "active") return;
+              const ref = grant.credentialSecretRefs.find(
+                (r) => r.configPath === "ai.credential",
+              );
+              if (!ref) return;
+              // Lock the referenced secret row for the rest of this
+              // transaction. The grant-row lock above does not cover it,
+              // so an authorized rotation of this secret could otherwise
+              // land between the read and the write below and be
+              // overwritten by this stale write-back.
+              await tx
+                .select({ id: companySecrets.id })
+                .from(companySecrets)
+                .where(
+                  and(
+                    eq(companySecrets.id, ref.secretId),
+                    eq(companySecrets.companyId, input.companyId),
+                  ),
+                )
+                .for("update");
+              const current = await aiConnectionService(
+                tx as unknown as Db,
+              ).credential({ ...selection, grant });
+              const destination = path.join(
+                providerHome,
+                "current-auth.json",
+              );
+              await writeFile(destination, current, { mode: 0o600 });
+              const decision =
+                input.binding.provider === "openai"
+                  ? await decideCodexAuthMerge(authFile, destination, {
+                      errorLabel: "AI account refresh",
+                    })
+                  : await decideGrokAuthMerge(authFile, destination, {
+                      errorLabel: "AI account refresh",
+                    });
+              if (decision !== 10) return;
+              await secretService(tx).rotate(
+                ref.secretId,
+                { value: refreshed },
+                { userId: grant.subjectUserId },
+              );
+              await tx
+                .update(connectionGrants)
+                .set({ updatedAt: new Date() })
+                .where(eq(connectionGrants.id, grant.id));
+            }));
+            for (let attempt = 0; ; attempt++) {
+              try { await writeBack(); break; }
+              catch (error) {
+                if (attempt >= 2 || (error as NodeJS.ErrnoException)?.code !== WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) throw error;
+              }
+            }
           }
-        } finally {
-          if (home) await rm(home, { recursive: true, force: true });
         }
+        // Preserve the private home on any failed write-back, including an
+        // exhausted lock retry or database failure. The caller can retry the
+        // same cleanup; only a committed save or an intentional discard
+        // (revoked grant, older/different identity) permits deletion.
+        if (home) await rm(home, { recursive: true, force: true });
       },
     };
   } catch (error) {