UI: Fix regression in select close handler focus

code/core/src/components/components/Select/Select.stories.tsx

@@ -1,6 +1,6 @@
 import React from 'react';
 
-import { Button } from 'storybook/internal/components';
+import { Button, Toolbar } from 'storybook/internal/components';
 
 import { LinuxIcon } from '@storybook/icons';
 
@@ -199,6 +199,113 @@ export const WithSiblings = meta.story({
       <Button ariaLabel={false}>After</Button>
     </Row>
   ),
+  play: async ({ canvas, step }) => {
+    const user = userEvent.setup();
+
+    await step('Open select and select an option', async () => {
+      const select = canvas.getByRole('button', { name: /Animal/i });
+      await user.click(select);
+
+      const listbox = await screen.findByRole('listbox');
+      expect(listbox).toBeInTheDocument();
+
+      const option = within(listbox).getByRole('option', { name: 'Frog' });
+      await user.click(option);
+    });
+
+    await step('Tab should land on sibling after select', async () => {
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+
+      await user.tab();
+
+      const afterButton = canvas.getByRole('button', { name: 'After' });
+      expect(afterButton).toHaveFocus();
+    });
+
+    await step('Navigate back and reopen select', async () => {
+      await user.tab({ shift: true });
+
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+
+      await user.keyboard('{Enter}');
+
+      const listbox = await screen.findByRole('listbox');
+      expect(listbox).toBeInTheDocument();
+    });
+
+    await step('Escape should return to select trigger', async () => {
+      await user.keyboard('{Escape}');
+
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+    });
+  },
+});
+
+export const WithSiblingsInToolbar = meta.story({
+  name: 'With Siblings in Toolbar',
+  render: (args) => (
+    <Toolbar aria-label="Test toolbar">
+      <Button ariaLabel="Before button">Before</Button>
+      <Select {...args} />
+      <Button ariaLabel="After button">After</Button>
+    </Toolbar>
+  ),
+  play: async ({ canvas, step }) => {
+    const user = userEvent.setup();
+
+    await step('Navigate to select with ArrowRight', async () => {
+      const beforeButton = canvas.getByRole('button', { name: 'Before button' });
+      beforeButton.focus();
+      expect(beforeButton).toHaveFocus();
+
+      await user.keyboard('{ArrowRight}');
+
+      const select = canvas.getByRole('button', { name: /Animal/i });
+      expect(select).toHaveFocus();
+    });
+
+    await step('Open select and select an option', async () => {
+      await user.keyboard('{Enter}');
+
+      const listbox = await screen.findByRole('listbox');
+      expect(listbox).toBeInTheDocument();
+
+      const option = within(listbox).getByRole('option', { name: 'Frog' });
+      await user.click(option);
+    });
+
+    await step('ArrowRight should land on sibling after select', async () => {
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+
+      await user.keyboard('{ArrowRight}');
+
+      const afterButton = canvas.getByRole('button', { name: 'After button' });
+      expect(afterButton).toHaveFocus();
+    });
+
+    await step('Navigate back with ArrowLeft and reopen select', async () => {
+      await user.keyboard('{ArrowLeft}');
+
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+
+      await user.keyboard('{Enter}');
+
+      const listbox = await screen.findByRole('listbox');
+      expect(listbox).toBeInTheDocument();
+    });
+
+    await step('Escape should return to select trigger', async () => {
+      await user.keyboard('{Escape}');
+
+      const select = canvas.getByRole('button', { name: /Frog/i });
+      expect(select).toHaveFocus();
+    });
+  },
 });
 
 export const DefaultOption = meta.story({

code/core/src/components/components/Select/Select.tsx

@@ -213,6 +213,7 @@ export const Select = forwardRef<HTMLButtonElement, SelectProps>(
     ref
   ) => {
     const [isOpen, setIsOpen] = useState(props.defaultOpen || false);
+    const [shouldRefocusTrigger, setShouldRefocusTrigger] = useState(false);
     const triggerRef = useObjectRef(ref);
 
     const id = useMemo(() => {
@@ -228,8 +229,17 @@ export const Select = forwardRef<HTMLButtonElement, SelectProps>(
 
     const handleClose = useCallback(() => {
       setIsOpen(false);
-      triggerRef.current?.focus();
-    }, [triggerRef]);
+      setShouldRefocusTrigger(true);
+    }, []);
+
+    // We must delay refocusing the trigger because we first need the listbox to close,
+    // and @react-aria/overlays to remove the inert attribute set up by MinimalistPopover.
+    useEffect(() => {
+      if (!otState.isOpen && shouldRefocusTrigger) {
+        triggerRef.current?.focus();
+        setShouldRefocusTrigger(false);
+      }
+    }, [otState.isOpen, shouldRefocusTrigger, triggerRef]);
 
     // The last selected option(s), which will be used by the app.
     const [selectedOptions, setSelectedOptions] = useState<InternalOption[]>(