From d5c170eaf48e916ec95c4ef71ce30d332ee40458 Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Sat, 13 Sep 2025 16:32:07 +0200
Subject: [PATCH] bui-themer: extract mui theme and then unmount theme provider

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 .../src/unified/UnifiedThemeProvider.tsx      | 14 +++-
 .../BuiThemerPage/BuiThemerPage.tsx           | 66 +++++++++++--------
 2 files changed, 49 insertions(+), 31 deletions(-)

diff --git a/packages/theme/src/unified/UnifiedThemeProvider.tsx b/packages/theme/src/unified/UnifiedThemeProvider.tsx
index a80e36cf35..006c45605e 100644
--- a/packages/theme/src/unified/UnifiedThemeProvider.tsx
+++ b/packages/theme/src/unified/UnifiedThemeProvider.tsx
@@ -74,12 +74,22 @@ export function UnifiedThemeProvider(
   const themeName = 'backstage';
 
   useEffect(() => {
+    const oldMode = document.body.getAttribute('data-theme-mode');
+    const oldName = document.body.getAttribute('data-theme-name');
     document.body.setAttribute('data-theme-mode', themeMode);
     document.body.setAttribute('data-theme-name', themeName);
 
     return () => {
-      document.body.removeAttribute('data-theme-mode');
-      document.body.removeAttribute('data-theme-name');
+      if (oldMode) {
+        document.body.setAttribute('data-theme-mode', oldMode);
+      } else {
+        document.body.removeAttribute('data-theme-mode');
+      }
+      if (oldName) {
+        document.body.setAttribute('data-theme-name', oldName);
+      } else {
+        document.body.removeAttribute('data-theme-name');
+      }
     };
   }, [themeMode, themeName]);
 
diff --git a/plugins/bui-themer/src/components/BuiThemerPage/BuiThemerPage.tsx b/plugins/bui-themer/src/components/BuiThemerPage/BuiThemerPage.tsx
index 4ae25b60ce..4df3aae006 100644
--- a/plugins/bui-themer/src/components/BuiThemerPage/BuiThemerPage.tsx
+++ b/plugins/bui-themer/src/components/BuiThemerPage/BuiThemerPage.tsx
@@ -15,9 +15,9 @@
  */
 
 import { useMemo, useState, useCallback, useEffect } from 'react';
-import { useApi } from '@backstage/core-plugin-api';
+import { AppTheme, useApi } from '@backstage/core-plugin-api';
 import { appThemeApiRef } from '@backstage/core-plugin-api';
-import { useTheme } from '@mui/material/styles';
+import { Theme, useTheme } from '@mui/material/styles';
 import {
   Box,
   Button,
@@ -33,13 +33,6 @@ import {
   ConvertMuiToBuiThemeOptions,
 } from './convertMuiToBuiTheme';
 
-interface ThemePreviewProps {
-  themeId: string;
-  themeTitle: string;
-  variant: 'light' | 'dark';
-  Provider: React.ComponentType<{ children: React.ReactNode }>;
-}
-
 // Memoization cache for generated CSS
 const cssCache = new Map<string, string>();
 
@@ -47,13 +40,18 @@ interface ThemeContentProps {
   themeId: string;
   themeTitle: string;
   variant: 'light' | 'dark';
+  muiTheme: Theme;
 }
 
-function ThemeContent({ themeId, themeTitle, variant }: ThemeContentProps) {
+function ThemeContent({
+  themeId,
+  themeTitle,
+  variant,
+  muiTheme,
+}: ThemeContentProps) {
   const [generatedCss, setGeneratedCss] = useState<string>('');
   const [isPreviewMode, setIsPreviewMode] = useState(false);
   const [includeThemeId, setIncludeThemeId] = useState(false);
-  const muiTheme = useTheme();
 
   const css = useMemo(() => {
     // Create cache key based on theme properties and options
@@ -243,23 +241,29 @@ function ThemeContent({ themeId, themeTitle, variant }: ThemeContentProps) {
   );
 }
 
-function ThemePreview({
-  themeId,
-  themeTitle,
-  variant,
-  Provider,
-}: ThemePreviewProps) {
+function MuiThemeExtractor(props: {
+  appTheme: AppTheme;
+  children: (theme: Theme) => JSX.Element;
+}): JSX.Element {
+  const { appTheme, children } = props;
+  const [theme, setTheme] = useState<Theme | null>(null);
+  const { Provider } = appTheme;
+  if (theme) {
+    return children(theme);
+  }
+
   return (
     <Provider>
-      <ThemeContent
-        themeId={themeId}
-        themeTitle={themeTitle}
-        variant={variant}
-      />
+      <MuiThemeExtractorInner setTheme={setTheme} />
     </Provider>
   );
 }
 
+function MuiThemeExtractorInner(props: { setTheme: (theme: Theme) => void }) {
+  props.setTheme(useTheme());
+  return null;
+}
+
 export function BuiThemerPage() {
   const appThemeApi = useApi(appThemeApiRef);
   const installedThemes = appThemeApi.getInstalledThemes();
@@ -287,13 +291,17 @@ export function BuiThemerPage() {
         ) : (
           <Flex direction="column" gap="4">
             {installedThemes.map(theme => (
-              <ThemePreview
-                key={theme.id}
-                themeId={theme.id}
-                themeTitle={theme.title}
-                variant={theme.variant}
-                Provider={theme.Provider}
-              />
+              <MuiThemeExtractor key={theme.id} appTheme={theme}>
+                {muiTheme => (
+                  <ThemeContent
+                    key={theme.id}
+                    themeId={theme.id}
+                    themeTitle={theme.title}
+                    variant={theme.variant}
+                    muiTheme={muiTheme}
+                  />
+                )}
+              </MuiThemeExtractor>
             ))}
           </Flex>
         )}