From 42164aef57a52f588bd52316c8136d21cc79615b Mon Sep 17 00:00:00 2001
From: Serg <serg.zhukovsky@gmail.com>
Date: Thu, 12 Sep 2024 14:21:51 -0400
Subject: [PATCH] Moves a focus check on the higher level

---
 .../page/components/input_box/index.tsx          | 16 +++++-----------
 .../resources/page/components/main/index.tsx     | 12 ++++++++++++
 .../page/components/BeginGeneration.tsx          |  1 -
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/components/ai_chat/resources/page/components/input_box/index.tsx b/components/ai_chat/resources/page/components/input_box/index.tsx
index 024d20b87d4f..44b4fddca210 100644
--- a/components/ai_chat/resources/page/components/input_box/index.tsx
+++ b/components/ai_chat/resources/page/components/input_box/index.tsx
@@ -27,14 +27,13 @@ type Props = Pick<AIChatContext,
   | 'isToolsMenuOpen'
   | 'setIsToolsMenuOpen'
   | 'isMobile'
-  | 'hasInitialHistory'
   | 'shouldDisableUserInput'
-  | 'hasAcceptedAgreement'
-  | 'handleShowSoftKeyboard'>
+  | 'hasAcceptedAgreement'>
 
 interface InputBoxProps {
     context: Props
     onFocusInputMobile?: () => unknown
+    maybeShowSoftKeyboard?: (querySubmitted: boolean) => unknown
 }
 
 function InputBox(props: InputBoxProps) {
@@ -86,16 +85,11 @@ function InputBox(props: InputBoxProps) {
       return
     }
 
-    if (props.context.selectedActionType) {
+    if (props.context.selectedActionType ||
+        props.maybeShowSoftKeyboard &&
+        props.maybeShowSoftKeyboard(querySubmitted.current)) {
       node.focus()
     }
-    if (props.context.isMobile && props.context.hasAcceptedAgreement &&
-      !props.context.hasInitialHistory && !querySubmitted.current) {
-        node.focus()
-        if (props.context.handleShowSoftKeyboard) {
-          props.context.handleShowSoftKeyboard();
-        }
-    }
   }
 
   return (
diff --git a/components/ai_chat/resources/page/components/main/index.tsx b/components/ai_chat/resources/page/components/main/index.tsx
index 808672d46e2a..063669833315 100644
--- a/components/ai_chat/resources/page/components/main/index.tsx
+++ b/components/ai_chat/resources/page/components/main/index.tsx
@@ -158,6 +158,17 @@ function Main() {
     }
   }
 
+  const maybeShowSoftKeyboard = (querySubmitted: boolean) => {
+    if (context.isMobile && context.hasAcceptedAgreement &&
+      !context.hasInitialHistory && !querySubmitted) {
+        if (context.handleShowSoftKeyboard) {
+          context.handleShowSoftKeyboard()
+          return true
+        }
+    }
+    return false
+  }
+
   return (
     <main className={styles.main}>
       {context.showAgreementModal && <PrivacyMessage />}
@@ -273,6 +284,7 @@ function Main() {
           <InputBox
             context={context}
             onFocusInputMobile={handleOnFocusInputMobile}
+            maybeShowSoftKeyboard={maybeShowSoftKeyboard}
           />
         </ToolsButtonMenu>
       </div>
diff --git a/components/ai_rewriter/resources/page/components/BeginGeneration.tsx b/components/ai_rewriter/resources/page/components/BeginGeneration.tsx
index 8dcdade0eec8..480a64b196f0 100644
--- a/components/ai_rewriter/resources/page/components/BeginGeneration.tsx
+++ b/components/ai_rewriter/resources/page/components/BeginGeneration.tsx
@@ -45,7 +45,6 @@ export default function BeginGeneration() {
             isCharLimitExceeded: context.isCharLimitExceeded,
             shouldDisableUserInput: context.isGenerating,
             isMobile: false,
-            hasInitialHistory: false,
             hasAcceptedAgreement: true
           }} />
       </ToolsButtonMenu>