diff --git a/doc/SPEC-implementation.md b/doc/SPEC-implementation.md index c52e309272f..f8a9ceaf7cb 100644 --- a/doc/SPEC-implementation.md +++ b/doc/SPEC-implementation.md @@ -2218,6 +2218,11 @@ Legacy agents retain their authentication until validated adoption. See [AI Connections](connections/AI-CONNECTIONS.md) for company isolation, compatible methods, lifecycle, runtime enforcement, and migration details. +AI connections support concurrent runs and account tests. Admission does not +take an account execution lock or wait for another runtime's credential cleanup. +Private runtime homes remain independent. Credential writeback still checks the +active grant and saved generation before an atomic update. + Missing personal AI credentials detected before adapter dispatch also produce the inline connection card. Every missing binding must belong to the same compatible AI provider. The responsible user connects their own account and diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 2a47847f46d..50457cee566 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -337,8 +337,10 @@ grant's credentials. Inherited credential variables are cleared. Conflicting project authentication and provider-routing overrides are rejected. Managed failure cannot reactivate host or legacy credentials. -A subscription invocation takes no lease. Two invocations of one grant, from -the same or a different provider account, run at the same time. At cleanup, +A subscription invocation takes no lease or account mutex. Runs and manual +account tests using one grant can start at the same time. Ordinary credential +reads take no company file lock or grant/secret write locks. Pending cleanup +belongs to its original runtime and does not block another invocation. At cleanup, each invocation re-reads the credential stored at that moment under a row lock on the grant, then compares it against its own refreshed copy using the provider's own freshness field: Codex compares `last_refresh` and bounds it @@ -553,15 +555,14 @@ existing account name. Connecting installs access for that agent and resumes the work automatically. Explicit incompatible bindings and shared-account permission denials still fail; hiring never expands a restricted shared account's audience. -Concurrent runs of one subscription do not hold a credential lease. A fresh -task can briefly wait when a credential rotation holds the company file lock or -a grant/secret database row lock. Lock timeouts become `ai_connection_busy`, -keeping the task on its automatic pre-provider scheduled retry path. It does -not request new credentials, and it does not consume the provider-failure retry -allowance. Each retry revalidates the account, and existing run-dispatch rules -still suppress cancelled, reassigned, or otherwise ineligible work. An assignee -retry must still keep execution-lock ownership at scheduling, promotion, and -dispatch. +Concurrent runs and manual tests do not hold an AI-connection execution lock. +An unfinished or failed cleanup from another runtime cannot cancel a new task +with `ai_connection_busy`. New invocations can request a bounded cleanup retry +without waiting for it. A home-specific cleanup check only joins work for that +same private home. Actual credential updates and durable replacement recovery +retain their atomic storage transactions and generation checks. Historical +scheduled retries still revalidate the account and task eligibility, including +assignee execution-lock ownership at scheduling, promotion, and dispatch. `server/src/__tests__/agent-hire-ai-connections.test.ts` covers both creation routes, both providers and methods, approval gates, native provider mapping, shared access diff --git a/server/src/__tests__/agent-hire-ai-connections.test.ts b/server/src/__tests__/agent-hire-ai-connections.test.ts index 2b397429aa8..024f133c9a0 100644 --- a/server/src/__tests__/agent-hire-ai-connections.test.ts +++ b/server/src/__tests__/agent-hire-ai-connections.test.ts @@ -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(); diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 9a02179171f..96324216d13 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -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> | 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 | 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> | 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]; diff --git a/server/src/__tests__/antigravity-ai-connections.integration.test.ts b/server/src/__tests__/antigravity-ai-connections.integration.test.ts index 58f0ab39b3b..b017327e390 100644 --- a/server/src/__tests__/antigravity-ai-connections.integration.test.ts +++ b/server/src/__tests__/antigravity-ai-connections.integration.test.ts @@ -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(resolve => { entered = resolve; }); + const hold = new Promise(resolve => { release = resolve; }); + const homes: string[] = []; + const first = withManagedAiProbe(db, f.runtime, async runtime => { + homes.push(runtime.home!); entered(); await hold; + }); + let second: Promise | 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>) => { - 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>) => + 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(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); }); diff --git a/server/src/__tests__/antigravity-connection-models.test.ts b/server/src/__tests__/antigravity-connection-models.test.ts index c96e981e339..2fe0cb219ea 100644 --- a/server/src/__tests__/antigravity-connection-models.test.ts +++ b/server/src/__tests__/antigravity-connection-models.test.ts @@ -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); }); diff --git a/server/src/__tests__/antigravity-managed-subscription.test.ts b/server/src/__tests__/antigravity-managed-subscription.test.ts index d2f54633310..0a60321d15e 100644 --- a/server/src/__tests__/antigravity-managed-subscription.test.ts +++ b/server/src/__tests__/antigravity-managed-subscription.test.ts @@ -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")) diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index d345002e451..593b6aeeb34 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -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 { - 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 | undefined; + const cleanupRuntime = (): Promise => cleanupRunning ??= (async () => { if (homeRemoved) return; const cleanup = async () => { if (removalOnly) { await removeHome(); - if (deferredWriteback) throw new RetiredCredentialRecoveryError(deferredWriteback.cause); return; } let googleRecovery: ReturnType | 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( db: Db, input: Parameters[1], probe: (runtime: Awaited>) => Promise, ): Promise { - 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 }).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 }).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; + } + } + } } diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 6314a5265d8..3622fc2c1c8 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -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[0], ownership?: { credentialLockHeld: true }) { + async function runtimeCredential(row: Parameters[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(); diff --git a/server/src/services/antigravity-connection-models.ts b/server/src/services/antigravity-connection-models.ts index d603327f01e..be787da0f3a 100644 --- a/server/src/services/antigravity-connection-models.ts +++ b/server/src/services/antigravity-connection-models.ts @@ -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" }, diff --git a/server/src/services/native-runtime/antigravity-runner-credentials.test.ts b/server/src/services/native-runtime/antigravity-runner-credentials.test.ts index b93aef47185..e5e11d9c580 100644 --- a/server/src/services/native-runtime/antigravity-runner-credentials.test.ts +++ b/server/src/services/native-runtime/antigravity-runner-credentials.test.ts @@ -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")); diff --git a/server/src/services/retained-ai-credential-cleanup.ts b/server/src/services/retained-ai-credential-cleanup.ts index 34a6dc6a204..78566c9ca6e 100644 --- a/server/src/services/retained-ai-credential-cleanup.ts +++ b/server/src/services/retained-ai-credential-cleanup.ts @@ -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; retryRetirement?: () => Promise; retired: "pending" | "confirmed" | "failed"; @@ -15,14 +16,13 @@ type Cleanup = { warned?: boolean; }; const retained = new Map>(); -const reconciling = new Map>(); const homes = new Map(); const active = new AsyncLocalStorage>(); 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 { 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(); 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 { 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 { 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; retired?: Promise; retryRetirement?: () => Promise; diff --git a/server/src/services/retained-google-credential-cleanup.test.ts b/server/src/services/retained-google-credential-cleanup.test.ts index ccdc1a68d43..06b0135a01f 100644 --- a/server/src/services/retained-google-credential-cleanup.test.ts +++ b/server/src/services/retained-google-credential-cleanup.test.ts @@ -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(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(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(); + } +}); diff --git a/server/src/services/retained-google-credential-cleanup.ts b/server/src/services/retained-google-credential-cleanup.ts index ddbc8721772..ce7964751ad 100644 --- a/server/src/services/retained-google-credential-cleanup.ts +++ b/server/src/services/retained-google-credential-cleanup.ts @@ -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,