diff --git a/tests/e2e/reasoning.test.ts b/tests/e2e/reasoning.test.ts
index 0ec4a9a..4491c83 100644
--- a/tests/e2e/reasoning.test.ts
+++ b/tests/e2e/reasoning.test.ts
@@ -48,8 +48,9 @@ test.describe("chat activity with reasoning", () => {
 
     const userMessage = await chatPage.getRecentUserMessage();
 
+    const generationCompletePromise = chatPage.isGenerationComplete();
     await userMessage.edit("Why is grass green?");
-    await chatPage.isGenerationComplete();
+    await generationCompletePromise;
 
     const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();