mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-24 01:41:42 +02:00
fix: retain and retry contended credential write-back
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
Devin FoleyandPaperclip committed
1 parent
8bd5846ba0
commit
fc97e24392
3 files changed
+149
-67
No files matched your search
@@ -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" });
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user