fix: Focus lock uses correct direction when initial focus is outside the lock

pages/focus-lock.page.tsx

@@ -2,28 +2,48 @@
 // SPDX-License-Identifier: Apache-2.0
 import React, { useRef, useState } from 'react';
 
-import Box from '~components/box';
-import Button from '~components/button';
-import FormField from '~components/form-field';
-import Input from '~components/input';
+import { Button, FormField, Input, SpaceBetween, Tabs } from '~components';
 import FocusLock, { FocusLockRef } from '~components/internal/components/focus-lock';
-import SpaceBetween from '~components/space-between';
+
+import { SimplePage } from './app/templates';
 
 export default function FocusLockPage() {
   const ref = useRef<FocusLockRef>(null);
   const [text, setText] = useState('');
 
   return (
-    <Box padding="l">
-      <h1>FocusLock component</h1>
-      <FocusLock ref={ref} autoFocus={true}>
-        <SpaceBetween size="m">
-          <FormField label="First input">
-            <Input value={text} onChange={event => setText(event.detail.value)} />
-          </FormField>
-          <Button onClick={() => ref.current?.focusFirst()}>Focus first</Button>
-        </SpaceBetween>
-      </FocusLock>
-    </Box>
+    <SimplePage title="FocusLock component">
+      <Tabs
+        tabs={[
+          {
+            id: 'with-auto-focus',
+            label: 'With auto focus',
+            content: (
+              <FocusLock ref={ref} autoFocus={true}>
+                <SpaceBetween size="m">
+                  <FormField label="First input">
+                    <Input value={text} onChange={event => setText(event.detail.value)} />
+                  </FormField>
+                  <Button onClick={() => ref.current?.focusFirst()}>Focus first</Button>
+                </SpaceBetween>
+              </FocusLock>
+            ),
+          },
+          {
+            id: 'without-auto-focus',
+            label: 'Without auto focus',
+            content: (
+              <FocusLock ref={ref}>
+                <SpaceBetween size="s" direction="horizontal">
+                  <Button>1</Button>
+                  <Button>2</Button>
+                  <Button>3</Button>
+                </SpaceBetween>
+              </FocusLock>
+            ),
+          },
+        ]}
+      />
+    </SimplePage>
   );
 }

src/internal/components/focus-lock/__tests__/focus-lock.test.tsx

@@ -1,20 +1,22 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 import React from 'react';
-import { render } from '@testing-library/react';
+import { fireEvent, render } from '@testing-library/react';
 
 import FocusLock, { FocusLockProps } from '../../../../../lib/components/internal/components/focus-lock';
 
 type TestFixtureProps = Omit<FocusLockProps, 'children'> & { unmount?: boolean };
 function TestFixture({ unmount = false, ...props }: TestFixtureProps) {
   return (
     <>
-      <button id="initial" />
+      <button id="before" />
       {!unmount && (
         <FocusLock {...props}>
-          <button id="inner" />
+          <button id="first" />
+          <button id="last" />
         </FocusLock>
       )}
+      <button id="after" />
     </>
   );
 }
@@ -23,13 +25,28 @@ function renderFocusLock(props: TestFixtureProps = {}) {
   const { rerender } = render(<TestFixture {...props} />);
   return {
     focusInitial() {
-      document.querySelector<HTMLButtonElement>('#initial')!.focus();
+      document.querySelector<HTMLButtonElement>('#before')!.focus();
     },
     getInitial() {
-      return document.querySelector<HTMLButtonElement>('#initial');
+      return document.querySelector<HTMLButtonElement>('#before');
     },
     getInner() {
-      return document.querySelector<HTMLButtonElement>('#inner');
+      return document.querySelector<HTMLButtonElement>('#first');
+    },
+    getBefore() {
+      return document.querySelector<HTMLButtonElement>('#before')!;
+    },
+    getAfter() {
+      return document.querySelector<HTMLButtonElement>('#after')!;
+    },
+    getFirst() {
+      return document.querySelector<HTMLButtonElement>('#first')!;
+    },
+    getLast() {
+      return document.querySelector<HTMLButtonElement>('#last')!;
+    },
+    getTabTraps() {
+      return document.querySelectorAll<HTMLDivElement>('[tabindex="0"]:not(button)');
     },
     rerender(props: TestFixtureProps = {}) {
       rerender(<TestFixture {...props} />);
@@ -100,3 +117,61 @@ describe('Focus lock', () => {
     });
   });
 });
