fix: defer onboarding-time restarts until setup is complete

apps/desktop/src/onboarding/final.tsx

@@ -1,11 +1,14 @@
 import { Icon } from "@iconify-icon/react";
+import { message } from "@tauri-apps/plugin-dialog";
 
 import { commands as analyticsCommands } from "@hypr/plugin-analytics";
 import { commands as openerCommands } from "@hypr/plugin-opener2";
 import { commands as sfxCommands } from "@hypr/plugin-sfx";
 
 import { OnboardingButton } from "./shared";
 
+import { consumePendingSystemAudioStatusChangedMessage } from "~/shared/hooks/usePermissions";
+import { flushAutomaticRelaunch } from "~/store/tinybase/store/save";
 import { commands } from "~/types/tauri.gen";
 
 const SOCIALS = [
@@ -61,5 +64,14 @@ export async function finishOnboarding(onContinue?: () => void) {
   await commands.setOnboardingNeeded(false).catch(console.error);
   await new Promise((resolve) => setTimeout(resolve, 100));
   await analyticsCommands.event({ event: "onboarding_completed" });
+  if (consumePendingSystemAudioStatusChangedMessage()) {
+    await message("The app will now restart to apply the changes", {
+      kind: "info",
+      title: "System Audio Status Changed",
+    });
+  }
+  if (await flushAutomaticRelaunch()) {
+    return;
+  }
   onContinue?.();
 }

apps/desktop/src/onboarding/folder-location.tsx

@@ -1,14 +1,28 @@
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
 import { homeDir } from "@tauri-apps/api/path";
-import { open as selectFolder } from "@tauri-apps/plugin-dialog";
+import { message, open as selectFolder } from "@tauri-apps/plugin-dialog";
 import { FolderIcon } from "lucide-react";
 
 import { commands as openerCommands } from "@hypr/plugin-opener2";
 import { commands as settingsCommands } from "@hypr/plugin-settings";
 
 import { ObsidianVaultList } from "~/settings/general/storage/obsidian-vault-list";
 import { displayPath } from "~/settings/general/storage/path-utils";
-import { relaunch } from "~/store/tinybase/store/save";
+import { scheduleAutomaticRelaunch } from "~/store/tinybase/store/save";
+
+async function handleStorageUpdate() {
+  const restartStatus = await scheduleAutomaticRelaunch();
+
+  if (restartStatus === "deferred") {
+    void message(
+      "The app will restart after onboarding to apply your storage changes",
+      {
+        kind: "info",
+        title: "Storage Updated",
+      },
+    );
+  }
+}
 
 export function FolderLocationSection({
   onContinue,
@@ -52,7 +66,7 @@ export function FolderLocationSection({
     },
     onSuccess: async () => {
       queryClient.invalidateQueries({ queryKey: ["vault-base-path"] });
-      await relaunch();
+      await handleStorageUpdate();
     },
   });
 
@@ -65,7 +79,7 @@ export function FolderLocationSection({
     },
     onSuccess: async () => {
       queryClient.invalidateQueries({ queryKey: ["vault-base-path"] });
-      await relaunch();
+      await handleStorageUpdate();
     },
   });
 

apps/desktop/src/settings/general/storage/use-storage-wizard.ts

@@ -4,7 +4,7 @@ import { useEffect, useState } from "react";
 
 import { commands as settingsCommands } from "@hypr/plugin-settings";
 
-import { relaunch } from "~/store/tinybase/store/save";
+import { scheduleAutomaticRelaunch } from "~/store/tinybase/store/save";
 
 export function useChangeContentPathWizard({
   open,
@@ -51,7 +51,7 @@ export function useChangeContentPathWizard({
     },
     onSuccess: async () => {
       onSuccess();
-      await relaunch();
+      await scheduleAutomaticRelaunch();
     },
   });
 

apps/desktop/src/shared/hooks/usePermissions.ts

@@ -1,15 +1,40 @@
 import { useMutation, useQuery } from "@tanstack/react-query";
 import { message } from "@tauri-apps/plugin-dialog";
+import { useState } from "react";
 
 import {
   type Permission,
   commands as permissionsCommands,
   type PermissionStatus,
 } from "@hypr/plugin-permissions";
 
-import { relaunch } from "~/store/tinybase/store/save";
+import { scheduleAutomaticRelaunch } from "~/store/tinybase/store/save";
+
+let pendingSystemAudioStatusChangedMessage = false;
+
+export function consumePendingSystemAudioStatusChangedMessage() {
+  const pending = pendingSystemAudioStatusChangedMessage;
+  pendingSystemAudioStatusChangedMessage = false;
+  return pending;
+}
+
+async function handleSystemAudioPermissionSuccess() {
+  const restartStatus = await scheduleAutomaticRelaunch(2000);
+
+  if (restartStatus === "deferred") {
+    pendingSystemAudioStatusChangedMessage = true;
+    return;
+  }
+
+  void message("The app will now restart to apply the changes", {
+    kind: "info",
+    title: "System Audio Status Changed",
+  });
+}
 
 export function usePermission(type: Permission) {
+  const [optimisticStatus, setOptimisticStatus] =
+    useState<PermissionStatus | null>(null);
   const status = useQuery({
     queryKey: [`${type}Permission`],
     queryFn: () => permissionsCommands.checkPermission(type),
@@ -24,22 +49,22 @@ export function usePermission(type: Permission) {
 
   const requestMutation = useMutation({
     mutationFn: () => permissionsCommands.requestPermission(type),
-    onSuccess: () => {
+    onSuccess: async () => {
       if (type === "systemAudio") {
-        void message("The app will now restart to apply the changes", {
-          kind: "info",
-          title: "System Audio Status Changed",
-        });
-        setTimeout(() => relaunch(), 2000);
+        setOptimisticStatus("authorized");
+        setTimeout(() => void status.refetch(), 1000);
+        await handleSystemAudioPermissionSuccess();
         return;
       }
+      setOptimisticStatus(null);
       setTimeout(() => status.refetch(), 1000);
     },
   });
 
   const resetMutation = useMutation({
     mutationFn: () => permissionsCommands.resetPermission(type),
     onSuccess: () => {
+      setOptimisticStatus(null);
       setTimeout(() => status.refetch(), 1000);
     },
   });
@@ -58,7 +83,13 @@ export function usePermission(type: Permission) {
     resetMutation.mutate();
   };
 
-  return { status: status.data, isPending, open, request, reset };
+  return {
+    status: optimisticStatus ?? status.data,
+    isPending,
+    open,
+    request,
+    reset,
+  };
 }
 
 export function usePermissions() {
@@ -112,12 +143,8 @@ export function usePermissions() {
 
   const systemAudioPermission = useMutation({
     mutationFn: () => permissionsCommands.requestPermission("systemAudio"),
-    onSuccess: () => {
-      void message("The app will now restart to apply the changes", {
-        kind: "info",
-        title: "System Audio Status Changed",
-      });
-      setTimeout(() => relaunch(), 2000);
+    onSuccess: async () => {
+      await handleSystemAudioPermissionSuccess();
     },
     onError: console.error,
   });

apps/desktop/src/store/tinybase/store/save.test.ts

@@ -0,0 +1,79 @@
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+
+const saveMock = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
+const relaunchMock = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
+const getOnboardingNeededMock = vi
+  .fn<() => Promise<{ status: "ok"; data: boolean }>>()
+  .mockResolvedValue({ status: "ok", data: false });
+
+vi.mock("@hypr/plugin-store2", () => ({
+  commands: {
+    save: saveMock,
+  },
+}));
+
+vi.mock("@tauri-apps/plugin-process", () => ({
+  relaunch: relaunchMock,
+}));
+
+vi.mock("~/types/tauri.gen", () => ({
+  commands: {
+    getOnboardingNeeded: getOnboardingNeededMock,
+  },
+}));
+
+describe("automatic relaunch", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+    vi.resetModules();
+    vi.clearAllMocks();
+    getOnboardingNeededMock.mockResolvedValue({ status: "ok", data: false });
+  });
+
+  afterEach(() => {
+    vi.runOnlyPendingTimers();
+    vi.useRealTimers();
+  });
+
+  test("schedules an immediate relaunch when onboarding is already done", async () => {
+    const { scheduleAutomaticRelaunch } = await import("./save");
+
+    await expect(scheduleAutomaticRelaunch(2000)).resolves.toBe("scheduled");
+
+    expect(saveMock).not.toHaveBeenCalled();
+    expect(relaunchMock).not.toHaveBeenCalled();
+
+    await vi.advanceTimersByTimeAsync(2000);
+
+    expect(saveMock).toHaveBeenCalledTimes(1);
+    expect(relaunchMock).toHaveBeenCalledTimes(1);
+  });
+
+  test("defers relaunch while onboarding is still required", async () => {
+    getOnboardingNeededMock.mockResolvedValue({ status: "ok", data: true });
+    const { scheduleAutomaticRelaunch } = await import("./save");
+
+    await expect(scheduleAutomaticRelaunch()).resolves.toBe("deferred");
+
+    expect(saveMock).not.toHaveBeenCalled();
+    expect(relaunchMock).not.toHaveBeenCalled();
+  });
+
+  test("flushes a deferred relaunch after onboarding completes", async () => {
+    getOnboardingNeededMock.mockResolvedValueOnce({
+      status: "ok",
+      data: true,
+    });
+    const { flushAutomaticRelaunch, scheduleAutomaticRelaunch } =
+      await import("./save");
+
+    await scheduleAutomaticRelaunch();
+
+    getOnboardingNeededMock.mockResolvedValue({ status: "ok", data: false });
+
+    await expect(flushAutomaticRelaunch()).resolves.toBe(true);
+
+    expect(saveMock).toHaveBeenCalledTimes(1);
+    expect(relaunchMock).toHaveBeenCalledTimes(1);
+  });
+});

apps/desktop/src/store/tinybase/store/save.ts

@@ -2,7 +2,11 @@ import { relaunch as tauriRelaunch } from "@tauri-apps/plugin-process";
 
 import { commands as store2Commands } from "@hypr/plugin-store2";
 
+import { commands } from "~/types/tauri.gen";
+
 const saveHandlers = new Map<string, () => Promise<void>>();
+let pendingAutomaticRelaunch = false;
+let automaticRelaunchTimeout: ReturnType<typeof setTimeout> | null = null;
 
 export function registerSaveHandler(id: string, handler: () => Promise<void>) {
   saveHandlers.set(id, handler);
@@ -22,3 +26,52 @@ export async function relaunch(): Promise<void> {
   await save();
   await tauriRelaunch();
 }
+
+async function getOnboardingNeeded() {
+  const result = await commands.getOnboardingNeeded().catch(() => null);
+  if (result?.status !== "ok") {
+    return false;
+  }
+  return result.data;
+}
+
+export async function scheduleAutomaticRelaunch(
+  delayMs = 0,
+): Promise<"scheduled" | "deferred"> {
+  if (await getOnboardingNeeded()) {
+    pendingAutomaticRelaunch = true;
+    return "deferred";
+  }
+
+  if (automaticRelaunchTimeout) {
+    return "scheduled";
+  }
+
+  automaticRelaunchTimeout = setTimeout(() => {
+    automaticRelaunchTimeout = null;
+    void relaunch().catch(console.error);
+  }, delayMs);
+
+  return "scheduled";
+}
+
+export async function flushAutomaticRelaunch(): Promise<boolean> {
+  if (!pendingAutomaticRelaunch || (await getOnboardingNeeded())) {
+    return false;
+  }
+
+  pendingAutomaticRelaunch = false;
+
+  if (automaticRelaunchTimeout) {
+    clearTimeout(automaticRelaunchTimeout);
+    automaticRelaunchTimeout = null;
+  }
+
+  try {
+    await relaunch();
+    return true;
+  } catch (error) {
+    pendingAutomaticRelaunch = true;
+    throw error;
+  }
+}