Improve custom theme hydration issues

Signed-off-by: Charles de Dreuille <charles.dedreuille@gmail.com>

docs-ui/src/components/CustomTheme/customTheme.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import { useEffect, useState, useCallback } from 'react';
+import { useEffect, useState, useCallback, useSyncExternalStore } from 'react';
 import CodeMirror from '@uiw/react-codemirror';
 import { sass } from '@codemirror/lang-sass';
 import styles from './styles.module.css';
@@ -14,6 +14,10 @@ const defaultTheme = `:root {
   --bui-bg-solid: #000;
 }`;
 
+// Stable server snapshots for useSyncExternalStore
+const serverIsClient = false;
+const serverDefaultTheme = defaultTheme;
+
 const myTheme = createTheme({
   theme: 'light',
   settings: {
@@ -46,12 +50,41 @@ const myTheme = createTheme({
 });
 
 export const CustomTheme = () => {
-  const [isClient, setIsClient] = useState(() => typeof window !== 'undefined');
   const [open, setOpen] = useState(true);
-  const [customTheme, setCustomTheme] = useState<string | undefined>(undefined);
   const { selectedThemeName } = usePlayground();
   const [savedMessage, setSavedMessage] = useState<string>('Save');
 
+  // SSR-safe client detection
+  const isClient = useSyncExternalStore(
+    () => () => {},
+    () => true,
+    () => serverIsClient,
+  );
+
+  // SSR-safe localStorage access for custom theme
+  const customThemeFromStorage = useSyncExternalStore(
+    callback => {
+      window.addEventListener('storage', callback);
+      return () => window.removeEventListener('storage', callback);
+    },
+    () => {
+      const stored = localStorage.getItem('customThemeCss');
+      if (!stored) {
+        localStorage.setItem('customThemeCss', defaultTheme);
+        return defaultTheme;
+      }
+      return stored;
+    },
+    () => serverDefaultTheme,
+  );
+
+  const [customTheme, setCustomTheme] = useState(customThemeFromStorage);
+
+  // Sync from storage when it changes
+  useEffect(() => {
+    setCustomTheme(customThemeFromStorage);
+  }, [customThemeFromStorage]);
+
   const updateStyleElement = (theme: string) => {
     let styleElement = document.getElementById(
       'custom-theme-style',
@@ -66,19 +99,11 @@ export const CustomTheme = () => {
     styleElement.textContent = theme;
   };
 
+  // Apply custom theme to DOM
   useEffect(() => {
-    if (selectedThemeName === 'custom') {
-      let storedTheme = localStorage.getItem('customThemeCss');
-      if (!storedTheme) {
-        storedTheme = defaultTheme;
-        localStorage.setItem('customThemeCss', storedTheme);
-      }
-      // This setState is intentional - we're syncing component state with localStorage
-      // when the user switches to the custom theme. This is a valid effect pattern.
-      // eslint-disable-next-line react-hooks/set-state-in-effect
-      setCustomTheme(storedTheme);
-      updateStyleElement(storedTheme);
-    } else {
+    if (selectedThemeName === 'custom' && customTheme && isClient) {
+      updateStyleElement(customTheme);
+    } else if (isClient) {
       const styleElement = document.getElementById(
         'custom-theme-style',
       ) as HTMLStyleElement;
@@ -86,7 +111,7 @@ export const CustomTheme = () => {
         styleElement.remove();
       }
     }
-  }, [selectedThemeName]);
+  }, [selectedThemeName, customTheme, isClient]);
 
   const handleSave = () => {
     if (customTheme) {

docs-ui/src/utils/playground-context.tsx

@@ -4,6 +4,7 @@ import {
   ReactNode,
   useState,
   useEffect,
+  useSyncExternalStore,
 } from 'react';
 import { components } from './data';
 
@@ -31,6 +32,14 @@ const PlaygroundContext = createContext<{
   setSelectedThemeName: () => {},
 });
 
+// Stable server snapshots (outside component to avoid recreating)
+const defaultThemeSet = new Set<Theme>(['light']);
+const defaultThemeName: ThemeName = 'backstage';
+
+// Cache for theme Sets to avoid creating new objects on every getSnapshot call
+let cachedThemeValue: string | null = null;
+let cachedThemeSet: Set<Theme> = defaultThemeSet;
+
 // Create a provider component
 export const PlaygroundProvider = ({ children }: { children: ReactNode }) => {
   // Check if running in a browser environment
@@ -41,49 +50,71 @@ export const PlaygroundProvider = ({ children }: { children: ReactNode }) => {
     components.map(component => component.slug),
   );
 
-  // Use lazy initialization to load from localStorage
-  const [selectedTheme, setSelectedTheme] = useState<Set<Theme>>(() => {
-    if (typeof window !== 'undefined') {
-      const savedThemeString = localStorage.getItem('theme-mode');
-      if (savedThemeString) {
-        const themeArray = savedThemeString
-          .split(',')
-          .filter(Boolean) as Theme[];
-        return new Set(themeArray);
-      }
-    }
-    return new Set(['light']);
-  });
+  // Use useSyncExternalStore for SSR-safe localStorage access
+  const selectedTheme = useSyncExternalStore(
+    callback => {
+      window.addEventListener('storage', callback);
+      return () => window.removeEventListener('storage', callback);
+    },
+    () => {
+      const saved = localStorage.getItem('theme-mode');
 
-  const [selectedThemeName, setSelectedThemeName] = useState<ThemeName>(() => {
-    if (typeof window !== 'undefined') {
-      const savedThemeName = localStorage.getItem('theme-name') as ThemeName;
-      if (savedThemeName) {
-        return savedThemeName;
+      // Return cached Set if value hasn't changed
+      if (saved === cachedThemeValue) {
+        return cachedThemeSet;
       }
-    }
-    return 'backstage';
-  });
+
+      // Update cache with new value
+      cachedThemeValue = saved;
+      if (saved) {
+        const themeArray = saved.split(',').filter(Boolean) as Theme[];
+        cachedThemeSet = new Set(themeArray);
+      } else {
+        cachedThemeSet = defaultThemeSet;
+      }
+
+      return cachedThemeSet;
+    },
+    () => defaultThemeSet, // Stable server snapshot
+  );
+
+  const selectedThemeName = useSyncExternalStore(
+    callback => {
+      window.addEventListener('storage', callback);
+      return () => window.removeEventListener('storage', callback);
+    },
+    () => {
+      const saved = localStorage.getItem('theme-name') as ThemeName;
+      return saved || defaultThemeName;
+    },
+    () => defaultThemeName, // Stable server snapshot
+  );
+
+  // Keep setter functions that update both state and localStorage
+  const setSelectedTheme = (keys: Set<Theme>) => {
+    const value = Array.from(keys).join(',');
+    localStorage.setItem('theme-mode', value);
+    // Invalidate cache
+    cachedThemeValue = null;
+    window.dispatchEvent(new Event('storage'));
+  };
+
+  const setSelectedThemeName = (name: ThemeName) => {
+    localStorage.setItem('theme-name', name);
+    window.dispatchEvent(new Event('storage'));
+  };
+
+  // Sync to DOM attributes when values change
+  useEffect(() => {
+    document.documentElement.setAttribute(
+      'data-theme-mode',
+      Array.from(selectedTheme).join(','),
+    );
+  }, [selectedTheme]);
 
   useEffect(() => {
-    if (isBrowser) {
-      document.documentElement.setAttribute(
-        'data-theme-mode',
-        Array.from(selectedTheme).join(','),
-      );
-      localStorage.setItem('theme-mode', Array.from(selectedTheme).join(','));
-    }
-  }, [selectedTheme, isBrowser]);
-
-  useEffect(() => {
-    if (isBrowser) {
-      document.documentElement.setAttribute(
-        'data-theme-name',
-        selectedThemeName || 'backstage',
-      );
-      localStorage.setItem('theme-name', selectedThemeName || 'backstage');
-    }
-  }, [selectedThemeName, isBrowser]);
+    document.documentElement.setAttribute('data-theme-name', selectedThemeName);
+  }, [selectedThemeName]);
 
   return (
     <PlaygroundContext.Provider