From 8f81d86d861554cc9ddc84a781d4658ff1a3322a Mon Sep 17 00:00:00 2001 From: Dotta Date: Sat, 5 Sep 2026 08:19:38 -0500 Subject: [PATCH] fix(connections): preserve access through enrollment --- .../connections/ConnectionSetupFlow.tsx | 82 +++++++++++++++-- ui/src/pages/apps/AppsConnect.test.tsx | 91 +++++++++++++++++++ 2 files changed, 167 insertions(+), 6 deletions(-) diff --git a/ui/src/features/connections/ConnectionSetupFlow.tsx b/ui/src/features/connections/ConnectionSetupFlow.tsx index 2b2471b36d..98432b1972 100644 --- a/ui/src/features/connections/ConnectionSetupFlow.tsx +++ b/ui/src/features/connections/ConnectionSetupFlow.tsx @@ -96,6 +96,56 @@ import { autoExtendNotice, INSTALL_ALL_WARNING, installInfoNotice, installPayloa type Step = "gallery" | "access" | "key" | "success"; export type OAuthConnectPhase = "entry" | "starting" | "redirecting" | "error"; +type EnrollmentAccessState = { + grantKind: ConnectionGrantKind; + installChoice: "specific" | "all"; + agentIds: string[]; +}; + +function enrollmentAccessStorageKey(companyId: string, appKey: string): string { + return `paperclip.connector-enrollment-access:${companyId}:${appKey}`; +} + +function validEnrollmentAccessState(value: unknown): value is EnrollmentAccessState { + if (!value || typeof value !== "object" || Array.isArray(value)) return false; + const candidate = value as Record; + if (!(candidate.grantKind === "user" || candidate.grantKind === "agent" || candidate.grantKind === "organization")) { + return false; + } + if (candidate.installChoice !== "specific" && candidate.installChoice !== "all") return false; + if (!Array.isArray(candidate.agentIds) || candidate.agentIds.some((id) => typeof id !== "string" || !id.trim())) { + return false; + } + const agentIds = new Set(candidate.agentIds); + if (agentIds.size !== candidate.agentIds.length) return false; + if (candidate.grantKind === "agent") { + return candidate.installChoice === "specific" && agentIds.size === 1; + } + return candidate.installChoice === "all" ? agentIds.size === 0 : agentIds.size > 0; +} + +function saveEnrollmentAccessState(companyId: string, appKey: string, state: EnrollmentAccessState): void { + try { + window.sessionStorage.setItem(enrollmentAccessStorageKey(companyId, appKey), JSON.stringify(state)); + } catch { + // Browser storage can be unavailable under restrictive privacy settings. + // The callback will safely use the provider's defaults in that case. + } +} + +function consumeEnrollmentAccessState(companyId: string, appKey: string): EnrollmentAccessState | null { + const key = enrollmentAccessStorageKey(companyId, appKey); + try { + const raw = window.sessionStorage.getItem(key); + window.sessionStorage.removeItem(key); + if (!raw) return null; + const parsed: unknown = JSON.parse(raw); + return validEnrollmentAccessState(parsed) ? parsed : null; + } catch { + return null; + } +} + function githubRecoveryUrl(value: string | null): string | null { if (!value) return null; try { @@ -489,6 +539,14 @@ export function ConnectionSetupFlow({ const zapierSource = (serviceSlug ?? sourceSlug ?? appKey) === "zapier"; const requestedAppKey = zapierSource ? undefined : routeAppKey; const byo = host === "page" && (byoOnly || searchParams.get("byo") === "1"); + const [restoredEnrollmentAccess] = useState(() => + host === "page" + && searchParams.get("cloud_connector") === "enrolled" + && selectedCompanyId + && requestedAppKey + ? consumeEnrollmentAccessState(selectedCompanyId, requestedAppKey) + : null, + ); // Prefill arrives from the app page for reconnects; read once so later // wizard navigation doesn't fight the URL. @@ -536,7 +594,7 @@ export function ConnectionSetupFlow({ const [access, setAccess] = useState<"all" | "specific">("all"); const [agentIds, setAgentIds] = useState>(new Set()); const [installAgentIds, setInstallAgentIds] = useState>( - () => new Set(requestedAgentId ? [requestedAgentId] : []), + () => new Set(restoredEnrollmentAccess?.agentIds ?? (requestedAgentId ? [requestedAgentId] : [])), ); /** * Access-step selections (PAP-17835). These are chosen before the credential @@ -544,10 +602,10 @@ export function ConnectionSetupFlow({ * backwards through the wizard. */ const [grantKind, setGrantKind] = useState( - reconnectGrantKindHint ?? "organization", + restoredEnrollmentAccess?.grantKind ?? reconnectGrantKindHint ?? "organization", ); const [installChoice, setInstallChoice] = useState<"specific" | "all">( - requestedAgentId ? "specific" : "all", + restoredEnrollmentAccess?.installChoice ?? (requestedAgentId ? "specific" : "all"), ); const resumingAfterOAuthFailure = Boolean( resumeConnectionId @@ -776,6 +834,14 @@ export function ConnectionSetupFlow({ ), }); const [connectorEnrollmentError, setConnectorEnrollmentError] = useState(null); + const preserveEnrollmentAccess = useCallback(() => { + if (!selectedCompanyId || !requestedAppKey) return; + saveEnrollmentAccessState(selectedCompanyId, requestedAppKey, { + grantKind, + installChoice, + agentIds: installChoice === "specific" ? [...installAgentIds] : [], + }); + }, [grantKind, installAgentIds, installChoice, requestedAppKey, selectedCompanyId]); const openConnectorEnrollment = useCallback((verificationUrl: string) => { const target = resolveAuthorizationTarget(verificationUrl); if (!target.ok) { @@ -1209,9 +1275,11 @@ export function ConnectionSetupFlow({ setGoogleSheetsLinks(""); setGoogleSheetsError(null); setConnectResult(null); - setGrantKind(reconnectGrantKind ?? defaultGrantKindFor(initialMethod)); - setInstallAgentIds(new Set(requestedAgentId ? [requestedAgentId] : [])); - setInstallChoice(requestedAgentId ? "specific" : "all"); + setGrantKind(reconnectGrantKind ?? restoredEnrollmentAccess?.grantKind ?? defaultGrantKindFor(initialMethod)); + setInstallAgentIds(new Set( + restoredEnrollmentAccess?.agentIds ?? (requestedAgentId ? [requestedAgentId] : []), + )); + setInstallChoice(restoredEnrollmentAccess?.installChoice ?? (requestedAgentId ? "specific" : "all")); // Route/service selection initializes the wizard once. Later renders must // preserve the user's current step in both hosts instead of snapping back // to Access after they continue. @@ -1254,6 +1322,7 @@ export function ConnectionSetupFlow({ fullRequestedDefinition, requestedAppKey, requestedAgentId, + restoredEnrollmentAccess, routeStage, zapierSource, ]); @@ -1868,6 +1937,7 @@ export function ConnectionSetupFlow({ disabled={connectorEnrollmentQuery.isLoading || startConnectorEnrollment.isPending} onClick={() => { setConnectorEnrollmentError(null); + preserveEnrollmentAccess(); const verificationUrl = connectorEnrollmentQuery.data?.verificationUrl; if (verificationUrl) openConnectorEnrollment(verificationUrl); else startConnectorEnrollment.mutate(); diff --git a/ui/src/pages/apps/AppsConnect.test.tsx b/ui/src/pages/apps/AppsConnect.test.tsx index 762b22099f..b6e48466ce 100644 --- a/ui/src/pages/apps/AppsConnect.test.tsx +++ b/ui/src/pages/apps/AppsConnect.test.tsx @@ -199,6 +199,7 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { beforeEach(() => { vi.resetAllMocks(); + window.sessionStorage.clear(); mockSearch.value = ""; mockParams.appKey = undefined; container = document.createElement("div"); @@ -768,6 +769,96 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { expect(container.textContent).not.toContain("Connect with Paperclip"); }); + it("preserves a dedicated agent identity across the full-page enrollment callback", async () => { + mockParams.appKey = "github"; + listGalleryMock.mockResolvedValue({ + apps: [{ + ...GITHUB, + methods: GITHUB.methods.filter((method) => !method.oauthStrategy), + ownershipAvailability: { platform_shared: false, customer: true, dcr: true }, + }], + capabilities: { + canCreateOrganizationGrant: true, + organizationGrantReason: null, + canSetCompanyInstall: true, + companyInstallReason: null, + }, + }); + getCloudConnectorEnrollmentMock.mockResolvedValue({ + configured: false, + status: "not_configured", + brokerBaseUrl: "https://my-staging.paperclip.app", + instanceId: null, + environment: "staging", + origins: [], + }); + + await render(); + await act(async () => { + radioContaining("A dedicated account for an agent")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + await act(async () => { + buttonByText("Select agents")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + await act(async () => { + document.body.querySelector('[aria-label="Allow Ada"]') + ?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + const accessContinue = buttonByText("Save and continue") ?? buttonByText("Continue to GitHub"); + expect(accessContinue?.disabled).toBe(false); + await act(async () => { + accessContinue?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + await act(async () => { + buttonByText("Connect with Paperclip")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + + expect(JSON.parse(window.sessionStorage.getItem( + "paperclip.connector-enrollment-access:company-1:github", + ) ?? "null")).toEqual({ + grantKind: "agent", + installChoice: "specific", + agentIds: ["agent-1"], + }); + + await act(async () => mountedRoot?.unmount()); + mountedRoot = null; + container.innerHTML = ""; + mockParams.appKey = undefined; + mockSearch.value = "source=github&stage=setup&cloud_connector=enrolled"; + listGalleryMock.mockResolvedValue({ apps: [GITHUB_MANAGED] }); + getCloudConnectorEnrollmentMock.mockResolvedValue({ + configured: true, + status: "active", + brokerBaseUrl: "https://my-staging.paperclip.app", + instanceId: "inst-test", + environment: "staging", + origins: ["https://paperclip.example.test"], + }); + + await render(); + + expect(container.textContent).toContain("Step 2 of 2"); + expect(window.sessionStorage.getItem( + "paperclip.connector-enrollment-access:company-1:github", + )).toBeNull(); + await act(async () => { + buttonByText("Continue to GitHub")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + + expect(connectAppMock).toHaveBeenCalledWith("company-1", expect.objectContaining({ + galleryKey: "github", + grantKind: "agent", + subjectAgentId: "agent-1", + })); + }); + it("never renders self-host enrollment when the connector identity is already active", async () => { mockSearch.value = "source=gmail&stage=setup"; listGalleryMock.mockResolvedValueOnce({