feat: add reasoning tests

.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/global.setup.ts → 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`;