From 30657cd42752d27cf9d101e566fcde163bdfc0d6 Mon Sep 17 00:00:00 2001
From: Tim <timothychang94@gmail.com>
Date: Fri, 4 Nov 2022 16:52:30 -0400
Subject: [PATCH] chore: ensure input focus when kbar is open (#251)

---
 src/InternalEvents.tsx | 26 +++++++++++++++++++++++---
 src/KBarSearch.tsx     |  6 ++----
 src/types.ts           |  2 ++
 src/useKBar.tsx        |  1 -
 src/useStore.tsx       | 13 +++++++++++++
 src/utils.ts           |  1 +
 6 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/InternalEvents.tsx b/src/InternalEvents.tsx
index 1d1252c..e48436b 100644
--- a/src/InternalEvents.tsx
+++ b/src/InternalEvents.tsx
@@ -169,11 +169,14 @@ function wrap(handler: (event: KeyboardEvent) => void) {
  * performs actions for patterns that match the user defined `shortcut`.
  */
 function useShortcuts() {
-  const { actions, query, options } = useKBar((state) => ({
+  const { actions, query, open, options } = useKBar((state) => ({
     actions: state.actions,
+    open: state.visualState === VisualState.showing,
   }));
 
   React.useEffect(() => {
+    if (open) return;
+
     const actionsList = Object.keys(actions).map((key) => actions[key]);
 
     let actionsWithShortcuts: ActionImpl[] = [];
@@ -214,7 +217,7 @@ function useShortcuts() {
     return () => {
       unsubscribe();
     };
-  }, [actions, options.callbacks, query]);
+  }, [actions, open, options.callbacks, query]);
 }
 
 /**
@@ -222,7 +225,7 @@ function useShortcuts() {
  * in focus prior to kbar being triggered.
  */
 function useFocusHandler() {
-  const { isShowing } = useKBar((state) => ({
+  const { isShowing, query } = useKBar((state) => ({
     isShowing:
       state.visualState === VisualState.showing ||
       state.visualState === VisualState.animatingIn,
@@ -249,4 +252,21 @@ function useFocusHandler() {
       activeElement.focus();
     }
   }, [isShowing]);
+
+  // When focus is blurred from the search input while kbar is still
+  // open, any keystroke should set focus back to the search input.
+  React.useEffect(() => {
+    function handler(event: KeyboardEvent) {
+      const input = query.getInput();
+      if (event.target !== input) {
+        input.focus();
+      }
+    }
+    if (isShowing) {
+      window.addEventListener("keydown", handler);
+      return () => {
+        window.removeEventListener("keydown", handler);
+      };
+    }
+  }, [isShowing, query]);
 }
diff --git a/src/KBarSearch.tsx b/src/KBarSearch.tsx
index 1322753..e1c1a4d 100644
--- a/src/KBarSearch.tsx
+++ b/src/KBarSearch.tsx
@@ -26,13 +26,11 @@ export function KBarSearch(
     showing: state.visualState === VisualState.showing,
   }));
 
-  const ownRef = React.useRef<HTMLInputElement>(null);
-
   const { defaultPlaceholder, ...rest } = props;
 
   React.useEffect(() => {
     query.setSearch("");
-    ownRef.current!.focus();
+    query.getInput().focus();
     return () => query.setSearch("");
   }, [currentRootActionId, query]);
 
@@ -46,7 +44,7 @@ export function KBarSearch(
   return (
     <input
       {...rest}
-      ref={ownRef}
+      ref={query.inputRefSetter}
       autoFocus
       autoComplete="off"
       role="combobox"
diff --git a/src/types.ts b/src/types.ts
index bf3f98c..943af69 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -91,6 +91,8 @@ export interface KBarQuery {
   registerActions: (actions: Action[]) => () => void;
   toggle: () => void;
   setActiveIndex: (cb: number | ((currIndex: number) => number)) => void;
+  inputRefSetter: (el: HTMLInputElement) => void;
+  getInput: () => HTMLInputElement;
 }
 
 export interface IKBarContext {
diff --git a/src/useKBar.tsx b/src/useKBar.tsx
index fffbc7f..19844de 100644
--- a/src/useKBar.tsx
+++ b/src/useKBar.tsx
@@ -47,4 +47,3 @@ export function useKBar<C = null>(
 
   return render;
 }
-
diff --git a/src/useStore.tsx b/src/useStore.tsx
index 492cd04..9957ee8 100644
--- a/src/useStore.tsx
+++ b/src/useStore.tsx
@@ -1,5 +1,6 @@
 import { deepEqual } from "fast-equals";
 import * as React from "react";
+import invariant from "tiny-invariant";
 import { ActionInterface } from "./action/ActionInterface";
 import { history } from "./action/HistoryImpl";
 import type {
@@ -72,6 +73,8 @@ export function useStore(props: useStoreProps) {
     [actionsInterface]
   );
 
+  const inputRef = React.useRef<HTMLInputElement | null>(null);
+
   return React.useMemo(() => {
     return {
       getState,
@@ -109,6 +112,16 @@ export function useStore(props: useStoreProps) {
             ...state,
             activeIndex: typeof cb === "number" ? cb : cb(state.activeIndex),
           })),
+        inputRefSetter: (el: HTMLInputElement) => {
+          inputRef.current = el;
+        },
+        getInput: () => {
+          invariant(
+            inputRef.current,
+            "Input is undefined, make sure you apple `query.inputRefSetter` to your search input."
+          );
+          return inputRef.current;
+        },
       },
       options: optionsRef.current,
       subscribe: (collector, cb) => publisher.subscribe(collector, cb),
diff --git a/src/utils.ts b/src/utils.ts
index 14b491f..9d3a85b 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -110,6 +110,7 @@ export function shouldRejectKeystrokes(
   );
 
   const activeElement = document.activeElement;
+
   const ignoreStrokes =
     activeElement &&
     (inputs.indexOf(activeElement.tagName.toLowerCase()) !== -1 ||