fix(connections): restore enrollment state before hydration

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<string, unknown>;
+  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<EnrollmentAccessState, "companyId">,
+): 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<EnrollmentAccessState | null>(() =>
     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.

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(
+        <QueryClientProvider client={coldLoadClient}>
+          <AppsConnect />
+        </QueryClientProvider>,
+      );
+    });
+    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 }));