🤖 fix: stabilize storybook theme by forcing context state (#712)

Fixes flaky visual diffs in Storybook where the theme would flash from
default (light) to selected (dark) on load.

## Changes
- Modified `ThemeProvider` to accept a `forcedTheme` prop that overrides
internal state and persistence.
- Updated Storybook `preview.tsx` to use `forcedTheme` instead of a
`useEffect`-based synchronization component.

This ensures the correct theme is applied during the first render pass,
eliminating the flash and preventing Storybook from polluting local
storage.

Author: ammar-agent

.storybook/preview.tsx

@@ -1,30 +1,13 @@
-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 { ThemeProvider, 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: {
       name: "Theme",
       description: "Choose between light and dark UI themes",
-      defaultValue: "dark",
       toolbar: {
         icon: "mirror",
         items: [
@@ -35,12 +18,22 @@ const preview: Preview = {
       },
     },
   },
+  initialGlobals: {
+    theme: "dark",
+  },
   decorators: [
     (Story, context) => {
-      const mode = (context.globals.theme ?? "dark") as ThemeMode;
+      // Default to dark if mode not set (e.g., Chromatic headless browser defaults to light)
+      const mode = (context.globals.theme as ThemeMode | undefined) ?? "dark";
+
+      // Apply theme synchronously before React renders - critical for Chromatic snapshots
+      if (typeof document !== "undefined") {
+        document.documentElement.dataset.theme = mode;
+        document.documentElement.style.colorScheme = mode;
+      }
+
       return (
-        <ThemeProvider>
-          <ThemeStorySync mode={mode} />
+        <ThemeProvider forcedTheme={mode}>
           <Story />
         </ThemeProvider>
       );
@@ -55,8 +48,8 @@ const preview: Preview = {
     },
     chromatic: {
       modes: {
-        dark: { globals: { theme: "dark" } },
-        light: { globals: { theme: "light" } },
+        dark: { theme: "dark" },
+        light: { theme: "light" },
       },
     },
   },

bun.lock

@@ -1,6 +1,5 @@
 {
   "lockfileVersion": 1,
-  "configVersion": 0,
   "workspaces": {
     "": {
       "name": "@coder/cmux",

src/browser/contexts/ThemeContext.test.tsx

@@ -0,0 +1,108 @@
+import { GlobalWindow } from "happy-dom";
+
+// Setup basic DOM environment for testing-library
+const dom = new GlobalWindow();
+/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
+(global as any).window = dom.window;
+(global as any).document = dom.window.document;
+// Polyfill console since happy-dom might interfere or we just want standard console
+(global as any).console = console;
+/* eslint-enable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
+
+import { afterEach, describe, expect, mock, test, beforeEach } from "bun:test";
+
+import { render, cleanup } from "@testing-library/react";
+import React from "react";
+import { ThemeProvider, useTheme } from "./ThemeContext";
+import { UI_THEME_KEY } from "@/common/constants/storage";
+
+// Helper to access internals
+const TestComponent = () => {
+  const { theme, toggleTheme } = useTheme();
+  return (
+    <div>
+      <span data-testid="theme-value">{theme}</span>
+      <button onClick={toggleTheme} data-testid="toggle-btn">
+        Toggle
+      </button>
+    </div>
+  );
+};
+
+describe("ThemeContext", () => {
+  // Mock matchMedia
+  const mockMatchMedia = mock(() => ({
+    matches: false,
+    media: "",
+    onchange: null,
+    addListener: () => {
+      // no-op
+    },
+    removeListener: () => {
+      // no-op
+    },
+    addEventListener: () => {
+      // no-op
+    },
+    removeEventListener: () => {
+      // no-op
+    },
+    dispatchEvent: () => true,
+  }));
+
+  beforeEach(() => {
+    // Ensure window exists (Bun test with happy-dom should provide it)
+    if (typeof window !== "undefined") {
+      window.matchMedia = mockMatchMedia;
+      window.localStorage.clear();
+    }
+  });
+
+  afterEach(() => {
+    cleanup();
+    if (typeof window !== "undefined") {
+      window.localStorage.clear();
+    }
+  });
+
+  test("uses persisted state by default", () => {
+    const { getByTestId } = render(
+      <ThemeProvider>
+        <TestComponent />
+      </ThemeProvider>
+    );
+    // If matchMedia matches is false (default mock), resolveSystemTheme returns 'dark' (since it checks prefers-color-scheme: light)
+    // resolveSystemTheme logic: window.matchMedia("(prefers-color-scheme: light)").matches ? "light" : "dark"
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+  });
+
+  test("respects forcedTheme prop", () => {
+    const { getByTestId, rerender } = render(
+      <ThemeProvider forcedTheme="light">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("light");
+
+    rerender(
+      <ThemeProvider forcedTheme="dark">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+  });
+
+  test("forcedTheme overrides persisted state", () => {
+    window.localStorage.setItem(UI_THEME_KEY, JSON.stringify("light"));
+
+    const { getByTestId } = render(
+      <ThemeProvider forcedTheme="dark">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+
+    // Check that localStorage is still light (since forcedTheme doesn't write to storage by itself)
+    expect(JSON.parse(window.localStorage.getItem(UI_THEME_KEY)!)).toBe("light");
+  });
+});

src/browser/contexts/ThemeContext.tsx

@@ -15,6 +15,8 @@ interface ThemeContextValue {
   theme: ThemeMode;
   setTheme: React.Dispatch<React.SetStateAction<ThemeMode>>;
   toggleTheme: () => void;
+  /** True if this provider has a forcedTheme - nested providers should not override */
+  isForced: boolean;
 }
 
 const ThemeContext = createContext<ThemeContextValue | null>(null);
@@ -51,29 +53,58 @@ function applyThemeToDocument(theme: ThemeMode) {
   }
 }
 
-export function ThemeProvider(props: { children: ReactNode }) {
-  const [theme, setTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
-    listener: true,
-  });
+export function ThemeProvider({
+  children,
+  forcedTheme,
+}: {
+  children: ReactNode;
+  forcedTheme?: ThemeMode;
+}) {
+  // Check if we're nested inside a forced theme provider
+  const parentContext = useContext(ThemeContext);
+  const isNestedUnderForcedProvider = parentContext?.isForced ?? false;
+
+  const [persistedTheme, setTheme] = usePersistedState<ThemeMode>(
+    UI_THEME_KEY,
+    resolveSystemTheme(),
+    {
+      listener: true,
+    }
+  );
+
+  // If nested under a forced provider, use parent's theme
+  // Otherwise, use forcedTheme (if provided) or persistedTheme
+  const theme =
+    isNestedUnderForcedProvider && parentContext
+      ? parentContext.theme
+      : (forcedTheme ?? persistedTheme);
+
+  const isForced = forcedTheme !== undefined || isNestedUnderForcedProvider;
 
+  // Only apply to document if we're the authoritative provider
   useLayoutEffect(() => {
-    applyThemeToDocument(theme);
-  }, [theme]);
+    if (!isNestedUnderForcedProvider) {
+      applyThemeToDocument(theme);
+    }
+  }, [theme, isNestedUnderForcedProvider]);
 
   const toggleTheme = useCallback(() => {
-    setTheme((current) => (current === "dark" ? "light" : "dark"));
-  }, [setTheme]);
+    if (!isNestedUnderForcedProvider) {
+      setTheme((current) => (current === "dark" ? "light" : "dark"));
+    }
+  }, [setTheme, isNestedUnderForcedProvider]);
 
   const value = useMemo<ThemeContextValue>(
     () => ({
       theme,
       setTheme,
       toggleTheme,
+      isForced,
     }),
-    [setTheme, theme, toggleTheme]
+    [setTheme, theme, toggleTheme, isForced]
   );
 
-  return <ThemeContext.Provider value={value}>{props.children}</ThemeContext.Provider>;
+  return <ThemeContext.Provider value={value}>{children}</ThemeContext.Provider>;
 }
 
 export function useTheme(): ThemeContextValue {