fix(runtime): re-attach removed style elements

Author: saudademjj

src/runtime/disconnected-callback.ts

@@ -3,7 +3,7 @@ import { getHostRef, plt } from '@platform';
 
 import type * as d from '../declarations';
 import { PLATFORM_FLAGS } from './runtime-constants';
-import { rootAppliedStyles } from './styles';
+import { rootAppliedStyleElms, rootAppliedStyles } from './styles';
 import { safeCall } from './update-component';
 
 const disconnectInstance = (instance: any, elm?: d.HostElement) => {
@@ -38,11 +38,17 @@ export const disconnectedCallback = async (elm: d.HostElement) => {
   if (rootAppliedStyles.has(elm)) {
     rootAppliedStyles.delete(elm);
   }
+  if (rootAppliedStyleElms.has(elm)) {
+    rootAppliedStyleElms.delete(elm);
+  }
 
   /**
    * Remove the shadow root from the `rootAppliedStyles` WeakMap
    */
   if (elm.shadowRoot && rootAppliedStyles.has(elm.shadowRoot as unknown as Element)) {
     rootAppliedStyles.delete(elm.shadowRoot as unknown as Element);
   }
+  if (elm.shadowRoot && rootAppliedStyleElms.has(elm.shadowRoot as unknown as Element)) {
+    rootAppliedStyleElms.delete(elm.shadowRoot as unknown as Element);
+  }
 };

src/runtime/styles.ts

@@ -16,6 +16,7 @@ import { createTime } from './profile';
 import { HYDRATED_STYLE_ID, NODE_TYPE, SLOT_FB_CSS } from './runtime-constants';
 
 export const rootAppliedStyles: d.RootAppliedStyleMap = /*@__PURE__*/ new WeakMap();
+export const rootAppliedStyleElms: WeakMap<any, Map<string, HTMLStyleElement>> = /*@__PURE__*/ new WeakMap();
 
 /**
  * Register the styles for a component by creating a stylesheet and then
@@ -71,24 +72,39 @@ export const addStyle = (styleContainerNode: any, cmpMeta: d.ComponentRuntimeMet
     if (typeof style === 'string') {
       styleContainerNode = styleContainerNode.head || (styleContainerNode as HTMLElement);
       let appliedStyles = rootAppliedStyles.get(styleContainerNode);
+      let appliedStyleElms = rootAppliedStyleElms.get(styleContainerNode);
       let styleElm;
       if (!appliedStyles) {
         rootAppliedStyles.set(styleContainerNode, (appliedStyles = new Set()));
       }
+      if (!appliedStyleElms) {
+        rootAppliedStyleElms.set(styleContainerNode, (appliedStyleElms = new Map()));
+      }
+
+      const trackedStyleElm = appliedStyleElms.get(scopeId);
+      if (trackedStyleElm && trackedStyleElm.parentNode !== styleContainerNode) {
+        appliedStyles.delete(scopeId);
+        appliedStyleElms.delete(scopeId);
+      }
 
       // Check if style element already exists (for HMR updates)
       // For shadow DOM components, directly update their dedicated style element
       // For scoped components, check if they have their own HMR-created style element
-      const existingStyleElm: HTMLStyleElement =
-        (BUILD.hydrateClientSide || BUILD.hotModuleReplacement) &&
-        styleContainerNode.querySelector(`[${HYDRATED_STYLE_ID}="${scopeId}"]`);
+      const existingStyleElm: HTMLStyleElement | undefined =
+        (trackedStyleElm?.parentNode === styleContainerNode && trackedStyleElm) ||
+        ((BUILD.hydrateClientSide || BUILD.hotModuleReplacement) &&
+          styleContainerNode.querySelector(`[${HYDRATED_STYLE_ID}="${scopeId}"]`)) ||
+        undefined;
 
       if (existingStyleElm) {
         // Update existing style element (for hydration or HMR)
         existingStyleElm.textContent = style;
+        appliedStyles.add(scopeId);
+        appliedStyleElms.set(scopeId, existingStyleElm);
       } else if (!appliedStyles.has(scopeId)) {
         styleElm = win.document.createElement('style');
         styleElm.textContent = style;
+        let appliedStyleElm = styleElm;
 
         // Apply CSP nonce to the style tag if it exists
         const nonce = plt.$nonce$ ?? queryNonceMetaTagContent(win.document);
@@ -165,6 +181,7 @@ export const addStyle = (styleContainerNode: any, cmpMeta: d.ComponentRuntimeMet
               const existingStyleContainer: HTMLStyleElement = styleContainerNode.querySelector('style');
               if (existingStyleContainer && !BUILD.hotModuleReplacement) {
                 existingStyleContainer.textContent = style + existingStyleContainer.textContent;
+                appliedStyleElm = existingStyleContainer;
               } else {
                 (styleContainerNode as HTMLElement).prepend(styleElm);
               }
@@ -189,6 +206,7 @@ export const addStyle = (styleContainerNode: any, cmpMeta: d.ComponentRuntimeMet
         if (appliedStyles) {
           appliedStyles.add(scopeId);
         }
+        appliedStyleElms.set(scopeId, appliedStyleElm);
       }
     } else if (BUILD.constructableCSS) {
       let appliedStyles = rootAppliedStyles.get(styleContainerNode);

src/runtime/test/style.spec.tsx

@@ -56,6 +56,45 @@ describe('style', () => {
     );
   });
 
+  it('re-attaches a removed style element when the component is rendered again', async () => {
+    @Component({
+      tag: 'cmp-a',
+      styles: `
+        cmp-a {
+          color: red;
+        }
+      `,
+    })
+    class CmpA {
+      render() {
+        return `innertext`;
+      }
+    }
+
+    const page = await newSpecPage({
+      components: [CmpA],
+      html: `<cmp-a></cmp-a>`,
+      attachStyles: true,
+    });
+
+    const findCmpStyle = () =>
+      Array.from(page.doc.head.querySelectorAll('style')).find((styleElm) =>
+        styleElm.textContent?.includes('color: red'),
+      );
+
+    const initialStyleElm = findCmpStyle();
+    expect(initialStyleElm).toBeDefined();
+
+    initialStyleElm!.remove();
+    expect(findCmpStyle()).toBeUndefined();
+
+    await page.setContent(`<cmp-a></cmp-a>`);
+
+    const reattachedStyleElm = findCmpStyle();
+    expect(reattachedStyleElm).toBeDefined();
+    expect(reattachedStyleElm!.isConnected).toBe(true);
+  });
+
   describe('mode', () => {
     it('md mode', async () => {
       setMode(() => 'md');