diff --git a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx
index 84c53620f..6d8d81720 100644
--- a/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx
@@ -113,7 +113,7 @@ export class ChatHistoryManager {
 
     updateMessageAtIndex(index: number, newContent: string, isAgentMessage: boolean = false): IOutgoingMessage {
         const activeCellID = getActiveCellID(this.notebookTracker)
-        const activeCellCode = getCellCodeByID(this.notebookTracker, activeCellID)
+        const activeCellCode = isAgentMessage ? undefined : getCellCodeByID(this.notebookTracker, activeCellID)
 
         const metadata: IChatMessageMetadata = {
             variables: this.variableManager.variables,
@@ -123,7 +123,11 @@ export class ChatHistoryManager {
         }
         
         this.displayOptimizedChatHistory[index] = { 
-            message: getDisplayedOptimizedUserMessage(newContent, activeCellCode),
+            message: getDisplayedOptimizedUserMessage(
+                newContent, 
+                activeCellCode,
+                isAgentMessage
+            ),
             type: isAgentMessage ? 'openai message:agent:planning' : 'openai message',
             codeCellID: activeCellID,
             promptType: isAgentMessage ? 'agent:planning' : 'chat'
@@ -288,10 +292,15 @@ export class ChatHistoryManager {
 }
 
 
-const getDisplayedOptimizedUserMessage = (input: string, activeCellCode?: string): OpenAI.Chat.ChatCompletionMessageParam => {
+const getDisplayedOptimizedUserMessage = (
+    input: string, 
+    activeCellCode?: string, 
+    isAgentPlanning: boolean = false
+): OpenAI.Chat.ChatCompletionMessageParam => {
     return {
         role: 'user',
-        content: activeCellCode ? `\`\`\`python
+        content: (!isAgentPlanning && activeCellCode) ? 
+`\`\`\`python
 ${activeCellCode}
 \`\`\`
 
diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx
index 66d27a1b4..ceab35af1 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx
@@ -20,6 +20,7 @@ interface ChatInputProps {
     variableManager?: IVariableManager;
     notebookTracker: INotebookTracker;
     renderMimeRegistry: IRenderMimeRegistry;
+    displayActiveCellCode?: boolean;
 }
 
 export interface ExpandedVariable extends Variable {
@@ -35,6 +36,7 @@ const ChatInput: React.FC<ChatInputProps> = ({
     variableManager,
     notebookTracker,
     renderMimeRegistry,
+    displayActiveCellCode = true,
 }) => {
 
     const [input, setInput] = useState(initialContent);
@@ -173,7 +175,7 @@ const ChatInput: React.FC<ChatInputProps> = ({
             }}
         >
             {/* Show the active cell preview if the text area has focus or the user has started typing */}
-            {activeCellCodePreview.length > 0 
+            {displayActiveCellCode && activeCellCodePreview.length > 0 
                 && (isFocused || input.length > 0)
                 && <div className='active-cell-preview-container'>
                     <div className='code-block-container'>
diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
index 45cd0f298..dc30fa720 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -94,6 +94,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                 variableManager={variableManager}
                 notebookTracker={notebookTracker}
                 renderMimeRegistry={renderMimeRegistry}
+                displayActiveCellCode={messageType !== 'openai message:agent:planning'}
             />
         );
     }
diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
index d1f4a0cc1..56049a5ac 100644
--- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
@@ -261,12 +261,14 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 3: Send the message to the AI
         await _sendMessageAndSaveResponse(outgoingMessage, newChatHistoryManager)
 
-        // Step 4: Scroll so that the top of the last AI message is visible
+        // Step 4: Scroll to the bottom of the chat smoothly
         setTimeout(() => {
-            const aiMessages = chatMessagesRef.current?.getElementsByClassName('message message-assistant');
-            if (aiMessages && aiMessages.length > 0) {
-                const lastAiMessage = aiMessages[aiMessages.length - 1];
-                lastAiMessage.scrollIntoView({ behavior: 'smooth' });
+            const chatContainer = chatMessagesRef.current;
+            if (chatContainer) {
+                chatContainer.scrollTo({
+                    top: chatContainer.scrollHeight,
+                    behavior: 'smooth'
+                });
             }
         }, 100);
 
@@ -298,11 +300,10 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 0: Reject the previous Ai generated code if they did not accept it
         rejectAICode()
 
-        // Step 1: Clear the chat history, and add the new error message
-        const newChatHistoryManager = clearChatHistory()
+        // Step 1: Add user message to chat history
+        const newChatHistoryManager = getDuplicateChatHistoryManager()
         const outgoingMessage = newChatHistoryManager.addAgentMessage(message)
         setChatHistoryManager(newChatHistoryManager)
-        console.log('outgoingMessage: ', outgoingMessage)
 
         // Step 2: Send the message to the AI
         await _sendMessageAndSaveResponse(outgoingMessage, newChatHistoryManager)
diff --git a/tests/mitoai_ui_tests/agent.spec.ts b/tests/mitoai_ui_tests/agent.spec.ts
index 1ab0770ac..913b82db4 100644
--- a/tests/mitoai_ui_tests/agent.spec.ts
+++ b/tests/mitoai_ui_tests/agent.spec.ts
@@ -64,6 +64,9 @@ test.describe("Agent mode integration tests", () => {
         const lastAgentMessage = await page.locator('.message-assistant-agent').last();
         await lastAgentMessage.locator('.message-edit-button').click();
 
+        // Make sure the active cell preview is not visible
+        await expect(page.locator('.active-cell-preview-container')).not.toBeVisible();
+
         // Edit the message
         await page.locator('.chat-input').fill(newMessage);
         await page.keyboard.press('Enter');