diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 3cefc7aaad..603c73380d 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -103,7 +103,13 @@ 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. -Subscription invocations take a grant-scoped database advisory lease. Two +Subscription invocations take a grant-scoped transaction advisory lease. The +reserved database client keeps one transaction open until cleanup, including on +transaction-pooling proxies such as PgBouncer. Session-level advisory locks must +not be used here: a pooled connection can return to a different backend for +cleanup and leave the original lock behind. The lease transaction disables its +idle timeout and contains no application data writes; cleanup rolls it back. +Two different users' grants can run concurrently; a second invocation of the same subscription receives a retryable busy response while it is in use. Refreshes are merged only into the originating active grant, with reconnect/revocation diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 5bf2b0c060..b4e14723b6 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -196,6 +196,21 @@ describe("managed AI connections", () => { it("serializes subscription refresh and releases the lease after execution", async () => { const subscription = { ...input, binding: { ...binding, method: "subscription" as const }, responsibleUserId: "alice", config: { model: "same-model" } }; const first = await prepareManagedAiRuntime(db, subscription); + const selected = await service.select({ ...subscription, userId: "alice" }); + const lockKey = `ai-runtime:${selected.grant.id}`; + const held = await db.execute(sql` + select activity.state, activity.xact_start + from pg_locks locks join pg_stat_activity activity on activity.pid = locks.pid + where locks.locktype = 'advisory' and locks.granted + and locks.classid = (hashtextextended(${lockKey}, 0) >> 32)::int::oid + and locks.objid = (hashtextextended(${lockKey}, 0) & 4294967295)::oid + and locks.objsubid = 1 + `); + // A transaction-pooling proxy may move an idle, unpinned client to a + // different backend. The lock must hold a transaction for its lifetime. + expect(held).toHaveLength(1); + expect(held[0].state).toBe("idle in transaction"); + expect(held[0].xact_start).not.toBeNull(); await expect(prepareManagedAiRuntime(db, subscription)).rejects.toThrow("in use"); await first.cleanup(); const next = await prepareManagedAiRuntime(db, subscription); diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 766a70d975..79489e7eb4 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -164,15 +164,24 @@ done`, async function acquireCredentialLease(db: Db, grantId: string) { const client = await db.$client.reserve(); try { + // A reserved client pins our connection to PgBouncer, not its backend. + // Keep the lease in one transaction so transaction-pooling deployments + // cannot acquire and release it on different PostgreSQL sessions. + await client`begin`; + await client`set local idle_in_transaction_session_timeout = 0`; const [result] = - await client`select pg_try_advisory_lock(hashtextextended(${`ai-runtime:${grantId}`}, 0)) as acquired`; + await client`select pg_try_advisory_xact_lock(hashtextextended(${`ai-runtime:${grantId}`}, 0)) as acquired`; if (!result.acquired) throw unprocessable( "This subscription is in use. Retry when its current execution finishes.", { code: "ai_connection_busy" }, ); } catch (error) { - client.release(); + try { + await client`rollback`; + } finally { + client.release(); + } throw error; } let released = false; @@ -180,7 +189,7 @@ async function acquireCredentialLease(db: Db, grantId: string) { if (released) return; released = true; try { - await client`select pg_advisory_unlock(hashtextextended(${`ai-runtime:${grantId}`}, 0))`; + await client`rollback`; } finally { client.release(); } diff --git a/ui/src/components/OnboardingWizard.test.tsx b/ui/src/components/OnboardingWizard.test.tsx index 234f2c2f10..1743d4e3f9 100644 --- a/ui/src/components/OnboardingWizard.test.tsx +++ b/ui/src/components/OnboardingWizard.test.tsx @@ -2012,6 +2012,52 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( return { root, queryClient }; } + it.each([ + ["claude_local", "anthropic", /Claude/, "claude-session-1", "claude-setup-token-status"], + ["codex_local", "openai", /OpenAI/, "codex-session-1", "adapter-login-status"], + ] as const)("finishes %s sign-in when its connection becomes visible before the completion poll", async (adapterType, provider, label, sessionId, statusKey) => { + mockAgentsApi.getAdapterAuthSignal.mockResolvedValue({ status: "absent" }); + mockAgentsApi.hire.mockRejectedValueOnce(new Error("Temporary hire failure")); + const { root, queryClient } = await openStep4({ adapterType }); + await pickSource(label); + for (let i = 0; i < 6; i++) await flushReact(); + try { + // The connection activity event arrives before the login poll. It must + // not replace/unmount the controller that still owns the completion. + await act(async () => { + queryClient.setQueryData(["ai-connections", "company-new"], { + currentUserId: "user-1", + connections: [{ id: "managed-connection", grantId: "managed-grant", companyId: "company-new", provider, method: "subscription", name: "My subscription", ownership: "personal", ownerUserId: "user-1", status: "connected", isDefault: true }], + }); + }); + for (let i = 0; i < 4; i++) await flushReact(); + expect(document.body.textContent).toContain(adapterType === "claude_local" ? "authorization code" : "Q2RJ-E1YIF"); + expect(mockAgentsApi.hire).not.toHaveBeenCalled(); + await act(async () => { + queryClient.setQueryData( + adapterType === "claude_local" ? [statusKey, "company-new", sessionId] : [statusKey, "company-new", adapterType, sessionId], + { sessionId, status: "authenticated", expiresAt: new Date(Date.now() + 600_000).toISOString() }, + ); + }); + for (let i = 0; i < 120 && !mockAgentsApi.hire.mock.calls.length; i++) { + await act(async () => { await new Promise(resolve => setTimeout(resolve, 25)); }); + } + expect(mockAgentsApi.hire).toHaveBeenCalledTimes(1); + expect(mockAgentsApi.hire).toHaveBeenCalledWith("company-new", expect.objectContaining({ runtimeConfig: expect.objectContaining({ aiConnection: { provider, method: "subscription", mode: "responsible_user" } }) })); + for (let i = 0; i < 4; i++) await flushReact(); + expect(document.body.textContent).toContain("Temporary hire failure"); + const retry = [...document.body.querySelectorAll("button")].find(button => button.textContent?.trim() === "Connect"); + expect(retry).toBeTruthy(); + expect(retry!.disabled).toBe(false); + await act(async () => { retry!.click(); }); + for (let i = 0; i < 6; i++) await flushReact(); + expect(mockAgentsApi.hire).toHaveBeenCalledTimes(2); + expect(mockAgentsApi.startClaudeSetupTokenLogin.mock.calls.length + mockAgentsApi.startAdapterAuthLogin.mock.calls.length).toBe(1); + } finally { + await act(async () => root.unmount()); + } + }); + it("names the tiles for the provider, not the adapter type", async () => { // `MODEL_SOURCE_NAMES` exists so this row says "Claude" and "OpenAI" — // which provider you are signing in to, the question the step's heading diff --git a/ui/src/components/OnboardingWizard.tsx b/ui/src/components/OnboardingWizard.tsx index 7641d62b83..5daf58cd44 100644 --- a/ui/src/components/OnboardingWizard.tsx +++ b/ui/src/components/OnboardingWizard.tsx @@ -1150,9 +1150,14 @@ function OnboardingWizardInner({ */ const connectStepNeedsLogin = Boolean( credentialMode !== "api" && - (showAdapterLoginPanel || (canShowAdapterLogin && adapterType === "codex_local" && subscriptionId?.companyId === createdCompanyId && subscriptionId.id === "")) && - !savedSubscription && - !(adapterType === "claude_local" && savedKeys.storedLogin.data) && + // Connection-list invalidation can arrive before the login's completion + // poll. Keep its controller mounted until it reports success; otherwise + // the saved account replaces the panel and "Connecting" never finishes. + (connectAuthUrl || ( + (showAdapterLoginPanel || (canShowAdapterLogin && adapterType === "codex_local" && subscriptionId?.companyId === createdCompanyId && subscriptionId.id === "")) && + !savedSubscription && + !(adapterType === "claude_local" && savedKeys.storedLogin.data) + )) && !savedKeys.loading && createdCompanyId && resolvedLoginEnvironmentId, @@ -2154,6 +2159,13 @@ function OnboardingWizardInner({ } finally { hiringAgentRef.current = false; setLoading(false); + // Authentication is already saved. A failed probe or hire must offer a + // retry with that account, rather than keep the completed login busy. + if (connectCredentialStored && stillTheSameCompany(createdCompanyId)) { + connectingSinceRef.current = null; + setConnectAuthUrl(null); + setConnectPhase((phase) => phase === "connecting" ? "ready" : phase); + } } } @@ -2740,6 +2752,7 @@ function OnboardingWizardInner({ }} onConnected={() => { if (managedProvider) managedSubscriptionRef.current = { companyId: createdCompanyId, binding: { provider: managedProvider, method: "subscription", mode: "responsible_user" } }; + setConnectAuthUrl(null); // Not into a card the customer has left. The panel is // still mounted through Back's exit, and a login that // finished there pulled the step back into "Connecting"