From 006d2a7ffc8bf093f9ffcc35f419dd0ab381274f Mon Sep 17 00:00:00 2001
From: Lu Wilson <l2wilson94@gmail.com>
Date: Thu, 25 Jan 2024 14:43:44 +0000
Subject: [PATCH] [draft] Keep editor focus after losing focus of an action
 button (#2630)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../tldraw/src/lib/ui/components/DuplicateButton.tsx  |  3 +++
 packages/tldraw/src/lib/ui/components/RedoButton.tsx  |  4 +++-
 packages/tldraw/src/lib/ui/components/TrashButton.tsx |  3 +++
 packages/tldraw/src/lib/ui/components/UndoButton.tsx  |  4 +++-
 .../src/lib/ui/components/primitives/Button.tsx       | 11 +++++++++++
 5 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/packages/tldraw/src/lib/ui/components/DuplicateButton.tsx b/packages/tldraw/src/lib/ui/components/DuplicateButton.tsx
index 4c69a9c59..81cadef13 100644
--- a/packages/tldraw/src/lib/ui/components/DuplicateButton.tsx
+++ b/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}
 		/>
 	)
 })
diff --git a/packages/tldraw/src/lib/ui/components/RedoButton.tsx b/packages/tldraw/src/lib/ui/components/RedoButton.tsx
index 2ff351ef6..d776fcd2f 100644
--- a/packages/tldraw/src/lib/ui/components/RedoButton.tsx
+++ b/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}
 		/>
 	)
 })
diff --git a/packages/tldraw/src/lib/ui/components/TrashButton.tsx b/packages/tldraw/src/lib/ui/components/TrashButton.tsx
index adc3b83cc..fcf2ada3d 100644
--- a/packages/tldraw/src/lib/ui/components/TrashButton.tsx
+++ b/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}
 		/>
 	)
 })
diff --git a/packages/tldraw/src/lib/ui/components/UndoButton.tsx b/packages/tldraw/src/lib/ui/components/UndoButton.tsx
index 2fc6e211d..39fca9b58 100644
--- a/packages/tldraw/src/lib/ui/components/UndoButton.tsx
+++ b/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}
 		/>
 	)
 })
diff --git a/packages/tldraw/src/lib/ui/components/primitives/Button.tsx b/packages/tldraw/src/lib/ui/components/primitives/Button.tsx
index 7f714e6ff..b3d405f03 100644
--- a/packages/tldraw/src/lib/ui/components/primitives/Button.tsx
+++ b/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)}