From b86dd0d7cb6e54318e08eaf71b63806b3beee0a4 Mon Sep 17 00:00:00 2001
From: Aaron Diamond-Reivich <aarondr77@gmail.com>
Date: Sun, 16 Feb 2025 15:23:56 -0500
Subject: [PATCH 1/5] mito-ai: move edit buttons to hover

---
 .../AiChat/ChatMessage/ChatMessage.tsx        | 24 +++++++++----------
 mito-ai/style/ChatMessage.css                 | 24 +++++++++++++++++++
 2 files changed, 35 insertions(+), 13 deletions(-)
 create mode 100644 mito-ai/style/ChatMessage.css

diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
index dc30fa720..af0999720 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -20,6 +20,7 @@ import CopyIcon from '../../../icons/CopyIcon';
 import copyToClipboard from '../../../utils/copyToClipboard';
 import TextButton from '../../../components/TextButton';
 import { IDisplayOptimizedChatHistory } from '../ChatHistoryManager';
+import '../../../../style/ChatMessage.css';
 
 interface IChatMessageProps {
     message: OpenAI.Chat.ChatCompletionMessageParam
@@ -178,7 +179,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                     }
                 } else {
                     return (
-                        <div className={classNames('markdown-message-part')} style={{ position: 'relative' }}>
+                        <div className={classNames('markdown-message-part')}>
                             <p 
                                 key={index + messagePart} 
                                 onDoubleClick={() => {
@@ -200,22 +201,19 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                     />
                                 )}
                             </p>
-                            {editable && (
-                                <div style={{ display: 'flex', justifyContent: 'flex-end', marginTop: '4px' }}>
-                                    <button
-                                        className="message-edit-button"
-                                        onClick={handleEditClick}
-                                        style={{ cursor: 'pointer' }}
-                                        title="Edit message"
-                                    >
-                                        <PencilIcon />
-                                    </button>
-                                </div>
-                            )}
                         </div>
                     )
                 }
             })}
+            {editable && (
+                <button
+                    className="message-edit-button"
+                    onClick={handleEditClick}
+                    title="Edit message"
+                >
+                    <PencilIcon />
+                </button>
+            )}
         </div>
     )
 }
diff --git a/mito-ai/style/ChatMessage.css b/mito-ai/style/ChatMessage.css
new file mode 100644
index 000000000..36f8c165c
--- /dev/null
+++ b/mito-ai/style/ChatMessage.css
@@ -0,0 +1,24 @@
+.message {
+    position: relative;
+}
+
+.message:hover .message-edit-button {
+    opacity: 1 !important;
+}
+
+.message-edit-button {
+    position: absolute !important;
+    bottom: 8px !important;
+    right: 8px !important;
+    cursor: pointer;
+    opacity: 0;
+    transition: opacity 0.2s ease-in-out;
+    background: none;
+    border: none;
+    padding: 4px;
+}
+
+.message-edit-button:hover {
+    background-color: rgba(0, 0, 0, 0.05) !important;
+    border-radius: 4px;
+} 
\ No newline at end of file

From 36f10b126acc01d4441e3db1cad1a7043d88627e Mon Sep 17 00:00:00 2001
From: Aaron Diamond-Reivich <aarondr77@gmail.com>
Date: Sun, 16 Feb 2025 15:32:07 -0500
Subject: [PATCH 2/5] mito-ai: update colors

---
 mito-ai/style/ChatMessage.css | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mito-ai/style/ChatMessage.css b/mito-ai/style/ChatMessage.css
index 36f8c165c..b242ed2fd 100644
--- a/mito-ai/style/ChatMessage.css
+++ b/mito-ai/style/ChatMessage.css
@@ -12,13 +12,13 @@
     right: 8px !important;
     cursor: pointer;
     opacity: 0;
-    transition: opacity 0.2s ease-in-out;
-    background: none;
+    transition: opacity 0.2s ease-in-out, background-color 0.2s ease-in-out;
+    background-color: var(--jp-layout-color3);
     border: none;
     padding: 4px;
+    border-radius: 4px;
 }
 
 .message-edit-button:hover {
-    background-color: rgba(0, 0, 0, 0.05) !important;
-    border-radius: 4px;
+    background-color: var(--jp-layout-color4);
 } 
\ No newline at end of file

From 719df2c46156ebd932a25b8ea3ea347842f53b42 Mon Sep 17 00:00:00 2001
From: Aaron Diamond-Reivich <aarondr77@gmail.com>
Date: Sun, 16 Feb 2025 15:41:02 -0500
Subject: [PATCH 3/5] mito-ai: polish

---
 mito-ai/style/ChatMessage.css | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mito-ai/style/ChatMessage.css b/mito-ai/style/ChatMessage.css
index b242ed2fd..caa0512af 100644
--- a/mito-ai/style/ChatMessage.css
+++ b/mito-ai/style/ChatMessage.css
@@ -17,6 +17,8 @@
     border: none;
     padding: 4px;
     border-radius: 4px;
