fix: enable removing tasklists and avatars (Chainlit#296)

- When splitting the frontend state for tasklists and avatars away from other elements, we introduced a but where it wouldn't be possible to remove the tasklist and avatars.
- Attempting to remove an element across the three existing element states (elements, tasklists, avatars) restores the feature.

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):

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");
   });
 });

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) => {