feat(ui): complete Registry provider onboarding for Private Cloud (#12494)

Co-authored-by: alejandrobailo <alejandrobailo94@gmail.com>

ui/store/task-watcher/store.test.ts

@@ -23,6 +23,7 @@ describe("task watcher store", () => {
   const onError = vi.fn();
 
   beforeEach(() => {
+    window.dispatchEvent(new PageTransitionEvent("pageshow"));
     vi.clearAllMocks();
     localStorage.clear();
     Object.defineProperty(navigator, "locks", {
@@ -33,6 +34,179 @@ describe("task watcher store", () => {
     registerTaskKindHandler("test-kind", { onReady, onError });
   });
 
+  it.each(["beforeunload", "pagehide"])(
+    "preserves pending work when %s aborts its RPC and resumes on return",
+    async (eventType) => {
+      let rejectPoll!: (error: Error) => void;
+      pollMock.mockImplementationOnce(
+        () =>
+          new Promise((_resolve, reject) => {
+            rejectPoll = reject;
+          }),
+      );
+      const tracking = trackAndPollTask({
+        taskId: "reload-task",
+        kind: "test-kind",
+        meta: {},
+      });
+      await vi.waitFor(() => expect(pollMock).toHaveBeenCalledOnce());
+      window.dispatchEvent(new Event(eventType));
+      rejectPoll(new Error("The document was unloaded"));
+      if (eventType === "beforeunload") {
+        window.dispatchEvent(new PageTransitionEvent("pagehide"));
+      }
+      expect(await tracking).toEqual({ status: TASK_WATCHER_STATUS.PENDING });
+      expect(useTaskWatcherStore.getState().tasks["reload-task"]?.status).toBe(
+        TASK_WATCHER_STATUS.PENDING,
+      );
+      expect(onError).not.toHaveBeenCalled();
+
+      if (eventType === "pagehide") {
+        await flush();
+        expect(pollMock).toHaveBeenCalledOnce();
+        expect(onReady).not.toHaveBeenCalled();
+      }
+
+      pollMock.mockResolvedValue({ ok: true, state: "completed" });
+      window.dispatchEvent(
+        new PageTransitionEvent("pageshow", { persisted: true }),
+      );
+      await vi.waitFor(() => expect(onReady).toHaveBeenCalledOnce());
+    },
+  );
+
+  it("keeps the caller's promise and notifications when beforeunload does not hide the page", async () => {
+    // Given: a download or cancelled navigation leaves this document alive.
+    let rejectPoll!: (error: Error) => void;
+    pollMock
+      .mockImplementationOnce(
+        () =>
+          new Promise((_resolve, reject) => {
+            rejectPoll = reject;
+          }),
+      )
+      .mockResolvedValue({ ok: true, state: "completed" });
+    const tracking = trackAndPollTask({
+      taskId: "download-task",
+      kind: "test-kind",
+      meta: {},
+      notifyHandler: false,
+    });
+    await vi.waitFor(() => expect(pollMock).toHaveBeenCalledOnce());
+
+    // When: only the RPC is interrupted; the caller remains in this document.
+    window.dispatchEvent(new Event("beforeunload"));
+    rejectPoll(new Error("The navigation interrupted the request"));
+    // Then: the original caller receives the final result and owns notification.
+    expect(await tracking).toEqual({ status: TASK_WATCHER_STATUS.READY });
+    expect(onReady).not.toHaveBeenCalled();
+    expect(onError).not.toHaveBeenCalled();
+    expect(useTaskWatcherStore.getState().tasks["download-task"]?.status).toBe(
+      TASK_WATCHER_STATUS.READY,
+    );
+  });
+
+  it("preserves a navigation abort delivered after the visible page has recovered", async () => {
+    // Given
+    let rejectPoll!: (error: Error) => void;
+    pollMock
+      .mockImplementationOnce(
+        () =>
+          new Promise((_resolve, reject) => {
+            rejectPoll = reject;
+          }),
+      )
+      .mockResolvedValue({ ok: true, state: "completed" });
+    const tracking = trackAndPollTask({
+      taskId: "delayed-abort",
+      kind: "test-kind",
+      meta: {},
+    });
+    await vi.waitFor(() => expect(pollMock).toHaveBeenCalledOnce());
+
+    // When: browser RPC cancellation arrives after the unload event's task.
+    window.dispatchEvent(new Event("beforeunload"));
+    await flush();
+    rejectPoll(new Error("The navigation interrupted the request"));
+
+    // Then: the original caller gets the final result, not a provisional failure.
+    expect(await tracking).toEqual({ status: TASK_WATCHER_STATUS.READY });
+    expect(onReady).toHaveBeenCalledOnce();
+    expect(onError).not.toHaveBeenCalled();
+  });
+
+  it("allows new tasks after a download without discarding existing results", async () => {
+    // Given
+    const existingResult = {
+      taskId: "previous-export",
+      kind: "export",
+      status: TASK_WATCHER_STATUS.READY,
+      meta: {},
+      startedAt: Date.now(),
+      result: { downloadUrl: "/download" },
+    };
+    useTaskWatcherStore.getState().upsertTask(existingResult);
+    pollMock.mockResolvedValue({ ok: true, state: "completed" });
+
+    // When: a download starts but leaves the document in place.
+    window.dispatchEvent(new Event("beforeunload"));
+    await flush();
+    const result = await trackAndPollTask({
+      taskId: "after-download",
+      kind: "test-kind",
+      meta: {},
+    });
+
+    // Then
+    expect(result.status).toBe(TASK_WATCHER_STATUS.READY);
+    expect(onReady).toHaveBeenCalledOnce();
+    expect(useTaskWatcherStore.getState().tasks["previous-export"]).toEqual(
+      existingResult,
+    );
+  });
+
+  it("keeps caller ownership when bfcache restores before its RPC rejects", async () => {
+    // Given
+    let rejectPoll!: (error: Error) => void;
+    pollMock
+      .mockImplementationOnce(
+        () =>
+          new Promise((_resolve, reject) => {
+            rejectPoll = reject;
+          }),
+      )
+      .mockResolvedValue({ ok: true, state: "completed" });
+    const tracking = trackAndPollTask({
+      taskId: "bfcache-task",
+      kind: "test-kind",
+      meta: {},
+      notifyHandler: false,
+    });
+    await vi.waitFor(() => expect(pollMock).toHaveBeenCalledOnce());
+
+    // When: the cached document returns before its interrupted RPC rejects.
+    window.dispatchEvent(
+      new PageTransitionEvent("pagehide", { persisted: true }),
+    );
+    window.dispatchEvent(
+      new PageTransitionEvent("pageshow", { persisted: true }),
+    );
+    rejectPoll(new Error("The cached document interrupted the request"));
+
+    // Then
+    expect(await tracking).toEqual({ status: TASK_WATCHER_STATUS.READY });
+    expect(onReady).not.toHaveBeenCalled();
+    window.dispatchEvent(
+      new PageTransitionEvent("pageshow", { persisted: true }),
+    );
+    await flush();
+    expect(onReady).not.toHaveBeenCalled();
+    expect(onError).not.toHaveBeenCalled();
+    expect(useTaskWatcherStore.getState().tasks["bfcache-task"]?.status).toBe(
+      TASK_WATCHER_STATUS.READY,
+    );
+  });
+
   it("tracks a task, polls it to completion and fires onReady once", async () => {
     pollMock.mockResolvedValue({ ok: true, state: "completed" });
 

ui/store/task-watcher/store.ts

@@ -117,12 +117,71 @@ export const useTaskWatcherStore = create<TaskWatcherState>()(
 const activePolls = new Map<string, Promise<TaskTrackingResult<unknown>>>();
 const suppressedHandlers = new Set<string>();
 
+// Navigation aborts outstanding Server Action requests. That is not a backend
+// task failure: keep its persisted identity for the next document to resume.
+let pageSuspended = false;
+let pageHidden = false;
+let navigationGeneration = 0;
+let pageRecoveryTimer: ReturnType<typeof setTimeout> | undefined;
+const pageRecoveryWaiters = new Set<(visible: boolean) => void>();
+
+const resolvePageRecovery = (visible: boolean) => {
+  pageRecoveryWaiters.forEach((resolve) => resolve(visible));
+  pageRecoveryWaiters.clear();
+};
+
+const resumeVisiblePage = () => {
+  clearTimeout(pageRecoveryTimer);
+  // Downloads and cancelled navigation emit beforeunload without pagehide.
+  // Let pagehide confirm navigation before retrying in the surviving document.
+  pageRecoveryTimer = setTimeout(() => {
+    pageRecoveryTimer = undefined;
+    if (pageHidden) return;
+    pageSuspended = false;
+    resolvePageRecovery(true);
+    void resumePendingTaskPolling();
+  }, 0);
+};
+
+const waitForVisiblePage = (): Promise<boolean> => {
+  if (pageHidden) return Promise.resolve(false);
+  return new Promise((resolve) => {
+    pageRecoveryWaiters.add(resolve);
+    resumeVisiblePage();
+  });
+};
+
+if (typeof window !== "undefined") {
+  // Browsers can abort a Server Action before pagehide is dispatched.
+  // Mark the navigation at its start so that abort cannot discard the task.
+  window.addEventListener("beforeunload", () => {
+    navigationGeneration++;
+    pageSuspended = true;
+    resumeVisiblePage();
+  });
+  window.addEventListener("pagehide", () => {
+    navigationGeneration++;
+    clearTimeout(pageRecoveryTimer);
+    pageHidden = true;
+    pageSuspended = true;
+    resolvePageRecovery(false);
+  });
+  window.addEventListener("pageshow", (event) => {
+    clearTimeout(pageRecoveryTimer);
+    pageHidden = false;
+    pageSuspended = false;
+    resolvePageRecovery(true);
+    if (event.persisted) void resumePendingTaskPolling();
+  });
+}
+
 const settleTask = (
   taskId: string,
   status: TaskWatcherStatus,
   error?: string,
   result?: unknown,
 ): TaskTrackingResult => {
+  if (pageSuspended) return { status: TASK_WATCHER_STATUS.PENDING };
   const store = useTaskWatcherStore.getState();
   const currentTask = store.tasks[taskId];
   if (!currentTask || currentTask.status !== TASK_WATCHER_STATUS.PENDING) {
@@ -224,23 +283,37 @@ const pollUntilDone = <R>(taskId: string): Promise<TaskTrackingResult<R>> => {
         return runPollLoop<R>(taskId);
       };
 
-      if (typeof navigator !== "undefined" && navigator.locks) {
-        return await navigator.locks.request(
-          `task-watcher:${taskId}`,
-          runIfPending,
-        );
-      }
+      for (;;) {
+        const pollNavigationGeneration = navigationGeneration;
+        try {
+          const result =
+            typeof navigator !== "undefined" && navigator.locks
+              ? await navigator.locks.request(
+                  `task-watcher:${taskId}`,
+                  runIfPending,
+                )
+              : await runIfPending();
+          if (result.status !== TASK_WATCHER_STATUS.PENDING) return result;
+        } catch {
+          if (
+            !pageSuspended &&
+            navigationGeneration === pollNavigationGeneration
+          ) {
+            return settleTask(
+              taskId,
+              TASK_WATCHER_STATUS.ERROR,
+              "Tracking the task failed unexpectedly. Try again later.",
+            ) as TaskTrackingResult<R>;
+          }
+        }
 
-      return await runIfPending();
-    } catch {
-      // A thrown poll (e.g. the server-action RPC failing on a network drop)
-      // must still settle the task, or it stays PENDING in the persisted
-      // store and blocks the UI until the staleness ceiling.
-      return settleTask(
-        taskId,
-        TASK_WATCHER_STATUS.ERROR,
-        "Tracking the task failed unexpectedly. Try again later.",
-      ) as TaskTrackingResult<R>;
+        // Downloads and cancelled navigation keep the original caller alive.
+        // Retry within its promise so it retains notification ownership and
+        // receives a terminal result. Only a hidden document hands off to resume.
+        if (!(await waitForVisiblePage())) {
+          return { status: TASK_WATCHER_STATUS.PENDING };
+        }
+      }
     } finally {
       activePolls.delete(taskId);
     }
@@ -295,9 +368,6 @@ export const trackAndPollTask = async <R = unknown>({
 export const resumePendingTasks = async (): Promise<void> => {
   const store = useTaskWatcherStore.getState();
   const persistedTasks = Object.values(store.tasks);
-  const pending = persistedTasks.filter(
-    (task) => task.status === TASK_WATCHER_STATUS.PENDING,
-  );
 
   // Settled entries already surfaced in the previous browser session. The
   // server-rendered feature UI resolves durable results again on reload, so
@@ -306,6 +376,13 @@ export const resumePendingTasks = async (): Promise<void> => {
     .filter((task) => task.status !== TASK_WATCHER_STATUS.PENDING)
     .forEach((task) => store.dismissTask(task.taskId));
 
+  await resumePendingTaskPolling();
+};
+
+async function resumePendingTaskPolling(): Promise<void> {
+  const pending = Object.values(useTaskWatcherStore.getState().tasks).filter(
+    (task) => task.status === TASK_WATCHER_STATUS.PENDING,
+  );
   await Promise.all(
     pending.map((task) => {
       if (Date.now() - task.startedAt > STALE_TASK_MS) {
@@ -319,4 +396,4 @@ export const resumePendingTasks = async (): Promise<void> => {
       return pollUntilDone(task.taskId);
     }),
   );
-};
+}

ui/store/ui/store-initializer.test.tsx

@@ -0,0 +1,36 @@
+import { render } from "@testing-library/react";
+import { beforeEach, describe, expect, it } from "vitest";
+
+import { useUIStore } from "./store";
+import { StoreInitializer } from "./store-initializer";
+
+describe("StoreInitializer", () => {
+  beforeEach(() => {
+    localStorage.clear();
+    useUIStore.setState({ hasProviders: false, registryEligible: false });
+  });
+
+  it("keeps Registry hidden when the server sends no eligibility decision", () => {
+    // Given / When
+    render(<StoreInitializer values={{ hasProviders: true }} />);
+
+    // Then
+    expect(useUIStore.getState().registryEligible).toBe(false);
+    expect(useUIStore.getState().hasProviders).toBe(true);
+  });
+
+  it("never persists Registry eligibility across sessions", () => {
+    // Given / When
+    render(
+      <StoreInitializer
+        values={{ hasProviders: true, registryEligible: true }}
+      />,
+    );
+
+    // Then
+    expect(useUIStore.getState().registryEligible).toBe(true);
+    const persisted = JSON.parse(localStorage.getItem("ui-store") ?? "{}");
+    expect(persisted.state?.hasProviders).toBe(true);
+    expect(persisted.state).not.toHaveProperty("registryEligible");
+  });
+});

ui/store/ui/store-initializer.tsx

@@ -7,6 +7,7 @@ import { useUIStore } from "@/store/ui/store";
 interface StoreInitializerProps {
   values: {
     hasProviders?: boolean;
+    registryEligible?: boolean;
     // Add more properties here as needed
     // otherProperty?: string;
   };
@@ -14,14 +15,23 @@ interface StoreInitializerProps {
 
 export function StoreInitializer({ values }: StoreInitializerProps) {
   const setHasProviders = useUIStore((state) => state.setHasProviders);
+  const setRegistryEligible = useUIStore((state) => state.setRegistryEligible);
 
   useEffect(() => {
     // Initialize store values from server
     if (values.hasProviders !== undefined) {
       setHasProviders(values.hasProviders);
     }
+    if (values.registryEligible !== undefined) {
+      setRegistryEligible(values.registryEligible);
+    }
     // Add more setters here as needed in the future
-  }, [values.hasProviders, setHasProviders]);
+  }, [
+    values.hasProviders,
+    values.registryEligible,
+    setHasProviders,
+    setRegistryEligible,
+  ]);
 
   return null;
 }

ui/store/ui/store.ts

@@ -4,10 +4,12 @@ import { persist } from "zustand/middleware";
 interface UIStoreState {
   isSideMenuOpen: boolean;
   hasProviders: boolean;
+  registryEligible: boolean;
 
   openSideMenu: () => void;
   closeSideMenu: () => void;
   setHasProviders: (value: boolean) => void;
+  setRegistryEligible: (value: boolean) => void;
 }
 
 export const useUIStore = create<UIStoreState>()(
@@ -15,12 +17,21 @@ export const useUIStore = create<UIStoreState>()(
     (set) => ({
       isSideMenuOpen: false,
       hasProviders: false,
+      registryEligible: false,
       openSideMenu: () => set({ isSideMenuOpen: true }),
       closeSideMenu: () => set({ isSideMenuOpen: false }),
       setHasProviders: (value: boolean) => set({ hasProviders: value }),
+      setRegistryEligible: (value: boolean) => set({ registryEligible: value }),
     }),
     {
       name: "ui-store",
+      // Registry eligibility is a per-request server decision; persisting it
+      // would resurface a stale entry on the next session before the server
+      // seed corrects it.
+      partialize: ({ isSideMenuOpen, hasProviders }) => ({
+        isSideMenuOpen,
+        hasProviders,
+      }),
     },
   ),
 );