From a86cfc6b37e5b4708bdab008f4b1fd291edc5f33 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sun, 29 Dec 2024 12:06:27 -0500 Subject: [PATCH 1/8] mito-ai: add overwrite active cell button to code block --- .../AiChat/ChatMessage/ChatMessage.tsx | 1 + .../Extensions/AiChat/ChatMessage/CodeBlock.tsx | 16 +++++++++++++++- mito-ai/style/IconButton.css | 3 +++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index acaad090e..29662ea65 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -106,6 +106,7 @@ const ChatMessage: React.FC = ({ code={messagePart} role={message.role} renderMimeRegistry={renderMimeRegistry} + previewAICode={previewAICode} /> {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..a56551003 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -3,18 +3,23 @@ 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'; interface ICodeBlockProps { code: string, role: 'user' | 'assistant' renderMimeRegistry: IRenderMimeRegistry + previewAICode: () => void } const CodeBlock: React.FC = ({ code, role, renderMimeRegistry, + previewAICode, }): JSX.Element => { if (role === 'user') { @@ -32,7 +37,16 @@ const CodeBlock: React.FC = ({ return (
- + } + title="Overwrite Active Cell" + onClick={() => {previewAICode()}} + /> + } + title="Copy" + onClick={() => {copyToClipboard(code)}} + />
Date: Sun, 29 Dec 2024 12:09:28 -0500 Subject: [PATCH 2/8] mito-ai: only show button if last message --- .../Extensions/AiChat/ChatMessage/CodeBlock.tsx | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index a56551003..ff50f20e3 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -13,6 +13,7 @@ interface ICodeBlockProps { role: 'user' | 'assistant' renderMimeRegistry: IRenderMimeRegistry previewAICode: () => void + isLastAiMessage: boolean } const CodeBlock: React.FC = ({ @@ -20,6 +21,7 @@ const CodeBlock: React.FC = ({ role, renderMimeRegistry, previewAICode, + isLastAiMessage }): JSX.Element => { if (role === 'user') { @@ -37,11 +39,13 @@ const CodeBlock: React.FC = ({ return (
- } - title="Overwrite Active Cell" - onClick={() => {previewAICode()}} - /> + {isLastAiMessage && + } + title="Overwrite Active Cell" + onClick={() => {previewAICode()}} + /> + } } title="Copy" From dfeeb656725e69ba9f64d19aee49949b39b40475 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sun, 29 Dec 2024 12:32:38 -0500 Subject: [PATCH 3/8] mito-ai: add accept and reject icons --- .../AiChat/ChatMessage/ChatMessage.tsx | 4 ++ .../AiChat/ChatMessage/CodeBlock.tsx | 39 +++++++++++++++---- mito-ai/src/icons/AcceptIcon.tsx | 12 ++++++ mito-ai/src/icons/CopyIcon.tsx | 5 --- mito-ai/src/icons/RejectIcon.tsx | 12 ++++++ mito-ai/style/ChatTaskpane.css | 1 + 6 files changed, 61 insertions(+), 12 deletions(-) create mode 100644 mito-ai/src/icons/AcceptIcon.tsx create mode 100644 mito-ai/src/icons/RejectIcon.tsx diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx index 29662ea65..f5d969e12 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx @@ -107,6 +107,10 @@ const ChatMessage: React.FC = ({ 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 ff50f20e3..b9526571f 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -6,6 +6,9 @@ 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 { @@ -13,7 +16,10 @@ interface ICodeBlockProps { role: 'user' | 'assistant' renderMimeRegistry: IRenderMimeRegistry previewAICode: () => void + acceptAICode: () => void + rejectAICode: () => void isLastAiMessage: boolean + codeReviewStatus: CodeReviewStatus } const CodeBlock: React.FC = ({ @@ -21,7 +27,10 @@ const CodeBlock: React.FC = ({ role, renderMimeRegistry, previewAICode, - isLastAiMessage + acceptAICode, + rejectAICode, + isLastAiMessage, + codeReviewStatus, }): JSX.Element => { if (role === 'user') { @@ -39,18 +48,34 @@ const CodeBlock: React.FC = ({ return (
- {isLastAiMessage && + {codeReviewStatus === 'chatPreview' && isLastAiMessage && } title="Overwrite Active Cell" onClick={() => {previewAICode()}} /> } - } - title="Copy" - onClick={() => {copyToClipboard(code)}} - /> + {codeReviewStatus === 'chatPreview' && + } + title="Copy" + onClick={() => {copyToClipboard(code)}} + /> + } + {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && + } + title="Accept AI Generated Code" + onClick={() => {acceptAICode()}} + /> + } + {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && + } + title="Reject AI Generated Code" + onClick={() => {rejectAICode()}} + /> + }
( + + + + + + +); + +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 = () => ( + + + + + + +); + +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..d9ef7c976 100644 --- a/mito-ai/style/ChatTaskpane.css +++ b/mito-ai/style/ChatTaskpane.css @@ -147,6 +147,7 @@ } .code-block-toolbar button { + height: 16px; font-size: 12px !important; } From 7b1c49529b9333d593d474f1e9311453e29cbd44 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sun, 29 Dec 2024 12:49:01 -0500 Subject: [PATCH 4/8] mito-ai: refactor css --- .../AiChat/ChatMessage/CodeBlock.tsx | 2 ++ .../src/Extensions/AiChat/ChatTaskpane.tsx | 4 +-- mito-ai/src/components/IconButton.tsx | 7 ++-- mito-ai/src/components/TextAndIconButton.tsx | 3 +- mito-ai/src/components/TextButton.tsx | 3 +- mito-ai/style/TextAndIconButton.css | 32 ------------------- mito-ai/style/button.css | 32 +++++++++++++++++++ 7 files changed, 45 insertions(+), 38 deletions(-) create mode 100644 mito-ai/style/button.css diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index b9526571f..e9454dbb4 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -67,6 +67,7 @@ const CodeBlock: React.FC = ({ icon={} title="Accept AI Generated Code" onClick={() => {acceptAICode()}} + variant='green' /> } {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && @@ -74,6 +75,7 @@ const CodeBlock: React.FC = ({ icon={} title="Reject AI Generated Code" onClick={() => {rejectAICode()}} + variant='red' /> }
diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 83f2833fe..c5b9af1fc 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -389,7 +389,7 @@ const ChatTaskpane: React.FC = ({ */ app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_ACCEPT_CODE, { label: `Accept ${operatingSystem === 'mac' ? '⌘Y' : 'Ctrl+Y'}`, - className: 'text-and-icon-button green small', + 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 @@ -405,7 +405,7 @@ const ChatTaskpane: React.FC = ({ app.commands.addCommand(COMMAND_MITO_AI_CELL_TOOLBAR_REJECT_CODE, { label: `Reject ${operatingSystem === 'mac' ? '⌘U' : 'Ctrl+U'}`, - className: 'text-and-icon-button red small', + 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..f13cb7713 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; + variant?: 'green' | 'red'; } -const IconButton: React.FC = ({ icon, onClick, title }) => { +const IconButton: React.FC = ({ icon, onClick, title, variant }) => { return ( - ) 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 = ({ text, icon: Icon, onClick, title, variant }) => { return ( - ) diff --git a/mito-ai/style/TextAndIconButton.css b/mito-ai/style/TextAndIconButton.css index 435a91dc2..e0055da21 100644 --- a/mito-ai/style/TextAndIconButton.css +++ b/mito-ai/style/TextAndIconButton.css @@ -22,36 +22,4 @@ .text-and-icon-button__icon > svg { height: 100%; width: auto; -} - -.text-and-icon-button.green { - background-color: var(--green-400); - color: var(--green-900); -} - -.text-and-icon-button.green:hover { - background-color: var(--green-500); -} - -.text-and-icon-button.red { - background-color: var(--red-400); - color: var(--red-900); -} - -.text-and-icon-button.red:hover { - background-color: var(--red-500); -} - -.text-and-icon-button.small { - padding: 2px; - font-size: 12px; -} - -.text-and-icon-button.gray { - background-color: var(--jp-layout-color2); - color: var(--jp-content-font-color1); -} - -.text-and-icon-button.gray:hover { - background-color: var(--jp-layout-color3); } \ No newline at end of file diff --git a/mito-ai/style/button.css b/mito-ai/style/button.css new file mode 100644 index 000000000..e1e01570a --- /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); +} + +.button-green:hover { + background-color: var(--green-500); +} + +.button-red { + background-color: var(--red-400); + color: var(--red-900); +} + +.button-red:hover { + background-color: var(--red-500); +} + +.button-gray { + background-color: var(--jp-layout-color2); + color: var(--jp-content-font-color1); +} + +.button-gray:hover { + background-color: var(--jp-layout-color3); +} \ No newline at end of file From f591dbedf07fcc4f80167189fa4cd8d3824d1bd4 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sun, 29 Dec 2024 12:57:37 -0500 Subject: [PATCH 5/8] mito-ai: add style to buttons --- .../AiChat/ChatMessage/CodeBlock.tsx | 4 ++-- mito-ai/src/components/IconButton.tsx | 6 +++--- mito-ai/style/ChatTaskpane.css | 20 +------------------ mito-ai/style/button.css | 6 +++--- 4 files changed, 9 insertions(+), 27 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index e9454dbb4..9e692669e 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -67,7 +67,7 @@ const CodeBlock: React.FC = ({ icon={} title="Accept AI Generated Code" onClick={() => {acceptAICode()}} - variant='green' + style={{color: 'var(--green-700)'}} /> } {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && @@ -75,7 +75,7 @@ const CodeBlock: React.FC = ({ icon={} title="Reject AI Generated Code" onClick={() => {rejectAICode()}} - variant='red' + style={{color: 'var(--red-700)'}} /> }
diff --git a/mito-ai/src/components/IconButton.tsx b/mito-ai/src/components/IconButton.tsx index f13cb7713..2b460964f 100644 --- a/mito-ai/src/components/IconButton.tsx +++ b/mito-ai/src/components/IconButton.tsx @@ -7,13 +7,13 @@ interface IconButtonProps { icon: React.ReactNode; onClick: () => void; title: string; - variant?: 'green' | 'red'; + style?: React.CSSProperties; } -const IconButton: React.FC = ({ icon, onClick, title, variant }) => { +const IconButton: React.FC = ({ icon, onClick, title, style }) => { return ( - ) diff --git a/mito-ai/style/ChatTaskpane.css b/mito-ai/style/ChatTaskpane.css index d9ef7c976..5cb25947b 100644 --- a/mito-ai/style/ChatTaskpane.css +++ b/mito-ai/style/ChatTaskpane.css @@ -149,22 +149,4 @@ .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/button.css b/mito-ai/style/button.css index e1e01570a..e650a6425 100644 --- a/mito-ai/style/button.css +++ b/mito-ai/style/button.css @@ -6,7 +6,7 @@ .button-green { background-color: var(--green-400); - color: var(--green-900); + color: var(--green-900) !important; } .button-green:hover { @@ -15,7 +15,7 @@ .button-red { background-color: var(--red-400); - color: var(--red-900); + color: var(--red-900) !important; } .button-red:hover { @@ -24,7 +24,7 @@ .button-gray { background-color: var(--jp-layout-color2); - color: var(--jp-content-font-color1); + color: var(--jp-content-font-color1) !important; } .button-gray:hover { From 4e6ea9e601856a1e38c99a1d49a9fe95b1495dfc Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Sun, 29 Dec 2024 13:14:14 -0500 Subject: [PATCH 6/8] mito-ai: cleanup code block toolbar logic" --- .../AiChat/ChatMessage/CodeBlock.tsx | 75 +++++++++++-------- 1 file changed, 45 insertions(+), 30 deletions(-) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx index 9e692669e..c25c9b030 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/CodeBlock.tsx @@ -47,38 +47,53 @@ const CodeBlock: React.FC = ({ if (role === 'assistant') { return (
-
- {codeReviewStatus === 'chatPreview' && isLastAiMessage && - } - title="Overwrite Active Cell" - onClick={() => {previewAICode()}} - /> + <> + {/* The code block toolbar for the last AI message */} + {isLastAiMessage && +
+ {codeReviewStatus === 'chatPreview' && + } + title="Overwrite Active Cell" + onClick={() => {previewAICode()}} + /> + } + {codeReviewStatus === 'codeCellPreview' && + } + title="Accept AI Generated Code" + onClick={() => {acceptAICode()}} + style={{color: 'var(--green-700)'}} + /> + } + {codeReviewStatus === 'codeCellPreview' && + } + title="Reject AI Generated Code" + onClick={() => {rejectAICode()}} + style={{color: 'var(--red-700)'}} + /> + } + {codeReviewStatus !== 'codeCellPreview' && + } + title="Copy" + onClick={() => {copyToClipboard(code)}} + /> + } +
} - {codeReviewStatus === 'chatPreview' && - } - title="Copy" - onClick={() => {copyToClipboard(code)}} - /> + {/* The code block toolbar for every other AI message */} + {!isLastAiMessage && +
+ } + title="Copy" + onClick={() => {copyToClipboard(code)}} + /> +
} - {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && - } - title="Accept AI Generated Code" - onClick={() => {acceptAICode()}} - style={{color: 'var(--green-700)'}} - /> - } - {codeReviewStatus === 'codeCellPreview' && isLastAiMessage && - } - title="Reject AI Generated Code" - onClick={() => {rejectAICode()}} - style={{color: 'var(--red-700)'}} - /> - } -
+ Date: Sun, 29 Dec 2024 14:13:40 -0500 Subject: [PATCH 7/8] tests: fix tests --- tests/mitoai_ui_tests/utils.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/mitoai_ui_tests/utils.ts b/tests/mitoai_ui_tests/utils.ts index f5e492e3d..a5b3dcbcd 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); } @@ -82,7 +82,7 @@ export const clickAcceptButton = async ( await page.getByLabel('notebook content').getByText('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); } @@ -99,7 +99,7 @@ export const clickDenyButton = async ( // up and down a bunch of times until it times out. So instead, we use this selector await page.getByLabel('notebook content').getByText('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); } From dab94d183e11ae31aea1ca0985807d66e5426292 Mon Sep 17 00:00:00 2001 From: Nawaz Gafar Date: Mon, 30 Dec 2024 12:05:46 -0500 Subject: [PATCH 8/8] Refactor btn click for useCellToolbar --- tests/mitoai_ui_tests/utils.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/mitoai_ui_tests/utils.ts b/tests/mitoai_ui_tests/utils.ts index 27bf7d918..be2d91d1f 100644 --- a/tests/mitoai_ui_tests/utils.ts +++ b/tests/mitoai_ui_tests/utils.ts @@ -78,7 +78,7 @@ 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-message-buttons').getByRole('button', { name: 'Accept code' }).click(); } @@ -92,7 +92,7 @@ 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-message-buttons').getByRole('button', { name: 'Reject code' }).click(); }