diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
index acaad090e..f5d969e12 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx
@@ -106,6 +106,11 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                     code={messagePart}
                                     role={message.role}
                                     renderMimeRegistry={renderMimeRegistry}
+                                    previewAICode={previewAICode}
+                                    acceptAICode={acceptAICode}
+                                    rejectAICode={rejectAICode}
+                                    isLastAiMessage={isLastAiMessage}
+                                    codeReviewStatus={codeReviewStatus}
                                 />
 
                                 {isLastAiMessage && codeReviewStatus === 'chatPreview' && 
diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx
index 9be20c938..c25c9b030 100644
--- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx
@@ -3,18 +3,34 @@ import PythonCode from './PythonCode';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import '../../../../style/CodeBlock.css'
 import copyToClipboard from '../../../utils/copyToClipboard';
+import IconButton from '../../../components/IconButton';
+import CopyIcon from '../../../icons/CopyIcon';
+import PlayButtonIcon from '../../../icons/PlayButtonIcon';
+import { CodeReviewStatus } from '../ChatTaskpane';
+import AcceptIcon from '../../../icons/AcceptIcon';
+import RejectIcon from '../../../icons/RejectIcon';
 
 
 interface ICodeBlockProps {
     code: string,
     role: 'user' | 'assistant'
     renderMimeRegistry: IRenderMimeRegistry
+    previewAICode: () => void
+    acceptAICode: () => void
+    rejectAICode: () => void
+    isLastAiMessage: boolean
+    codeReviewStatus: CodeReviewStatus
 }
 
 const CodeBlock: React.FC<ICodeBlockProps> = ({
     code,
     role,
     renderMimeRegistry,
+    previewAICode,
+    acceptAICode,
+    rejectAICode,
+    isLastAiMessage,
+    codeReviewStatus,
 }): JSX.Element => {
 
     if (role === 'user') {
@@ -31,9 +47,53 @@ const CodeBlock: React.FC<ICodeBlockProps> = ({
     if (role === 'assistant') {
         return (
             <div className='code-block-container'>
-                <div className='code-block-toolbar'>
-                    <button onClick={() => {copyToClipboard(code)}}>Copy</button>
-                </div>
+                <>
+                    {/* The code block toolbar for the last AI message */}
+                    {isLastAiMessage && 
+                        <div className='code-block-toolbar'>
+                            {codeReviewStatus === 'chatPreview' && 
+                                <IconButton 
+                                    icon={<PlayButtonIcon />}
+                                    title="Overwrite Active Cell"
+                                    onClick={() => {previewAICode()}}
+                                />
+                            }
+                            {codeReviewStatus === 'codeCellPreview' && 
+                                <IconButton 
+                                    icon={<AcceptIcon />}
+                                    title="Accept AI Generated Code"
+                                    onClick={() => {acceptAICode()}}
+                                    style={{color: 'var(--green-700)'}}
+                                />
+                            }
+                            {codeReviewStatus === 'codeCellPreview' && 
+                                <IconButton 
+                                    icon={<RejectIcon />}
+                                    title="Reject AI Generated Code"
+                                    onClick={() => {rejectAICode()}}
+                                    style={{color: 'var(--red-700)'}}
+                                />
+                            }
+                            {codeReviewStatus !== 'codeCellPreview' && 
+                                <IconButton
+                                    icon={<CopyIcon />}
+                                    title="Copy"
+                                    onClick={() => {copyToClipboard(code)}}
+                                />
+                            }
+                        </div>
+                    }
+                    {/* The code block toolbar for every other AI message */}
+                    {!isLastAiMessage && 
+                        <div className='code-block-toolbar'>
+                            <IconButton
+                                icon={<CopyIcon />}
+                                title="Copy"
+                                onClick={() => {copyToClipboard(code)}}
+                            />
+                        </div>
+                    }
+                </>
                 <PythonCode
                     code={code}
                     renderMimeRegistry={renderMimeRegistry}
diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
index 18d86d1bf..c5b9af1fc 100644
--- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
+++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx
@@ -388,8 +388,8 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
             Register the code cell toolbar buttons for accepting and rejecting code.
         */
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_ACCEPT_CODE, {
-            label: `Accept AI edits ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`,
-            className: 'text-and-icon-button green',
+            label: `Accept ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`,
+            className: 'text-and-icon-button button-green small',
             caption: 'Accept Code',
             execute: () => {acceptAICode()},
             // We use the cellStateBeforeDiff because it contains the code cell ID that we want to write to
@@ -404,8 +404,8 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         });
 
         app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_REJECT_CODE, {
-            label: `Reject AI edits ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`,
-            className: 'text-and-icon-button red',
+            label: `Reject ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`,
+            className: 'text-and-icon-button button-red small',
             caption: 'Reject Code',
             execute: () => {rejectAICode()},
             isVisible: () => {
diff --git a/mito-ai/src/components/IconButton.tsx b/mito-ai/src/components/IconButton.tsx
index 24fd890cd..2b460964f 100644
--- a/mito-ai/src/components/IconButton.tsx
+++ b/mito-ai/src/components/IconButton.tsx
@@ -1,16 +1,19 @@
 import React from 'react';
 import '../../style/IconButton.css';
+import '../../style/button.css';
+import { classNames } from '../utils/classNames';
 
 interface IconButtonProps {
     icon: React.ReactNode;
     onClick: () => void;
     title: string;
+    style?: React.CSSProperties;
 }
 
-const IconButton: React.FC<IconButtonProps> = ({ icon, onClick, title }) => {
+const IconButton: React.FC<IconButtonProps> = ({ icon, onClick, title, style }) => {
 
     return (
-        <button className="icon-button" onClick={onClick} title={title}>
+        <button className={classNames("icon-button")} onClick={onClick} title={title} style={style}>
             {icon}
         </button>
     )
diff --git a/mito-ai/src/components/TextAndIconButton.tsx b/mito-ai/src/components/TextAndIconButton.tsx
index 042fb5385..0719daa16 100644
--- a/mito-ai/src/components/TextAndIconButton.tsx
+++ b/mito-ai/src/components/TextAndIconButton.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import '../../style/TextAndIconButton.css';
+import '../../style/button.css';
 import { ButtonProps } from './TextButton';
 import { classNames } from '../utils/classNames';
 
@@ -11,7 +12,7 @@ interface TextAndIconButtonProps extends ButtonProps {
 const TextAndIconButton: React.FC<TextAndIconButtonProps> = ({ text, icon: Icon, onClick, title, variant }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", variant)} onClick={onClick} title={title}>
+        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
             {text}
             <span className="text-and-icon-button__icon">
                 <Icon />
diff --git a/mito-ai/src/components/TextButton.tsx b/mito-ai/src/components/TextButton.tsx
index a3f248a03..2125a258d 100644
--- a/mito-ai/src/components/TextButton.tsx
+++ b/mito-ai/src/components/TextButton.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import '../../style/TextAndIconButton.css';
+import '../../style/button.css';
 import { classNames } from '../utils/classNames';
 
 
@@ -16,7 +17,7 @@ interface TextButtonProps extends ButtonProps {}
 const TextButton: React.FC<TextButtonProps> = ({ text, onClick, title, variant }) => {
 
     return (
-        <button className={classNames("text-and-icon-button", variant)} onClick={onClick} title={title}>
+        <button className={classNames("text-and-icon-button", `button-${variant}`)} onClick={onClick} title={title}>
             {text}
         </button>
     )
diff --git a/mito-ai/src/icons/AcceptIcon.tsx b/mito-ai/src/icons/AcceptIcon.tsx
new file mode 100644
index 000000000..7735761b5
--- /dev/null
+++ b/mito-ai/src/icons/AcceptIcon.tsx
@@ -0,0 +1,12 @@
+import React from 'react';
+
+const AcceptIcon: React.FC = () => (
+    <svg width="1em" height="1em" viewBox="0 0 21 20" fill="none" xmlns="http://www.w3.org/2000/svg">
+        <g fill="currentColor">
+            <path d="M10.2236 0C4.70642 0 0.224487 4.48351 0.224487 10.0009C0.224487 15.5182 4.70642 20 10.2236 20C15.7408 20 20.2245 15.5181 20.2245 10.0009C20.2245 4.48368 15.741 0 10.2236 0ZM14.6192 5.46298C14.986 5.46298 15.3511 5.60114 15.6283 5.87746C16.1817 6.43091 16.1817 7.34456 15.6283 7.89886L9.56828 13.9589C9.01482 14.5123 8.10035 14.5123 7.54687 13.9589L4.85354 11.2648C4.30008 10.7114 4.30008 9.79937 4.85354 9.24518C5.407 8.69172 6.319 8.69172 6.87319 9.24518L8.55657 10.9294L13.6083 5.87851C13.8855 5.60054 14.2525 5.46298 14.6192 5.46298Z" />
+        </g>
+    </svg>
+
+);
+
+export default AcceptIcon;
\ No newline at end of file
diff --git a/mito-ai/src/icons/CopyIcon.tsx b/mito-ai/src/icons/CopyIcon.tsx
index eba9efa41..f6b770134 100644
--- a/mito-ai/src/icons/CopyIcon.tsx
+++ b/mito-ai/src/icons/CopyIcon.tsx
@@ -1,8 +1,3 @@
-
-
-
-
-
 import React from 'react';
 
 const CopyIcon: React.FC = () => (
diff --git a/mito-ai/src/icons/RejectIcon.tsx b/mito-ai/src/icons/RejectIcon.tsx
new file mode 100644
index 000000000..e77027170
--- /dev/null
+++ b/mito-ai/src/icons/RejectIcon.tsx
@@ -0,0 +1,12 @@
+import React from 'react';
+
+const RejectIcon: React.FC = () => (
+
+    <svg width="1em" height="1em" viewBox="0 0 21 20" fill="none" xmlns="http://www.w3.org/2000/svg">
+        <g fill="currentColor">
+            <path d="M10.2245 0C7.57229 0 5.02869 1.0536 3.15349 2.929C1.27817 4.8044 0.224487 7.348 0.224487 10C0.224487 12.652 1.27809 15.1958 3.15349 17.071C5.02889 18.9463 7.57249 20 10.2245 20C12.8765 20 15.4203 18.9464 17.2955 17.071C19.1708 15.1956 20.2245 12.652 20.2245 10C20.2245 8.2446 19.7624 6.5202 18.8847 5C18.0071 3.47984 16.7447 2.2174 15.2245 1.3398C13.7044 0.46206 11.9799 3.9978e-05 10.2245 3.9978e-05L10.2245 0ZM13.9315 12.293C14.1244 12.4794 14.2344 12.7354 14.2367 13.0036C14.2391 13.2718 14.1335 13.5298 13.9438 13.7194C13.7542 13.9091 13.4962 14.0146 13.228 14.0123C12.9598 14.01 12.7038 13.9 12.5174 13.7071L10.2244 11.4141L7.93143 13.7071C7.67737 13.9524 7.31283 14.0455 6.97229 13.9521C6.63175 13.8588 6.36565 13.5927 6.27229 13.2521C6.17893 12.9116 6.27205 12.5471 6.51737 12.293L8.81037 10L6.51737 7.707C6.27205 7.45294 6.17893 7.0884 6.27229 6.74786C6.36565 6.40732 6.63175 6.14122 6.97229 6.04786C7.31283 5.9545 7.67737 6.04763 7.93143 6.29294L10.2244 8.58594L12.5174 6.29294V6.29286C12.7715 6.04754 13.136 5.95442 13.4766 6.04778C13.8171 6.14114 14.0832 6.40724 14.1766 6.74778C14.2699 7.08832 14.1768 7.45286 13.9315 7.70692L11.6385 9.99992L13.9315 12.293Z" />
+        </g>
+    </svg>
+);
+
+export default RejectIcon;
\ No newline at end of file
diff --git a/mito-ai/style/ChatTaskpane.css b/mito-ai/style/ChatTaskpane.css
index 77b54f2c2..5cb25947b 100644
--- a/mito-ai/style/ChatTaskpane.css
+++ b/mito-ai/style/ChatTaskpane.css
@@ -147,23 +147,6 @@
 }
 
 .code-block-toolbar button {
+    height: 16px;
     font-size: 12px !important;
-}
-
-.code-block-accept-button {
-    background-color: var(--green-600) !important;
-    color: white !important;
-}
-
-.code-block-accept-button:hover {
-    background-color: var(--green-700) !important;
-}
-
-.code-block-deny-button {
-    background-color: var(--red-600) !important;
-    color: white !important;
-}
-
-.code-block-deny-button:hover {
-    background-color: var(--red-700) !important;
-}
+}
\ No newline at end of file
diff --git a/mito-ai/style/IconButton.css b/mito-ai/style/IconButton.css
index 47fa61fe7..9a7a683ed 100644
--- a/mito-ai/style/IconButton.css
+++ b/mito-ai/style/IconButton.css
@@ -1,4 +1,7 @@
 .icon-button {
+    display: flex;
+    align-items: center;
+    justify-content: center;
     background-color: transparent;
     border: none;
     cursor: pointer;
diff --git a/mito-ai/style/button.css b/mito-ai/style/button.css
new file mode 100644
index 000000000..e650a6425
--- /dev/null
+++ b/mito-ai/style/button.css
@@ -0,0 +1,32 @@
+/* 
+    Classes that can be used for any button, making it easier to keep 
+    the theme consistent. For example, use them to make a textButton or
+    textAndIconButton green.
+*/
+
+.button-green {
+    background-color: var(--green-400);
+    color: var(--green-900) !important;
+}
+
+.button-green:hover {
+    background-color: var(--green-500);
+}
+
+.button-red {
+    background-color: var(--red-400);
+    color: var(--red-900) !important;
+}
+
+.button-red:hover {
+    background-color: var(--red-500);
+}
+
+.button-gray {
+    background-color: var(--jp-layout-color2);
+    color: var(--jp-content-font-color1) !important;
+}
+
+.button-gray:hover {
+    background-color: var(--jp-layout-color3);
+}
\ No newline at end of file
diff --git a/tests/mitoai_ui_tests/utils.ts b/tests/mitoai_ui_tests/utils.ts
index 3a454b41b..be2d91d1f 100644
--- a/tests/mitoai_ui_tests/utils.ts
+++ b/tests/mitoai_ui_tests/utils.ts
@@ -67,7 +67,7 @@ export const editMitoAIMessage = async (
 }
 
 export const clickPreviewButton = async (page: IJupyterLabPageFixture) => {
-    await page.getByRole('button', { name: 'Overwrite Active Cell' }).click();
+    await page.locator('.chat-message-buttons').getByRole('button', { name: 'Overwrite Active Cell' }).click();
     await waitForIdle(page);
 }
 
@@ -78,9 +78,9 @@ export const clickAcceptButton = async (
     { useCellToolbar = false }: { useCellToolbar?: boolean } = {useCellToolbar: false}
 ) => {
     if (useCellToolbar) {
-        await page.getByText('Accept AI edits').click();
+        await page.locator('.jp-ToolbarButtonComponent-label').filter({ hasText: 'Accept' }).click();
     } else {
-        await page.locator('.chat-taskpane').getByRole('button', { name: 'Accept code' }).click();
+        await page.locator('.chat-message-buttons').getByRole('button', { name: 'Accept code' }).click();
     }
     await waitForIdle(page);
 }
@@ -92,9 +92,9 @@ export const clickDenyButton = async (
     { useCellToolbar = false }: { useCellToolbar?: boolean } = {useCellToolbar: false}
 ) => {
     if (useCellToolbar) {
-        await page.getByText('Reject AI edits').click();
+        await page.locator('.jp-ToolbarButtonComponent-label').filter({ hasText: 'Reject' }).click();
     } else {
-        await page.locator('.chat-taskpane').getByRole('button', { name: 'Reject code' }).click();
+        await page.locator('.chat-message-buttons').getByRole('button', { name: 'Reject code' }).click();
     }
     await waitForIdle(page);
 }