original but with a preview-head.html

Author: ethanndickson

.storybook/preview.tsx

@@ -1,13 +1,30 @@
-import React from "react";
+import React, { useLayoutEffect } from "react";
 import type { Preview } from "@storybook/react-vite";
-import { ThemeProvider, type ThemeMode } from "../src/browser/contexts/ThemeContext";
+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();
+
+  useLayoutEffect(() => {
+    if (theme !== mode) {
+      setTheme(mode);
+    }
+  }, [mode, setTheme, theme]);
+
+  return null;
+};
+
 const preview: Preview = {
   globalTypes: {
     theme: {
       name: "Theme",
       description: "Choose between light and dark UI themes",
+      defaultValue: "dark",
       toolbar: {
         icon: "mirror",
         items: [
@@ -20,10 +37,10 @@ const preview: Preview = {
   },
   decorators: [
     (Story, context) => {
-      const mode = context.globals.theme as ThemeMode | undefined;
-
+      const mode = (context.globals.theme ?? "dark") as ThemeMode;
       return (
-        <ThemeProvider forcedTheme={mode}>
+        <ThemeProvider>
+          <ThemeStorySync mode={mode} />
           <Story />
         </ThemeProvider>
       );

src/browser/contexts/ThemeContext.tsx

@@ -51,35 +51,26 @@ function applyThemeToDocument(theme: ThemeMode) {
   }
 }
 
-export function ThemeProvider(props: { children: ReactNode; forcedTheme?: ThemeMode }) {
-  const [persistedTheme, setPersistedTheme] = usePersistedState<ThemeMode>(
-    UI_THEME_KEY,
-    resolveSystemTheme(),
-    {
-      listener: !props.forcedTheme, // Disable listener when theme is forced
-    }
-  );
-
-  const activeTheme = props.forcedTheme ?? persistedTheme;
+export function ThemeProvider(props: { children: ReactNode }) {
+  const [theme, setTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
+    listener: true,
+  });
 
   useLayoutEffect(() => {
-    applyThemeToDocument(activeTheme);
-  }, [activeTheme]);
+    applyThemeToDocument(theme);
+  }, [theme]);
 
   const toggleTheme = useCallback(() => {
-    // Only allow toggling when not forced
-    if (!props.forcedTheme) {
-      setPersistedTheme((current) => (current === "dark" ? "light" : "dark"));
-    }
-  }, [props.forcedTheme, setPersistedTheme]);
+    setTheme((current) => (current === "dark" ? "light" : "dark"));
+  }, [setTheme]);
 
   const value = useMemo<ThemeContextValue>(
     () => ({
-      theme: activeTheme,
-      setTheme: setPersistedTheme,
+      theme,
+      setTheme,
       toggleTheme,
     }),
-    [activeTheme, setPersistedTheme, toggleTheme]
+    [setTheme, theme, toggleTheme]
   );
 
   return <ThemeContext.Provider value={value}>{props.children}</ThemeContext.Provider>;