From 643118d3de80f5e457cf49796605d400cf12c600 Mon Sep 17 00:00:00 2001
From: Pavel Tarnopolsky <pavelta@wix.com>
Date: Mon, 29 Sep 2025 11:26:36 +0300
Subject: [PATCH 1/2] AsChildSlot: fix data tag propogation for fragments #pr

---
 .../utils/src/react/AsChildSlot.test.tsx      | 27 ++++++++++++++
 .../utils/src/react/AsChildSlot.tsx           | 35 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/packages/headless-components/utils/src/react/AsChildSlot.test.tsx b/packages/headless-components/utils/src/react/AsChildSlot.test.tsx
index 49bb111e6..0eefd5113 100644
--- a/packages/headless-components/utils/src/react/AsChildSlot.test.tsx
+++ b/packages/headless-components/utils/src/react/AsChildSlot.test.tsx
@@ -1003,6 +1003,33 @@ describe('AsChildSlot', () => {
         );
         expect(secondElement).not.toHaveAttribute('data-component-tag');
       });
+
+      it('should inject data-component-tag into first element of fragment with asChild=false', () => {
+        render(
+          <TestComponent
+            asChild={false}
+            data-component-tag="fragment-no-asChild"
+          >
+            <>
+              <div data-testid="fragment-first-no-asChild">
+                First element in fragment
+              </div>
+              <span data-testid="fragment-second-no-asChild">
+                Second element in fragment
+              </span>
+            </>
+          </TestComponent>,
+        );
+
+        const firstElement = screen.getByTestId('fragment-first-no-asChild');
+        const secondElement = screen.getByTestId('fragment-second-no-asChild');
+
+        expect(firstElement).toHaveAttribute(
+          'data-component-tag',
+          'fragment-no-asChild',
+        );
+        expect(secondElement).not.toHaveAttribute('data-component-tag');
+      });
     });
 
     describe('edge cases', () => {
diff --git a/packages/headless-components/utils/src/react/AsChildSlot.tsx b/packages/headless-components/utils/src/react/AsChildSlot.tsx
index edf35e710..f66a9a16b 100644
--- a/packages/headless-components/utils/src/react/AsChildSlot.tsx
+++ b/packages/headless-components/utils/src/react/AsChildSlot.tsx
@@ -219,6 +219,41 @@ export const AsChildSlot = React.forwardRef<HTMLElement, AsChildSlot>(
       return <Slot {...restProps}>{renderedElement}</Slot>;
     }
 
+    // Handle React Fragment children specifically
+    if (React.isValidElement(children) && children.type === React.Fragment) {
+      const fragmentChildren = React.Children.toArray(
+        (children.props as any).children,
+      );
+
+      if (fragmentChildren.length > 0) {
+        const firstChild = fragmentChildren[0];
+        const restChildren = fragmentChildren.slice(1);
+
+        if (React.isValidElement(firstChild)) {
+          // Only inject data-component-tag if the child doesn't already have one
+          const existingTag = (firstChild.props as any)['data-component-tag'];
+          const enhancedFirstChild = React.cloneElement(firstChild, {
+            ...getConditionalDataComponentTagProps(dataComponentTag, existingTag),
+          });
+
+          return (
+            <Slot ref={ref} {...restProps}>
+              <>
+                {enhancedFirstChild}
+                {restChildren}
+              </>
+            </Slot>
+          );
+        }
+      }
+
+      return (
+        <Slot ref={ref} {...restProps}>
+          {children}
+        </Slot>
+      );
+    }
+
     if (React.Children.count(children) > 1) {
       const childrenArray = React.Children.toArray(children);
       const firstChild = childrenArray[0];

From 4ff18e07476bcd686d1fb6ac516c69a4ca830a34 Mon Sep 17 00:00:00 2001
From: Pavel Tarnopolsky <pavelta@wix.com>
Date: Mon, 29 Sep 2025 14:16:27 +0300
Subject: [PATCH 2/2] wip

---
 .../blog/src/react/Categories.tsx             |   1 +
 .../utils/src/react/AsChildSlot.tsx           | 296 +++++++++---------
 2 files changed, 153 insertions(+), 144 deletions(-)

diff --git a/packages/headless-components/blog/src/react/Categories.tsx b/packages/headless-components/blog/src/react/Categories.tsx
index ca890e064..378349016 100644
--- a/packages/headless-components/blog/src/react/Categories.tsx
+++ b/packages/headless-components/blog/src/react/Categories.tsx
@@ -114,6 +114,7 @@ export const Root = React.forwardRef<HTMLElement, BlogCategoriesRootProps>((prop
           {...attributes}
           customElement={children}
           customElementProps={{ hasCategories }}
+          data-component-tag={"wow"}
         >
           <div>{isValidChildren(children) ? children : null}</div>
         </AsChildSlot>
diff --git a/packages/headless-components/utils/src/react/AsChildSlot.tsx b/packages/headless-components/utils/src/react/AsChildSlot.tsx
index f66a9a16b..e2c88faa1 100644
--- a/packages/headless-components/utils/src/react/AsChildSlot.tsx
+++ b/packages/headless-components/utils/src/react/AsChildSlot.tsx
@@ -44,6 +44,20 @@ export interface AsChildSlot {
   [key: string]: any;
 }
 
+// Helper functions for data-component-tag prop merging
+function getDataComponentTagProps(dataComponentTag?: string) {
+  return dataComponentTag ? { 'data-component-tag': dataComponentTag } : {};
+}
+
+function getConditionalDataComponentTagProps(
+  dataComponentTag?: string,
+  existingTag?: string,
+) {
+  return dataComponentTag && !existingTag
+    ? { 'data-component-tag': dataComponentTag }
+    : {};
+}
+
 // Helper function to inject data-component-tag into rendered elements
 function injectDataComponentTag(
   renderedElement: React.ReactNode,
@@ -97,18 +111,133 @@ function injectDataComponentTag(
   return <Slot {...restProps}>{renderedElement}</Slot>;
 }
 
-// Helper functions for data-component-tag prop merging
-function getDataComponentTagProps(dataComponentTag?: string) {
-  return dataComponentTag ? { 'data-component-tag': dataComponentTag } : {};
+// Helper to handle string elements
+function handleStringElement(
+  element: string,
+  dataComponentTag: string | undefined,
+  ref: React.Ref<HTMLElement>,
+  restProps: any,
+  WrapperComponent: 'span' | 'div' = 'span',
+): React.ReactElement {
+  return (
+    <Slot ref={ref} {...restProps}>
+      <WrapperComponent {...getDataComponentTagProps(dataComponentTag)}>
+        {element}
+      </WrapperComponent>
+    </Slot>
+  );
 }
 
-function getConditionalDataComponentTagProps(
-  dataComponentTag?: string,
-  existingTag?: string,
-) {
-  return dataComponentTag && !existingTag
-    ? { 'data-component-tag': dataComponentTag }
-    : {};
+// Helper to handle render functions and objects
+function handleRenderElement(
+  element: AsChildRenderFunction | AsChildRenderObject,
+  props: any,
+  ref: React.Ref<HTMLElement>,
+  dataComponentTag: string | undefined,
+  restProps: any,
+): React.ReactElement {
+  const renderFn = typeof element === 'function' ? element : element.render;
+  const renderedElement = renderFn(props, ref);
+
+  if (dataComponentTag) {
+    return injectDataComponentTag(renderedElement, dataComponentTag, restProps);
+  }
+
+  return <Slot {...restProps}>{renderedElement}</Slot>;
+}
+
+// Helper to handle React elements (for customElement case)
+function handleReactElement(
+  element: React.ReactElement,
+  content: React.ReactNode | undefined,
+  dataComponentTag: string | undefined,
+  ref: React.Ref<HTMLElement>,
+  restProps: any,
+): React.ReactElement {
+  return (
+    <Slot {...restProps}>
+      {React.cloneElement(element, {
+        ref,
+        ...(content !== undefined ? { children: content } : {}),
+        ...getDataComponentTagProps(dataComponentTag),
+      })}
+    </Slot>
+  );
+}
+
+// Helper to handle fragments
+function handleFragmentChildren(
+  children: React.ReactElement,
+  dataComponentTag: string | undefined,
+  ref: React.Ref<HTMLElement>,
+  restProps: any,
+): React.ReactElement {
+  const fragmentChildren = React.Children.toArray(
+    (children.props as any).children,
+  );
+
+  if (fragmentChildren.length > 0) {
+    const firstChild = fragmentChildren[0];
+    const restChildren = fragmentChildren.slice(1);
+
+    if (React.isValidElement(firstChild)) {
+      // Only inject data-component-tag if the child doesn't already have one
+      const existingTag = (firstChild.props as any)['data-component-tag'];
+      const enhancedFirstChild = React.cloneElement(firstChild, {
+        ...getConditionalDataComponentTagProps(dataComponentTag, existingTag),
+      });
+
+      return (
+        <Slot ref={ref} {...restProps}>
+          <>
+            {enhancedFirstChild}
+            {restChildren}
+          </>
+        </Slot>
+      );
+    }
+  }
+
+  return (
+    <Slot ref={ref} {...restProps}>
+      {children}
+    </Slot>
+  );
+}
+
+// Helper to handle multiple children
+function handleMultipleChildren(
+  children: React.ReactNode,
+  dataComponentTag: string | undefined,
+  ref: React.Ref<HTMLElement>,
+  restProps: any,
+): React.ReactElement {
+  const childrenArray = React.Children.toArray(children);
+  const firstChild = childrenArray[0];
+  const restChildren = childrenArray.slice(1);
+
+  if (React.isValidElement(firstChild)) {
+    // Only inject data-component-tag if the child doesn't already have one
+    const existingTag = (firstChild.props as any)['data-component-tag'];
+    const enhancedFirstChild = React.cloneElement(firstChild, {
+      ...getConditionalDataComponentTagProps(dataComponentTag, existingTag),
+    });
+
+    return (
+      <Slot ref={ref} {...restProps}>
+        <>
+          {enhancedFirstChild}
+          {restChildren}
+        </>
+      </Slot>
+    );
+  }
+
+  return (
+    <Slot ref={ref} {...restProps}>
+      <>{children}</>
+    </Slot>
+  );
 }
 
 export const AsChildSlot = React.forwardRef<HTMLElement, AsChildSlot>(
@@ -123,166 +252,45 @@ export const AsChildSlot = React.forwardRef<HTMLElement, AsChildSlot>(
       ...restProps
     } = props;
 
+    // Handle customElement when asChild is true
     if (asChild && customElement) {
-      // Handle string
       if (typeof customElement === 'string') {
-        return (
-          <Slot ref={ref} {...restProps}>
-            <span {...getDataComponentTagProps(dataComponentTag)}>
-              {customElement}
-            </span>
-          </Slot>
-        );
+        return handleStringElement(customElement, dataComponentTag, ref, restProps);
       }
 
-      // Handle React element pattern
       if (React.isValidElement(customElement)) {
-        return (
-          <Slot {...restProps}>
-            {React.cloneElement(customElement as React.ReactElement, {
-              ref,
-              ...(content !== undefined ? { children: content } : {}),
-              ...getDataComponentTagProps(dataComponentTag),
-            })}
-          </Slot>
-        );
+        return handleReactElement(customElement, content, dataComponentTag, ref, restProps);
       }
 
-      // Handle render function pattern
-      if (typeof customElement === 'function') {
-        const renderedElement = customElement(customElementProps, ref);
-
-        if (dataComponentTag) {
-          return injectDataComponentTag(
-            renderedElement,
-            dataComponentTag,
-            restProps,
-          );
-        }
-
-        return <Slot {...restProps}>{renderedElement}</Slot>;
-      }
-
-      // Handle render object pattern
-      if (typeof customElement === 'object' && 'render' in customElement) {
-        const renderedElement = customElement.render(customElementProps, ref);
-
-        if (dataComponentTag) {
-          return injectDataComponentTag(
-            renderedElement,
-            dataComponentTag,
-            restProps,
-          );
-        }
-
-        return <Slot {...restProps}>{renderedElement}</Slot>;
+      if (typeof customElement === 'function' || (typeof customElement === 'object' && 'render' in customElement)) {
+        return handleRenderElement(customElement, customElementProps, ref, dataComponentTag, restProps);
       }
     }
 
+    // Handle children cases
     if (!children) {
       return null;
     }
 
     if (typeof children === 'string') {
-      return (
-        <Slot ref={ref} {...restProps}>
-          <div {...getDataComponentTagProps(dataComponentTag)}>{children}</div>
-        </Slot>
-      );
+      return handleStringElement(children, dataComponentTag, ref, restProps, 'div');
     }
 
-    if (typeof children === 'function') {
-      const renderedElement = children(customElementProps, ref);
-
-      if (dataComponentTag) {
-        return injectDataComponentTag(
-          renderedElement,
-          dataComponentTag,
-          restProps,
-        );
-      }
-
-      return <Slot {...restProps}>{renderedElement}</Slot>;
-    }
-
-    if (typeof children === 'object' && 'render' in children) {
-      const renderedElement = children.render(customElementProps, ref);
-
-      if (dataComponentTag) {
-        return injectDataComponentTag(
-          renderedElement,
-          dataComponentTag,
-          restProps,
-        );
-      }
-
-      return <Slot {...restProps}>{renderedElement}</Slot>;
+    if (typeof children === 'function' || (typeof children === 'object' && 'render' in children)) {
+      return handleRenderElement(children, customElementProps, ref, dataComponentTag, restProps);
     }
 
     // Handle React Fragment children specifically
     if (React.isValidElement(children) && children.type === React.Fragment) {
-      const fragmentChildren = React.Children.toArray(
-        (children.props as any).children,
-      );
-
-      if (fragmentChildren.length > 0) {
-        const firstChild = fragmentChildren[0];
-        const restChildren = fragmentChildren.slice(1);
-
-        if (React.isValidElement(firstChild)) {
-          // Only inject data-component-tag if the child doesn't already have one
-          const existingTag = (firstChild.props as any)['data-component-tag'];
-          const enhancedFirstChild = React.cloneElement(firstChild, {
-            ...getConditionalDataComponentTagProps(dataComponentTag, existingTag),
-          });
-
-          return (
-            <Slot ref={ref} {...restProps}>
-              <>
-                {enhancedFirstChild}
-                {restChildren}
-              </>
-            </Slot>
-          );
-        }
-      }
-
-      return (
-        <Slot ref={ref} {...restProps}>
-          {children}
-        </Slot>
-      );
+      return handleFragmentChildren(children, dataComponentTag, ref, restProps);
     }
 
+    // Handle multiple children
     if (React.Children.count(children) > 1) {
-      const childrenArray = React.Children.toArray(children);
-      const firstChild = childrenArray[0];
-      const restChildren = childrenArray.slice(1);
-
-      if (React.isValidElement(firstChild)) {
-        // Only inject data-component-tag if the child doesn't already have one
-        const existingTag = (firstChild.props as any)['data-component-tag'];
-        const enhancedFirstChild = React.cloneElement(firstChild, {
-          ...getConditionalDataComponentTagProps(dataComponentTag, existingTag),
-        });
-
-        return (
-          <Slot ref={ref} {...restProps}>
-            <>
-              {enhancedFirstChild}
-              {restChildren}
-            </>
-          </Slot>
-        );
-      }
-
-      return (
-        <Slot ref={ref} {...restProps}>
-          <>{children}</>
-        </Slot>
-      );
+      return handleMultipleChildren(children, dataComponentTag, ref, restProps);
     }
 
+    // Handle single child (default case)
     return (
       <Slot
         ref={ref}