+    display: flex;
+    align-items: center;
 }
 
 .message-edit-button:hover {

From 63dfa8aacbb61f073ec000d5bae4e4c4c0e30f44 Mon Sep 17 00:00:00 2001
From: Aaron Diamond-Reivich <aarondr77@gmail.com>
Date: Sun, 16 Feb 2025 15:52:29 -0500
Subject: [PATCH 4/5] mito-ai: cleanup

---
 .../AiChat/ChatMessage/ChatMessage.tsx        |  3 +--
 mito-ai/style/ChatMessage.css                 | 26 -------------------
 mito-ai/style/ChatTaskpane.css                | 23 ++++++++++++++--
 3 files changed, 22 insertions(+), 30 deletions(-)
 delete mode 100644 mito-ai/style/ChatMessage.css

diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
index af0999720..535d6b786 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -20,7 +20,6 @@ import CopyIcon from '../../../icons/CopyIcon';
 import copyToClipboard from '../../../utils/copyToClipboard';
 import TextButton from '../../../components/TextButton';
 import { IDisplayOptimizedChatHistory } from '../ChatHistoryManager';
-import '../../../../style/ChatMessage.css';
 
 interface IChatMessageProps {
     message: OpenAI.Chat.ChatCompletionMessageParam
@@ -207,7 +206,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
             })}
             {editable && (
                 <button
-                    className="message-edit-button"
+                    className="message-start-editing-button"
                     onClick={handleEditClick}
                     title="Edit message"
                 >
diff --git a/mito-ai/style/ChatMessage.css b/mito-ai/style/ChatMessage.css
deleted file mode 100644
index caa0512af..000000000
--- a/mito-ai/style/ChatMessage.css
+++ /dev/null
@@ -1,26 +0,0 @@
-.message {
-    position: relative;
-}
-
-.message:hover .message-edit-button {
-    opacity: 1 !important;
-}
-
-.message-edit-button {
-    position: absolute !important;
-    bottom: 8px !important;
-    right: 8px !important;
-    cursor: pointer;
-    opacity: 0;
-    transition: opacity 0.2s ease-in-out, background-color 0.2s ease-in-out;
-    background-color: var(--jp-layout-color3);
-    border: none;
-    padding: 4px;
-    border-radius: 4px;
-    display: flex;
-    align-items: center;
-}
-
-.message-edit-button:hover {
-    background-color: var(--jp-layout-color4);
-} 
\ No newline at end of file
diff --git a/mito-ai/style/ChatTaskpane.css b/mito-ai/style/ChatTaskpane.css
index b3e7cec05..b6675618d 100644
--- a/mito-ai/style/ChatTaskpane.css
+++ b/mito-ai/style/ChatTaskpane.css
@@ -48,6 +48,7 @@
 }
 
 .message {
+    position: relative;
     height: min-content;
     box-sizing: border-box;
     padding: 10px;
@@ -60,6 +61,10 @@
     margin-bottom: 20px;
 }
 
+.message:hover .message-start-editing-button {
+    opacity: 1 !important;
+}
+
 .message-user {
     background-color: var(--chat-user-message-background-color);
     color: var(--chat-user-message-font-color);
@@ -143,12 +148,26 @@
     align-items: center;
 }
 
-.message-edit-button {
-    background-color: rgba(255, 255, 255, 0);
+.message-start-editing-button {
+    position: absolute !important;
+    bottom: 8px !important;
+    right: 8px !important;
+    cursor: pointer;
+    opacity: 0;
+    transition: opacity 0.2s ease-in-out, background-color 0.2s ease-in-out;
+    background-color: var(--jp-layout-color3);
     border: none;
+    padding: 4px;
+    border-radius: 4px;
+    display: flex;
+    align-items: center;
     color: var(--grey-900) !important;
 }
 
+.message-start-editing-button:hover {
+    background-color: var(--jp-layout-color4);
+}
+
 .message-edit-buttons {
     display: flex;
     gap: 8px;

From 1a30d790701bca95ef0a1091c5c69a096296e29b Mon Sep 17 00:00:00 2001
From: Aaron Diamond-Reivich <aarondr77@gmail.com>
Date: Sun, 16 Feb 2025 15:52:46 -0500
Subject: [PATCH 5/5] tests: update tests

---
 tests/mitoai_ui_tests/agent.spec.ts  | 2 +-
 tests/mitoai_ui_tests/mitoai.spec.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/mitoai_ui_tests/agent.spec.ts b/tests/mitoai_ui_tests/agent.spec.ts
index 913b82db4..e92efa909 100644
--- a/tests/mitoai_ui_tests/agent.spec.ts
+++ b/tests/mitoai_ui_tests/agent.spec.ts
@@ -62,7 +62,7 @@ test.describe("Agent mode integration tests", () => {
 
         // Get the last agent message, and click on it's edit button 
         const lastAgentMessage = await page.locator('.message-assistant-agent').last();
-        await lastAgentMessage.locator('.message-edit-button').click();
+        await lastAgentMessage.locator('.message-start-editing-button').click();
 
         // Make sure the active cell preview is not visible
         await expect(page.locator('.active-cell-preview-container')).not.toBeVisible();
diff --git a/tests/mitoai_ui_tests/mitoai.spec.ts b/tests/mitoai_ui_tests/mitoai.spec.ts
index 7c3478b67..e3d376d5e 100644
--- a/tests/mitoai_ui_tests/mitoai.spec.ts
+++ b/tests/mitoai_ui_tests/mitoai.spec.ts
@@ -132,7 +132,7 @@ test.describe('Mito AI Chat', () => {
     expect(messageAssistantDivs).toBe(1);
 
     // Ensure you cannot edit the AI's messages by clicking the pencil icon or double clicking the message
-    await expect(page.locator('.message-assistant-chat .message-edit-button')).not.toBeVisible();
+    await expect(page.locator('.message-assistant-chat .message-start-editing-button')).not.toBeVisible();
     await page.locator('.message-assistant-chat p').last().dblclick();
     await expect(page.locator('.message-assistant-chat').getByRole('button', { name: 'Save' })).not.toBeVisible();
   });