TextInputWithTokens bugfix: focus next token when removing first token with keyboard

.changeset/eighty-beds-sneeze.md

@@ -0,0 +1,5 @@
+---
+'@primer/components': minor
+---
+
+fixes a bug in TextInputWithToken where the next token would not receive focus after the user deleted the first token using the keyboard

src/TextInputWithTokens.tsx

@@ -11,6 +11,7 @@ import {useProvidedRefOrCreate} from './hooks'
 import UnstyledTextInput from './_UnstyledTextInput'
 import TextInputWrapper from './_TextInputWrapper'
 import Box from './Box'
+import {isFocusable} from './utils/iterateFocusableElements'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type AnyReactComponent = React.ComponentType<any>
@@ -117,10 +118,24 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
   const handleTokenRemove = (tokenId: string | number) => {
     onTokenRemove(tokenId)
 
-    if (selectedTokenIndex) {
-      const nextElementToFocus = containerRef.current?.children[selectedTokenIndex] as HTMLElement
-      nextElementToFocus.focus()
-    }
+    // HACK: wait a tick for the the token node to be removed from the DOM
+    setTimeout(() => {
+      const nextElementToFocus = containerRef.current?.children[selectedTokenIndex || 0] as HTMLElement | undefined
+
+      // when removing the first token by keying "Backspace" or "Delete",
+      // `nextFocusableElement` is the div that wraps the input
+      const firstFocusable =
+        nextElementToFocus && isFocusable(nextElementToFocus)
+          ? nextElementToFocus
+          : (Array.from(containerRef.current?.children || []) as HTMLElement[]).find(el => isFocusable(el))
+
+      if (firstFocusable) {
+        firstFocusable.focus()
+      } else {
+        // if there are no tokens left, focus the input
+        ref.current?.focus()
+      }
+    }, 0)
   }
 
   const handleTokenFocus: (tokenIndex: number) => FocusEventHandler = tokenIndex => () => {

src/__tests__/TextInputWithTokens.test.tsx

@@ -238,6 +238,44 @@ describe('TextInputWithTokens', () => {
     expect(onRemoveMock).toHaveBeenCalledWith(mockTokens[4].id)
   })
 
+  it('moves focus to the next token when removing the first token', () => {
+    jest.useFakeTimers()
+    const onRemoveMock = jest.fn()
+    const {getByText} = HTMLRender(
+      <TextInputWithTokens tokens={[...mockTokens].slice(0, 2)} onTokenRemove={onRemoveMock} />
+    )
+    const tokenNode = getByText(mockTokens[0].text)
+
+    fireEvent.focus(tokenNode)
+    fireEvent.keyDown(tokenNode, {key: 'Backspace'})
+
+    jest.runAllTimers()
+    setTimeout(() => {
+      expect(document.activeElement?.textContent).toBe(mockTokens[1].text)
+    }, 0)
+
+    jest.useRealTimers()
+  })
+
+  it('moves focus to the input when the last token is removed', () => {
+    jest.useFakeTimers()
+    const onRemoveMock = jest.fn()
+    const {getByText, getByLabelText} = HTMLRender(
+      <LabelledTextInputWithTokens tokens={[mockTokens[0]]} onTokenRemove={onRemoveMock} />
+    )
+    const tokenNode = getByText(mockTokens[0].text)
+    const inputNode = getByLabelText('Tokens')
+
+    fireEvent.focus(tokenNode)
+    fireEvent.keyDown(tokenNode, {key: 'Backspace'})
+
+    jest.runAllTimers()
+    setTimeout(() => {
+      expect(document.activeElement?.id).toBe(inputNode.id)
+    }, 0)
+    jest.useRealTimers()
+  })
+
   it('calls onKeyDown', () => {
     const onRemoveMock = jest.fn()
     const onKeyDownMock = jest.fn()

src/stories/TextInputWithTokens.stories.tsx

@@ -47,7 +47,7 @@ const mockTokens = [
 ]
 
 export const Default = () => {
-  const [tokens, setTokens] = useState([...mockTokens].slice(0, 2))
+  const [tokens, setTokens] = useState([...mockTokens].slice(0, 3))
   const onTokenRemove: (tokenId: string | number) => void = tokenId => {
     setTokens(tokens.filter(token => token.id !== tokenId))
   }