[draft] Keep editor focus after losing focus of an action button (#2630)

This PR fixes a bug where you could lose focus of the editor, which
caused keyboard shortcuts to stop working.

The problem was this:
- The duplicate button can become disabled while you have it focused.
- This would shove focus back to the document body, and the editor would
lose focus.
- When we disable the button, we should keep focus in the editor
instead.
- This shouldn't interfere with a developer manually handling focus of
the editor themselves.

I applied the same fix to the undo, redo, delete and duplicate buttons.

**Is this is a bit hacky? Not sure if I'm handling those `ref`s
correctly? WDYT?**

![2024-01-25 at 12 14 50 - Gold
Nightingale](https://github.com/tldraw/tldraw/assets/15892272/5ca71f92-45fa-48f6-9039-f6c01c495ce7)


### Change Type

- [x] `patch` — Bug fix

[^1]: publishes a `patch` release, for devDependencies use `internal`
[^2]: will not publish a new version

### Test Plan

1. Create a shape.
2. Select it.
3. Click the duplicate button at the top of the screen.
4. Press the 'd' key.
5. Press the 'a' key.
6. You should have the Arrow tool selected.

- [ ] Unit Tests
- [ ] End to end tests

### Release Notes

- Fixed a bug where keyboard shortcuts could stop working after using an
action button.

packages/tldraw/src/lib/ui/components/DuplicateButton.tsx

@@ -1,4 +1,5 @@
 import { track, useEditor } from '@tldraw/editor'
+import { useRef } from 'react'
 import { useActions } from '../hooks/useActions'
 import { useTranslation } from '../hooks/useTranslation/useTranslation'
 import { Button } from './primitives/Button'
@@ -9,6 +10,7 @@ export const DuplicateButton = track(function DuplicateButton() {
 	const actions = useActions()
 	const msg = useTranslation()
 	const action = actions['duplicate']
+	const ref = useRef<HTMLButtonElement>(null)
 
 	return (
 		<Button
@@ -18,6 +20,7 @@ export const DuplicateButton = track(function DuplicateButton() {
 			disabled={!(editor.isIn('select') && editor.getSelectedShapeIds().length > 0)}
 			title={`${msg(action.label!)} ${kbdStr(action.kbd!)}`}
 			smallIcon
+			ref={ref}
 		/>
 	)
 })

packages/tldraw/src/lib/ui/components/RedoButton.tsx

@@ -1,4 +1,4 @@
-import { memo } from 'react'
+import { memo, useRef } from 'react'
 import { useActions } from '../hooks/useActions'
 import { useCanRedo } from '../hooks/useCanRedo'
 import { useTranslation } from '../hooks/useTranslation/useTranslation'
@@ -11,6 +11,7 @@ export const RedoButton = memo(function RedoButton() {
 	const actions = useActions()
 
 	const redo = actions['redo']
+	const ref = useRef<HTMLButtonElement>(null)
 
 	return (
 		<Button
@@ -21,6 +22,7 @@ export const RedoButton = memo(function RedoButton() {
 			disabled={!canRedo}
 			onClick={() => redo.onSelect('quick-actions')}
 			smallIcon
+			ref={ref}
 		/>
 	)
 })

packages/tldraw/src/lib/ui/components/TrashButton.tsx

@@ -1,4 +1,5 @@
 import { track, useEditor } from '@tldraw/editor'
+import { useRef } from 'react'
 import { useActions } from '../hooks/useActions'
 import { useReadonly } from '../hooks/useReadonly'
 import { useTranslation } from '../hooks/useTranslation/useTranslation'
@@ -12,6 +13,7 @@ export const TrashButton = track(function TrashButton() {
 	const action = actions['delete']
 
 	const isReadonly = useReadonly()
+	const ref = useRef<HTMLButtonElement>(null)
 
 	if (isReadonly) return null
 
@@ -27,6 +29,7 @@ export const TrashButton = track(function TrashButton() {
 			disabled={!(editor.isIn('select') && editor.getSelectedShapeIds().length > 0)}
 			title={`${msg(action.label!)} ${kbdStr(action.kbd!)}`}
 			smallIcon
+			ref={ref}
 		/>
 	)
 })

packages/tldraw/src/lib/ui/components/UndoButton.tsx

@@ -1,4 +1,4 @@
-import { memo } from 'react'
+import { memo, useRef } from 'react'
 import { useActions } from '../hooks/useActions'
 import { useCanUndo } from '../hooks/useCanUndo'
 import { useTranslation } from '../hooks/useTranslation/useTranslation'
@@ -9,6 +9,7 @@ export const UndoButton = memo(function UndoButton() {
 	const msg = useTranslation()
 	const canUndo = useCanUndo()
 	const actions = useActions()
+	const ref = useRef<HTMLButtonElement>(null)
 
 	const undo = actions['undo']
 
@@ -21,6 +22,7 @@ export const UndoButton = memo(function UndoButton() {
 			disabled={!canUndo}
 			onClick={() => undo.onSelect('quick-actions')}
 			smallIcon
+			ref={ref}
 		/>
 	)
 })

packages/tldraw/src/lib/ui/components/primitives/Button.tsx

@@ -1,3 +1,4 @@
+import { useEditor } from '@tldraw/editor'
 import classnames from 'classnames'
 import * as React from 'react'
 import { TLUiTranslationKey } from '../../hooks/useTranslation/TLUiTranslationKey'
@@ -35,18 +36,28 @@ export const Button = React.forwardRef<HTMLButtonElement, TLUiButtonProps>(funct
 		type,
 		children,
 		spinner,
+		disabled,
 		...props
 	},
 	ref
 ) {
 	const msg = useTranslation()
 	const labelStr = label ? msg(label) : ''
+	const editor = useEditor()
+
+	// If the button is getting disabled while it's focused, move focus to the editor
+	// so that the user can continue using keyboard shortcuts
+	const current = (ref as React.MutableRefObject<HTMLButtonElement | null>)?.current
+	if (disabled && current === document.activeElement) {
+		editor.getContainer().focus()
+	}
 
 	return (
 		<button
 			ref={ref}
 			draggable={false}
 			type="button"
+			disabled={disabled}
 			{...props}
 			title={props.title ?? labelStr}
 			className={classnames('tlui-button', `tlui-button__${type}`, props.className)}