diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/components/edit-message.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/components/edit-message.tsx
index 42ec8117f..3f0ee6552 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/components/edit-message.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/components/edit-message.tsx
@@ -39,7 +39,7 @@ export const MarkdownField = ({
         linkTarget="_blank"
         rehypePlugins={[rehypeMathjax, rehypeRaw]}
         className={cn(
-          "markdown prose flex w-fit max-w-full flex-col items-baseline text-sm font-normal word-break-break-word dark:prose-invert",
+          "markdown prose flex w-full max-w-full flex-col items-baseline text-sm font-normal word-break-break-word dark:prose-invert",
           isEmpty ? "text-muted-foreground" : "text-primary",
         )}
         components={{
diff --git a/src/frontend/tests/core/features/freeze.spec.ts b/src/frontend/tests/core/features/freeze.spec.ts
index f5a1a6dc5..7cc56b460 100644
--- a/src/frontend/tests/core/features/freeze.spec.ts
+++ b/src/frontend/tests/core/features/freeze.spec.ts
@@ -2,6 +2,7 @@ import { expect, test } from "@playwright/test";
 import { addFlowToTestOnEmptyLangflow } from "../../utils/add-flow-to-test-on-empty-langflow";
 import { addLegacyComponents } from "../../utils/add-legacy-components";
 import { awaitBootstrapTest } from "../../utils/await-bootstrap-test";
+import { runChatOutput } from "../../utils/run-chat-output";
 import { zoomOut } from "../../utils/zoom-out";
 
 test(
@@ -162,7 +163,7 @@ test(
       .getByTestId("inputlist_str_urls_0")
       .fill("https://www.lipsum.com/");
 
-    await page.getByTestId("button_run_chat output").click();
+    await runChatOutput(page);
 
     await page.waitForSelector("text=built successfully", {
       timeout: 30000 * 3,
@@ -188,7 +189,7 @@ test(
 
     await page.getByTestId("textarea_str_input_value").first().fill(",");
 
-    await page.getByTestId("button_run_chat output").click();
+    await runChatOutput(page);
 
     await page.waitForSelector("text=built successfully", {
       timeout: 30000 * 3,
@@ -237,13 +238,9 @@ test(
       .first()
       .fill("lorem ipsum");
 
-    await page.waitForSelector('[data-testid="button_run_chat output"]', {
-      timeout: 1000,
-    });
-
     await page.waitForTimeout(2000);
 
-    await page.getByTestId("button_run_chat output").click();
+    await runChatOutput(page);
 
     await page.waitForSelector("text=built successfully", {
       timeout: 30000 * 3,
@@ -273,7 +270,7 @@ test(
 
     await page.keyboard.press("Escape");
 
-    await page.getByTestId("button_run_chat output").click();
+    await runChatOutput(page);
 
     await page.waitForSelector("text=built successfully", {
       timeout: 30000 * 3,
diff --git a/src/frontend/tests/utils/run-chat-output.ts b/src/frontend/tests/utils/run-chat-output.ts
new file mode 100644
index 000000000..7f436e752
--- /dev/null
+++ b/src/frontend/tests/utils/run-chat-output.ts
@@ -0,0 +1,14 @@
+import { Page } from "@playwright/test";
+
+export async function runChatOutput(page: Page) {
+  try {
+    await page.getByTestId("button_run_chat output").click({
+      timeout: 1000,
+    });
+  } catch (error) {
+    await page.getByTestId("generic-node-title-arrangement").last().click();
+    await page.getByTestId("more-options-modal").last().click();
+    await page.getByTestId("expand-button-modal").last().click();
+    await page.getByTestId("button_run_chat output").click();
+  }
+}