+
+describe('Focus lock — tab trap direction', () => {
+  function focusTrap(trap: HTMLElement, relatedTarget: HTMLElement | null) {
+    fireEvent.focus(trap, { relatedTarget });
+  }
+
+  it('wraps to last element when tabbing backward from first element (hits first trap)', () => {
+    const { getFirst, getLast, getTabTraps } = renderFocusLock();
+    const firstTrap = getTabTraps()[0];
+
+    // Simulate: focus was on first inner element, user tabs backward → hits first trap
+    focusTrap(firstTrap, getFirst());
+
+    expect(getLast()).toHaveFocus();
+  });
+
+  it('wraps to first element when tabbing forward from last element (hits last trap)', () => {
+    const { getFirst, getLast, getTabTraps } = renderFocusLock();
+    const traps = getTabTraps();
+    const lastTrap = traps[traps.length - 1];
+
+    // Simulate: focus was on last inner element, user tabs forward → hits last trap
+    focusTrap(lastTrap, getLast());
+
+    expect(getFirst()).toHaveFocus();
+  });
+
+  it('focuses first element when entering from before (hits first trap from outside)', () => {
+    const { getBefore, getFirst, getTabTraps } = renderFocusLock();
+    const firstTrap = getTabTraps()[0];
+
+    // Simulate: focus was on element before the lock, user tabs forward → hits first trap
+    focusTrap(firstTrap, getBefore());
+
+    expect(getFirst()).toHaveFocus();
+  });
+
+  it('focuses last element when entering from after (hits last trap from outside)', () => {
+    const { getAfter, getLast, getTabTraps } = renderFocusLock();
+    const traps = getTabTraps();
+    const lastTrap = traps[traps.length - 1];
+
+    // Simulate: focus was on element after the lock, user tabs backward → hits last trap
+    focusTrap(lastTrap, getAfter());
+
+    expect(getLast()).toHaveFocus();
+  });
+
+  it('does not trap focus when disabled', () => {
+    const { getFirst, getLast } = renderFocusLock({ disabled: true });
+    // When disabled, tab traps have tabIndex=-1 and won't receive focus naturally.
+    // Verify that inner elements are not affected.
+    getFirst().focus();
+    expect(getFirst()).toHaveFocus();
+    getLast().focus();
+    expect(getLast()).toHaveFocus();
+  });
+});

src/internal/components/focus-lock/index.tsx

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 import React, { useCallback, useEffect, useImperativeHandle, useRef, useState } from 'react';
 
+import { nodeBelongs } from '@cloudscape-design/component-toolkit/dom';
 import { useMergeRefs } from '@cloudscape-design/component-toolkit/internal';
 
 import TabTrap from '../tab-trap/index';
@@ -82,13 +83,14 @@ function FocusLock(
   useImperativeHandle(ref, () => ({ focusFirst }));
   const mergedRef = useMergeRefs(containerRef, restoreFocusHandler);
 
+  const isFocusInside = (event: React.FocusEvent) => nodeBelongs(containerRef.current, event.relatedTarget);
   return (
     <>
-      <TabTrap disabled={disabled} focusNextCallback={focusLast} />
+      <TabTrap disabled={disabled} focusNextCallback={event => (isFocusInside(event) ? focusLast() : focusFirst())} />
       <div className={className} ref={mergedRef}>
         {children}
       </div>
-      <TabTrap disabled={disabled} focusNextCallback={focusFirst} />
+      <TabTrap disabled={disabled} focusNextCallback={event => (isFocusInside(event) ? focusFirst() : focusLast())} />
     </>
   );
 }