From 848fc8be1f6f63e64eb8bfa06daf3e5f249ae951 Mon Sep 17 00:00:00 2001
From: Yihui Liao <44729383+yihuiliao@users.noreply.github.com>
Date: Mon, 8 Dec 2025 16:52:16 -0800
Subject: [PATCH] fix: update action button to use grid area to position
 notification badge (#9303)

* Revert "Revert "fix: update action button to use grid area to position notification badge (#8409)" (#8552)"

This reverts commit 0cdf701003aa586d820221591bba6e137d7db97d.

* add chromatic story

* fix lint
---
 .../s2/chromatic/ActionButton.stories.tsx     | 14 ++++++++-
 .../@react-spectrum/s2/src/ActionButton.tsx   | 30 ++++++++++++-------
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/packages/@react-spectrum/s2/chromatic/ActionButton.stories.tsx b/packages/@react-spectrum/s2/chromatic/ActionButton.stories.tsx
index 9aa0e2149eb..da80c53f128 100644
--- a/packages/@react-spectrum/s2/chromatic/ActionButton.stories.tsx
+++ b/packages/@react-spectrum/s2/chromatic/ActionButton.stories.tsx
@@ -10,7 +10,9 @@
  * governing permissions and limitations under the License.
  */
 
-import {ActionButton, ActionButtonProps, Avatar, Text} from '../src';
+import {ActionButton, ActionButtonProps, Avatar, NotificationBadge, Text} from '../src';
+import BellIcon from '../s2wf-icons/S2_Icon_Bell_20_N.svg';
+import CommentIcon from '../s2wf-icons/S2_Icon_Comment_20_N.svg';
 import {Fonts, NotificationBadges, UnsafeClassName} from '../stories/ActionButton.stories';
 import {generatePowerset} from '@react-spectrum/story-utils';
 import type {Meta, StoryObj} from '@storybook/react';
@@ -100,3 +102,13 @@ export const AvatarOnly: ActionButtonStory = {
 };
 
 export {Fonts, UnsafeClassName, NotificationBadges};
+
+export const NotificationBadgesCustomWidth: ActionButtonStory = {
+  render: (args) => (
+    <div className={style({display: 'flex', flexDirection: 'column', gap: 8})}>
+      <ActionButton aria-label="Messages has new activity" styles={style({width: 200})} {...args}><CommentIcon /><NotificationBadge /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><BellIcon /><NotificationBadge value={10} /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><CommentIcon /><Text>Messages</Text><NotificationBadge value={5} /></ActionButton>
+      <ActionButton styles={style({width: 200})} {...args}><Text>Notifications</Text><NotificationBadge value={105} /></ActionButton>
+    </div>)
+};
diff --git a/packages/@react-spectrum/s2/src/ActionButton.tsx b/packages/@react-spectrum/s2/src/ActionButton.tsx
index 9d28e68100f..be70cb4396f 100644
--- a/packages/@react-spectrum/s2/src/ActionButton.tsx
+++ b/packages/@react-spectrum/s2/src/ActionButton.tsx
@@ -73,7 +73,7 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
   ...focusRing(),
   ...staticColor(),
   ...controlStyle,
-  position: 'relative',
+  display: 'grid',
   justifyContent: 'center',
   flexShrink: {
     default: 1,
@@ -86,9 +86,21 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
     isJustified: 0
   },
   fontWeight: 'medium',
+  width: 'fit',
   userSelect: 'none',
   transition: 'default',
   forcedColorAdjust: 'none',
+  position: 'relative',
+  gridTemplateAreas: {
+    default: ['icon text'],
+    [iconOnly]: ['icon'],
+    [textOnly]: ['text']
+  },
+  gridTemplateColumns: {
+    default: ['auto', 'auto'],
+    [iconOnly]: ['auto'],
+    [textOnly]: ['auto']
+  },
   backgroundColor: {
     default: {
       ...baseColor('gray-100'),
@@ -230,8 +242,7 @@ export const btnStyles = style<ButtonRenderProps & ActionButtonStyleProps & Togg
   '--badgePosition': {
     type: 'width',
     value: {
-      default: 'calc(self(paddingStart) + var(--iconWidth))',
-      [iconOnly]: 'calc(self(minWidth)/2 + var(--iconWidth)/2)',
+      default: '--iconWidth',
       [textOnly]: 'full'
     }
   },
@@ -311,7 +322,7 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
               [SkeletonContext, null],
               [TextContext, {styles:
                 style({
-                  order: 1,
+                  gridArea: 'text',
                   truncate: true,
                   visibility: {
                     isProgressVisible: 'hidden'
@@ -319,11 +330,10 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
                 })({isProgressVisible})
               }],
               [IconContext, {
-                render: centerBaseline({slot: 'icon', styles: style({order: 0})}),
+                render: centerBaseline({slot: 'icon', styles: style({gridArea: 'icon'})}),
                 styles: style({
                   size: fontRelative(20),
                   marginStart: '--iconMargin',
-                  flexShrink: 0,
                   visibility: {
                     isProgressVisible: 'hidden'
                   }
@@ -333,8 +343,7 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
                 size: avatarSize[size],
                 styles: style({
                   marginStart: '--iconMargin',
-                  flexShrink: 0,
-                  order: 0
+                  gridArea: 'icon'
                 })
               }],
               [ImageContext, {
@@ -351,9 +360,10 @@ export const ActionButton = forwardRef(function ActionButton(props: ActionButton
                 styles: style({
                   position: 'absolute',
                   top: '--badgeTop',
-                  insetStart: '--badgePosition',
                   marginTop: 'calc((self(height) * -1)/2)',
-                  marginStart: 'calc((self(height) * -1)/2)',
+                  marginStart: 'calc(var(--iconMargin) * 2 + (self(height) * -1)/4)',
+                  gridColumnStart: 1,
+                  insetStart: '--badgePosition',
                   visibility: {
                     isProgressVisible: 'hidden'
                   }