🤖 refactor: simplify Storybook theme injection and remove race conditions

Author: ammar-agent

.storybook/preview.tsx

@@ -1,11 +1,17 @@
 import React from "react";
 import type { Preview } from "@storybook/react-vite";
-import { ThemeProvider, type ThemeMode } from "../src/browser/contexts/ThemeContext";
-import { getStorageChangeEvent } from "@/common/constants/events";
+import {
+  ThemeProvider,
+  type ThemeMode,
+  useTheme,
+} from "../src/browser/contexts/ThemeContext";
 import isChromatic from "chromatic/isChromatic";
 import { UI_THEME_KEY } from "@/common/constants/storage";
+import { updatePersistedState } from "@/browser/hooks/usePersistedState";
 import "../src/browser/styles/globals.css";
 
+const DEFAULT_THEME: ThemeMode = "light";
+
 const isAutomatedChromaticRun = () => {
   if (typeof window === "undefined" || !isChromatic()) {
     return false;
@@ -30,65 +36,45 @@ const isE2ETestEnv = () => {
   return false;
 };
 
-const applyStorybookTheme = (mode: ThemeMode) => {
+const getPersistedTheme = (): ThemeMode | undefined => {
   if (typeof window === "undefined") {
-    return;
+    return undefined;
   }
 
   try {
-    const serialized = JSON.stringify(mode);
-    const prev = window.localStorage.getItem(UI_THEME_KEY);
-    const root = document.documentElement;
-    const current = root?.dataset?.theme as ThemeMode | undefined;
-    if (prev === serialized && current === mode) {
-      return; // no-op if already applied
-    }
-
-    window.localStorage.setItem(UI_THEME_KEY, serialized);
-    root.dataset.theme = mode;
-    root.style.colorScheme = mode;
-
-    const event = new CustomEvent(getStorageChangeEvent(UI_THEME_KEY), {
-      detail: { key: UI_THEME_KEY, newValue: mode, origin: "storybook" },
-    });
-    window.dispatchEvent(event);
+    const stored = window.localStorage.getItem(UI_THEME_KEY);
+    return stored ? (JSON.parse(stored) as ThemeMode) : undefined;
   } catch (error) {
-    console.warn("Failed to write Storybook theme:", error);
+    console.warn("Failed to read Storybook theme:", error);
+    return undefined;
   }
 };
 
 const syncStorybookTheme = (mode?: ThemeMode): ThemeMode => {
   if (typeof window === "undefined") {
-    return mode ?? "light"; // Keep default aligned with baseline to avoid Chromatic diffs
+    return mode ?? DEFAULT_THEME; // Keep default aligned with baseline to avoid Chromatic diffs
   }
 
-  const stored = window.localStorage.getItem(UI_THEME_KEY);
-  const persisted = stored ? (JSON.parse(stored) as ThemeMode) : undefined;
-  const resolved = mode ?? persisted ?? "light"; // Default to light for Latest when unset
-  applyStorybookTheme(resolved);
-  return resolved;
-};
+  const persisted = getPersistedTheme();
+  const resolved = mode ?? persisted ?? DEFAULT_THEME; // Default to light for Latest when unset
 
-const StorybookThemeToggle: React.FC<{ initialTheme: ThemeMode }> = ({ initialTheme }) => {
-  const [theme, setTheme] = React.useState(initialTheme);
+  if (resolved !== persisted) {
+    updatePersistedState(UI_THEME_KEY, resolved);
+  }
 
-  React.useEffect(() => {
-    setTheme(initialTheme);
-  }, [initialTheme]);
+  return resolved;
+};
 
-  const handleToggle = () => {
-    const next = theme === "dark" ? "light" : "dark";
-    applyStorybookTheme(next);
-    setTheme(next);
-  };
+const StorybookThemeToggle: React.FC = () => {
+  const { theme, toggleTheme } = useTheme();
 
   if (typeof window === "undefined") {
     return null;
   }
 
   return (
     <button
-      onClick={handleToggle}
+      onClick={toggleTheme}
       style={{
         position: "fixed",
         bottom: "1rem",
@@ -128,13 +114,11 @@ const preview: Preview = {
   decorators: [
     (Story, context) => {
       const mode = context.globals.theme as ThemeMode | undefined;
-      const resolved = syncStorybookTheme(mode);
+      syncStorybookTheme(mode);
       return (
         <ThemeProvider>
           <Story />
-          {!mode && !isE2ETestEnv() && (
-            <StorybookThemeToggle initialTheme={resolved} />
-          )}
+          {!mode && !isE2ETestEnv() && <StorybookThemeToggle />}
         </ThemeProvider>
       );
     },