diff --git a/ui/src/features/connections/ConnectionSetupFlow.tsx b/ui/src/features/connections/ConnectionSetupFlow.tsx index 98432b1972..dab0571889 100644 --- a/ui/src/features/connections/ConnectionSetupFlow.tsx +++ b/ui/src/features/connections/ConnectionSetupFlow.tsx @@ -97,18 +97,20 @@ type Step = "gallery" | "access" | "key" | "success"; export type OAuthConnectPhase = "entry" | "starting" | "redirecting" | "error"; type EnrollmentAccessState = { + companyId: string; grantKind: ConnectionGrantKind; installChoice: "specific" | "all"; agentIds: string[]; }; -function enrollmentAccessStorageKey(companyId: string, appKey: string): string { - return `paperclip.connector-enrollment-access:${companyId}:${appKey}`; +function enrollmentAccessStorageKey(appKey: string): string { + return `paperclip.connector-enrollment-access:${appKey}`; } function validEnrollmentAccessState(value: unknown): value is EnrollmentAccessState { if (!value || typeof value !== "object" || Array.isArray(value)) return false; const candidate = value as Record; + if (typeof candidate.companyId !== "string" || !candidate.companyId.trim()) return false; if (!(candidate.grantKind === "user" || candidate.grantKind === "agent" || candidate.grantKind === "organization")) { return false; } @@ -124,17 +126,21 @@ function validEnrollmentAccessState(value: unknown): value is EnrollmentAccessSt return candidate.installChoice === "all" ? agentIds.size === 0 : agentIds.size > 0; } -function saveEnrollmentAccessState(companyId: string, appKey: string, state: EnrollmentAccessState): void { +function saveEnrollmentAccessState( + companyId: string, + appKey: string, + state: Omit, +): void { try { - window.sessionStorage.setItem(enrollmentAccessStorageKey(companyId, appKey), JSON.stringify(state)); + window.sessionStorage.setItem(enrollmentAccessStorageKey(appKey), JSON.stringify({ ...state, companyId })); } 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); +function consumeEnrollmentAccessState(appKey: string): EnrollmentAccessState | null { + const key = enrollmentAccessStorageKey(appKey); try { const raw = window.sessionStorage.getItem(key); window.sessionStorage.removeItem(key); @@ -542,9 +548,8 @@ export function ConnectionSetupFlow({ const [restoredEnrollmentAccess] = useState(() => host === "page" && searchParams.get("cloud_connector") === "enrolled" - && selectedCompanyId && requestedAppKey - ? consumeEnrollmentAccessState(selectedCompanyId, requestedAppKey) + ? consumeEnrollmentAccessState(requestedAppKey) : null, ); @@ -1275,11 +1280,14 @@ export function ConnectionSetupFlow({ setGoogleSheetsLinks(""); setGoogleSheetsError(null); setConnectResult(null); - setGrantKind(reconnectGrantKind ?? restoredEnrollmentAccess?.grantKind ?? defaultGrantKindFor(initialMethod)); + const matchingEnrollmentAccess = restoredEnrollmentAccess?.companyId === selectedCompanyId + ? restoredEnrollmentAccess + : null; + setGrantKind(reconnectGrantKind ?? matchingEnrollmentAccess?.grantKind ?? defaultGrantKindFor(initialMethod)); setInstallAgentIds(new Set( - restoredEnrollmentAccess?.agentIds ?? (requestedAgentId ? [requestedAgentId] : []), + matchingEnrollmentAccess?.agentIds ?? (requestedAgentId ? [requestedAgentId] : []), )); - setInstallChoice(restoredEnrollmentAccess?.installChoice ?? (requestedAgentId ? "specific" : "all")); + setInstallChoice(matchingEnrollmentAccess?.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. diff --git a/ui/src/pages/apps/AppsConnect.test.tsx b/ui/src/pages/apps/AppsConnect.test.tsx index b6e48466ce..a81915e19a 100644 --- a/ui/src/pages/apps/AppsConnect.test.tsx +++ b/ui/src/pages/apps/AppsConnect.test.tsx @@ -24,6 +24,12 @@ const mockNavigate = vi.hoisted(() => vi.fn()); const navigateTopLevelMock = vi.hoisted(() => vi.fn()); const mockSearch = vi.hoisted(() => ({ value: "" })); const mockParams = vi.hoisted(() => ({ appKey: undefined as string | undefined })); +const mockCompany = vi.hoisted(() => ({ + value: { + selectedCompanyId: "company-1" as string | undefined, + selectedCompany: { id: "company-1", name: "Paperclip" } as { id: string; name: string } | null, + }, +})); const ZAPIER = CONNECTABLE_APP_DEFINITIONS.find((app) => app.slug === "zapier")!; const GITHUB = CONNECTABLE_APP_DEFINITIONS.find((app) => app.slug === "github")!; @@ -78,10 +84,7 @@ vi.mock("@/lib/router", () => ({ })); vi.mock("@/context/CompanyContext", () => ({ - useCompany: () => ({ - selectedCompanyId: "company-1", - selectedCompany: { id: "company-1", name: "Paperclip" }, - }), + useCompany: () => mockCompany.value, })); vi.mock("@/context/BreadcrumbContext", () => ({ @@ -200,6 +203,10 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { beforeEach(() => { vi.resetAllMocks(); window.sessionStorage.clear(); + mockCompany.value = { + selectedCompanyId: "company-1", + selectedCompany: { id: "company-1", name: "Paperclip" }, + }; mockSearch.value = ""; mockParams.appKey = undefined; container = document.createElement("div"); @@ -819,8 +826,9 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { await flushReact(); expect(JSON.parse(window.sessionStorage.getItem( - "paperclip.connector-enrollment-access:company-1:github", + "paperclip.connector-enrollment-access:github", ) ?? "null")).toEqual({ + companyId: "company-1", grantKind: "agent", installChoice: "specific", agentIds: ["agent-1"], @@ -831,6 +839,7 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { container.innerHTML = ""; mockParams.appKey = undefined; mockSearch.value = "source=github&stage=setup&cloud_connector=enrolled"; + mockCompany.value = { selectedCompanyId: undefined, selectedCompany: null }; listGalleryMock.mockResolvedValue({ apps: [GITHUB_MANAGED] }); getCloudConnectorEnrollmentMock.mockResolvedValue({ configured: true, @@ -841,11 +850,27 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { origins: ["https://paperclip.example.test"], }); - await render(); + const coldLoadClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + await render(coldLoadClient); + + expect(window.sessionStorage.getItem("paperclip.connector-enrollment-access:github")).toBeNull(); + mockCompany.value = { + selectedCompanyId: "company-1", + selectedCompany: { id: "company-1", name: "Paperclip" }, + }; + await act(async () => { + mountedRoot?.render( + + + , + ); + }); + await flushReact(); + await flushReact(); expect(container.textContent).toContain("Step 2 of 2"); expect(window.sessionStorage.getItem( - "paperclip.connector-enrollment-access:company-1:github", + "paperclip.connector-enrollment-access:github", )).toBeNull(); await act(async () => { buttonByText("Continue to GitHub")?.dispatchEvent(new MouseEvent("click", { bubbles: true }));