diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx
index 435f231b1..08e466b10 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx
@@ -166,18 +166,20 @@ export default function ChatInput({
 
   const send = async () => {
     const storedChatValue = chatValue;
+    const filesToSend = files
+      .map((file) => file.path ?? "")
+      .filter((file) => file !== "");
+    const storedFiles = [...files];
+    setFiles([]);
     try {
       await sendMessage({
         repeat: 1,
-        files: files
-          .map((file) => file.path ?? "")
-          .filter((file) => file !== ""),
+        files: filesToSend,
       });
     } catch (error) {
       setChatValueStore(storedChatValue);
+      setFiles(storedFiles);
     }
-
-    setFiles([]);
   };
 
   const checkSendingOk = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/no-input.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/no-input.tsx
index a05add60a..bf690eb16 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/no-input.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/no-input.tsx
@@ -7,7 +7,7 @@ import { cn } from "../../../../../../utils/utils";
 
 interface NoInputViewProps {
   isBuilding: boolean;
-  sendMessage: (args: { repeat: number }) => void;
+  sendMessage: (args: { repeat: number }) => Promise<void>;
   stopBuilding: () => void;
 }
 
@@ -23,8 +23,8 @@ const NoInputView: React.FC<NoInputViewProps> = ({
           <Button
             data-testid="button-send"
             className="font-semibold"
-            onClick={() => {
-              sendMessage({
+            onClick={async () => {
+              await sendMessage({
                 repeat: 1,
               });
             }}
diff --git a/src/frontend/src/modals/IOModal/playground-modal.tsx b/src/frontend/src/modals/IOModal/playground-modal.tsx
index 728b7bc38..9ef1b2ec1 100644
--- a/src/frontend/src/modals/IOModal/playground-modal.tsx
+++ b/src/frontend/src/modals/IOModal/playground-modal.tsx
@@ -356,7 +356,7 @@ export default function IOModal({
       setOpen={setOpen}
       disable={disable}
       type={isPlayground ? "full-screen" : undefined}
-      onSubmit={() => sendMessage({ repeat: 1 })}
+      onSubmit={async () => await sendMessage({ repeat: 1 })}
       size="x-large"
       className="!rounded-[12px] p-0"
     >
diff --git a/src/frontend/tests/core/unit/fileUploadComponent.spec.ts b/src/frontend/tests/core/unit/fileUploadComponent.spec.ts
index 354e6c6d7..9df144a59 100644
--- a/src/frontend/tests/core/unit/fileUploadComponent.spec.ts
+++ b/src/frontend/tests/core/unit/fileUploadComponent.spec.ts
@@ -304,19 +304,18 @@ test(
       await page.dispatchEvent('[data-testid="drag-files-component"]', "drop", {
         dataTransfer,
       });
-
       await expect(page.getByText(`${newTxtFile}.txt`).last()).toBeVisible({
         timeout: 1000,
       });
 
       await expect(
         page.getByTestId(`checkbox-${newTxtFile}`).last(),
-      ).toHaveAttribute("data-state", "checked", { timeout: 1000 });
+      ).toHaveAttribute("data-state", "checked", { timeout: 3000 });
 
       await page.getByTestId("select-files-modal-button").click();
       await expect(page.getByText(`${renamedJsonFile}.txt`).first()).toBeHidden(
         {
-          timeout: 1000,
+          timeout: 3000,
         },
       );
       await expect(page.getByText(`${newTxtFile}.txt`).first()).toBeVisible({
@@ -324,16 +323,6 @@ test(
       });
       await page.getByTestId(`remove-file-button-${renamedTxtFile}`).click();
 
-      await page
-        .getByTestId("handle-file-shownode-raw content-right")
-        .first()
-        .click();
-
-      await page
-        .getByTestId("handle-chatoutput-noshownode-inputs-target")
-        .first()
-        .click();
-
       await page
         .getByRole("button", { name: "Playground", exact: true })
         .click();