diff --git a/components/artifact-messages.tsx b/components/artifact-messages.tsx
index b7cddaa..e4097d4 100644
--- a/components/artifact-messages.tsx
+++ b/components/artifact-messages.tsx
@@ -1,29 +1,26 @@
 import { PreviewMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Vote } from '@/lib/db/schema';
-import { ChatRequestOptions, Message } from 'ai';
+import { Message } from 'ai';
 import { memo } from 'react';
 import equal from 'fast-deep-equal';
 import { UIArtifact } from './artifact';
+import { UseChatHelpers } from '@ai-sdk/react';
 
 interface ArtifactMessagesProps {
   chatId: string;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
   votes: Array<Vote> | undefined;
   messages: Array<Message>;
-  setMessages: (
-    messages: Message[] | ((messages: Message[]) => Message[]),
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  setMessages: UseChatHelpers['setMessages'];
+  reload: UseChatHelpers['reload'];
   isReadonly: boolean;
   artifactStatus: UIArtifact['status'];
 }
 
 function PureArtifactMessages({
   chatId,
-  isLoading,
+  status,
   votes,
   messages,
   setMessages,
@@ -43,8 +40,7 @@ function PureArtifactMessages({
           chatId={chatId}
           key={message.id}
           message={message}
-          isLoading={isLoading && index === messages.length - 1}
-          index={index}
+          isLoading={status === 'streaming' && index === messages.length - 1}
           vote={
             votes
               ? votes.find((vote) => vote.messageId === message.id)
@@ -74,8 +70,8 @@ function areEqual(
   )
     return true;
 
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
-  if (prevProps.isLoading && nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
+  if (prevProps.status && nextProps.status) return false;
   if (prevProps.messages.length !== nextProps.messages.length) return false;
   if (!equal(prevProps.votes, nextProps.votes)) return false;
 
diff --git a/components/artifact.tsx b/components/artifact.tsx
index 5be9f1a..ec7a7b5 100644
--- a/components/artifact.tsx
+++ b/components/artifact.tsx
@@ -1,9 +1,4 @@
-import type {
-  Attachment,
-  ChatRequestOptions,
-  CreateMessage,
-  Message,
-} from 'ai';
+import type { Attachment, Message } from 'ai';
 import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
@@ -31,6 +26,7 @@ import { codeArtifact } from '@/artifacts/code/client';
 import { sheetArtifact } from '@/artifacts/sheet/client';
 import { textArtifact } from '@/artifacts/text/client';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers } from '@ai-sdk/react';
 
 export const artifactDefinitions = [
   textArtifact,
@@ -60,7 +56,7 @@ function PureArtifact({
   input,
   setInput,
   handleSubmit,
-  isLoading,
+  status,
   stop,
   attachments,
   setAttachments,
@@ -73,27 +69,17 @@ function PureArtifact({
 }: {
   chatId: string;
   input: string;
-  setInput: (input: string) => void;
-  isLoading: boolean;
-  stop: () => void;
+  setInput: UseChatHelpers['setInput'];
+  status: UseChatHelpers['status'];
+  stop: UseChatHelpers['stop'];
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
   messages: Array<Message>;
   setMessages: Dispatch<SetStateAction<Array<Message>>>;
   votes: Array<Vote> | undefined;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  handleSubmit: (
-    event?: {
-      preventDefault?: () => void;
-    },
-    chatRequestOptions?: ChatRequestOptions,
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  append: UseChatHelpers['append'];
+  handleSubmit: UseChatHelpers['handleSubmit'];
+  reload: UseChatHelpers['reload'];
   isReadonly: boolean;
 }) {
   const { artifact, setArtifact, metadata, setMetadata } = useArtifact();
@@ -326,7 +312,7 @@ function PureArtifact({
               <div className="flex flex-col h-full justify-between items-center gap-4">
                 <ArtifactMessages
                   chatId={chatId}
-                  isLoading={isLoading}
+                  status={status}
                   votes={votes}
                   messages={messages}
                   setMessages={setMessages}
@@ -341,7 +327,7 @@ function PureArtifact({
                     input={input}
                     setInput={setInput}
                     handleSubmit={handleSubmit}
-                    isLoading={isLoading}
+                    status={status}
                     stop={stop}
                     attachments={attachments}
                     setAttachments={setAttachments}
@@ -487,7 +473,7 @@ function PureArtifact({
                     isToolbarVisible={isToolbarVisible}
                     setIsToolbarVisible={setIsToolbarVisible}
                     append={append}
-                    isLoading={isLoading}
+                    status={status}
                     stop={stop}
                     setMessages={setMessages}
                     artifactKind={artifact.kind}
@@ -513,7 +499,7 @@ function PureArtifact({
 }
 
 export const Artifact = memo(PureArtifact, (prevProps, nextProps) => {
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
   if (!equal(prevProps.votes, nextProps.votes)) return false;
   if (prevProps.input !== nextProps.input) return false;
   if (!equal(prevProps.messages, nextProps.messages.length)) return false;
diff --git a/components/chat.tsx b/components/chat.tsx
index 4a5f236..2f40593 100644
--- a/components/chat.tsx
+++ b/components/chat.tsx
@@ -36,7 +36,7 @@ export function Chat({
     input,
     setInput,
     append,
-    isLoading,
+    status,
     stop,
     reload,
   } = useChat({
@@ -74,7 +74,7 @@ export function Chat({
 
         <Messages
           chatId={id}
-          isLoading={isLoading}
+          status={status}
           votes={votes}
           messages={messages}
           setMessages={setMessages}
@@ -90,7 +90,7 @@ export function Chat({
               input={input}
               setInput={setInput}
               handleSubmit={handleSubmit}
-              isLoading={isLoading}
+              status={status}
               stop={stop}
               attachments={attachments}
               setAttachments={setAttachments}
@@ -107,7 +107,7 @@ export function Chat({
         input={input}
         setInput={setInput}
         handleSubmit={handleSubmit}
-        isLoading={isLoading}
+        status={status}
         stop={stop}
         attachments={attachments}
         setAttachments={setAttachments}
diff --git a/components/message.tsx b/components/message.tsx
index 7b42222..ddbbf68 100644
--- a/components/message.tsx
+++ b/components/message.tsx
@@ -27,7 +27,6 @@ const PurePreviewMessage = ({
   setMessages,
   reload,
   isReadonly,
-  index,
 }: {
   chatId: string;
   message: Message;
@@ -40,7 +39,6 @@ const PurePreviewMessage = ({
     chatRequestOptions?: ChatRequestOptions,
   ) => Promise<string | null | undefined>;
   isReadonly: boolean;
-  index: number;
 }) => {
   const [mode, setMode] = useState<'view' | 'edit'>('view');
 
diff --git a/components/messages.tsx b/components/messages.tsx
index 5e801d0..18c7785 100644
--- a/components/messages.tsx
+++ b/components/messages.tsx
@@ -1,29 +1,26 @@
-import { ChatRequestOptions, Message } from 'ai';
+import { Message } from 'ai';
 import { PreviewMessage, ThinkingMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Overview } from './overview';
 import { memo } from 'react';
 import { Vote } from '@/lib/db/schema';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers } from '@ai-sdk/react';
 
 interface MessagesProps {
   chatId: string;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
   votes: Array<Vote> | undefined;
   messages: Array<Message>;
-  setMessages: (
-    messages: Message[] | ((messages: Message[]) => Message[]),
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  setMessages: UseChatHelpers['setMessages'];
+  reload: UseChatHelpers['reload'];
   isReadonly: boolean;
   isArtifactVisible: boolean;
 }
 
 function PureMessages({
   chatId,
-  isLoading,
+  status,
   votes,
   messages,
   setMessages,
@@ -43,10 +40,9 @@ function PureMessages({
       {messages.map((message, index) => (
         <PreviewMessage
           key={message.id}
-          index={index}
           chatId={chatId}
           message={message}
-          isLoading={isLoading && messages.length - 1 === index}
+          isLoading={status === 'streaming' && messages.length - 1 === index}
           vote={
             votes
               ? votes.find((vote) => vote.messageId === message.id)
@@ -58,7 +54,7 @@ function PureMessages({
         />
       ))}
 
-      {isLoading &&
+      {status === 'submitted' &&
         messages.length > 0 &&
         messages[messages.length - 1].role === 'user' && <ThinkingMessage />}
 
@@ -73,8 +69,8 @@ function PureMessages({
 export const Messages = memo(PureMessages, (prevProps, nextProps) => {
   if (prevProps.isArtifactVisible && nextProps.isArtifactVisible) return true;
 
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
-  if (prevProps.isLoading && nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
+  if (prevProps.status && nextProps.status) return false;
   if (prevProps.messages.length !== nextProps.messages.length) return false;
   if (!equal(prevProps.messages, nextProps.messages)) return false;
   if (!equal(prevProps.votes, nextProps.votes)) return false;
diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx
index 2eeb93c..74c2a80 100644
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@@ -29,12 +29,13 @@ import { Button } from './ui/button';
 import { Textarea } from './ui/textarea';
 import { SuggestedActions } from './suggested-actions';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers, UseChatOptions } from '@ai-sdk/react';
 
 function PureMultimodalInput({
   chatId,
   input,
   setInput,
-  isLoading,
+  status,
   stop,
   attachments,
   setAttachments,
@@ -45,24 +46,16 @@ function PureMultimodalInput({
   className,
 }: {
   chatId: string;
-  input: string;
-  setInput: (value: string) => void;
-  isLoading: boolean;
+  input: UseChatHelpers['input'];
+  setInput: UseChatHelpers['setInput'];
+  status: UseChatHelpers['status'];
   stop: () => void;
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
   messages: Array<Message>;
   setMessages: Dispatch<SetStateAction<Array<Message>>>;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  handleSubmit: (
-    event?: {
-      preventDefault?: () => void;
-    },
-    chatRequestOptions?: ChatRequestOptions,
-  ) => void;
+  append: UseChatHelpers['append'];
+  handleSubmit: UseChatHelpers['handleSubmit'];
   className?: string;
 }) {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -253,7 +246,7 @@ function PureMultimodalInput({
           ) {
             event.preventDefault();
 
-            if (isLoading) {
+            if (status !== 'ready') {
               toast.error('Please wait for the model to finish its response!');
             } else {
               submitForm();
@@ -263,11 +256,11 @@ function PureMultimodalInput({
       />
 
       <div className="absolute bottom-0 p-2 w-fit flex flex-row justify-start">
-        <AttachmentsButton fileInputRef={fileInputRef} isLoading={isLoading} />
+        <AttachmentsButton fileInputRef={fileInputRef} status={status} />
       </div>
 
       <div className="absolute bottom-0 right-0 p-2 w-fit flex flex-row justify-end">
-        {isLoading ? (
+        {status === 'submitted' ? (
           <StopButton stop={stop} setMessages={setMessages} />
         ) : (
           <SendButton
@@ -285,7 +278,7 @@ export const MultimodalInput = memo(
   PureMultimodalInput,
   (prevProps, nextProps) => {
     if (prevProps.input !== nextProps.input) return false;
-    if (prevProps.isLoading !== nextProps.isLoading) return false;
+    if (prevProps.status !== nextProps.status) return false;
     if (!equal(prevProps.attachments, nextProps.attachments)) return false;
 
     return true;
@@ -294,10 +287,10 @@ export const MultimodalInput = memo(
 
 function PureAttachmentsButton({
   fileInputRef,
-  isLoading,
+  status,
 }: {
   fileInputRef: React.MutableRefObject<HTMLInputElement | null>;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
 }) {
   return (
     <Button
@@ -307,7 +300,7 @@ function PureAttachmentsButton({
         event.preventDefault();
         fileInputRef.current?.click();
       }}
-      disabled={isLoading}
+      disabled={status !== 'ready'}
       variant="ghost"
     >
       <PaperclipIcon size={14} />
diff --git a/components/toolbar.tsx b/components/toolbar.tsx
index 2a10926..768f28d 100644
--- a/components/toolbar.tsx
+++ b/components/toolbar.tsx
@@ -319,19 +319,16 @@ const PureToolbar = ({
   isToolbarVisible,
   setIsToolbarVisible,
   append,
-  isLoading,
+  status,
   stop,
   setMessages,
   artifactKind,
 }: {
   isToolbarVisible: boolean;
   setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
-  isLoading: boolean;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  stop: () => void;
+  status: UseChatHelpers['status'];
+  append: UseChatHelpers['append'];
+  stop: UseChatHelpers['stop'];
   setMessages: Dispatch<SetStateAction<Message[]>>;
   artifactKind: ArtifactKind;
 }) => {
@@ -372,10 +369,10 @@ const PureToolbar = ({
   }, []);
 
   useEffect(() => {
-    if (isLoading) {
+    if (status === 'streaming') {
       setIsToolbarVisible(false);
     }
-  }, [isLoading, setIsToolbarVisible]);
+  }, [status, setIsToolbarVisible]);
 
   const artifactDefinition = artifactDefinitions.find(
     (definition) => definition.kind === artifactKind,
@@ -418,13 +415,13 @@ const PureToolbar = ({
         exit={{ opacity: 0, y: -20, transition: { duration: 0.1 } }}
         transition={{ type: 'spring', stiffness: 300, damping: 25 }}
         onHoverStart={() => {
-          if (isLoading) return;
+          if (status === 'streaming') return;
 
           cancelCloseTimer();
           setIsToolbarVisible(true);
         }}
         onHoverEnd={() => {
-          if (isLoading) return;
+          if (status === 'streaming') return;
 
           startCloseTimer();
         }}
@@ -436,7 +433,7 @@ const PureToolbar = ({
         }}
         ref={toolbarRef}
       >
-        {isLoading ? (
+        {status === 'streaming' ? (
           <motion.div
             key="stop-icon"
             initial={{ scale: 1 }}
@@ -475,7 +472,7 @@ const PureToolbar = ({
 };
 
 export const Toolbar = memo(PureToolbar, (prevProps, nextProps) => {
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
   if (prevProps.isToolbarVisible !== nextProps.isToolbarVisible) return false;
   if (prevProps.artifactKind !== nextProps.artifactKind) return false;
 
diff --git a/lib/ai/models.test.ts b/lib/ai/models.test.ts
index a487134..5e4bb98 100644
--- a/lib/ai/models.test.ts
+++ b/lib/ai/models.test.ts
@@ -11,6 +11,8 @@ export const chatModel = new MockLanguageModelV1({
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
+      chunkDelayInMs: 50,
+      initialDelayInMs: 100,
       chunks: getResponseChunksByPrompt(prompt),
     }),
     rawCall: { rawPrompt: null, rawSettings: {} },
@@ -26,6 +28,8 @@ export const reasoningModel = new MockLanguageModelV1({
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
+      chunkDelayInMs: 50,
+      initialDelayInMs: 500,
       chunks: getResponseChunksByPrompt(prompt, true),
     }),
     rawCall: { rawPrompt: null, rawSettings: {} },
@@ -41,6 +45,8 @@ export const titleModel = new MockLanguageModelV1({
   }),
   doStream: async () => ({
     stream: simulateReadableStream({
+      chunkDelayInMs: 50,
+      initialDelayInMs: 100,
       chunks: [
         { type: 'text-delta', textDelta: 'This is a test title' },
         {
@@ -64,6 +70,8 @@ export const artifactModel = new MockLanguageModelV1({
   }),
   doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
+      chunkDelayInMs: 50,
+      initialDelayInMs: 100,
       chunks: getResponseChunksByPrompt(prompt),
     }),
     rawCall: { rawPrompt: null, rawSettings: {} },
diff --git a/tests/pages/chat.ts b/tests/pages/chat.ts
index 7a99875..85ea3d8 100644
--- a/tests/pages/chat.ts
+++ b/tests/pages/chat.ts
@@ -179,15 +179,4 @@ export class ChatPage {
       },
     };
   }
-
-  async waitForMessageGeneration(timeout = 10000) {
-    await this.page.waitForFunction(
-      () => {
-        return document.querySelector('[data-testid="send-button"]') !== null;
-      },
-      { timeout },
-    );
-
-    await this.page.waitForTimeout(500);
-  }
 }
diff --git a/tests/reasoning.test.ts b/tests/reasoning.test.ts
index 9000285..1e395cc 100644
--- a/tests/reasoning.test.ts
+++ b/tests/reasoning.test.ts
@@ -11,7 +11,7 @@ test.describe('chat activity with reasoning', () => {
 
   test('send user message and generate response with reasoning', async () => {
     await chatPage.sendUserMessage('Why is the sky blue?');
-    await chatPage.waitForMessageGeneration();
+    await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
     expect(assistantMessage.content).toBe("It's just blue duh!");
@@ -23,7 +23,7 @@ test.describe('chat activity with reasoning', () => {
 
   test('toggle reasoning visibility', async () => {
     await chatPage.sendUserMessage('Why is the sky blue?');
-    await chatPage.waitForMessageGeneration();
+    await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
     const reasoningElement =
@@ -39,7 +39,7 @@ test.describe('chat activity with reasoning', () => {
 
   test('edit message and resubmit', async () => {
     await chatPage.sendUserMessage('Why is the sky blue?');
-    await chatPage.waitForMessageGeneration();
+    await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
     const reasoningElement =
@@ -49,7 +49,7 @@ test.describe('chat activity with reasoning', () => {
     const userMessage = await chatPage.getRecentUserMessage();
 
     await userMessage.edit('Why is grass green?');
-    await chatPage.waitForMessageGeneration();
+    await chatPage.isGenerationComplete();
 
     const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();