diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
index fc012f6b7..21ef182c3 100644
--- a/src/frontend/src/modals/codeAreaModal/index.tsx
+++ b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -5,7 +5,7 @@ import "ace-builds/src-noconflict/mode-python";
 import "ace-builds/src-noconflict/theme-github";
 import "ace-builds/src-noconflict/theme-twilight";
 import { TerminalSquare } from "lucide-react";
-import { useContext, useEffect, useRef, useState } from "react";
+import { useContext, useRef, useState } from "react";
 import AceEditor from "react-ace";
 import { Button } from "../../components/ui/button";
 import { CODE_PROMPT_DIALOG_SUBTITLE } from "../../constants";
@@ -37,9 +37,6 @@ export default function CodeAreaModal({
     detail: { error: string; traceback: string };
   }>(null);
   const ref = useRef();
-  useEffect(() => {
-    setValue(code);
-  }, [code, setValue]);
 
   function setModalOpen(x: boolean) {
     if (x === false) {
diff --git a/src/frontend/src/modals/formModal/chatMessage/codeBlock/index.tsx b/src/frontend/src/modals/formModal/chatMessage/codeBlock/index.tsx
index c78974a3c..884215ebb 100644
--- a/src/frontend/src/modals/formModal/chatMessage/codeBlock/index.tsx
+++ b/src/frontend/src/modals/formModal/chatMessage/codeBlock/index.tsx
@@ -1,5 +1,5 @@
 import { IconCheck, IconClipboard, IconDownload } from "@tabler/icons-react";
-import { FC, memo, useState } from "react";
+import { useState } from "react";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { oneDark } from "react-syntax-highlighter/dist/cjs/styles/prism";
 import { programmingLanguages } from "../../../../utils";
@@ -9,7 +9,7 @@ interface Props {
   value: string;
 }
 
-export const CodeBlock: FC<Props> = memo(({ language, value }) => {
+export function CodeBlock({ language, value }) {
   const [isCopied, setIsCopied] = useState<Boolean>(false);
 
   const copyToClipboard = () => {
@@ -63,7 +63,7 @@ export const CodeBlock: FC<Props> = memo(({ language, value }) => {
       </div>
 
       <SyntaxHighlighter
-        className=" w-[47vw]"
+        className="overflow-auto"
         language={language}
         style={oneDark}
         customStyle={{ margin: 0 }}
@@ -72,5 +72,5 @@ export const CodeBlock: FC<Props> = memo(({ language, value }) => {
       </SyntaxHighlighter>
     </div>
   );
-});
+}
 CodeBlock.displayName = "CodeBlock";
diff --git a/src/frontend/src/modals/formModal/chatMessage/index.tsx b/src/frontend/src/modals/formModal/chatMessage/index.tsx
index 6e7457549..4e941e278 100644
--- a/src/frontend/src/modals/formModal/chatMessage/index.tsx
+++ b/src/frontend/src/modals/formModal/chatMessage/index.tsx
@@ -1,6 +1,6 @@
 import Convert from "ansi-to-html";
 import { ChevronDown } from "lucide-react";
-import { useState } from "react";
+import { useMemo, useState } from "react";
 import ReactMarkdown from "react-markdown";
 import rehypeMathjax from "rehype-mathjax";
 import remarkGfm from "remark-gfm";
@@ -75,47 +75,60 @@ export default function ChatMessage({
             <div className="w-full">
               <div className="w-full dark:text-white">
                 <div className="w-full">
-                  <ReactMarkdown
-                    remarkPlugins={[remarkGfm, remarkMath]}
-                    rehypePlugins={[rehypeMathjax]}
-                    className="markdown prose inline-block break-words text-primary
+                  {useMemo(
+                    () => (
+                      <ReactMarkdown
+                        remarkPlugins={[remarkGfm, remarkMath]}
+                        rehypePlugins={[rehypeMathjax]}
+                        className="markdown prose inline-block break-words text-primary
                      dark:prose-invert sm:max-w-[30vw] lg:max-w-[40vw]"
-                    components={{
-                      code({ node, inline, className, children, ...props }) {
-                        if (children.length) {
-                          if (children[0] === "▍") {
-                            return (
-                              <span className="form-modal-markdown-span">
-                                ▍
-                              </span>
+                        components={{
+                          code: ({
+                            node,
+                            inline,
+                            className,
+                            children,
+                            ...props
+                          }) => {
+                            if (children.length) {
+                              if (children[0] === "▍") {
+                                return (
+                                  <span className="form-modal-markdown-span">
+                                    ▍
+                                  </span>
+                                );
+                              }
+
+                              children[0] = (children[0] as string).replace(
+                                "`▍`",
+                                "▍"
+                              );
+                            }
+
+                            const match = /language-(\w+)/.exec(
+                              className || ""
                             );
-                          }
 
-                          children[0] = (children[0] as string).replace(
-                            "`▍`",
-                            "▍"
-                          );
-                        }
-
-                        const match = /language-(\w+)/.exec(className || "");
-
-                        return !inline ? (
-                          <CodeBlock
-                            key={Math.random()}
-                            language={(match && match[1]) || ""}
-                            value={String(children).replace(/\n$/, "")}
-                            {...props}
-                          />
-                        ) : (
-                          <code className={className} {...props}>
-                            {children}
-                          </code>
-                        );
-                      },
-                    }}
-                  >
-                    {chat.message.toString()}
-                  </ReactMarkdown>
+                            return !inline ? (
+                              <CodeBlock
+                                key={Math.random()}
+                                language={(match && match[1]) || ""}
+                                value={String(children).replace(/\n$/, "")}
+                                {...props}
+                              />
+                            ) : (
+                              <code className={className} {...props}>
+                                {children}
+                              </code>
+                            );
+                          },
+                        }}
+                      >
+                        {chat.message.toString()}
+                      </ReactMarkdown>
+                    ),
+                    [chat.message.toString()]
+                  )}
                 </div>
                 {chat.files && (
                   <div className="my-2 w-full">
@@ -138,51 +151,56 @@ export default function ChatMessage({
         </div>
       ) : (
         <div>
-          <button
-            className="form-modal-initial-prompt-btn"
-            onClick={() => {
-              setPromptOpen((old) => !old);
-            }}
-          >
-            Display Prompt
-            <ChevronDown
-              className={
-                "h-3 w-3 transition-all " + (promptOpen ? "rotate-180" : "")
-              }
-            />
-          </button>
-          <span className="prose inline-block break-words text-primary dark:prose-invert">
-            {promptOpen
-              ? template?.split("\n")?.map((line, index) => {
-                  const regex = /{([^}]+)}/g;
-                  let match;
-                  let parts = [];
-                  let lastIndex = 0;
-                  while ((match = regex.exec(line)) !== null) {
-                    // Push text up to the match
-                    if (match.index !== lastIndex) {
-                      parts.push(line.substring(lastIndex, match.index));
-                    }
-                    // Push div with matched text
-                    if (chat.message[match[1]]) {
-                      parts.push(
-                        <span className="my-1 rounded-md bg-indigo-100">
-                          {chat.message[match[1]]}
-                        </span>
-                      );
-                    }
+          {template && (
+            <>
+              <button
+                className="form-modal-initial-prompt-btn"
+                onClick={() => {
+                  setPromptOpen((old) => !old);
+                }}
+              >
+                Display Prompt
+                <ChevronDown
+                  className={
+                    "h-3 w-3 transition-all " + (promptOpen ? "rotate-180" : "")
+                  }
+                />
+              </button>
+              <span className="prose inline-block break-words text-primary dark:prose-invert">
+                {promptOpen
+                  ? template?.split("\n")?.map((line, index) => {
+                      const regex = /{([^}]+)}/g;
+                      let match;
+                      let parts = [];
+                      let lastIndex = 0;
+                      while ((match = regex.exec(line)) !== null) {
+                        // Push text up to the match
+                        if (match.index !== lastIndex) {
+                          parts.push(line.substring(lastIndex, match.index));
+                        }
+                        // Push div with matched text
+                        if (chat.message[match[1]]) {
+                          parts.push(
+                            <span className="my-1 rounded-md bg-indigo-100">
+                              {chat.message[match[1]]}
+                            </span>
+                          );
+                        }
 
-                    // Update last index
-                    lastIndex = regex.lastIndex;
-                  }
-                  // Push text after the last match
-                  if (lastIndex !== line.length) {
-                    parts.push(line.substring(lastIndex));
-                  }
-                  return <p>{parts}</p>;
-                })
-              : chat.message[chat.chatKey]}
-          </span>
+                        // Update last index
+                        lastIndex = regex.lastIndex;
+                      }
+                      // Push text after the last match
+                      if (lastIndex !== line.length) {
+                        parts.push(line.substring(lastIndex));
+                      }
+                      return <p>{parts}</p>;
+                    })
+                  : chat.message[chat.chatKey]}
+              </span>
+            </>
+          )}
+          {chat.message[chat.chatKey]}
         </div>
       )}
     </div>