From 8d02ae28f5022b7cb2d8738c7b638b70d786a2e0 Mon Sep 17 00:00:00 2001
From: PathGao <42336971+PathGao@users.noreply.github.com>
Date: Thu, 30 Jul 2026 08:59:35 +0800
Subject: [PATCH] fix(frontend): preserve theme body classes (#6157)

* fix(storybook): preserve preview body classes

* fix(frontend): retain theme body classes

* fix(storybook): mirror panel theme attributes

* test(storybook): cover theme switches

* test(storybook): strengthen theme DOM coverage

* fix(frontend): preserve message container classes

---------

Co-authored-by: PathGao <gaoyanbo@gaoyanbodeMacBook-Air.local>
---
 frontend/.storybook/preview.tsx            | 11 ++---
 frontend/src/hooks/useTheme.tsx            | 12 +++--
 frontend/src/test/storybook-theme.test.tsx | 51 ++++++++++++++++++++++
 3 files changed, 65 insertions(+), 9 deletions(-)
 create mode 100644 frontend/src/test/storybook-theme.test.tsx

diff --git a/frontend/.storybook/preview.tsx b/frontend/.storybook/preview.tsx
index 2e1e48f63..c5281cb16 100644
--- a/frontend/.storybook/preview.tsx
+++ b/frontend/.storybook/preview.tsx
@@ -1,4 +1,4 @@
-import { useEffect } from 'react';
+import { useLayoutEffect } from 'react';
 import type { Decorator, Preview } from '@storybook/react-vite';
 import { ConfigProvider } from 'antd';
 import i18next from 'i18next';
@@ -17,11 +17,12 @@ if (!i18next.isInitialized) {
   });
 }
 
-const withTheme: Decorator = (Story, context) => {
+export const withTheme: Decorator = (Story, context) => {
   const dark = context.globals.theme === 'dark';
-  useEffect(() => {
-    document.body.setAttribute('class', dark ? 'dark' : 'light');
-    document.documentElement.setAttribute('data-theme', dark ? 'dark' : 'light');
+  useLayoutEffect(() => {
+    document.body.classList.remove('dark', 'light');
+    document.body.classList.add(dark ? 'dark' : 'light');
+    document.documentElement.removeAttribute('data-theme');
   }, [dark]);
   return (
     <ConfigProvider theme={buildAntdThemeConfig(dark, false)}>
diff --git a/frontend/src/hooks/useTheme.tsx b/frontend/src/hooks/useTheme.tsx
index a3620e035..c64be0427 100644
--- a/frontend/src/hooks/useTheme.tsx
+++ b/frontend/src/hooks/useTheme.tsx
@@ -1,4 +1,4 @@
-import { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
+import { createContext, useCallback, useContext, useLayoutEffect, useMemo, useState } from 'react';
 import type { ReactNode } from 'react';
 import { theme as antdTheme } from 'antd';
 import type { ThemeConfig } from 'antd';
@@ -13,14 +13,18 @@ function readBool(key: string, fallback: boolean): boolean {
 }
 
 function applyDom(isDark: boolean, isUltra: boolean) {
-  document.body.setAttribute('class', isDark ? 'dark' : 'light');
+  document.body.classList.remove('dark', 'light');
+  document.body.classList.add(isDark ? 'dark' : 'light');
   if (isUltra) {
     document.documentElement.setAttribute('data-theme', 'ultra-dark');
   } else {
     document.documentElement.removeAttribute('data-theme');
   }
   const msg = document.getElementById('message');
-  if (msg) msg.className = isDark ? 'dark' : 'light';
+  if (msg) {
+    msg.classList.remove('dark', 'light');
+    msg.classList.add(isDark ? 'dark' : 'light');
+  }
 }
 
 // module load so the document is in the right theme before React mounts.
@@ -158,7 +162,7 @@ export function ThemeProvider({ children }: { children: ReactNode }) {
   const [isDark, setIsDark] = useState<boolean>(initialDark);
   const [isUltra, setIsUltra] = useState<boolean>(initialUltra);
 
-  useEffect(() => {
+  useLayoutEffect(() => {
     applyDom(isDark, isUltra);
     localStorage.setItem(STORAGE_DARK, String(isDark));
     localStorage.setItem(STORAGE_ULTRA, String(isUltra));
diff --git a/frontend/src/test/storybook-theme.test.tsx b/frontend/src/test/storybook-theme.test.tsx
new file mode 100644
index 000000000..6795a97c1
--- /dev/null
+++ b/frontend/src/test/storybook-theme.test.tsx
@@ -0,0 +1,51 @@
+import { render } from '@testing-library/react';
+import { afterEach, expect, test } from 'vitest';
+
+import { withTheme } from '../../.storybook/preview';
+import { ThemeProvider } from '@/hooks/useTheme';
+
+function Story() {
+  return <div>Story</div>;
+}
+
+function StorybookTheme({ theme }: { theme: 'light' | 'dark' }) {
+  return withTheme(Story, { globals: { theme } } as Partial<Parameters<typeof withTheme>[1]> as Parameters<typeof withTheme>[1]);
+}
+
+afterEach(() => {
+  document.body.className = '';
+  document.documentElement.removeAttribute('data-theme');
+});
+
+test('preserves unrelated body classes when applying the Storybook theme', () => {
+  document.body.className = 'storybook-fixture dark';
+  document.documentElement.setAttribute('data-theme', 'ultra-dark');
+  const { rerender } = render(<StorybookTheme theme="light" />);
+
+  expect(document.body.classList.contains('storybook-fixture')).toBe(true);
+  expect(document.body.classList.contains('light')).toBe(true);
+  expect(document.body.classList.contains('dark')).toBe(false);
+  expect(document.documentElement.hasAttribute('data-theme')).toBe(false);
+
+  rerender(<StorybookTheme theme="dark" />);
+
+  expect(document.body.classList.contains('storybook-fixture')).toBe(true);
+  expect(document.body.classList.contains('dark')).toBe(true);
+  expect(document.body.classList.contains('light')).toBe(false);
+});
+
+test('preserves unrelated body classes when applying the panel theme', () => {
+  document.body.className = 'panel-fixture';
+  const message = document.createElement('div');
+  message.id = 'message';
+  message.className = 'message-fixture';
+  document.body.append(message);
+
+  render(<ThemeProvider><div>Panel</div></ThemeProvider>);
+
+  expect(document.body.classList.contains('panel-fixture')).toBe(true);
+  expect(document.body.classList.contains('dark')).toBe(true);
+  expect(document.body.classList.contains('light')).toBe(false);
+  expect(message.classList.contains('message-fixture')).toBe(true);
+  expect(message.classList.contains('dark')).toBe(true);
+});