feat: add reasoning tests (#856)

.gitignore

@@ -41,5 +41,4 @@ yarn-error.log*
 /test-results/
 /playwright-report/
 /blob-report/
-/playwright/.cache/
-/playwright/.auth/
+/playwright/*

components/message-reasoning.tsx

@@ -44,6 +44,7 @@ export function MessageReasoning({
         <div className="flex flex-row gap-2 items-center">
           <div className="font-medium">Reasoned for a few seconds</div>
           <button
+            data-testid="message-reasoning-toggle"
             type="button"
             className="cursor-pointer"
             onClick={() => {
@@ -58,6 +59,7 @@ export function MessageReasoning({
       <AnimatePresence initial={false}>
         {isExpanded && (
           <motion.div
+            data-testid="message-reasoning"
             key="content"
             initial="collapsed"
             animate="expanded"

components/message.tsx

@@ -47,7 +47,7 @@ const PurePreviewMessage = ({
   return (
     <AnimatePresence>
       <motion.div
-        data-testid={`message-${message.role}-${index}`}
+        data-testid={`message-${message.role}`}
         className="w-full mx-auto max-w-3xl px-4 group/message"
         initial={{ y: 5, opacity: 0 }}
         animate={{ y: 0, opacity: 1 }}
@@ -73,7 +73,7 @@ const PurePreviewMessage = ({
           <div className="flex flex-col gap-4 w-full">
             {message.experimental_attachments && (
               <div
-                data-testid={`message-attachments-${index}`}
+                data-testid={`message-attachments`}
                 className="flex flex-row justify-end gap-2"
               >
                 {message.experimental_attachments.map((attachment) => (
@@ -93,12 +93,15 @@ const PurePreviewMessage = ({
             )}
 
             {(message.content || message.reasoning) && mode === 'view' && (
-              <div className="flex flex-row gap-2 items-start">
+              <div
+                data-testid="message-content"
+                className="flex flex-row gap-2 items-start"
+              >
                 {message.role === 'user' && !isReadonly && (
                   <Tooltip>
                     <TooltipTrigger asChild>
                       <Button
-                        data-testid={`edit-${message.role}-${index}`}
+                        data-testid={`message-edit`}
                         variant="ghost"
                         className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
                         onClick={() => {
@@ -243,6 +246,7 @@ export const ThinkingMessage = () => {
 
   return (
     <motion.div
+      data-testid="message-assistant-loading"
       className="w-full mx-auto max-w-3xl px-4 group/message "
       initial={{ y: 5, opacity: 0 }}
       animate={{ y: 0, opacity: 1, transition: { delay: 1 } }}

components/model-selector.tsx

@@ -39,7 +39,11 @@ export function ModelSelector({
           className,
         )}
       >
-        <Button variant="outline" className="md:px-2 md:h-[34px]">
+        <Button
+          data-testid="model-selector"
+          variant="outline"
+          className="md:px-2 md:h-[34px]"
+        >
           {selectedChatModel?.name}
           <ChevronDownIcon />
         </Button>
@@ -50,6 +54,7 @@ export function ModelSelector({
 
           return (
             <DropdownMenuItem
+              data-testid={`model-selector-item-${id}`}
               key={id}
               onSelect={() => {
                 setOpen(false);
@@ -59,19 +64,24 @@ export function ModelSelector({
                   saveChatModelAsCookie(id);
                 });
               }}
-              className="gap-4 group/item flex flex-row justify-between items-center"
               data-active={id === optimisticModelId}
+              asChild
             >
-              <div className="flex flex-col gap-1 items-start">
-                <div>{chatModel.name}</div>
-                <div className="text-xs text-muted-foreground">
-                  {chatModel.description}
+              <button
+                type="button"
+                className="gap-4 group/item flex flex-row justify-between items-center w-full"
+              >
+                <div className="flex flex-col gap-1 items-start">
+                  <div>{chatModel.name}</div>
+                  <div className="text-xs text-muted-foreground">
+                    {chatModel.description}
+                  </div>
                 </div>
-              </div>
 
-              <div className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
-                <CheckCircleFillIcon />
-              </div>
+                <div className="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100">
+                  <CheckCircleFillIcon />
+                </div>
+              </button>
             </DropdownMenuItem>
           );
         })}

lib/ai/models.test.ts

@@ -1,6 +1,11 @@
 import { CoreMessage, FinishReason, simulateReadableStream } from 'ai';
 import { MockLanguageModelV1 } from 'ai/test';
 
+interface ReasoningChunk {
+  type: 'reasoning';
+  textDelta: string;
+}
+
 interface TextDeltaChunk {
   type: 'text-delta';
   textDelta: string;
@@ -13,15 +18,73 @@ interface FinishChunk {
   usage: { completionTokens: number; promptTokens: number };
 }
 
-type Chunk = TextDeltaChunk | FinishChunk;
+type Chunk = TextDeltaChunk | ReasoningChunk | FinishChunk;
 
-const getResponseChunksByPrompt = (prompt: CoreMessage[]): Array<Chunk> => {
+const getResponseChunksByPrompt = (
+  prompt: CoreMessage[],
+  isReasoningEnabled: boolean = false,
+): Array<Chunk> => {
   const userMessage = prompt.at(-1);
 
   if (!userMessage) {
     throw new Error('No user message found');
   }
 
+  if (isReasoningEnabled) {
+    if (
+      compareMessages(userMessage, {
+        role: 'user',
+        content: [{ type: 'text', text: 'why is the sky blue?' }],
+      })
+    ) {
+      return [
+        { type: 'reasoning', textDelta: 'the ' },
+        { type: 'reasoning', textDelta: 'sky ' },
+        { type: 'reasoning', textDelta: 'is ' },
+        { type: 'reasoning', textDelta: 'blue ' },
+        { type: 'reasoning', textDelta: 'because ' },
+        { type: 'reasoning', textDelta: 'of ' },
+        { type: 'reasoning', textDelta: 'rayleigh ' },
+        { type: 'reasoning', textDelta: 'scattering! ' },
+        { type: 'text-delta', textDelta: "it's " },
+        { type: 'text-delta', textDelta: 'just ' },
+        { type: 'text-delta', textDelta: 'blue ' },
+        { type: 'text-delta', textDelta: 'duh!' },
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ];
+    } else if (
+      compareMessages(userMessage, {
+        role: 'user',
+        content: [{ type: 'text', text: 'why is grass green?' }],
+      })
+    ) {
+      return [
+        { type: 'reasoning', textDelta: 'grass ' },
+        { type: 'reasoning', textDelta: 'is ' },
+        { type: 'reasoning', textDelta: 'green ' },
+        { type: 'reasoning', textDelta: 'because ' },
+        { type: 'reasoning', textDelta: 'of ' },
+        { type: 'reasoning', textDelta: 'chlorophyll ' },
+        { type: 'reasoning', textDelta: 'absorption! ' },
+        { type: 'text-delta', textDelta: "it's " },
+        { type: 'text-delta', textDelta: 'just ' },
+        { type: 'text-delta', textDelta: 'green ' },
+        { type: 'text-delta', textDelta: 'duh!' },
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ];
+    }
+  }
+
   if (
     compareMessages(userMessage, {
       role: 'user',
@@ -135,17 +198,9 @@ export const reasoningModel = new MockLanguageModelV1({
     usage: { promptTokens: 10, completionTokens: 20 },
     text: `Hello, world!`,
   }),
-  doStream: async () => ({
+  doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
-      chunks: [
-        { type: 'text-delta', textDelta: 'test' },
-        {
-          type: 'finish',
-          finishReason: 'stop',
-          logprobs: undefined,
-          usage: { completionTokens: 10, promptTokens: 3 },
-        },
-      ],
+      chunks: getResponseChunksByPrompt(prompt, true),
     }),
     rawCall: { rawPrompt: null, rawSettings: {} },
   }),

playwright.config.ts

@@ -49,19 +49,38 @@ export default defineConfig({
     timeout: 30000,
   },
 
-  /* Configure projects for major browsers */
+  /* Configure projects */
   projects: [
     {
-      name: 'setup',
-      testMatch: /.*\.setup\.ts/,
+      name: 'setup:auth',
+      testMatch: /auth.setup.ts/,
     },
     {
-      name: 'chromium',
+      name: 'setup:reasoning',
+      testMatch: /reasoning.setup.ts/,
+      dependencies: ['setup:auth'],
       use: {
         ...devices['Desktop Chrome'],
-        storageState: 'playwright/.auth/user.json',
+        storageState: 'playwright/.auth/session.json',
+      },
+    },
+    {
+      name: 'chat',
+      testMatch: /chat.test.ts/,
+      dependencies: ['setup:auth'],
+      use: {
+        ...devices['Desktop Chrome'],
+        storageState: 'playwright/.auth/session.json',
+      },
+    },
+    {
+      name: 'reasoning',
+      testMatch: /reasoning.test.ts/,
+      dependencies: ['setup:reasoning'],
+      use: {
+        ...devices['Desktop Chrome'],
+        storageState: 'playwright/.reasoning/session.json',
       },
-      dependencies: ['setup'],
     },
 
     // {

tests/auth.setup.ts

@@ -3,7 +3,7 @@ import { generateId } from 'ai';
 import { getUnixTime } from 'date-fns';
 import { expect, test as setup } from '@playwright/test';
 
-const authFile = path.join(__dirname, '../playwright/.auth/user.json');
+const authFile = path.join(__dirname, '../playwright/.auth/session.json');
 
 setup('authenticate', async ({ page }) => {
   const testEmail = `test-${getUnixTime(new Date())}@playwright.com`;

tests/chat.test.ts

@@ -1,132 +1,37 @@
-import fs from 'fs';
-import path from 'path';
-import { test, expect, Page } from '@playwright/test';
-
-class ChatPage {
-  constructor(private page: Page) {}
-
-  async goto() {
-    await this.page.goto('/');
-  }
-
-  public getCurrentURL(): string {
-    return this.page.url();
-  }
-
-  async sendUserMessage(message: string) {
-    await this.page.getByTestId('multimodal-input').click();
-    await this.page.getByTestId('multimodal-input').fill(message);
-    await this.page.getByTestId('send-button').click();
-    await this.page.getByTestId('message-user-0').isVisible();
-    expect(await this.page.getByTestId('message-user-0').innerText()).toContain(
-      message,
-    );
-  }
-
-  async getAssistantResponse() {
-    return this.page.getByTestId('message-assistant-1').innerText();
-  }
-
-  async isGenerationComplete() {
-    await expect(this.page.getByTestId('send-button')).toBeVisible();
-    await this.page.waitForTimeout(2000);
-  }
-
-  async hasChatIdInUrl() {
-    await expect(this.page).toHaveURL(
-      /^http:\/\/localhost:3000\/chat\/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/,
-    );
-  }
-
-  async sendUserMessageFromSuggestion() {
-    await this.page
-      .getByRole('button', { name: 'What are the advantages of' })
-      .click();
-  }
-
-  async editAndResendUserMessage(message: string) {
-    await this.page.getByTestId('edit-user-0').click();
-    await this.page.getByTestId('message-editor').fill(message);
-    await this.page.getByTestId('message-editor-send-button').click();
-    await this.page.getByTestId('message-user-0').isVisible();
-    await expect(
-      this.page.getByTestId('message-editor-send-button'),
-    ).not.toBeVisible();
-    expect(await this.page.getByTestId('message-user-0').innerText()).toContain(
-      message,
-    );
-  }
-
-  async isElementVisible(elementId: string) {
-    await expect(this.page.getByTestId(elementId)).toBeVisible();
-  }
-
-  async isElementNotVisible(elementId: string) {
-    await expect(this.page.getByTestId(elementId)).not.toBeVisible();
-  }
-
-  async addImageAttachment() {
-    this.page.on('filechooser', async (fileChooser) => {
-      const filePath = path.join(
-        process.cwd(),
-        'public',
-        'images',
-        'mouth of the seine, monet.jpg',
-      );
-      const imageBuffer = fs.readFileSync(filePath);
-
-      await fileChooser.setFiles({
-        name: 'mouth of the seine, monet.jpg',
-        mimeType: 'image/jpeg',
-        buffer: imageBuffer,
-      });
-    });
-
-    await this.page.getByTestId('attachments-button').click();
-  }
-
-  public get sendButton() {
-    return this.page.getByTestId('send-button');
-  }
-
-  public get stopButton() {
-    return this.page.getByTestId('stop-button');
-  }
-
-  public get multimodalInput() {
-    return this.page.getByTestId('multimodal-input');
-  }
-}
+import { test, expect } from '@playwright/test';
+import { ChatPage } from './pages/chat';
 
 test.describe('chat activity', () => {
   let chatPage: ChatPage;
 
   test.beforeEach(async ({ page }) => {
     chatPage = new ChatPage(page);
-    await chatPage.goto();
+    await chatPage.createNewChat();
   });
 
   test('send a user message and receive response', async () => {
     await chatPage.sendUserMessage('why is grass green?');
     await chatPage.isGenerationComplete();
-    expect(await chatPage.getAssistantResponse()).toContain(
-      "it's just green duh!",
-    );
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toContain("it's just green duh!");
   });
 
   test('redirect to /chat/:id after submitting message', async () => {
     await chatPage.sendUserMessage('why is grass green?');
     await chatPage.isGenerationComplete();
-    expect(await chatPage.getAssistantResponse()).toContain(
-      "it's just green duh!",
-    );
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toContain("it's just green duh!");
     await chatPage.hasChatIdInUrl();
   });
 
   test('send a user message from suggestion', async () => {
     await chatPage.sendUserMessageFromSuggestion();
     await chatPage.isGenerationComplete();
-    expect(await chatPage.getAssistantResponse()).toContain(
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toContain(
       'with next.js you can ship fast!',
     );
   });
@@ -148,6 +53,7 @@ test.describe('chat activity', () => {
 
   test('stop generation during submission', async () => {
     await chatPage.sendUserMessage('why is grass green?');
+    await expect(chatPage.stopButton).toBeVisible();
     await chatPage.stopButton.click();
     await expect(chatPage.sendButton).toBeVisible();
   });
@@ -156,16 +62,16 @@ test.describe('chat activity', () => {
     await chatPage.sendUserMessage('why is grass green?');
     await chatPage.isGenerationComplete();
 
-    expect(await chatPage.getAssistantResponse()).toContain(
-      "it's just green duh!",
-    );
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toContain("it's just green duh!");
+
+    const userMessage = await chatPage.getRecentUserMessage();
+    await userMessage.edit('why is the sky blue?');
 
-    await chatPage.editAndResendUserMessage('why is the sky blue?');
     await chatPage.isGenerationComplete();
 
-    expect(await chatPage.getAssistantResponse()).toContain(
-      "it's just blue duh!",
-    );
+    const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(updatedAssistantMessage.content).toContain("it's just blue duh!");
   });
 
   test('hide suggested actions after sending message', async () => {
@@ -174,7 +80,7 @@ test.describe('chat activity', () => {
     await chatPage.isElementNotVisible('suggested-actions');
   });
 
-  test('handle file upload and send image attachment with message', async () => {
+  test('upload file and send image attachment with message', async () => {
     await chatPage.addImageAttachment();
 
     await chatPage.isElementVisible('attachments-preview');
@@ -182,13 +88,13 @@ test.describe('chat activity', () => {
     await chatPage.isElementNotVisible('input-attachment-loader');
 
     await chatPage.sendUserMessage('who painted this?');
-    await chatPage.isElementVisible('message-attachments-0');
+
+    const userMessage = await chatPage.getRecentUserMessage();
+    expect(userMessage.attachments).toHaveLength(1);
 
     await chatPage.isGenerationComplete();
 
-    await chatPage.isElementVisible('message-assistant-1');
-    expect(await chatPage.getAssistantResponse()).toContain(
-      'this painting is by monet!',
-    );
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toBe('this painting is by monet!');
   });
 });

tests/pages/chat.ts

@@ -0,0 +1,186 @@
+import fs from 'fs';
+import path from 'path';
+import { chatModels } from '@/lib/ai/models';
+import { expect, Page } from '@playwright/test';
+
+export class ChatPage {
+  constructor(private page: Page) {}
+
+  public get sendButton() {
+    return this.page.getByTestId('send-button');
+  }
+
+  public get stopButton() {
+    return this.page.getByTestId('stop-button');
+  }
+
+  public get multimodalInput() {
+    return this.page.getByTestId('multimodal-input');
+  }
+
+  async createNewChat() {
+    await this.page.goto('/');
+  }
+
+  public getCurrentURL(): string {
+    return this.page.url();
+  }
+
+  async sendUserMessage(message: string) {
+    await this.page.getByTestId('multimodal-input').click();
+    await this.page.getByTestId('multimodal-input').fill(message);
+    await this.page.getByTestId('send-button').click();
+    await this.page.getByTestId('message-user').isVisible();
+    expect(await this.page.getByTestId('message-user').innerText()).toContain(
+      message,
+    );
+  }
+
+  async isGenerationComplete() {
+    await expect(this.page.getByTestId('send-button')).toBeVisible();
+    await this.page.waitForTimeout(2000);
+  }
+
+  async hasChatIdInUrl() {
+    await expect(this.page).toHaveURL(
+      /^http:\/\/localhost:3000\/chat\/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/,
+    );
+  }
+
+  async sendUserMessageFromSuggestion() {
+    await this.page
+      .getByRole('button', { name: 'What are the advantages of' })
+      .click();
+  }
+
+  async isElementVisible(elementId: string) {
+    await expect(this.page.getByTestId(elementId)).toBeVisible();
+  }
+
+  async isElementNotVisible(elementId: string) {
+    await expect(this.page.getByTestId(elementId)).not.toBeVisible();
+  }
+
+  async addImageAttachment() {
+    this.page.on('filechooser', async (fileChooser) => {
+      const filePath = path.join(
+        process.cwd(),
+        'public',
+        'images',
+        'mouth of the seine, monet.jpg',
+      );
+      const imageBuffer = fs.readFileSync(filePath);
+
+      await fileChooser.setFiles({
+        name: 'mouth of the seine, monet.jpg',
+        mimeType: 'image/jpeg',
+        buffer: imageBuffer,
+      });
+    });
+
+    await this.page.getByTestId('attachments-button').click();
+  }
+
+  public async getSelectedModel() {
+    const modelId = await this.page.getByTestId('model-selector').innerText();
+    return modelId;
+  }
+
+  public async chooseModelFromSelector(chatModelId: string) {
+    const chatModel = chatModels.find(
+      (chatModel) => chatModel.id === chatModelId,
+    );
+
+    if (!chatModel) {
+      throw new Error(`Model with id ${chatModelId} not found`);
+    }
+
+    await this.page.getByTestId('model-selector').click();
+    await this.page.getByTestId(`model-selector-item-${chatModelId}`).click();
+    expect(await this.getSelectedModel()).toBe(chatModel.name);
+  }
+
+  async getRecentAssistantMessage() {
+    const messageElements = await this.page
+      .getByTestId('message-assistant')
+      .all();
+    const lastMessageElement = messageElements[messageElements.length - 1];
+
+    const content = await lastMessageElement
+      .getByTestId('message-content')
+      .innerText()
+      .catch(() => null);
+
+    const reasoningElement = await lastMessageElement
+      .getByTestId('message-reasoning')
+      .isVisible()
+      .then(async (visible) =>
+        visible
+          ? await lastMessageElement
+              .getByTestId('message-reasoning')
+              .innerText()
+          : null,
+      )
+      .catch(() => null);
+
+    const page = this.page;
+
+    return {
+      element: lastMessageElement,
+      content,
+      reasoning: reasoningElement,
+      async waitForGenerationComplete() {
+        await expect(page.getByTestId('send-button')).toBeVisible();
+        await page.waitForTimeout(2000);
+      },
+      async toggleReasoningVisibility() {
+        await lastMessageElement
+          .getByTestId('message-reasoning-toggle')
+          .click();
+      },
+    };
+  }
+
+  async getRecentUserMessage() {
+    const messageElements = await this.page.getByTestId('message-user').all();
+    const lastMessageElement = messageElements[messageElements.length - 1];
+
+    const content = await lastMessageElement.innerText();
+
+    const hasAttachments = await lastMessageElement
+      .getByTestId('message-attachments')
+      .isVisible()
+      .catch(() => false);
+
+    const attachments = hasAttachments
+      ? await lastMessageElement.getByTestId('message-attachments').all()
+      : [];
+
+    const page = this.page;
+
+    return {
+      element: lastMessageElement,
+      content,
+      attachments,
+      async edit(newMessage: string) {
+        await page.getByTestId('message-edit').click();
+        await page.getByTestId('message-editor').fill(newMessage);
+        await page.getByTestId('message-editor-send-button').click();
+        await expect(
+          page.getByTestId('message-editor-send-button'),
+        ).not.toBeVisible();
+      },
+    };
+  }
+
+  async waitForMessageGeneration(timeout = 10000) {
+    await this.page.waitForFunction(
+      () => {
+        return document.querySelector('[data-testid="send-button"]') !== null;
+      },
+      { timeout },
+    );
+
+    await this.page.waitForTimeout(500);
+  }
+}

tests/reasoning.setup.ts

@@ -0,0 +1,20 @@
+import path from 'path';
+import { expect, test as setup } from '@playwright/test';
+import { ChatPage } from './pages/chat';
+
+const reasoningFile = path.join(
+  __dirname,
+  '../playwright/.reasoning/session.json',
+);
+
+setup('switch to reasoning model', async ({ page }) => {
+  const chatPage = new ChatPage(page);
+  await chatPage.createNewChat();
+
+  await chatPage.chooseModelFromSelector('chat-model-reasoning');
+
+  await expect(chatPage.getSelectedModel()).resolves.toEqual('Reasoning model');
+
+  await page.waitForTimeout(1000);
+  await page.context().storageState({ path: reasoningFile });
+});

tests/reasoning.test.ts

@@ -0,0 +1,64 @@
+import { useAssistant } from '@ai-sdk/react';
+import { ChatPage } from './pages/chat';
+import { test, expect } from '@playwright/test';
+
+test.describe('chat activity with reasoning', () => {
+  let chatPage: ChatPage;
+
+  test.beforeEach(async ({ page }) => {
+    chatPage = new ChatPage(page);
+    await chatPage.createNewChat();
+  });
+
+  test('send user message and generate response with reasoning', async () => {
+    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.waitForMessageGeneration();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+
+    expect(assistantMessage.content).toBe("it's just blue duh!");
+
+    expect(assistantMessage.reasoning).toBe(
+      'the sky is blue because of rayleigh scattering!',
+    );
+  });
+
+  test('toggle reasoning visibility', async () => {
+    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.waitForMessageGeneration();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    const reasoningElement =
+      assistantMessage.element.getByTestId('message-reasoning');
+    expect(reasoningElement).toBeVisible();
+
+    await assistantMessage.toggleReasoningVisibility();
+    await expect(reasoningElement).not.toBeVisible();
+
+    await assistantMessage.toggleReasoningVisibility();
+    await expect(reasoningElement).toBeVisible();
+  });
+
+  test('edit message and resubmit', async () => {
+    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.waitForMessageGeneration();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    const reasoningElement =
+      assistantMessage.element.getByTestId('message-reasoning');
+    expect(reasoningElement).toBeVisible();
+
+    const userMessage = await chatPage.getRecentUserMessage();
+
+    await userMessage.edit('why is grass green?');
+    await chatPage.waitForMessageGeneration();
+
+    const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();
+
+    expect(updatedAssistantMessage.content).toBe("it's just green duh!");
+
+    expect(updatedAssistantMessage.reasoning).toBe(
+      'grass is green because of chlorophyll absorption!',
+    );
+  });
+});