diff --git a/cypress/e2e/tasklist/main.py b/cypress/e2e/tasklist/main.py
index 149f56ba1a..2765f423c4 100644
--- a/cypress/e2e/tasklist/main.py
+++ b/cypress/e2e/tasklist/main.py
@@ -21,9 +21,22 @@
     "Preparing for shutdown",
 ]
 
+# Not a good practice in a normal chainlit server as it's global to all users
+# However it work in a testing scenario where we have just one user
+task_list = None
+
+
+@cl.on_message
+async def on_message():
+    # Waiting on a message to remove the tasklist to make sure
+    # all checks are successful before we remove it
+    await task_list.remove()
+
 
 @cl.on_chat_start
 async def main():
+    global task_list
+
     task_list = cl.TaskList()
     task_list.status = "Running..."
     for i in range(17):
diff --git a/cypress/e2e/tasklist/spec.cy.ts b/cypress/e2e/tasklist/spec.cy.ts
index 6d8a83b157..bd48e27831 100644
--- a/cypress/e2e/tasklist/spec.cy.ts
+++ b/cypress/e2e/tasklist/spec.cy.ts
@@ -1,8 +1,8 @@
-import { runTestServer } from "../../support/testUtils";
+import { runTestServer, submitMessage } from "../../support/testUtils";
 
 describe("tasklist", () => {
   before(() => {
-    runTestServer()
+    runTestServer();
   });
 
   it("should display the tasklist ", () => {
@@ -32,5 +32,9 @@ describe("tasklist", () => {
       "have.length",
       9
     );
+
+    submitMessage("ok");
+
+    cy.get(".tasklist").should("not.exist");
   });
 });
diff --git a/src/chainlit/frontend/src/components/socket.tsx b/src/chainlit/frontend/src/components/socket.tsx
index da8ebb0aa9..0107c79454 100644
--- a/src/chainlit/frontend/src/components/socket.tsx
+++ b/src/chainlit/frontend/src/components/socket.tsx
@@ -218,6 +218,12 @@ export default memo(function Socket() {
       setElements((old) => {
         return old.filter((e) => e.id !== remove.id);
       });
+      setTasklists((old) => {
+        return old.filter((e) => e.id !== remove.id);
+      });
+      setAvatars((old) => {
+        return old.filter((e) => e.id !== remove.id);
+      });
     });
 
     socket.on('action', (action: IAction) => {