From 06fe8484b159e55ab1d5e91583e016ae026f74c7 Mon Sep 17 00:00:00 2001
From: Connor Lanigan <connlani@amazon.com>
Date: Tue, 12 Sep 2023 16:46:32 +0000
Subject: [PATCH] chore: Fix funnel support for modals

---
 src/container/index.tsx                       | 12 ++++++++++--
 src/container/internal.tsx                    | 12 ++++++++----
 .../expandable-section-container.tsx          | 10 ++++++++--
 src/expandable-section/index.tsx              | 19 ++++++++++++++-----
 src/expandable-section/internal.tsx           | 10 ++++++++--
 .../analytics/components/analytics-funnel.tsx | 17 ++++++++++++++---
 src/internal/analytics/hooks/use-funnel.ts    |  3 ++-
 src/modal/internal.tsx                        | 14 ++++++++++----
 8 files changed, 74 insertions(+), 23 deletions(-)

diff --git a/src/container/index.tsx b/src/container/index.tsx
index 2235c80157..b7792800bb 100644
--- a/src/container/index.tsx
+++ b/src/container/index.tsx
@@ -7,6 +7,7 @@ import { getExternalProps } from '../internal/utils/external-props';
 import { applyDisplayName } from '../internal/utils/apply-display-name';
 import useBaseComponent from '../internal/hooks/use-base-component';
 import { AnalyticsFunnelSubStep } from '../internal/analytics/components/analytics-funnel';
+import { useFunnelSubStep } from '../internal/analytics/hooks/use-funnel';
 
 export { ContainerProps };
 
