🤖 fix: stabilize storybook theme by forcing context state

Author: ammar-agent

.storybook/preview.tsx

@@ -1,24 +1,11 @@
-import React, { useEffect } from "react";
+import React from "react";
 import type { Preview } from "@storybook/react-vite";
 import {
   ThemeProvider,
-  useTheme,
   type ThemeMode,
 } from "../src/browser/contexts/ThemeContext";
 import "../src/browser/styles/globals.css";
 
-const ThemeStorySync: React.FC<{ mode: ThemeMode }> = ({ mode }) => {
-  const { theme, setTheme } = useTheme();
-
-  useEffect(() => {
-    if (theme !== mode) {
-      setTheme(mode);
-    }
-  }, [mode, setTheme, theme]);
-
-  return null;
-};
-
 const preview: Preview = {
   globalTypes: {
     theme: {
@@ -39,8 +26,7 @@ const preview: Preview = {
     (Story, context) => {
       const mode = (context.globals.theme ?? "dark") as ThemeMode;
       return (
-        <ThemeProvider>
-          <ThemeStorySync mode={mode} />
+        <ThemeProvider forcedTheme={mode}>
           <Story />
         </ThemeProvider>
       );

src/browser/contexts/ThemeContext.tsx

@@ -51,18 +51,29 @@ function applyThemeToDocument(theme: ThemeMode) {
   }
 }
 
-export function ThemeProvider(props: { children: ReactNode }) {
-  const [theme, setTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
+export function ThemeProvider(props: { children: ReactNode; forcedTheme?: ThemeMode }) {
+  const [persistedTheme, setPersistedTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
     listener: true,
   });
 
+  const theme = props.forcedTheme ?? persistedTheme;
+
   useLayoutEffect(() => {
     applyThemeToDocument(theme);
   }, [theme]);
 
+  const setTheme = useCallback(
+    (value: React.SetStateAction<ThemeMode>) => {
+      if (props.forcedTheme) return;
+      setPersistedTheme(value);
+    },
+    [props.forcedTheme, setPersistedTheme]
+  );
+
   const toggleTheme = useCallback(() => {
-    setTheme((current) => (current === "dark" ? "light" : "dark"));
-  }, [setTheme]);
+    if (props.forcedTheme) return;
+    setPersistedTheme((current) => (current === "dark" ? "light" : "dark"));
+  }, [props.forcedTheme, setPersistedTheme]);
 
   const value = useMemo<ThemeContextValue>(
     () => ({