Move logic from useDefinition to useBg

Signed-off-by: Charles de Dreuille <charles.dedreuille@gmail.com>

packages/ui/report.api.md

@@ -2181,13 +2181,10 @@ export const TooltipTrigger: (
 export const useBg: (options?: UseBgOptions) => BgContextValue;
 
 // @public (undocumented)
-export type UseBgOptions =
-  | {
-      bg: Responsive<Bg> | undefined;
-    }
-  | {
-      leaf: true;
-    };
+export interface UseBgOptions {
+  bg?: Responsive<Bg>;
+  mode: 'container' | 'leaf';
+}
 
 // @public (undocumented)
 export const useBreakpoint: () => {

packages/ui/src/components/Flex/Flex.tsx

@@ -25,7 +25,7 @@ import { BgProvider, useBg } from '../../hooks/useBg';
 /** @public */
 export const Flex = forwardRef<HTMLDivElement, FlexProps>((props, ref) => {
   // Resolve the bg this Flex creates for its children
-  const { bg: resolvedBg } = useBg({ bg: props.bg });
+  const { bg: resolvedBg } = useBg({ mode: 'container', bg: props.bg });
 
   const { classNames, dataAttributes, utilityClasses, style, cleanedProps } =
     useStyles(FlexDefinition, {

packages/ui/src/components/Grid/Grid.tsx

@@ -24,7 +24,7 @@ import { BgProvider, useBg } from '../../hooks/useBg';
 
 const GridRoot = forwardRef<HTMLDivElement, GridProps>((props, ref) => {
   // Resolve the bg this Grid creates for its children
-  const { bg: resolvedBg } = useBg({ bg: props.bg });
+  const { bg: resolvedBg } = useBg({ mode: 'container', bg: props.bg });
 
   const { classNames, dataAttributes, utilityClasses, style, cleanedProps } =
     useStyles(GridDefinition, {
@@ -60,7 +60,7 @@ const GridRoot = forwardRef<HTMLDivElement, GridProps>((props, ref) => {
 
 const GridItem = forwardRef<HTMLDivElement, GridItemProps>((props, ref) => {
   // Resolve the bg this GridItem creates for its children
-  const { bg: resolvedBg } = useBg({ bg: props.bg });
+  const { bg: resolvedBg } = useBg({ mode: 'container', bg: props.bg });
 
   const { classNames, dataAttributes, utilityClasses, style, cleanedProps } =
     useStyles(GridItemDefinition, {

packages/ui/src/components/ToggleButton/ToggleButton.tsx

@@ -36,7 +36,7 @@ export const ToggleButton = forwardRef(
 
     const { children, className, iconStart, iconEnd, ...rest } = cleanedProps;
 
-    const { bg } = useBg({ leaf: true });
+    const { bg } = useBg({ mode: 'leaf' });
 
     return (
       <AriaToggleButton

packages/ui/src/hooks/useBg.tsx

@@ -33,21 +33,22 @@ export interface BgProviderProps {
 }
 
 /** @public */
-export type UseBgOptions =
-  | {
-      /**
-       * Container mode: the explicit bg value from the component's prop.
-       * If undefined, the container auto-increments from parent context.
-       */
-      bg: Responsive<Bg> | undefined;
-    }
-  | {
-      /**
-       * Leaf mode: automatically reads bg from context and increments by 1.
-       * No prop is needed on the component.
-       */
-      leaf: true;
-    };
+export interface UseBgOptions {
+  /**
+   * The bg mode of the component.
+   *
+   * - `'container'` — for components like Box, Card, Flex that establish bg context.
+   *   If `bg` prop is provided, uses that value. Otherwise auto-increments from parent.
+   * - `'leaf'` — for components like Button that consume bg context.
+   *   Always auto-increments from parent context. The `bg` prop is ignored.
+   */
+  mode: 'container' | 'leaf';
+  /**
+   * The explicit bg value from the component's prop.
+   * Only used in container mode — leaf mode ignores this.
+   */
+  bg?: Responsive<Bg>;
+}
 
 const BgContext = createVersionedContext<{
   1: BgContextValue;
@@ -140,12 +141,13 @@ export const BgProvider = ({ bg, children }: BgProviderProps) => {
 /**
  * Hook to access and resolve the current bg context.
  *
- * Supports two modes:
- * - **Container mode** (`{ bg }`) - for components like Box that establish bg context.
- *   If bg prop is provided, uses that value. If bg is undefined but parent context exists,
- *   auto-increments from parent.
- * - **Leaf mode** (`{ leaf: true }`) - for components like Button that consume bg context.
- *   Always auto-increments from parent context. No prop needed.
+ * All bg resolution logic lives here — callers only need to specify the mode
+ * and (for containers) the explicit bg prop value.
+ *
+ * - **Container mode** — uses explicit `bg` if provided, otherwise auto-increments
+ *   from parent context. Caps at `neutral-4`.
+ * - **Leaf mode** — always auto-increments from parent context. No prop needed.
+ * - **No options** — returns the raw context value without resolution.
  *
  * @param options - Configuration for bg resolution
  * @public
@@ -158,7 +160,7 @@ export const useBg = (options?: UseBgOptions): BgContextValue => {
     return context;
   }
 
-  if ('leaf' in options) {
+  if (options.mode === 'leaf') {
     return { bg: resolveBgForLeaf(context.bg) };
   }
 

packages/ui/src/hooks/useDefinition/useDefinition.tsx

@@ -17,7 +17,7 @@
 import { ReactNode } from 'react';
 import clsx from 'clsx';
 import { useBreakpoint } from '../useBreakpoint';
-import { useBg, BgProvider, UseBgOptions } from '../useBg';
+import { useBg, BgProvider } from '../useBg';
 import { resolveResponsiveValue, processUtilityProps } from './helpers';
 import type {
   ComponentConfig,
@@ -36,14 +36,9 @@ export function useDefinition<
 ): UseDefinitionResult<D, P> {
   const { breakpoint } = useBreakpoint();
 
-  const bgOptions: UseBgOptions | undefined =
-    definition.bg === 'container'
-      ? { bg: props.bg }
-      : definition.bg === 'leaf'
-      ? { leaf: true }
-      : undefined;
-
-  const { bg: resolvedBg } = useBg(bgOptions);
+  const { bg: resolvedBg } = useBg(
+    definition.bg ? { mode: definition.bg, bg: props.bg } : undefined,
+  );
 
   const ownPropKeys = new Set(Object.keys(definition.propDefs));
   const utilityPropKeys = new Set(definition.utilityProps ?? []);
@@ -77,20 +72,8 @@ export function useDefinition<
     }
   }
 
-  // Set data-bg for bg container components with the resolved value
-  // This covers both explicit bg props and auto-incremented values
-  if (definition.bg === 'container' && resolvedBg !== undefined) {
-    const bgValue =
-      typeof resolvedBg === 'object'
-        ? resolveResponsiveValue(resolvedBg as any, breakpoint)
-        : resolvedBg;
-    if (bgValue !== undefined) {
-      dataAttributes['data-bg'] = String(bgValue);
-    }
-  }
-
-  // Add data-bg for bg leaf components (auto-incremented from context)
-  if (definition.bg === 'leaf' && resolvedBg !== undefined) {
+  // Set data-bg from the resolved bg value (works for both container and leaf)
+  if (definition.bg && resolvedBg !== undefined) {
     const bgValue =
       typeof resolvedBg === 'object'
         ? resolveResponsiveValue(resolvedBg as any, breakpoint)