@@ -18,12 +19,14 @@ export default function Container({
 }: ContainerProps) {
   const baseComponentProps = useBaseComponent('Container');
   const externalProps = getExternalProps(props);
+
   return (
     <AnalyticsFunnelSubStep>
-      <InternalContainer
+      <InternalContainerAsSubstep
         variant={variant}
-        disableHeaderPaddings={disableHeaderPaddings}
         disableContentPaddings={disableContentPaddings}
+        disableHeaderPaddings={disableHeaderPaddings}
+        {...props}
         {...externalProps}
         {...baseComponentProps}
       />
@@ -31,4 +34,9 @@ export default function Container({
   );
 }
 
+function InternalContainerAsSubstep(props: ContainerProps) {
+  const { subStepRef, funnelSubStepProps } = useFunnelSubStep();
+  return <InternalContainer {...props} __subStepRef={subStepRef} __funnelSubStepProps={funnelSubStepProps} />;
+}
+
 applyDisplayName(Container, 'Container');
diff --git a/src/container/internal.tsx b/src/container/internal.tsx
index 0a2b6bd63c..26cd45d4bc 100644
--- a/src/container/internal.tsx
+++ b/src/container/internal.tsx
@@ -12,7 +12,7 @@ import { useMergeRefs } from '../internal/hooks/use-merge-refs';
 import { useMobile } from '../internal/hooks/use-mobile';
 import { useVisualRefresh } from '../internal/hooks/use-visual-mode';
 import styles from './styles.css.js';
-import { useFunnelSubStep } from '../internal/analytics/hooks/use-funnel';
+import type { useFunnelSubStep } from '../internal/analytics/hooks/use-funnel';
 
 export interface InternalContainerProps extends Omit<ContainerProps, 'variant'>, InternalBaseComponentProps {
   __stickyHeader?: boolean;
@@ -31,6 +31,9 @@ export interface InternalContainerProps extends Omit<ContainerProps, 'variant'>,
    * * `full-page` – Only for internal use in table, cards and other components
    */
   variant?: ContainerProps['variant'] | 'embedded' | 'full-page' | 'cards';
+
+  __funnelSubStepProps?: ReturnType<typeof useFunnelSubStep>['funnelSubStepProps'];
+  __subStepRef?: ReturnType<typeof useFunnelSubStep>['subStepRef'];
 }
 
 export default function InternalContainer({
@@ -52,6 +55,8 @@ export default function InternalContainer({
   __headerRef,
   __darkHeader = false,
   __disableStickyMobile = true,
+  __funnelSubStepProps,
+  __subStepRef,
   ...restProps
 }: InternalContainerProps) {
   const isMobile = useMobile();
@@ -68,12 +73,11 @@ export default function InternalContainer({
   );
   const { setHasStickyBackground } = useAppLayoutContext();
   const isRefresh = useVisualRefresh();
-  const { subStepRef, funnelSubStepProps } = useFunnelSubStep();
 
   const hasDynamicHeight = isRefresh && variant === 'full-page';
   const overlapElement = useDynamicOverlap({ disabled: !hasDynamicHeight || !__darkHeader });
 
-  const mergedRef = useMergeRefs(rootRef, subStepRef, __internalRootRef);
+  const mergedRef = useMergeRefs(rootRef, __subStepRef, __internalRootRef);
   const headerMergedRef = useMergeRefs(headerRef, overlapElement, __headerRef);
 
   /**
@@ -104,7 +108,7 @@ export default function InternalContainer({
   return (
     <div
       {...baseProps}
-      {...funnelSubStepProps}
+      {...__funnelSubStepProps}
       className={clsx(
         baseProps.className,
         styles.root,
diff --git a/src/expandable-section/expandable-section-container.tsx b/src/expandable-section/expandable-section-container.tsx
index ced421bb14..fd11c6e45c 100644
--- a/src/expandable-section/expandable-section-container.tsx
+++ b/src/expandable-section/expandable-section-container.tsx
@@ -1,17 +1,19 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import InternalContainer from '../container/internal';
+import InternalContainer, { InternalContainerProps } from '../container/internal';
 import React from 'react';
 import { ExpandableSectionProps } from './interfaces';
 import { InternalBaseComponentProps } from '../internal/hooks/use-base-component';
 
-interface ExpandableSectionContainerProps extends InternalBaseComponentProps {
+export interface ExpandableSectionContainerProps extends InternalBaseComponentProps {
   className?: string;
   header: React.ReactNode;
   children?: React.ReactNode;
   variant: ExpandableSectionProps.Variant;
   expanded: boolean | undefined;
   disableContentPaddings: boolean | undefined;
+  __funnelSubStepProps?: InternalContainerProps['__funnelSubStepProps'];
+  __subStepRef?: InternalContainerProps['__subStepRef'];
 }
 
 export const ExpandableSectionContainer = ({
@@ -22,6 +24,8 @@ export const ExpandableSectionContainer = ({
   expanded,
   disableContentPaddings,
   __internalRootRef,
+  __funnelSubStepProps,
+  __subStepRef,
   ...rest
 }: ExpandableSectionContainerProps) => {
   if (variant === 'container' || variant === 'stacked') {
@@ -35,6 +39,8 @@ export const ExpandableSectionContainer = ({
         disableHeaderPaddings={true}
         __hiddenContent={!expanded}
         __internalRootRef={__internalRootRef}
+        __funnelSubStepProps={__funnelSubStepProps}
+        __subStepRef={__subStepRef}
       >
         {children}
       </InternalContainer>
diff --git a/src/expandable-section/index.tsx b/src/expandable-section/index.tsx
index 38dddeb2d2..19d1ec58fe 100644
--- a/src/expandable-section/index.tsx
+++ b/src/expandable-section/index.tsx
@@ -4,22 +4,31 @@ import React from 'react';
 
 import { ExpandableSectionProps } from './interfaces';
 import { applyDisplayName } from '../internal/utils/apply-display-name';
-import InternalExpandableSection from './internal';
+import InternalExpandableSection, { InternalExpandableSectionProps } from './internal';
 import useBaseComponent from '../internal/hooks/use-base-component';
 import { AnalyticsFunnelSubStep } from '../internal/analytics/components/analytics-funnel';
+import { useFunnelSubStep } from '../internal/analytics/hooks/use-funnel';
 
 export { ExpandableSectionProps };
 
 export default function ExpandableSection({ variant = 'default', ...props }: ExpandableSectionProps) {
   const baseComponentProps = useBaseComponent('ExpandableSection');
 
-  const expandableSection = <InternalExpandableSection variant={variant} {...props} {...baseComponentProps} />;
-
   if (variant === 'container' || variant === 'stacked') {
-    return <AnalyticsFunnelSubStep>{expandableSection}</AnalyticsFunnelSubStep>;
+    return (
+      <AnalyticsFunnelSubStep>
+        <InternalExpandableSectionAsSubstep variant={variant} {...props} {...baseComponentProps} />
+      </AnalyticsFunnelSubStep>
+    );
   } else {
-    return expandableSection;
+    return <InternalExpandableSection variant={variant} {...props} {...baseComponentProps} />;
   }
 }
 
+function InternalExpandableSectionAsSubstep(props: InternalExpandableSectionProps) {
+  const { subStepRef, funnelSubStepProps } = useFunnelSubStep();
+
+  return <InternalExpandableSection {...props} __subStepRef={subStepRef} __funnelSubStepProps={funnelSubStepProps} />;
+}
+
 applyDisplayName(ExpandableSection, 'ExpandableSection');
diff --git a/src/expandable-section/internal.tsx b/src/expandable-section/internal.tsx
index 353780fcae..34ac689202 100644
--- a/src/expandable-section/internal.tsx
+++ b/src/expandable-section/internal.tsx
@@ -13,12 +13,14 @@ import { fireNonCancelableEvent } from '../internal/events';
 import { ExpandableSectionProps } from './interfaces';
 
 import styles from './styles.css.js';
-import { ExpandableSectionContainer } from './expandable-section-container';
+import { ExpandableSectionContainer, ExpandableSectionContainerProps } from './expandable-section-container';
 import { ExpandableSectionHeader } from './expandable-section-header';
 import { InternalBaseComponentProps } from '../internal/hooks/use-base-component';
 import { variantSupportsDescription } from './utils';
 
-type InternalExpandableSectionProps = ExpandableSectionProps & InternalBaseComponentProps;
+export type InternalExpandableSectionProps = ExpandableSectionProps &
+  InternalBaseComponentProps &
+  Pick<ExpandableSectionContainerProps, '__funnelSubStepProps' | '__subStepRef'>;
 
 export default function InternalExpandableSection({
   expanded: controlledExpanded,
@@ -36,6 +38,8 @@ export default function InternalExpandableSection({
   disableContentPaddings,
   headerAriaLabel,
   __internalRootRef,
+  __funnelSubStepProps,
+  __subStepRef,
   ...props
 }: InternalExpandableSectionProps) {
   const ref = useRef<HTMLDivElement>(null);
@@ -98,6 +102,8 @@ export default function InternalExpandableSection({
       expanded={expanded}
       className={clsx(baseProps.className, styles.root)}
       variant={variant}
+      __funnelSubStepProps={__funnelSubStepProps}
+      __subStepRef={__subStepRef}
       disableContentPaddings={disableContentPaddings}
       header={
         <ExpandableSectionHeader
diff --git a/src/internal/analytics/components/analytics-funnel.tsx b/src/internal/analytics/components/analytics-funnel.tsx
index 8a775a84e8..97b4802952 100644
--- a/src/internal/analytics/components/analytics-funnel.tsx
+++ b/src/internal/analytics/components/analytics-funnel.tsx
@@ -29,6 +29,7 @@ import {
   getSubStepSelector,
 } from '../selectors';
 import { useDebounceCallback } from '../../hooks/use-debounce-callback';
+import { nodeBelongs } from '../../utils/node-belongs';
 
 export const FUNNEL_VERSION = '1.1';
 
@@ -349,7 +350,7 @@ const InnerAnalyticsFunnelStep = ({ children, stepNumber, stepNameSelector }: An
   );
 };
 interface AnalyticsFunnelSubStepProps {
-  children?: React.ReactNode;
+  children?: React.ReactNode | ((props: FunnelSubStepContextValue) => React.ReactNode);
 }
 
 export const AnalyticsFunnelSubStep = ({ children }: AnalyticsFunnelSubStepProps) => {
@@ -395,6 +396,10 @@ export const AnalyticsFunnelSubStep = ({ children }: AnalyticsFunnelSubStepProps
   const context = isNested ? inheritedContext : newContext;
 
   useEffect(() => {
+    if (isNested || !subStepRef.current) {
+      return;
+    }
+
     const onMouseDown = () => (mousePressed.current = true);
 
     const onMouseUp = async () => {
@@ -412,7 +417,7 @@ export const AnalyticsFunnelSubStep = ({ children }: AnalyticsFunnelSubStepProps
       */
       await new Promise(r => setTimeout(r, 1));
 
-      if (!subStepRef.current || !subStepRef.current.contains(document.activeElement)) {
+      if (!subStepRef.current || !document.activeElement || !nodeBelongs(subStepRef.current, document.activeElement)) {
         isFocusedSubStep.current = false;
 
         /*
@@ -436,7 +441,13 @@ export const AnalyticsFunnelSubStep = ({ children }: AnalyticsFunnelSubStepProps
     subStepNameSelector,
     subStepSelector,
     focusCleanupFunction,
+    isNested,
+    subStepRef,
   ]);
 
-  return <FunnelSubStepContext.Provider value={context}>{children}</FunnelSubStepContext.Provider>;
+  return (
+    <FunnelSubStepContext.Provider value={context}>
+      {typeof children === 'function' ? children(context) : children}
+    </FunnelSubStepContext.Provider>
+  );
 };
diff --git a/src/internal/analytics/hooks/use-funnel.ts b/src/internal/analytics/hooks/use-funnel.ts
index 442c8b06ce..3972437f47 100644
--- a/src/internal/analytics/hooks/use-funnel.ts
+++ b/src/internal/analytics/hooks/use-funnel.ts
@@ -10,6 +10,7 @@ import {
   getSubStepAllSelector,
 } from '../selectors';
 import { FunnelMetrics } from '../';
+import { nodeBelongs } from '../../utils/node-belongs';
 
 /**
  * Custom React Hook to manage and interact with FunnelSubStep.
@@ -121,7 +122,7 @@ export const useFunnelSubStep = () => {
       return;
     }
 
-    if (!subStepRef.current || !subStepRef.current.contains(event.relatedTarget) || !event.relatedTarget) {
+    if (!subStepRef.current || !event.relatedTarget || !nodeBelongs(subStepRef.current, event.relatedTarget)) {
       isFocusedSubStep.current = false;
 
       if (funnelInteractionId && subStepId && funnelState.current !== 'cancelled') {
diff --git a/src/modal/internal.tsx b/src/modal/internal.tsx
index 4cda58cc65..e67f19ffc5 100644
--- a/src/modal/internal.tsx
+++ b/src/modal/internal.tsx
@@ -29,10 +29,14 @@ import { ButtonContext } from '../internal/context/button-context';
 type InternalModalProps = SomeRequired<ModalProps, 'size'> & InternalBaseComponentProps;
 
 export default function InternalModal({ modalRoot, ...rest }: InternalModalProps) {
+  const referrerId = useUniqueId('modal');
+
   return (
-    <Portal container={modalRoot}>
-      <InnerModal {...rest} />
-    </Portal>
+    <div id={referrerId}>
+      <Portal container={modalRoot}>
+        <InnerModal {...rest} referrerId={referrerId} />
+      </Portal>
+    </div>
   );
 }
 
@@ -47,8 +51,9 @@ function InnerModal({
   disableContentPaddings,
   onDismiss,
   __internalRootRef = null,
+  referrerId,
   ...rest
-}: InternalModalProps) {
+}: InternalModalProps & { referrerId: string }) {
   const instanceUniqueId = useUniqueId();
   const headerId = `${rest.id || instanceUniqueId}-header`;
   const lastMouseDownElementRef = useRef<HTMLElement | null>(null);
@@ -128,6 +133,7 @@ function InnerModal({
           onClick={onOverlayClick}
           ref={mergedRef}
           style={footerHeight ? { scrollPaddingBottom: footerHeight } : undefined}
+          data-awsui-referrer-id={referrerId}
         >
           <FocusLock disabled={!visible} autoFocus={true} restoreFocus={true} className={styles['focus-lock']}>
             <div