From 4e6df280df0e0493809db41bc07fb5f9838d284c Mon Sep 17 00:00:00 2001
From: Jocelyn Liu <yrliou@gmail.com>
Date: Fri, 13 Sep 2024 14:14:46 -0700
Subject: [PATCH] Move edit answer from quick action to menu

---
 components/ai_chat/core/browser/constants.cc           |  1 +
 .../page/components/context_menu_assistant/index.tsx   | 10 ++++++++++
 .../page/components/conversation_list/index.tsx        | 10 ++++++----
 .../resources/page/components/edit_button/index.tsx    |  2 --
 components/resources/ai_chat_ui_strings.grdp           |  3 +++
 5 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/components/ai_chat/core/browser/constants.cc b/components/ai_chat/core/browser/constants.cc
index e948640993f8..8ec2866559b5 100644
--- a/components/ai_chat/core/browser/constants.cc
+++ b/components/ai_chat/core/browser/constants.cc
@@ -84,6 +84,7 @@ base::span<const webui::LocalizedString> GetLocalizedStrings() {
       {"saveButtonLabel", IDS_CHAT_UI_SAVE_BUTTON_LABEL},
       {"editedLabel", IDS_CHAT_UI_EDITED_LABEL},
       {"editButtonLabel", IDS_CHAT_UI_EDIT_BUTTON_LABEL},
+      {"editAnswerLabel", IDS_CHAT_UI_EDIT_ANSWER_LABEL},
       {"optionNotHelpful", IDS_CHAT_UI_OPTION_NOT_HELPFUL},
       {"optionIncorrect", IDS_CHAT_UI_OPTION_INCORRECT},
       {"optionUnsafeHarmful", IDS_CHAT_UI_OPTION_UNSAFE_HARMFUL},
diff --git a/components/ai_chat/resources/page/components/context_menu_assistant/index.tsx b/components/ai_chat/resources/page/components/context_menu_assistant/index.tsx
index c9f8b98e1e22..e22879466881 100644
--- a/components/ai_chat/resources/page/components/context_menu_assistant/index.tsx
+++ b/components/ai_chat/resources/page/components/context_menu_assistant/index.tsx
@@ -22,6 +22,7 @@ interface ContextMenuAssistantProps {
   isOpen: boolean
   onClick: () => void
   onClose: () => void
+  onEditAnswerClicked: () => void
 }
 
 enum RatingStatus {
@@ -35,6 +36,7 @@ function ContextMenuAssistant_(
   ref: React.RefObject<Map<number, Element>>
 ) {
   const context = React.useContext(DataContext)
+  const {shouldDisableUserInput} = context
   const [feedbackId, setFeedbackId] = React.useState<string | null>()
   const [isFormVisible, setIsFormVisible] = React.useState(false)
   const [currentRatingStatus, setCurrentRatingStatus] =
@@ -157,6 +159,14 @@ function ContextMenuAssistant_(
         >
           <Icon name='more-vertical' />
         </Button>
+        {!shouldDisableUserInput && (
+          <leo-menu-item
+            onClick={props.onEditAnswerClicked}
+          >
+            <Icon name='edit-pencil' />
+            <span>{getLocale('editAnswerLabel')}</span>
+          </leo-menu-item>
+        )}
         <leo-menu-item
           class={classnames({
             [styles.liked]: currentRatingStatus === RatingStatus.Liked
diff --git a/components/ai_chat/resources/page/components/conversation_list/index.tsx b/components/ai_chat/resources/page/components/conversation_list/index.tsx
index 2baefa067e19..6ae1790ffc35 100644
--- a/components/ai_chat/resources/page/components/conversation_list/index.tsx
+++ b/components/ai_chat/resources/page/components/conversation_list/index.tsx
@@ -186,10 +186,11 @@ function ConversationList(props: ConversationListProps) {
                         )}
                         <div className={styles.turnActions}>
                           <CopyButton onClick={handleCopyText} />
-                          <EditButton
-                            onClick={() => setEditInputId(id)}
-                            isDisabled={isAIAssistant && shouldDisableUserInput}
-                          />
+                          {!isAIAssistant &&
+                            <EditButton
+                              onClick={() => setEditInputId(id)}
+                            />
+                          }
                           {isAIAssistant &&
                             context.currentModel?.options.leoModelOptions && (
                             <ContextMenuAssistant
@@ -198,6 +199,7 @@ function ConversationList(props: ConversationListProps) {
                               isOpen={activeMenuId === id}
                               onClick={() => showAssistantMenu(id)}
                               onClose={hideAssistantMenu}
+                              onEditAnswerClicked={() => setEditInputId(id)}
                             />
                           )}
                         </div>
diff --git a/components/ai_chat/resources/page/components/edit_button/index.tsx b/components/ai_chat/resources/page/components/edit_button/index.tsx
index 51db56c80bc0..d3ad939cb8c5 100644
--- a/components/ai_chat/resources/page/components/edit_button/index.tsx
+++ b/components/ai_chat/resources/page/components/edit_button/index.tsx
@@ -11,7 +11,6 @@ import styles from './style.module.scss'
 
 interface Props {
   onClick: () => void
-  isDisabled: boolean
 }
 
 function EditButton (props: Props) {
@@ -20,7 +19,6 @@ function EditButton (props: Props) {
       fab
       size='tiny'
       kind='plain-faint'
-      isDisabled={props.isDisabled}
       title={getLocale('editButtonLabel')}
     >
       <Icon name='edit-pencil' />
diff --git a/components/resources/ai_chat_ui_strings.grdp b/components/resources/ai_chat_ui_strings.grdp
index 7e88ba96eac1..fcaed04ab6c4 100644
--- a/components/resources/ai_chat_ui_strings.grdp
+++ b/components/resources/ai_chat_ui_strings.grdp
@@ -199,6 +199,9 @@
   <message name="IDS_CHAT_UI_EDIT_BUTTON_LABEL" desc="Button title to edit content">
     Edit
   </message>
+  <message name="IDS_CHAT_UI_EDIT_ANSWER_LABEL" desc="Button title to edit AI agent answer">
+    Edit answer
+  </message>
   <message name="IDS_CHAT_UI_OPTION_NOT_HELPFUL" desc="Select option for feedback form if the answer is not helpful">
     Not helpful
   </message>