fix(ai-connections): remove account execution locks (#18778)

## Thinking Path

> - Paperclip manages AI agents and the accounts that supply their model
credentials.
> - Several runs can use the same AI connection at the same time.
> - Runtime preparation treated unfinished credential cleanup as
exclusive account ownership.
> - Manual account tests also held a company credential mutex for the
entire provider test.
> - A failed cleanup could therefore cancel fresh runs before the
provider started and keep them in automatic retries.
> - This pull request removes that execution gate and those read and
probe locks.
> - Each invocation keeps its own credential home. Credential updates
still use the existing freshness and authorization checks.

## Linked Issues or Issue Description

**What happened?**

Fresh runs stopped with `ai_connection_busy` and the message “The
selected AI account is still completing credential cleanup. Retry
shortly.” Pending cleanup from an earlier invocation blocked every new
invocation of that account. Concurrent manual account tests also waited
on a company mutex.

**Expected behavior**

AI connections support concurrent runs and account tests. Cleanup for
one private runtime must not prevent another private runtime from using
the same account.

**Steps to reproduce**

1. Prepare a managed OpenAI account and retain an earlier runtime whose
cleanup has not finished.
2. Prepare another run with that account. The old implementation rejects
it before provider work starts.
3. Alternatively, keep one manual account test open and start a second
test with the same account. The old implementation prevents the second
provider callback from starting.

**Paperclip version or commit**

Reproduced with two PostgreSQL regressions on
`17e8bd773543d62ab3f444c55255acfd82575efd`. Both fail before the fix and
pass after it.

**Deployment mode**

The reported deployment used authenticated Cloud. Local reproduction
uses disposable PostgreSQL and fixture credentials. No live provider
credentials are required for these regressions.

Related search found #15952, which concerns account removal and provider
application archival. It does not remove this runtime cleanup gate. This
fix does not introduce the account concurrency cap requested in #14564.
ROADMAP.md was checked; this change restores concurrent account use.

## What Changed

- Remove the account-wide cleanup assertion and cleanup mutex from run
admission.
- Read ordinary subscription credentials without acquiring the company
file mutex or grant and secret write locks. Only an actual durable
replacement recovery enters the existing write transaction.
- Remove the mutex held across manual account tests.
- Retry retained cleanup without making new runtime preparation wait for
it. A home-specific cleanup check considers only that exact private
home.
- Keep cleanup attempts for one runtime idempotent. Keep physical
retirement, credential freshness, reconnect, revocation and company
checks.
- Update concurrency, refresh and cleanup regressions, and the
implementation contract.

## Verification

- Before/after: the pending-cleanup and simultaneous-probe regressions
both fail on the original code and pass with the fix.
- Affected suites: 186 tests passed across `ai-connections.test.ts`,
`antigravity-ai-connections.integration.test.ts`, and
`retained-google-credential-cleanup.test.ts`. They use real PostgreSQL,
encrypted fixture credentials and controlled provider callbacks.
- `pnpm -r typecheck` passed.
- `node scripts/check-maps.mjs` and `git diff --check` passed.
- `pnpm build` passed on the rebased branch. The server typecheck and
production server build both passed again on final commit
`03bd3973bdf85108bbd82355abdb217f0e1c83bd` after the runtime-home
ownership changes.
- Caller coverage: all 166 cases pass across native credential handoff,
managed Google subscription writeback, model discovery, heartbeat
configuration and hired-agent execution. The full hiring file passed all
79 cases after its admission assertion was corrected. Together with the
affected suites, 352 targeted cases pass.
- The final home-binding check passed all 188 cases across the account,
managed Google writeback, model discovery and retained-cleanup suites.
- Review follow-up: two real-PostgreSQL cases pass for quota refresh and
concurrent Google account tests. The quota case observes a credential
row update in progress, verifies that the runtime takes no account file
mutex, and confirms that the private home receives the committed
replacement token. The Google case completes a second provider callback
while the first callback remains open in a separate private home.
- Greptile rated the runtime change 4/5 and found only misplaced map
patch notes. Those additions were removed. The implementation
documentation remains in `doc/connections/AI-CONNECTIONS.md` and the
implementation contract.
- `pnpm test:run` was attempted. Chat partition 1 passed 351 expanded
cases. Partition 2 passed 324 cases and failed one Teams file-receipt
API case at its 15-second deadline: `refuses generic resolution with
stale version transfer preconditions without changing receipts`. That
case exercises publication resolution, not AI account execution. Its
isolated rerun passed. The remaining local partitions did not run after
that failure. This is not a claim that the full local suite passed.
- Wider caller checks initially identified old serialization
expectations and two cleanup retry cases. Updated assertions now require
concurrent account use. The partial-removal correction lets a retained
home finish removal after a durable credential handoff. All affected
caller cases pass, including the full hiring file.
- Live provider concurrency and deployment recovery have not been
qualified. This PR does not claim that the reported Cloud deployment
contains the fix.

## Risks

- A new invocation can use the currently saved credential while another
runtime still owns a refreshed copy. The existing writeback checks
decide which credential can replace the saved value.
- Actual credential updates and durable replacement recovery retain
atomic storage transactions. This change removes execution
serialization, not the integrity checks on stored credentials.
- Unfinished cleanup retains its private home and its existing idle
hold. This change does not certify a previous provider process as
stopped.
- There is no database migration, credential format change, permission
expansion or ambient-account fallback.

## Model Used

- OpenAI GPT-6 through Codex desktop, with source editing, shell
execution, PostgreSQL fixture tests and read-only browser diagnosis. The
session does not expose the exact deployment ID or context-window size.

## Checklist

- [x] I have included a thinking path that traces from project context
to this change
- [x] I have specified the model used (with version and capability
details)
- [x] I have checked ROADMAP.md and confirmed this PR does not duplicate
planned core work
- [x] I have searched GitHub for duplicate or related PRs and linked
them above
- [x] I have either (a) linked existing issues with `Fixes: #` / `Closes
#` / `Refs #` OR (b) described the issue in-PR following the relevant
issue template
- [x] I have not referenced internal/instance-local Paperclip issues or
links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip`
URLs)
- [x] My branch name describes the change (e.g. `docs/...`, `fix/...`)
and contains no internal Paperclip ticket id or instance-derived details
- [x] I have run tests locally and they pass
- [x] I have added or updated tests where applicable
- [x] I have updated relevant documentation to reflect my changes
- [x] I have considered and documented any risks above
- [ ] All Paperclip CI gates are green
- [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups
- [x] I will address all Greptile and reviewer comments before
requesting merge

---------

Co-authored-by: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip authored and GitHub committed 2026-11-03 07:56:37 -06:00
1 parent 86b365f992
commit 2b14ec03a4
14 files changed
+256 -219

No files matched your search

@@ -575,14 +575,17 @@ describe("hired agents sharing a subscription", () => {
const worker = configureAgentLifecycle(db, { requiredPluginIds: async () => [], runPlugin: async () => "complete", runHost: async () => "complete" });
try { await worker.process(agentId); } finally { await worker.stop(); }
}
it("keeps a credential-lock timeout on automatic retry without blocking the task or starting a provider", async () => {
it("starts a hired agent without acquiring an AI connection execution lock", 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, runner: "legacy", reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } }));
const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Concurrent teammate", role: "engineer", adapterType: f.adapterType, runner: "legacy", reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } }));
// Only an operator may set host execution paths after the agent is hired.
await db.update(agents).set({ adapterConfig: { ...agent.adapterConfig, cwd: home } }).where(eq(agents.id, agent.id));
await finishSetup(agent.id);
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: {} }));
const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Use the account concurrently", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning();
const execute = vi.fn(async () => {
await db.update(issues).set({ status: "done", completedAt: new Date() }).where(eq(issues.id, issue.id));
return { 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 }));
@@ -590,27 +593,17 @@ describe("hired agents sharing a subscription", () => {
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 expect.poll(async () => (await heartbeat.getRun(run!.id))?.status, { timeout: 20_000 }).toBe("succeeded");
await heartbeat.drainActiveRunExecutions();
const cancelled = await heartbeat.getRun(run!.id);
expect(cancelled).toMatchObject({ errorCode: "ai_connection_busy", resultJson: {
executionRecovery: { kind: "ai_connection_wait", providerWorkStarted: false },
cancellation: {
source: "control_plane", expected: true, initiator: { type: "system" },
reason: "Waiting for shared AI credentials",
recordedAt: cancelled!.finishedAt!.toISOString(),
},
} });
const completed = await heartbeat.getRun(run!.id);
expect(completed?.errorCode).toBeNull();
await waitForPendingRunFailureReports();
expect(captureRunFailure.mock.calls.filter(([report]) => report.runId === run!.id)).toHaveLength(0);
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" });
expect(retries[0].resultJson?.cancellation).toBeUndefined();
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 db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.retryOfRunId, run!.id), eq(heartbeatRuns.scheduledRetryReason, "ai_connection_busy"),
))).toEqual([]);
expect(execute).toHaveBeenCalledOnce();
expect(lock).not.toHaveBeenCalled();
expect(await aiConnectionService(db).list(f.companyId, f.userId)).toEqual([expect.objectContaining({ status: "connected" })]);
} finally {
lock.mockRestore();
+97 -40
View File
@@ -425,7 +425,63 @@ describe("managed AI connections", () => {
} finally { open.mockRestore(); mkdir.mockRestore(); request.mockRestore(); }
});
it.each(["timeout", "unexpected"] as const)("classifies a manual probe's %s lock failure before starting the provider", async (kind) => {
it("starts a second run while the same account has unfinished credential cleanup", async () => {
const owner = "concurrent-retained-cleanup";
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
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 scope = { companyId, accountKey: row.grant.id }, retired = deferredSignal(), cleanup = vi.fn(async () => {});
const input = { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: {} };
retainAiCredentialCleanup({ scope, retired: retired.promise, cleanup });
let runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
try {
runtime = await prepareManagedAiRuntime(db, input);
expect(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth);
expect(cleanup).not.toHaveBeenCalled();
await runtime.cleanup();
expect(cleanup).not.toHaveBeenCalled();
} finally { retired.resolve(); await reconcileRetainedAiCleanup(scope); await runtime?.cleanup(); }
});
it("runs concurrent probes of the same account without holding the company credential mutex", async () => {
const owner = "concurrent-account-probes";
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth);
const input = { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: {} };
const entered = deferredSignal(), release = deferredSignal(), secondEntered = vi.fn();
const first = withManagedAiProbe(db, input, async () => { entered.resolve(); await release.promise; });
let second: Promise<void> | undefined;
try {
await entered.promise;
second = withManagedAiProbe(db, input, async () => { secondEntered(); });
await vi.waitFor(() => expect(secondEntered).toHaveBeenCalledOnce(), { timeout: 1000 });
await second;
} finally { release.resolve(); await Promise.allSettled([first, second]); }
});
it("prepares a run without waiting for another runtime's active cleanup", async () => {
const owner = "concurrent-active-cleanup";
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
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 scope = { companyId, accountKey: row.grant.id }, entered = deferredSignal(), release = deferredSignal();
retainAiCredentialCleanup({ scope, retired: Promise.resolve(), cleanup: async () => { entered.resolve(); await release.promise; } });
await entered.promise;
let runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }).then(value => { runtime = value; });
try {
await vi.waitFor(() => expect(runtime).toBeDefined(), { timeout: 1000 });
expect(await readFile(path.join(String(runtime!.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth);
} finally { release.resolve(); await preparing; await reconcileRetainedAiCleanup(scope); await runtime?.cleanup(); }
});
it.each(["timeout", "unexpected"] as const)("does not acquire a company credential lock for an ordinary probe (%s)", async (kind) => {
const owner = `manual-probe-lock-${kind}`;
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
@@ -437,20 +493,17 @@ describe("managed AI connections", () => {
const failure = Object.assign(new Error("injected private lock diagnostic"), { code: kind === "timeout" ? WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE : "EACCES" });
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValueOnce(failure);
try {
const error = await withManagedAiProbe(db, input, probe).then(() => { throw new Error("Probe must not start"); }, error => error);
if (kind === "timeout") {
expect(isAiConnectionBusy(error)).toBe(true);
expect(error).toMatchObject({ status: 422, message: "AI credentials are being updated. Retry shortly.", details: { code: "ai_connection_busy" } });
} else expect(error).toBe(failure);
expect(probe).not.toHaveBeenCalled();
expect(await withManagedAiProbe(db, input, probe)).toBe("ok");
expect(lock).not.toHaveBeenCalled();
expect(probe).toHaveBeenCalledOnce();
expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
} finally { lock.mockRestore(); }
expect(await withManagedAiProbe(db, input, probe)).toBe("ok");
expect(probe).toHaveBeenCalledTimes(1);
expect(probe).toHaveBeenCalledTimes(2);
expect(await service.credential(row)).toBe(auth);
});
it("fences retained OpenAI credentials before reads, including callers holding the mutation lock", async () => {
it("reads OpenAI credentials independently of another runtime cleanup or mutation lock", async () => {
const owner = "retained-openai-read-fence";
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
@@ -459,9 +512,9 @@ describe("managed AI connections", () => {
const scope = { companyId, accountKey: row.grant.id }, retired = deferredSignal(), cleanup = vi.fn(async () => {});
retainAiCredentialCleanup({ scope, retired: retired.promise, cleanup });
try {
await expect(service.runtimeCredential(row)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(service.runtimeCredential(row)).resolves.toBe(auth);
await codexAdapter.withAccountHomeSecretMutationLock(undefined, companyId, async () => {
await expect(service.runtimeCredential(row, { credentialLockHeld: true })).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(service.runtimeCredential(row)).resolves.toBe(auth);
});
expect(cleanup).not.toHaveBeenCalled();
retired.resolve(); await reconcileRetainedAiCleanup(scope);
@@ -470,7 +523,7 @@ describe("managed AI connections", () => {
} finally { retired.resolve(); await reconcileRetainedAiCleanup(scope).catch(() => {}); }
});
it("retries retired Codex probe cleanup before the same account can launch again", async () => {
it("allows the same account to run while retired Codex probe cleanup is pending or fails", async () => {
const owner = "retired-probe-retry", otherOwner = "retired-probe-other";
const auth = (user: string, hour: string) => JSON.stringify({ tokens: { account_id: user, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` });
for (const user of [owner, otherOwner]) {
@@ -490,29 +543,30 @@ describe("managed AI connections", () => {
await writeFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), auth(owner, "11"));
throw new runnerProbe.NativeEnvironmentProbeCleanupError(stopped.promise, new Error("fixture probe interruption"), []);
})).rejects.toMatchObject({ details: { code: "native_probe_cleanup_pending" } });
await expect(withManagedAiProbe(db, input, provider)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
expect(provider).not.toHaveBeenCalled(); expect(merge).not.toHaveBeenCalled();
await expect(withManagedAiProbe(db, input, provider)).resolves.toBe(auth(owner, "10"));
expect(provider).toHaveBeenCalledOnce(); expect(merge).not.toHaveBeenCalled();
stopped.resolve();
await vi.waitFor(() => expect(merge).toHaveBeenCalledOnce());
await expect(withManagedAiProbe(db, input, provider)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
expect(provider).not.toHaveBeenCalled(); expect(merge).toHaveBeenCalledTimes(2);
await expect(withManagedAiProbe(db, input, provider)).resolves.toBe(auth(owner, "10"));
await vi.waitFor(() => expect(merge).toHaveBeenCalledTimes(2));
await expect(withManagedAiProbe(db, { ...input, responsibleUserId: otherOwner }, provider)).resolves.toBe(auth(otherOwner, "10"));
expect(await readFile(path.join(String(retained!.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth(owner, "11"));
merge.mockRestore();
await reconcileRetainedAiCleanup({ companyId, accountKey: row.grant.id });
await expect(withManagedAiProbe(db, input, provider)).resolves.toBe(auth(owner, "11"));
expect(await service.credential(row)).toBe(auth(owner, "11"));
await expect(access(retained!.home!)).rejects.toMatchObject({ code: "ENOENT" });
} finally { stopped.resolve(); merge.mockRestore(); await retained?.cleanup().catch(() => {}); }
});
it("keeps quota refresh outside a manual probe until its rotated credential is saved", async () => {
it("allows quota refresh alongside a manual probe and keeps the freshest saved credential", async () => {
const owner = "manual-probe-race";
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 [row] = await service.quotaAccounts(companyId, owner);
const entered = deferredSignal(), release = deferredSignal();
const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Quota must not exchange the probe's token"));
const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: "quota-access", refresh_token: "quota-refresh" }));
const probing = withManagedAiProbe(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }, async runtime => {
entered.resolve(); await release.promise;
@@ -525,12 +579,12 @@ describe("managed AI connections", () => {
lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
// Start outside the probe's async scope; it is a separate HTTP request.
refreshing = service.refreshQuotaCredential(row, auth("10"), new AbortController().signal);
await vi.waitFor(() => expect(lock).toHaveBeenCalled());
expect(request).not.toHaveBeenCalled();
const refreshed = await refreshing;
expect(request).toHaveBeenCalledOnce();
expect(JSON.parse(refreshed)).toMatchObject({ tokens: { access_token: "quota-access", refresh_token: "quota-refresh" } });
release.resolve(); await probing;
expect(await refreshing).toBe(auth("11"));
expect(await service.credential(row)).toBe(auth("11"));
expect(request).not.toHaveBeenCalled();
expect(await service.credential(row)).toBe(refreshed);
expect(request).toHaveBeenCalledOnce();
} finally { release.resolve(); await Promise.allSettled([probing, refreshing]); lock?.mockRestore(); request.mockRestore(); }
});
@@ -570,7 +624,7 @@ describe("managed AI connections", () => {
} finally { probe.mockRestore(); resolveBinary.mockRestore(); }
});
it("keeps the real Codex setup cleanup error and retries its original owner before another probe", async () => {
it("preserves the original Codex setup cleanup while another probe can run", async () => {
const owner = "native-codex-setup-retirement-retry";
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 db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
@@ -596,7 +650,7 @@ describe("managed AI connections", () => {
throw new runnerProbe.NativeEnvironmentProbeCleanupError(retired, new Error("fixture interrupted probe"), [], retry);
}).mockImplementation(async options => {
directories.push(options.runtimeDirectory);
expect(await readFile(path.join(options.transportOptions!.sourceCodexHome!, "auth.json"), "utf8")).toBe(auth("11"));
expect(await readFile(path.join(options.transportOptions!.sourceCodexHome!, "auth.json"), "utf8")).toBe(auth(physical ? "11" : "10"));
await options.onCleanupConfirmed!();
return { provider: "codex", providerDriver: "codex_app_server", effectiveModel: "gpt-5.6-sol", helloProbePassed: true };
});
@@ -610,12 +664,13 @@ describe("managed AI connections", () => {
await expect(access(retainedHome!)).resolves.toBeUndefined();
rejectRetirement(new Error("fixture initial retirement proof failed"));
await Promise.resolve();
await expect(invoke()).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
expect(probe).toHaveBeenCalledOnce(); expect(retry).toHaveBeenCalledOnce();
expect((await invoke()).status).toBe("pass");
expect(probe).toHaveBeenCalledTimes(2); expect(retry).toHaveBeenCalledOnce();
await expect(access(retainedHome!)).resolves.toBeUndefined();
physical = true;
await reconcileRetainedAiCleanup({ companyId, accountKey: row.grant.id });
expect((await invoke()).status).toBe("pass");
expect(retry).toHaveBeenCalledTimes(2); expect(probe).toHaveBeenCalledTimes(2);
expect(retry).toHaveBeenCalledTimes(2); expect(probe).toHaveBeenCalledTimes(3);
expect(await service.credential(row)).toBe(auth("11"));
await expect(access(retainedHome!)).rejects.toMatchObject({ code: "ENOENT" });
} finally {
@@ -636,11 +691,11 @@ describe("managed AI connections", () => {
const replacement = await secretService(db).createCurrentUserSecretValue(companyId, owner, { definitionId: definition.id, value: auth("new-account") }, { userId: owner });
const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
const original = codexAdapter.withAccountHomeSecretMutationLock;
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockImplementationOnce(async (env, company, callback) => {
const original = localEncryptedProvider.resolveVersion;
const lock = vi.spyOn(localEncryptedProvider, "resolveVersion").mockImplementationOnce(async options => {
await db.update(connectionGrants).set({ credentialSecretRefs: row.grant.credentialSecretRefs.map(ref => ref.configPath === "ai.credential" ? { ...ref, secretId: replacement.id } : ref) }).where(eq(connectionGrants.id, row.grant.id));
await syncConnectionCredentialBindings(db, row.connection);
return original(env, company, callback);
return original(options);
});
try { await expect(prepare()).rejects.toThrow("AI credential changed during preparation"); }
finally { lock.mockRestore(); }
@@ -752,7 +807,7 @@ 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) => {
it.each(["company_file", "grant_row"] as const)("reads runtime credentials without an execution lock while a %s writer is active", 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 } });
@@ -771,9 +826,8 @@ describe("managed AI connections", () => {
});
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 service.runtimeCredential(row)).toBe(auth);
expect(lock?.mock.calls ?? []).toHaveLength(0);
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();
@@ -1200,7 +1254,7 @@ describe("managed AI connections", () => {
await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" });
});
it("serializes quota exchange with a concurrent secret edit and runtime credential read", async () => {
it("rereads committed quota rotation without taking an account execution mutex", 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" } });
@@ -1218,15 +1272,18 @@ describe("managed AI connections", () => {
const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal);
try {
await exchangeStarted.promise;
const edit = secretService(db).update(secretId, { description: "Changed during refresh" });
lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
// Runtime reads now wait at the company file lock before touching the
// old credential. The returned home must contain the committed rotation.
let prepared = false;
void preparing.then(() => { prepared = true; });
await vi.waitFor(() => expect(lock).toHaveBeenCalled());
await vi.waitFor(async () => {
const waiting = await db.execute(sql`select pid from pg_stat_activity
where datname = current_database() and wait_event_type = 'Lock' and query like '%company_secrets%'`);
expect(waiting.length).toBeGreaterThan(0);
});
expect(prepared).toBe(false);
expect(lock).not.toHaveBeenCalled();
const edit = secretService(db).update(secretId, { description: "Changed during refresh" });
releaseExchange.resolve();
const results = await Promise.all([refreshing, edit, preparing]);
runtime = results[2];
@@ -138,6 +138,29 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
expect(request).not.toHaveBeenCalled();
});
it("runs two probes of the same Google subscription concurrently in separate homes", async () => {
const f = await fixture("subscription");
let entered!: () => void, release!: () => void;
const started = new Promise<void>(resolve => { entered = resolve; });
const hold = new Promise<void>(resolve => { release = resolve; });
const homes: string[] = [];
const first = withManagedAiProbe(db, f.runtime, async runtime => {
homes.push(runtime.home!); entered(); await hold;
});
let second: Promise<void> | undefined;
try {
await started;
second = withManagedAiProbe(db, f.runtime, async runtime => {
homes.push(runtime.home!);
expect(await readFile(join(String(runtime.config.env.GEMINI_HOME), "antigravity-acp", "acp_token.json"), "utf8")).toBe(f.credential);
});
await vi.waitFor(() => expect(homes).toHaveLength(2), { timeout: 1000 });
await second;
expect(homes[0]).not.toBe(homes[1]);
expect(request).not.toHaveBeenCalled();
} finally { release(); await Promise.allSettled([first, second]); }
});
it.each(["api_key", "subscription"] as const)("refuses revoked %s defaults without falling back to another saved credential", async method => {
const f = await fixture(method);
await f.service.save(f.companyId, f.owner, { ...f.intent, name: "Unselected fallback", connectionId: undefined }, f.credential);
@@ -240,7 +263,7 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
return journal;
}
it("settles a failed probe journal before another provider starts for the same Google account", async () => {
it("runs another provider while the same Google account retains failed probe cleanup", async () => {
const f = await fixture("subscription");
const replacement = googleCredential("recovered-probe-project");
const journalPath = quotaCredentialRecoveryPath(f.companyId, f.saved.grantId);
@@ -253,9 +276,8 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
return realOpen(...args);
});
let originalHome: string | undefined;
const provider = vi.fn(async (runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>>) => {
expect(await readFile(join(String(runtime.config.env.GEMINI_HOME), "antigravity-acp", "acp_token.json"), "utf8")).toBe(replacement);
});
const provider = vi.fn(async (runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>>) =>
readFile(join(String(runtime.config.env.GEMINI_HOME), "antigravity-acp", "acp_token.json"), "utf8"));
try {
await expect(withManagedAiProbe(db, f.runtime, async runtime => {
originalHome = runtime.home;
@@ -263,12 +285,13 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
unavailable = true;
})).rejects.toThrow("fixture Google journal unavailable");
await expect(access(originalHome!)).resolves.toBeUndefined();
await expect(withManagedAiProbe(db, f.runtime, provider)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
expect(provider).not.toHaveBeenCalled();
await expect(withManagedAiProbe(db, f.runtime, provider)).resolves.toBe(f.credential);
expect(provider).toHaveBeenCalledOnce();
expect(await f.service.credential(await f.service.select(f.select))).toBe(f.credential);
unavailable = false;
await withManagedAiProbe(db, f.runtime, provider);
expect(provider).toHaveBeenCalledOnce();
await reconcileRetainedGoogleCleanup({ companyId: f.companyId, accountKey: f.saved.grantId });
await expect(withManagedAiProbe(db, f.runtime, provider)).resolves.toBe(replacement);
expect(provider).toHaveBeenCalledTimes(2);
await expect(access(originalHome!)).rejects.toMatchObject({ code: "ENOENT" });
expect(await f.service.credential(await f.service.select(f.select))).toBe(replacement);
expect(await quotaCredentialRecovery(f.companyId, f.saved.grantId).read()).toBeNull();
@@ -276,7 +299,7 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
} finally { unavailable = false; open.mockRestore(); }
});
it("lets retirement cleanup acquire the credential lock before another account use joins it", async () => {
it("reads Google credentials without joining an unfinished retirement cleanup", async () => {
const f = await fixture("subscription"), runtime = await prepareManagedAiRuntime(db, f.runtime);
const replacement = googleCredential("retirement-concurrency-project");
await writeFile(join(String(runtime.config.env.GEMINI_HOME), "antigravity-acp", "acp_token.json"), replacement);
@@ -303,9 +326,10 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
const using = f.service.runtimeCredential(f.selection);
void using.catch(() => undefined);
try {
if (requestedLock) await credentialLockEntered;
await expect(using).resolves.toBe(f.credential);
expect(requestedLock).toBe(false);
release();
await expect(using).resolves.toBe(replacement);
await reconcileRetainedGoogleCleanup(scope);
await expect(access(runtime.home!)).rejects.toMatchObject({ code: "ENOENT" });
expect(await f.service.credential(await f.service.select(f.select))).toBe(replacement);
expect(request).not.toHaveBeenCalled();
@@ -318,18 +342,16 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
}
});
it("defers credential use when another retained owner arrives before lock acquisition", async () => {
it("reads Google credentials while another retained owner is still live", async () => {
const f = await fixture("subscription"), scope = { companyId: f.companyId, accountKey: f.saved.grantId };
let retire!: () => void;
const retired = new Promise<void>(resolve => { retire = resolve; });
const cleanup = vi.fn(async () => {});
const originalLock = codexAdapter.withAccountHomeSecretMutationLock;
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockImplementationOnce((env, company, callback) => {
retainGoogleCredentialCleanup({ scope, cleanup, retired });
return originalLock(env, company, callback);
});
retainGoogleCredentialCleanup({ scope, cleanup, retired });
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
try {
await expect(f.service.runtimeCredential(f.selection)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(f.service.runtimeCredential(f.selection)).resolves.toBe(f.credential);
expect(lock).not.toHaveBeenCalled();
expect(cleanup).not.toHaveBeenCalled();
expect(request).not.toHaveBeenCalled();
} finally { lock.mockRestore(); retire(); await reconcileRetainedGoogleCleanup(scope); }
@@ -352,8 +374,7 @@ describeDatabase("Google AI connection authority with real PostgreSQL", () => {
try {
await codexAdapter.withAccountHomeSecretMutationLock(undefined, f.companyId, async () => {
retire(); await cleanupEntered;
await expect(f.service.runtimeCredential(f.selection, { credentialLockHeld: true }))
.rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(f.service.runtimeCredential(f.selection)).resolves.toBe(f.credential);
expect(failedLock).toBe(false);
expect(finished).toBe(false);
});
@@ -19,10 +19,10 @@ const input = { companyId: "company", userId: "alice", binding: { provider: "goo
let sequence = 0;
let scope = { companyId: "company", accountKey: "" };
beforeEach(() => { vi.clearAllMocks(); scope = { companyId: "company", accountKey: `catalog-${++sequence}` }; mocks.cleanup.mockResolvedValue(undefined); mocks.managed.mockImplementation(async () => {
await reconcileRetainedGoogleCleanup(scope);
void reconcileRetainedGoogleCleanup(scope).catch(() => undefined);
return {
attribution: { method: "subscription" }, config: { env: { GEMINI_HOME: "/private/selected-google-home" } },
cleanup: mocks.cleanup, retainedCleanupScope: scope,
cleanup: () => mocks.cleanup(), retainedCleanupScope: scope,
}; }); });
afterEach(async () => { vi.useRealTimers(); mocks.cleanup.mockResolvedValue(undefined); await reconcileRetainedGoogleCleanup(scope); });
@@ -105,21 +105,22 @@ it("waits for native retirement, then lets the durable journal own failed writeb
expect(JSON.stringify(mocks.warn.mock.calls)).not.toContain("private-token");
});
it("settles the original failed journal cleanup before retrying model discovery", async () => {
it("retries the original failed journal cleanup independently of model discovery", async () => {
const order: string[] = [];
mocks.discover.mockImplementation(async () => { order.push("provider"); return { models: [] }; });
mocks.cleanup.mockImplementationOnce(async () => { order.push("journal-failed"); throw new Error("storage unavailable"); })
.mockImplementation(async () => { order.push("journal-saved"); });
await expect(antigravityConnectionModels({} as never, input)).rejects.toThrow("Could not complete Google account credential cleanup");
await expect(antigravityConnectionModels({} as never, input)).resolves.toEqual([]);
expect(order).toEqual(["provider", "journal-failed", "journal-saved", "provider", "journal-saved"]);
expect(order.filter(step => step === "provider")).toHaveLength(2);
expect(order.filter(step => step === "journal-saved")).toHaveLength(2);
});
it("blocks another provider launch while the original cleanup still cannot journal", async () => {
it("allows another provider launch while the original cleanup still cannot journal", async () => {
mocks.discover.mockResolvedValue({ models: [] });
mocks.cleanup.mockRejectedValue(new Error("storage unavailable"));
await expect(antigravityConnectionModels({} as never, input)).rejects.toThrow("Could not complete Google account credential cleanup");
await expect(antigravityConnectionModels({} as never, input)).rejects.toThrow("still completing credential cleanup");
expect(mocks.discover).toHaveBeenCalledTimes(1);
expect(mocks.cleanup).toHaveBeenCalledTimes(2);
await expect(antigravityConnectionModels({} as never, input)).rejects.toThrow("Could not complete Google account credential cleanup");
expect(mocks.discover).toHaveBeenCalledTimes(2);
expect(mocks.cleanup).toHaveBeenCalledTimes(3);
});
@@ -195,6 +195,7 @@ describe("managed Antigravity subscription projection and writeback", () => {
account.rotation.mockImplementationOnce(async (_secretId, input) => { account.credential = input.value; account.version += 1; });
request.mockResolvedValueOnce(Response.json({ access_token: "fixture-access", refresh_token: "fixture-retry-rotation" }))
.mockResolvedValueOnce(Response.json({ email_verified: true, email: "alice@example.test" }));
await runtime.cleanup();
const next = await prepare();
expect(new URLSearchParams(request.mock.calls[2][1].body).get("refresh_token")).toBe("fixture-latest-rotation");
expect(await readFile(join(String(next.config.env.GEMINI_HOME), "antigravity-acp", "acp_token.json"), "utf8"))
+33 -67
View File
@@ -1,4 +1,3 @@
import { AsyncLocalStorage } from "node:async_hooks";
import { environmentTaskResultSchema } from "@paperclipai/plugin-sdk";
import type { PluginTaskCredentialWritebackBinding } from "./plugin-task-credential-writeback.js";
import { OPENROUTER_ATTRIBUTION_HEADERS, openRouterClaudeHeaders } from "@paperclipai/shared/openrouter-attribution";
@@ -28,7 +27,7 @@ import { NativeEnvironmentProbeCleanupError, sanitizeAntigravityCredential } fro
import { readLocalAiCredentialFile } from "./local-ai-credential-file.js";
import { prepareAntigravityCredentialWriteback } from "./native-runtime/antigravity-runner-credentials.js";
import { quotaCredentialHash, quotaCredentialRecovery, RetiredCredentialRecoveryError } from "./quota-credential-recovery.js";
import { assertAiCredentialCleanupSettled, bindRetainedAiCredentialHome, reconcileRetainedAiCleanup, retainAiCredentialCleanup, type RetainedAiCleanupScope } from "./retained-ai-credential-cleanup.js";
import { bindRetainedAiCredentialHome, reconcileRetainedAiCleanup, retainAiCredentialCleanup } from "./retained-ai-credential-cleanup.js";
export function isAiConnectionBusy(error: unknown): error is HttpError {
return error instanceof HttpError && error.status === 422 &&
@@ -230,13 +229,6 @@ export function managedAiSessionFingerprintConfig(
return { ...config, env, ...(managed?.sessionIdentity ? { managedAiConnection: { ...managed, identity: managed.sessionIdentity } } : {}) };
}
const managedProbeCredentialLocks = new AsyncLocalStorage<{ companyId: string; active: boolean }>();
async function reconcileManagedCredentialCleanup(scope: RetainedAiCleanupScope): Promise<void> {
const ownership = managedProbeCredentialLocks.getStore();
if (ownership?.active && ownership.companyId === scope.companyId) assertAiCredentialCleanupSettled(scope);
else await reconcileRetainedAiCleanup(scope);
}
export async function prepareManagedAiRuntime(
db: Db,
input: {
@@ -282,7 +274,7 @@ export async function prepareManagedAiRuntime(
acpxAgent: input.config.acpxAgent,
});
const retainedCleanupScope = { companyId: input.companyId, accountKey: selection.grant.id };
await reconcileManagedCredentialCleanup(retainedCleanupScope);
void reconcileRetainedAiCleanup(retainedCleanupScope).catch(() => undefined);
const subscriptionFile =
selection.attribution.method === "subscription" &&
input.binding.provider !== "anthropic";
@@ -321,9 +313,7 @@ export async function prepareManagedAiRuntime(
// against a concurrent reconnect or explicit rotation.
for (let attempt = 0; attempt < 2; attempt += 1) {
const before = await readFreshness();
const ownership = managedProbeCredentialLocks.getStore();
const value = await service.runtimeCredential(selection,
ownership?.active && ownership.companyId === input.companyId ? { credentialLockHeld: true } : undefined);
const value = await service.runtimeCredential(selection);
const after = await readFreshness();
if (before && after && before.version === after.version && before.epoch === after.epoch) {
return { value, freshness: after };
@@ -441,14 +431,12 @@ export async function prepareManagedAiRuntime(
homeRemoved = true;
releaseCredentialHome?.();
};
const cleanupRuntime = async () => {
if (homeRemoved) return;
await reconcileManagedCredentialCleanup(retainedCleanupScope);
let cleanupRunning: Promise<void> | undefined;
const cleanupRuntime = (): Promise<void> => cleanupRunning ??= (async () => {
if (homeRemoved) return;
const cleanup = async () => {
if (removalOnly) {
await removeHome();
if (deferredWriteback) throw new RetiredCredentialRecoveryError(deferredWriteback.cause);
return;
}
let googleRecovery: ReturnType<typeof quotaCredentialRecovery> | undefined;
@@ -465,7 +453,7 @@ export async function prepareManagedAiRuntime(
if (prior && prior.value !== incoming && !ownsGoogleRecovery) {
// A previous retired owner keeps its first recoverable refresh.
// Settle it before this owner can stage another generation.
await aiConnectionService(db).runtimeCredential(selection, { credentialLockHeld: true });
await aiConnectionService(db).runtimeCredential(selection);
}
// Physical retirement is the caller's cleanup precondition. Save
// before network verification or a database write can fail; only
@@ -600,10 +588,10 @@ export async function prepareManagedAiRuntime(
return await (antigravitySubscription
? withAccountHomeSecretMutationLock(undefined, input.companyId, cleanup) : cleanup());
} catch (error) {
if (retainedCleanupScope && !homeRemoved) retainAiCredentialCleanup({ scope: retainedCleanupScope, cleanup: cleanupRuntime });
if (retainedCleanupScope && !homeRemoved) retainAiCredentialCleanup({ scope: retainedCleanupScope, home: providerHome, cleanup: cleanupRuntime });
throw error;
}
};
})().finally(() => { cleanupRunning = undefined; });
return {
sessionIdentity,
...(retainedCleanupScope ? { retainedCleanupScope } : {}),
@@ -632,58 +620,36 @@ export async function prepareManagedAiRuntime(
}
}
/** Manual tests/adoption have no heartbeat row for quota's active-run guard.
* Hold the shared credential lock from resolution through the provider probe
* and write-back. Nested credential writes reuse this lock's async ownership. */
export async function withManagedAiProbe<T>(
db: Db,
input: Parameters<typeof prepareManagedAiRuntime>[1],
probe: (runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>>) => Promise<T>,
): Promise<T> {
const run = async () => {
const runtime = await prepareManagedAiRuntime(db, input);
let retained = false;
try {
assertAiCredentialCleanupSettled(runtime.retainedCleanupScope!);
return await probe(runtime);
}
catch (error) {
if (error instanceof NativeEnvironmentProbeCleanupError) {
// Keep the selected credential generation until the native owner has
// retired and completed its refresh handoff. HTTP need not wait for it.
retained = true;
const retryCleanup = (error as { retryCleanup?: () => Promise<void> }).retryCleanup;
retainAiCredentialCleanup({ scope: runtime.retainedCleanupScope!, cleanup: runtime.cleanup, retired: error.retired,
...(typeof retryCleanup === "function" ? { retryRetirement: () => retryCleanup.call(error) } : {}) });
throw unprocessable("The native account test is still cleaning up. Retry after runtime cleanup completes.", { code: "native_probe_cleanup_pending" });
}
throw error;
} finally {
if (!retained) {
try { await runtime.cleanup(); }
catch (error) {
if (runtime.retainedCleanupScope && !(error instanceof RetiredCredentialRecoveryError)) {
retainAiCredentialCleanup({ scope: runtime.retainedCleanupScope, cleanup: runtime.cleanup });
}
throw error;
}
}
}
};
if (input.binding.provider !== "openai") return run();
// A retirement callback may need the company lock. Join it before taking
// the probe's outer lock; preparation rechecks this scope synchronously.
const selected = await aiConnectionService(db).select({ ...input, userId: input.responsibleUserId,
model: input.config.model, runnerProvider: input.config.provider, acpxAgent: input.config.acpxAgent });
await reconcileRetainedAiCleanup({ companyId: input.companyId, accountKey: selected.grant.id });
return withAccountHomeSecretMutationLock(undefined, input.companyId, async () => {
const ownership = { companyId: input.companyId, active: true };
try { return await managedProbeCredentialLocks.run(ownership, run); }
finally { ownership.active = false; }
}).catch(error => {
if ((error as NodeJS.ErrnoException)?.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) {
throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" });
const runtime = await prepareManagedAiRuntime(db, input);
let retained = false;
try {
return await probe(runtime);
}
catch (error) {
if (error instanceof NativeEnvironmentProbeCleanupError) {
// Keep the selected credential generation until the native owner has
// retired and completed its refresh handoff. HTTP need not wait for it.
retained = true;
const retryCleanup = (error as { retryCleanup?: () => Promise<void> }).retryCleanup;
retainAiCredentialCleanup({ scope: runtime.retainedCleanupScope!, home: runtime.home && path.join(runtime.home, "provider"), cleanup: runtime.cleanup, retired: error.retired,
...(typeof retryCleanup === "function" ? { retryRetirement: () => retryCleanup.call(error) } : {}) });
throw unprocessable("The native account test is still cleaning up. Retry after runtime cleanup completes.", { code: "native_probe_cleanup_pending" });
}
throw error;
});
} finally {
if (!retained) {
try { await runtime.cleanup(); }
catch (error) {
if (runtime.retainedCleanupScope && !(error instanceof RetiredCredentialRecoveryError)) {
retainAiCredentialCleanup({ scope: runtime.retainedCleanupScope, home: runtime.home && path.join(runtime.home, "provider"), cleanup: runtime.cleanup });
}
throw error;
}
}
}
}
+2 -11
View File
@@ -45,7 +45,6 @@ import { forbidden, notFound, unprocessable } from "../errors.js";
import { aiConnectionConfigurationFailure, aiConnectionCredentialNotSharedFailure } from "./ai-connection-configuration-failure.js";
import { quotaCredentialRecovery, quotaCredentialHash } from "./quota-credential-recovery.js";
import { prepareAntigravityCredentialRecovery } from "./native-runtime/antigravity-runner-credentials.js";
import { assertAiCredentialCleanupSettled, reconcileRetainedAiCleanup } from "./retained-ai-credential-cleanup.js";
import { logger } from "../middleware/logger.js";
import { logActivity } from "./activity-log.js";
import { secretService } from "./secrets.js";
@@ -987,20 +986,12 @@ export function aiConnectionService(db: Db) {
return pending.value;
}
/** Never hand a consumed token to a new provider while its replacement awaits commit. */
async function runtimeCredential(row: Parameters<typeof credential>[0], ownership?: { credentialLockHeld: true }) {
async function runtimeCredential(row: Parameters<typeof credential>[0]) {
const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai);
if (!metadata.success || !["openai", "google"].includes(metadata.data.provider) || metadata.data.method !== "subscription") return credential(row);
// Retirement callbacks acquire this same mutation lock to persist their
// final generation. Join them before taking the lock they need. A caller
// already holding it must only check, for both rotated-credential lanes.
const scope = { companyId: row.connection.companyId, accountKey: row.grant.id };
if (ownership?.credentialLockHeld) assertAiCredentialCleanupSettled(scope);
else await reconcileRetainedAiCleanup(scope);
if (!(await quotaCredentialRecovery(row.connection.companyId, row.grant.id).read())) return credential(row);
let savingReplacement = false;
return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => {
// An owner can arrive while lock acquisition waits. Never join under it.
assertAiCredentialCleanupSettled(scope);
if (metadata.data.provider === "google") {
const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id);
const pending = await recovery.read();
@@ -40,7 +40,7 @@ export async function antigravityConnectionModels(db: Db, input: {
// Keep this selected credential home until that owner proves retirement.
retained = true;
if (runtime.retainedCleanupScope) retainGoogleCredentialCleanup({
scope: runtime.retainedCleanupScope, cleanup: runtime.cleanup, retired: error.retired,
scope: runtime.retainedCleanupScope, home, cleanup: runtime.cleanup, retired: error.retired,
});
}
throw unprocessable("Could not load models for this Google account. Check the account's Antigravity access, then refresh models.");
@@ -53,7 +53,9 @@ export async function antigravityConnectionModels(db: Db, input: {
// Storage failures still retain the private home and block sleep.
// No provider or repeated database work is started by this hold.
if (runtime.retainedCleanupScope) retainGoogleCredentialCleanup({
scope: runtime.retainedCleanupScope, cleanup: runtime.cleanup,
scope: runtime.retainedCleanupScope,
home: typeof runtime.config.env.GEMINI_HOME === "string" ? runtime.config.env.GEMINI_HOME : undefined,
cleanup: runtime.cleanup,
});
}
logger.warn({ companyId: input.companyId, code: "ANTIGRAVITY_CATALOG_CREDENTIAL_CLEANUP_PENDING" },
@@ -169,7 +169,7 @@ describe("Antigravity managed Runner credentials", () => {
expect(await readFile(destination, "utf8")).toBe(credential("fixture-refreshed"));
await expect(access(source)).rejects.toMatchObject({ code: "ENOENT" });
});
it("settles the retained runtime refresh before preparing another launch for the same account", async () => {
it("settles the retained runtime refresh before reusing its original home", async () => {
const managedHome = await home(), runtimeHome = await home();
const release = bindRetainedGoogleCredentialHome(managedHome, { companyId: "company", accountKey: managedHome });
const destination = join(managedHome, "antigravity-acp", "acp_token.json");
@@ -182,7 +182,7 @@ describe("Antigravity managed Runner credentials", () => {
try {
await expect(completeManagedNativeCredentialTurn(session)).rejects.toThrow("destination is invalid");
await expect(prepareAntigravityRunnerCredentials({ managedHome, remote: false, environment: {} }))
.rejects.toThrow("still completing credential cleanup");
.rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
expect(await readFile(source, "utf8")).toBe(credential("fixture-refreshed"));
await chmod(destination, 0o600);
expect((await prepareAntigravityRunnerCredentials({ managedHome, remote: false, environment: {} })).credential)
@@ -191,7 +191,7 @@ describe("Antigravity managed Runner credentials", () => {
await expect(access(source)).rejects.toMatchObject({ code: "ENOENT" });
} finally { release(); }
});
it.each(["turn completion", "direct close"])("retries original process proof after rejected %s before account reuse, without clearing session failure", async origin => {
it.each(["turn completion", "direct close"])("retries original process proof after rejected %s before reusing the original home, without clearing session failure", async origin => {
const managedHome = await home(), runtimeHome = await home();
const release = bindRetainedGoogleCredentialHome(managedHome, { companyId: "company-retirement", accountKey: managedHome });
const destination = join(managedHome, "antigravity-acp", "acp_token.json"), source = join(runtimeHome, "antigravity-acp", "auth-refresh.json");
@@ -212,7 +212,7 @@ describe("Antigravity managed Runner credentials", () => {
try {
await expect(origin === "turn completion" ? completeManagedNativeCredentialTurn(session) : session.close({ reason: "direct close" })).rejects.toBe(terminal);
await expect(prepareAntigravityRunnerCredentials({ managedHome, remote: false, environment: {} }))
.rejects.toThrow("still completing credential cleanup");
.rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
expect(proof).toHaveBeenCalledOnce();
expect(await readFile(source, "utf8")).toBe(credential("fixture-refreshed"));
expect(await readFile(destination, "utf8")).toBe(credential("fixture-original"));
@@ -7,6 +7,7 @@ import { RetiredCredentialRecoveryError } from "./quota-credential-recovery.js";
export type RetainedAiCleanupScope = { companyId: string; accountKey: string };
type Cleanup = {
scope: RetainedAiCleanupScope;
home?: string;
cleanup: () => Promise<void>;
retryRetirement?: () => Promise<void>;
retired: "pending" | "confirmed" | "failed";
@@ -15,14 +16,13 @@ type Cleanup = {
warned?: boolean;
};
const retained = new Map<string, Set<Cleanup>>();
const reconciling = new Map<string, Promise<void>>();
const homes = new Map<string, RetainedAiCleanupScope>();
const active = new AsyncLocalStorage<ReadonlySet<Cleanup>>();
const key = (scope: RetainedAiCleanupScope) => JSON.stringify([scope.companyId, scope.accountKey]);
registerIdleLocalWorkProbe(() => retained.size ? "present" : "none");
function pending() {
return unprocessable("The selected AI account is still completing credential cleanup. Retry shortly.", { code: "ai_connection_busy" });
return unprocessable("This runtime still has unfinished credential cleanup.", { code: "ai_credential_cleanup_pending" });
}
function remove(entry: Cleanup) {
const entries = retained.get(key(entry.scope));
@@ -38,11 +38,6 @@ function cleaningScope(scope: RetainedAiCleanupScope) {
return [...(active.getStore() ?? [])].some(entry => entry.executing && key(entry.scope) === key(scope));
}
/** Check without joining work that may need the mutation lock already held by
* the caller. A cleanup's own nested journal recovery shares its ownership. */
export function assertAiCredentialCleanupSettled(scope: RetainedAiCleanupScope): void {
if (!cleaningScope(scope) && retained.get(key(scope))?.size) throw pending();
}
function attempt(entry: Cleanup): Promise<void> {
if (!retained.get(key(entry.scope))?.has(entry)) return Promise.resolve();
if (entry.retired === "pending" || (entry.retired === "failed" && !entry.retryRetirement)) return Promise.reject(pending());
@@ -75,34 +70,20 @@ export function retainAiCredentialCleanup(input: {
const scopeKey = key(input.scope);
const entries = retained.get(scopeKey) ?? new Set<Cleanup>();
if ([...entries].some(entry => entry.cleanup === input.cleanup)) return;
const entry: Cleanup = { scope: { ...input.scope }, cleanup: input.cleanup, retryRetirement: input.retryRetirement,
const entry: Cleanup = { scope: { ...input.scope }, home: input.home, cleanup: input.cleanup, retryRetirement: input.retryRetirement,
retired: input.retired ? "pending" : "confirmed" };
entries.add(entry); retained.set(scopeKey, entries);
if (input.retired) void input.retired.then(() => {
entry.retired = "confirmed";
return reconcileRetainedAiCleanup(entry.scope).catch(() => undefined);
return attempt(entry).catch(() => undefined);
}, () => { entry.retired = "failed"; warn(entry); });
}
/** Fail closed before another provider starts. Concurrent attempts join the
* original cleanup. Nested credential reads skip this account's peer join to
* avoid recursing into the current callback. The outer bounded pass and final
* fence still settle every peer before any external provider admission. */
export async function reconcileRetainedAiCleanup(scope: RetainedAiCleanupScope): Promise<void> {
await Promise.resolve();
const scopeKey = key(scope);
if (cleaningScope(scope)) return;
const existing = reconciling.get(scopeKey);
if (existing) return existing;
const operation = (async () => {
for (const entry of [...(retained.get(scopeKey) ?? [])]) await attempt(entry);
// A new obligation can arrive while a callback awaits storage. This
// request attempts its original bounded set and keeps admission closed.
if (retained.get(scopeKey)?.size) throw pending();
})();
reconciling.set(scopeKey, operation);
try { await operation; }
finally { if (reconciling.get(scopeKey) === operation) reconciling.delete(scopeKey); }
for (const entry of [...(retained.get(scopeKey) ?? [])]) await attempt(entry);
}
export function bindRetainedAiCredentialHome(home: string, scope: RetainedAiCleanupScope): () => void {
@@ -112,7 +93,10 @@ export function bindRetainedAiCredentialHome(home: string, scope: RetainedAiClea
}
export async function reconcileRetainedAiCleanupForHome(home: string): Promise<void> {
const scope = homes.get(home);
if (scope) await reconcileRetainedAiCleanup(scope);
if (!scope || cleaningScope(scope)) return;
for (const entry of [...(retained.get(key(scope)) ?? [])]) {
if (entry.home === home) await attempt(entry);
}
}
export function retainAiCredentialCleanupForHome(home: string, input: {
cleanup: () => Promise<void>; retired?: Promise<void>; retryRetirement?: () => Promise<void>;
@@ -1,6 +1,6 @@
import { expect, it, vi } from "vitest";
import {
assertGoogleCredentialCleanupSettled, bindRetainedGoogleCredentialHome, reconcileRetainedGoogleCleanup,
bindRetainedGoogleCredentialHome, reconcileRetainedGoogleCleanup,
reconcileRetainedGoogleCleanupForHome, retainGoogleCredentialCleanup,
retainGoogleCredentialCleanupForHome,
} from "./retained-google-credential-cleanup.js";
@@ -33,7 +33,7 @@ it("preserves failed cleanup for an event-driven retry without launching a provi
const provider = vi.fn();
retainGoogleCredentialCleanup({ scope: owner, cleanup });
const useAccount = async () => { await reconcileRetainedGoogleCleanup(owner); provider(); };
await expect(useAccount()).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(useAccount()).rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
expect(provider).not.toHaveBeenCalled();
await useAccount();
expect(cleanup).toHaveBeenCalledTimes(2);
@@ -46,7 +46,7 @@ it("requires confirmed physical retirement before retrying the retained callback
const retired = new Promise<void>(resolve => { finish = resolve; });
const cleanup = vi.fn(async () => {});
retainGoogleCredentialCleanup({ scope: owner, cleanup, retired });
await expect(reconcileRetainedGoogleCleanup(owner)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(reconcileRetainedGoogleCleanup(owner)).rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
expect(cleanup).not.toHaveBeenCalled();
finish();
await vi.waitFor(() => expect(cleanup).toHaveBeenCalledOnce());
@@ -72,7 +72,7 @@ it("retries a rejected physical proof only on the next account action and coales
const cleanup = vi.fn(async () => {});
retainGoogleCredentialCleanup({ scope: owner, cleanup, retired, retryRetirement });
await retired.catch(() => {});
await expect(reconcileRetainedGoogleCleanup(owner)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await expect(reconcileRetainedGoogleCleanup(owner)).rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
expect(retryRetirement).toHaveBeenCalledOnce(); expect(cleanup).not.toHaveBeenCalled();
const first = reconcileRetainedGoogleCleanup(owner), second = reconcileRetainedGoogleCleanup(owner);
await vi.waitFor(() => expect(retryRetirement).toHaveBeenCalledTimes(2));
@@ -85,23 +85,22 @@ it("retries a rejected physical proof only on the next account action and coales
it("does not recurse through cleanup wrappers and settles peers once", async () => {
const owner = scope();
const first = vi.fn(async () => { assertGoogleCredentialCleanupSettled(owner); await reconcileRetainedGoogleCleanup(owner); });
const first = vi.fn(async () => { await reconcileRetainedGoogleCleanup(owner); });
const second = vi.fn(async () => { await reconcileRetainedGoogleCleanup(owner); });
retainGoogleCredentialCleanup({ scope: owner, cleanup: first });
retainGoogleCredentialCleanup({ scope: owner, cleanup: second });
expect(() => assertGoogleCredentialCleanupSettled(owner)).toThrow("still completing credential cleanup");
await reconcileRetainedGoogleCleanup(owner);
expect(first).toHaveBeenCalledOnce();
expect(second).toHaveBeenCalledOnce();
});
it("keeps admission closed when a new obligation arrives during bounded cleanup", async () => {
it("leaves a new obligation for the next bounded cleanup attempt", async () => {
const owner = scope();
const second = vi.fn(async () => {});
retainGoogleCredentialCleanup({ scope: owner, cleanup: async () => {
retainGoogleCredentialCleanup({ scope: owner, cleanup: second });
} });
await expect(reconcileRetainedGoogleCleanup(owner)).rejects.toMatchObject({ details: { code: "ai_connection_busy" } });
await reconcileRetainedGoogleCleanup(owner);
expect(second).not.toHaveBeenCalled();
await reconcileRetainedGoogleCleanup(owner);
expect(second).toHaveBeenCalledOnce();
@@ -119,3 +118,20 @@ it("retains the original home's cleanup scope until its owner releases it", asyn
await reconcileRetainedGoogleCleanup(owner);
expect(cleanup).toHaveBeenCalledOnce();
});
it("does not join cleanup from another home that uses the same account", async () => {
const owner = scope(), originalHome = `/fixture/original-${owner.accountKey}`, nextHome = `/fixture/next-${owner.accountKey}`;
const releaseOriginal = bindRetainedGoogleCredentialHome(originalHome, owner);
const releaseNext = bindRetainedGoogleCredentialHome(nextHome, owner);
let finish!: () => void;
const retired = new Promise<void>(resolve => { finish = resolve; });
const cleanup = vi.fn(async () => {});
retainGoogleCredentialCleanupForHome(originalHome, { cleanup, retired });
try {
await reconcileRetainedGoogleCleanupForHome(nextHome);
expect(cleanup).not.toHaveBeenCalled();
await expect(reconcileRetainedGoogleCleanupForHome(originalHome)).rejects.toMatchObject({ details: { code: "ai_credential_cleanup_pending" } });
} finally {
finish(); await reconcileRetainedGoogleCleanup(owner); releaseOriginal(); releaseNext();
}
});
@@ -1,7 +1,6 @@
// Compatibility names for Google-specific callers; all managed accounts share
// the same company/account-scoped owner and bounded reconciliation registry.
export {
assertAiCredentialCleanupSettled as assertGoogleCredentialCleanupSettled,
bindRetainedAiCredentialHome as bindRetainedGoogleCredentialHome,
reconcileRetainedAiCleanup as reconcileRetainedGoogleCleanup,
reconcileRetainedAiCleanupForHome as reconcileRetainedGoogleCleanupForHome,