From d6c7793147edf67398f27a2b0793b159dcec78b4 Mon Sep 17 00:00:00 2001
From: Kevin On <40454531+kevin-on@users.noreply.github.com>
Date: Tue, 29 Oct 2024 13:29:48 +0900
Subject: [PATCH] Autofocus input (#63)

* Implement custom AutoFocusPlugin

* Update comment
---
 .../chat-view/chat-input/ChatUserInput.tsx    |  4 +--
 .../plugins/auto-focus/AutoFocusPlugin.tsx    | 29 +++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 src/components/chat-view/chat-input/plugins/auto-focus/AutoFocusPlugin.tsx

diff --git a/src/components/chat-view/chat-input/ChatUserInput.tsx b/src/components/chat-view/chat-input/ChatUserInput.tsx
index e9a626a..c34c8ff 100644
--- a/src/components/chat-view/chat-input/ChatUserInput.tsx
+++ b/src/components/chat-view/chat-input/ChatUserInput.tsx
@@ -1,4 +1,3 @@
-import { AutoFocusPlugin } from '@lexical/react/LexicalAutoFocusPlugin'
 import {
   InitialConfigType,
   LexicalComposer,
@@ -34,6 +33,7 @@ import { MemoizedSyntaxHighlighterWrapper } from '../SyntaxHighlighterWrapper'
 
 import MentionableBadge from './MentionableBadge'
 import { ModelSelect } from './ModelSelect'
+import AutoFocusPlugin from './plugins/auto-focus/AutoFocusPlugin'
 import AutoLinkMentionPlugin from './plugins/mention/AutoLinkMentionPlugin'
 import { MentionNode } from './plugins/mention/MentionNode'
 import MentionPlugin from './plugins/mention/MentionPlugin'
@@ -318,7 +318,7 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
             ErrorBoundary={LexicalErrorBoundary}
           />
           <HistoryPlugin />
-          {autoFocus && <AutoFocusPlugin />}
+          {autoFocus && <AutoFocusPlugin defaultSelection="rootEnd" />}
           <MentionPlugin searchResultByQuery={searchResultByQuery} />
           <OnChangePlugin
             onChange={(editorState) => {
diff --git a/src/components/chat-view/chat-input/plugins/auto-focus/AutoFocusPlugin.tsx b/src/components/chat-view/chat-input/plugins/auto-focus/AutoFocusPlugin.tsx
new file mode 100644
index 0000000..714a990
--- /dev/null
+++ b/src/components/chat-view/chat-input/plugins/auto-focus/AutoFocusPlugin.tsx
@@ -0,0 +1,29 @@
+import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
+import { useEffect } from 'react'
+
+export type AutoFocusPluginProps = {
+  defaultSelection?: 'rootStart' | 'rootEnd'
+}
+
+export default function AutoFocusPlugin({
+  defaultSelection,
+}: AutoFocusPluginProps) {
+  const [editor] = useLexicalComposerContext()
+
+  useEffect(() => {
+    editor.focus(
+      () => {
+        const rootElement = editor.getRootElement()
+        if (rootElement) {
+          // requestAnimationFrame is required here for unknown reasons, possibly related to the Obsidian plugin environment.
+          requestAnimationFrame(() => {
+            rootElement.focus()
+          })
+        }
+      },
+      { defaultSelection },
+    )
+  }, [defaultSelection, editor])
+
+  return null
+}