Ensure loading button with asChild renders the correct element

apps/playground/app/test-as-child/page.tsx

@@ -4,6 +4,7 @@ import {
   Avatar,
   Badge,
   Box,
+  Button,
   Card,
   Code,
   Container,
@@ -17,6 +18,8 @@ import {
   TabNav,
   Text,
   Theme,
+  IconButton,
+  Heading,
 } from '@radix-ui/themes';
 import { NextThemeProvider } from '../next-theme-provider';
 import NextLink from 'next/link';
@@ -133,6 +136,33 @@ export default function Test() {
                     <Container asChild>
                       <section>Container as child</section>
                     </Container>
+
+                    <Flex direction="column" gap="2">
+                      <Flex gap="2">
+                        <Button asChild>
+                          <button>Button as child</button>
+                        </Button>
+                        <Button asChild>
+                          <a href="#">Button as child (link)</a>
+                        </Button>
+                      </Flex>
+                      <Flex gap="2">
+                        <Button asChild disabled>
+                          <button>Button as child</button>
+                        </Button>
+                        <Button asChild disabled>
+                          <a href="#">Button as child (link)</a>
+                        </Button>
+                      </Flex>
+                      <Flex gap="2">
+                        <Button asChild loading>
+                          <button>Button as child</button>
+                        </Button>
+                        <Button asChild loading>
+                          <a href="#">Button as child (link)</a>
+                        </Button>
+                      </Flex>
+                    </Flex>
                   </Flex>
                 </Section>
               </Container>

packages/radix-ui-themes/src/components/_internal/base-button.tsx

@@ -1,3 +1,5 @@
+'use client';
+
 import * as React from 'react';
 import classNames from 'classnames';
 import { Slot } from 'radix-ui';
@@ -32,6 +34,25 @@ const BaseButton = React.forwardRef<BaseButtonElement, BaseButtonProps>((props,
     ...baseButtonProps
   } = extractProps(props, baseButtonPropDefs, marginPropDefs);
   const Comp = asChild ? Slot.Root : 'button';
+  let child = children;
+  if (props.loading) {
+    // Loading buttons will wrap the contents of the button for hiding them
+    // visually while retaining the button's size. This does not work with the
+    // Radix Slot since the slot root expects the slottable content to be one of
+    // its direct descendants. To get around this we need to clone the child
+    // with its wrapped inner children.
+    if (asChild && React.isValidElement(children)) {
+      const props = children.props as { children?: React.ReactNode };
+      const childNode = props.children;
+      child = React.cloneElement<any>(children, {
+        ...props,
+        children: <LoadingButtonContents size={size}>{childNode}</LoadingButtonContents>,
+      });
+    } else {
+      child = <LoadingButtonContents size={size}>{children}</LoadingButtonContents>;
+    }
+  }
+
   return (
     <Comp
       // The `data-disabled` attribute enables correct styles when doing `<Button asChild disabled>`
@@ -43,33 +64,40 @@ const BaseButton = React.forwardRef<BaseButtonElement, BaseButtonProps>((props,
       className={classNames('rt-reset', 'rt-BaseButton', className)}
       disabled={disabled}
     >
-      {props.loading ? (
-        <>
-          {/**
-           * We need a wrapper to set `visibility: hidden` to hide the button content whilst we show the `Spinner`.
-           * The button is a flex container with a `gap`, so we use `display: contents` to ensure the correct flex layout.
-           *
-           * However, `display: contents` removes the content from the accessibility tree in some browsers,
-           * so we force remove it with `aria-hidden` and re-add it in the tree with `VisuallyHidden`
-           */}
-          <span style={{ display: 'contents', visibility: 'hidden' }} aria-hidden>
-            {children}
-          </span>
-          <VisuallyHidden>{children}</VisuallyHidden>
-
-          <Flex asChild align="center" justify="center" position="absolute" inset="0">
-            <span>
-              <Spinner size={mapResponsiveProp(size, mapButtonSizeToSpinnerSize)} />
-            </span>
-          </Flex>
-        </>
-      ) : (
-        children
-      )}
+      {child}
     </Comp>
   );
 });
 BaseButton.displayName = 'BaseButton';
 
+const LoadingButtonContents: React.FC<{ children: React.ReactNode; size: any }> = ({
+  children,
+  size,
+}) => {
+  return (
+    <>
+      {/*
+       * We need a wrapper to set `visibility: hidden` to hide the button content
+       * whilst we show the `Spinner`. The button is a flex container with a `gap`,
+       * so we use `display: contents` to ensure the correct flex layout.
+       *
+       * However, `display: contents` removes the content from the accessibility
+       * tree in some browsers, so we force remove it with `aria-hidden` and
+       * re-add it in the tree with `VisuallyHidden`
+       */}
+      <span style={{ display: 'contents', visibility: 'hidden' }} aria-hidden>
+        {children}
+      </span>
+      <VisuallyHidden>{children}</VisuallyHidden>
+      <Flex asChild align="center" justify="center" position="absolute" inset="0">
+        <span>
+          <Spinner size={mapResponsiveProp(size, mapButtonSizeToSpinnerSize)} />
+        </span>
+      </Flex>
+    </>
+  );
+};
+LoadingButtonContents.displayName = 'LoadingButtonContents';
+
 export { BaseButton };
 export type { BaseButtonProps };