fix(e2e): de-flake thread_resume by waiting on the resume replay

cypress/e2e/thread_resume/main.py

@@ -3,6 +3,7 @@
 
 import chainlit as cl
 import chainlit.data as cl_data
+from chainlit.data.utils import queue_until_user_message
 from chainlit.element import Element, ElementDict
 from chainlit.step import StepDict
 from chainlit.types import (
@@ -23,6 +24,42 @@
 THREADS: Dict[str, List[ThreadDict]] = {}
 
 
+def find_thread(thread_id: str) -> Optional[ThreadDict]:
+    """Return the stored thread with this id, whichever user bucket holds it."""
+    for threads in THREADS.values():
+        for thread in threads:
+            if thread["id"] == thread_id:
+                return thread
+    return None
+
+
+def ensure_thread(thread_id: str, user_id: Optional[str] = None) -> ThreadDict:
+    """Return the stored thread with this id, creating the row if needed.
+
+    Steps can be persisted before the first `update_thread` call, so a row may
+    have to be created without an owner. Re-home it when the owner shows up
+    instead of appending a second, empty row under the user.
+    """
+    thread = find_thread(thread_id)
+    if thread is None:
+        thread = {
+            "id": thread_id,
+            "createdAt": utc_now(),
+            "userId": user_id,
+            "userIdentifier": user_id,
+            "name": thread_id,
+            "steps": [],
+        }
+        THREADS.setdefault(user_id or "", []).append(thread)
+    elif user_id and thread.get("userIdentifier") != user_id:
+        for uid, threads in THREADS.items():
+            THREADS[uid] = [t for t in threads if t is not thread]
+        thread["userId"] = user_id
+        thread["userIdentifier"] = user_id
+        THREADS.setdefault(user_id, []).append(thread)
+    return thread
+
+
 class MemoryDataLayer(cl_data.BaseDataLayer):
     async def get_user(self, identifier: str):
         return cl.PersistedUser(id=identifier, createdAt=now, identifier=identifier)
@@ -55,14 +92,26 @@ async def get_element(
     async def delete_element(self, element_id: str, thread_id: Optional[str] = None):
         pass
 
+    @queue_until_user_message()
     async def create_step(self, step_dict: "StepDict"):
-        pass
-
+        thread_id = step_dict.get("threadId")
+        if not thread_id:
+            return
+        steps = ensure_thread(thread_id)["steps"]
+        for index, existing in enumerate(steps):
+            if existing["id"] == step_dict["id"]:
+                steps[index] = step_dict
+                return
+        steps.append(step_dict)
+
+    @queue_until_user_message()
     async def update_step(self, step_dict: "StepDict"):
-        pass
+        await self.create_step(step_dict)
 
+    @queue_until_user_message()
     async def delete_step(self, step_id: str):
-        pass
+        for thread in (t for threads in THREADS.values() for t in threads):
+            thread["steps"] = [s for s in thread["steps"] if s["id"] != step_id]
 
     async def get_thread_author(self, thread_id: str) -> str:
         return (await self.get_thread(thread_id))["userIdentifier"]
@@ -96,18 +145,7 @@ async def update_thread(
         metadata: Optional[Dict] = None,
         tags: Optional[List[str]] = None,
     ):
-        user_threads = THREADS.setdefault(user_id or "", [])
-        thr = next((t for t in user_threads if t["id"] == thread_id), None)
-        if not thr:
-            thr = {
-                "id": thread_id,
-                "createdAt": utc_now(),
-                "userId": user_id,
-                "userIdentifier": user_id,
-                "name": name or thread_id,
-                "steps": [],
-            }
-            user_threads.append(thr)
+        thr = ensure_thread(thread_id, user_id)
         if name:
             thr["name"] = name
         if metadata is not None:

cypress/e2e/thread_resume/spec.cy.ts

@@ -21,9 +21,24 @@ describe('Thread resume (author)', () => {
   it('resumes own thread, composer visible, can continue chatting', () => {
     login('alice');
 
-    // Start a thread
+    // `on_chat_start` sends its welcome message asynchronously. While the
+    // message list is empty `WelcomeScreen` renders the composer and `Footer`
+    // renders nothing; the first message flips `hasMessage()` and swaps which
+    // of the two owns it, remounting the input. Wait for that swap before
+    // typing, or the typed value is dropped mid-`type()`.
+    cy.get("[data-step-type='assistant_message']").should(
+      'contain',
+      'Welcome, say hi to start!'
+    );
+
+    // Start a thread and let the first turn finish, so the data layer has
+    // persisted it and there is a history for resume to replay.
     submitMessage('hi');
     cy.location('pathname').should('match', /\/thread\//);
+    cy.get("[data-step-type='assistant_message']").should(
+      'contain',
+      'Echo: hi'
+    );
 
     // Reload to trigger resume
     cy.reload();
@@ -32,6 +47,18 @@ describe('Thread resume (author)', () => {
     cy.get('#message-composer').should('be.visible');
     cy.get('[data-testid="read-only-banner"]').should('not.exist');
 
+    // `#message-composer` is visible well before the thread has hydrated. The
+    // backend's resume path emits, in order: the `@cl.on_chat_resume` message,
+    // then `resume_thread` with the persisted history -- and the client's
+    // `resume_thread` handler *replaces* the whole message list with it. So
+    // `Echo: hi` can only come from that replay, which is the last write to
+    // the list; once it is on screen every resume event has landed and the
+    // composer has stopped remounting.
+    cy.get("[data-step-type='assistant_message']").should(
+      'contain',
+      'Echo: hi'
+    );
+
     // Continue chatting
     submitMessage('still here');
     cy.get("[data-step-type='assistant_message']").contains('Echo: still here');

cypress/support/testUtils.ts

@@ -8,16 +8,25 @@ Cypress.on('uncaught:exception', (err) => {
   }
 });
 
+// The composer is re-rendered by React while a message round-trip is in flight,
+// so the input and the send button can be replaced between the two actions
+// below. Cypress retries queries but never re-runs an action command, so the
+// docs recommend a separate `cy.get()` per action rather than one chain:
+// "Call cy.get() separately for each action if the DOM node might be replaced
+// between steps." -- https://docs.cypress.io/app/core-concepts/retry-ability
+// ("Use separate queries for re-rendering elements"); see also
+// https://docs.cypress.io/app/core-concepts/interacting-with-elements
+// ("Detached").
 export function submitMessage(message: string) {
-  cy.get('#chat-input')
-    .should('be.visible')
-    .should('not.be.disabled')
-    .type(message);
-  cy.get('#chat-submit').should('not.be.disabled').click();
+  cy.get('#chat-input').should('be.visible').should('not.be.disabled');
+  cy.get('#chat-input').type(message);
+  cy.get('#chat-submit').should('not.be.disabled');
+  cy.get('#chat-submit').click();
 }
 
 export function openHistory() {
-  cy.get(`#chat-input`).should('not.be.disabled').type(`{upArrow}`);
+  cy.get(`#chat-input`).should('not.be.disabled');
+  cy.get(`#chat-input`).type(`{upArrow}`);
 }
 
 export function closeHistory() {