vercel · jeremyphilemon · May 20, 2025
diff --git a/app/(chat)/chat/[id]/page.tsx b/app/(chat)/chat/[id]/page.tsx
@@ -5,10 +5,10 @@ import { Chat } from '@/components/chat';
 import { getChatById, getMessagesByChatId } from '@/lib/db/queries';
 import { DEFAULT_CHAT_MODEL } from '@/lib/ai/models';
 import type { DBMessage } from '@/lib/db/schema';
-import type { UIMessage } from 'ai';
-import type { Attachment } from '@/lib/types';
+import type { InferUIDataTypes, UIMessage } from 'ai';
 import { DataStreamHandler } from '@/components/data-stream-handler';
 import { ChatStoreProvider } from '@/components/chat-store';
+import { dataPartSchemas, type MessageMetadata } from '@/lib/types';
 
 export default async function Page(props: { params: Promise<{ id: string }> }) {
   const params = await props.params;
@@ -39,16 +39,31 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
     id,
   });
 
-  function convertToUIMessages(messages: Array<DBMessage>): Array<UIMessage> {
+  function convertToUIMessages(
+    messages: Array<DBMessage>,
+  ): Array<
+    UIMessage<MessageMetadata, InferUIDataTypes<typeof dataPartSchemas>>
+  > {
     return messages.map((message) => ({
       id: message.id,
-      parts: message.parts as UIMessage['parts'],
       role: message.role as UIMessage['role'],
-      // Note: content will soon be deprecated in @ai-sdk/react
-      content: '',
       createdAt: message.createdAt,
-      experimental_attachments:
-        (message.attachments as Array<Attachment>) ?? [],
+      parts: (message.parts as any[]).map((part: any) => {
+        const schema =
+          dataPartSchemas[part.type as keyof typeof dataPartSchemas];
+
+        if (!schema) {
+          throw new Error(
+            `Unknown message part type '${part.type}' in message ${message.id}`,
+          );
+        }
+
+        return {
+          type: part.type,
+          id: part.id,
+          data: schema.parse(part.data),
+        };
+      }),
     }));
   }
 

diff --git a/artifacts/code/client.tsx b/artifacts/code/client.tsx
@@ -79,7 +79,7 @@ export const codeArtifact = new Artifact<'code', Metadata>({
     if (streamPart.type === 'data-artifacts-code-delta') {
       setArtifact((draftArtifact) => ({
         ...draftArtifact,
-        content: streamPart.value as string,
+        content: streamPart.data,
         isVisible:
           draftArtifact.status === 'streaming' &&
           draftArtifact.content.length > 300 &&

diff --git a/artifacts/image/client.tsx b/artifacts/image/client.tsx
@@ -10,7 +10,7 @@ export const imageArtifact = new Artifact({
     if (streamPart.type === 'data-artifacts-image-delta') {
       setArtifact((draftArtifact) => ({
         ...draftArtifact,
-        content: streamPart.value as string,
+        content: streamPart.data,
         isVisible: true,
         status: 'streaming',
       }));

diff --git a/artifacts/sheet/client.tsx b/artifacts/sheet/client.tsx
@@ -20,7 +20,7 @@ export const sheetArtifact = new Artifact<'sheet', Metadata>({
     if (streamPart.type === 'data-artifacts-sheet-delta') {
       setArtifact((draftArtifact) => ({
         ...draftArtifact,
-        content: streamPart.value as string,
+        content: streamPart.data,
         isVisible: true,
         status: 'streaming',
       }));

diff --git a/artifacts/text/client.tsx b/artifacts/text/client.tsx
@@ -32,10 +32,7 @@ export const textArtifact = new Artifact<'text', TextArtifactMetadata>({
     if (streamPart.type === 'data-artifacts-suggestion') {
       setMetadata((metadata) => {
         return {
-          suggestions: [
-            ...metadata.suggestions,
-            streamPart.value as Suggestion,
-          ],
+          suggestions: [...metadata.suggestions, streamPart.data],
         };
       });
     }
@@ -44,7 +41,7 @@ export const textArtifact = new Artifact<'text', TextArtifactMetadata>({
       setArtifact((draftArtifact) => {
         return {
           ...draftArtifact,
-          content: draftArtifact.content + (streamPart.value as string),
+          content: draftArtifact.content + streamPart.data,
           isVisible:
             draftArtifact.status === 'streaming' &&
             draftArtifact.content.length > 400 &&

diff --git a/components/artifact-messages.tsx b/components/artifact-messages.tsx
@@ -1,19 +1,17 @@
 import { PreviewMessage, ThinkingMessage } from './message';
 import type { Vote } from '@/lib/db/schema';
-import type { UIMessage } from 'ai';
 import { memo } from 'react';
 import equal from 'fast-deep-equal';
 import type { UIArtifact } from './artifact';
 import type { UseChatHelpers } from '@ai-sdk/react';
 import { motion } from 'framer-motion';
 import { useMessages } from '@/hooks/use-messages';
+import { useChatStore } from './chat-store';
 
 interface ArtifactMessagesProps {
   chatId: string;
   status: UseChatHelpers['status'];
   votes: Array<Vote> | undefined;
-  messages: Array<UIMessage>;
-  setMessages: UseChatHelpers['setMessages'];
   reload: UseChatHelpers['reload'];
   isReadonly: boolean;
   artifactStatus: UIArtifact['status'];
@@ -23,11 +21,12 @@ function PureArtifactMessages({
   chatId,
   status,
   votes,
-  messages,
-  setMessages,
   reload,
   isReadonly,
 }: ArtifactMessagesProps) {
+  const chatStore = useChatStore();
+  const messages = chatStore.getMessages(chatId);
+
   const {
     containerRef: messagesContainerRef,
     endRef: messagesEndRef,
@@ -55,7 +54,6 @@ function PureArtifactMessages({
               ? votes.find((vote) => vote.messageId === message.id)
               : undefined
           }
-          setMessages={setMessages}
           reload={reload}
           isReadonly={isReadonly}
           requiresScrollPadding={
@@ -90,7 +88,6 @@ function areEqual(
 
   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;
 
   return true;

diff --git a/components/artifact.tsx b/components/artifact.tsx
@@ -1,4 +1,3 @@
-import type { UIMessage } from 'ai';
 import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
@@ -62,9 +61,6 @@ function PureArtifact({
   stop,
   attachments,
   setAttachments,
-  append,
-  messages,
-  setMessages,
   reload,
   votes,
   isReadonly,
@@ -77,10 +73,7 @@ function PureArtifact({
   stop: UseChatHelpers['stop'];
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
-  messages: Array<UIMessage>;
-  setMessages: UseChatHelpers['setMessages'];
   votes: Array<Vote> | undefined;
-  append: UseChatHelpers['append'];
   handleSubmit: UseChatHelpers['handleSubmit'];
   reload: UseChatHelpers['reload'];
   isReadonly: boolean;
@@ -318,8 +311,6 @@ function PureArtifact({
                   chatId={chatId}
                   status={status}
                   votes={votes}
-                  messages={messages}
-                  setMessages={setMessages}
                   reload={reload}
                   isReadonly={isReadonly}
                   artifactStatus={artifact.status}
@@ -330,15 +321,11 @@ function PureArtifact({
                     chatId={chatId}
                     input={input}
                     setInput={setInput}
-                    handleSubmit={handleSubmit}
                     status={status}
                     stop={stop}
                     attachments={attachments}
                     setAttachments={setAttachments}
-                    messages={messages}
-                    append={append}
                     className="bg-background dark:bg-muted"
-                    setMessages={setMessages}
                     selectedVisibilityType={selectedVisibilityType}
                   />
                 </form>
@@ -477,10 +464,8 @@ function PureArtifact({
                   <Toolbar
                     isToolbarVisible={isToolbarVisible}
                     setIsToolbarVisible={setIsToolbarVisible}
-                    append={append}
                     status={status}
                     stop={stop}
-                    setMessages={setMessages}
                     artifactKind={artifact.kind}
                   />
                 )}
@@ -507,7 +492,6 @@ export const Artifact = memo(PureArtifact, (prevProps, nextProps) => {
   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;
   if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
     return false;
 

diff --git a/components/chat-store.tsx b/components/chat-store.tsx
@@ -1,19 +1,27 @@
 'use client';
 
-import React, {
-  createContext,
-  type ReactNode,
-  useContext,
-  useMemo,
-} from 'react';
-import { defaultChatStore, type UIMessage } from 'ai';
+import React, { createContext, type ReactNode, useContext } from 'react';
+import { defaultChatStore, type InferUIDataTypes, type UIMessage } from 'ai';
 import { fetchWithErrorHandlers, generateUUID } from '@/lib/utils';
-import { zodSchema } from '@ai-sdk/provider-utils';
-import { messageMetadataSchema } from '@/lib/types';
+import type { dataPartSchemas, MessageMetadata } from '@/lib/types';
 
-type ChatStoreType = ReturnType<typeof defaultChatStore>;
+const chatStore = defaultChatStore<MessageMetadata, typeof dataPartSchemas>({
+  api: '/api/chat',
+  fetch: fetchWithErrorHandlers,
+  generateId: generateUUID,
+  prepareRequestBody: (body) => ({
+    id: body.chatId,
+    message: body.messages.at(-1),
+    selectedChatModel: 'chat-model',
+    selectedVisibilityType: 'private',
+  }),
 
-const ChatStoreContext = createContext<ChatStoreType | null>(null);
+  chats: {},
+});
+
+export type ChatStore = typeof chatStore;
+
+const ChatStoreContext = createContext(chatStore);
 
 export function useChatStore() {
   const context = useContext(ChatStoreContext);
@@ -30,7 +38,10 @@ type Props = {
   id: string;
   initialChatModel: string;
   visibilityType: string;
-  initialMessages: UIMessage[];
+  initialMessages: UIMessage<
+    MessageMetadata,
+    InferUIDataTypes<typeof dataPartSchemas>
+  >[];
 };
 
 export function ChatStoreProvider({
@@ -40,29 +51,10 @@ export function ChatStoreProvider({
   visibilityType,
   initialMessages,
 }: Props) {
-  const store = useMemo(
-    () =>
-      defaultChatStore({
-        api: '/api/chat',
-        fetch: fetchWithErrorHandlers,
-        messageMetadataSchema: zodSchema(messageMetadataSchema),
-        generateId: generateUUID,
-        prepareRequestBody: (body) => ({
-          id,
-          message: body.messages.at(-1),
-          selectedChatModel: initialChatModel,
-          selectedVisibilityType: visibilityType,
-        }),
-        chats: {
-          [id]: {
-            messages: initialMessages,
-          },
-        },
-      }),
-    [id, initialChatModel, visibilityType, initialMessages],
-  );
+  chatStore.addChat(id, initialMessages);
+
   return (
-    <ChatStoreContext.Provider value={store}>
+    <ChatStoreContext.Provider value={chatStore}>
       {children}
     </ChatStoreContext.Provider>
   );

diff --git a/components/chat.tsx b/components/chat.tsx
@@ -58,7 +58,7 @@ export function Chat({
     reload,
     experimental_resume,
   } = useChat({
-    id,
+    chatId: id,
     experimental_throttle: 100,
     generateId: generateUUID,
     chatStore,
@@ -82,15 +82,18 @@ export function Chat({
 
   useEffect(() => {
     if (query && !hasAppendedQuery) {
-      append({
-        role: 'user',
-        parts: [{ type: 'text', text: query }],
+      chatStore.submitMessage({
+        chatId: id,
+        message: {
+          role: 'user',
+          parts: [{ type: 'text', text: query }],
+        },
       });
 
       setHasAppendedQuery(true);
       window.history.replaceState({}, '', `/chat/${id}`);
     }
-  }, [query, append, hasAppendedQuery, id]);
+  }, [query, append, hasAppendedQuery, id, chatStore]);
 
   const { data: votes } = useSWR<Array<Vote>>(
     messages.length >= 2 ? `/api/vote?chatId=${id}` : null,
@@ -101,10 +104,9 @@ export function Chat({
   const isArtifactVisible = useArtifactSelector((state) => state.isVisible);
 
   useAutoResume({
+    chatId: id,
     autoResume,
     experimental_resume,
-    messages,
-    setMessages,
   });
 
   return (
@@ -123,7 +125,6 @@ export function Chat({
           status={status}
           votes={votes}
           messages={messages}
-          setMessages={setMessages}
           reload={reload}
           isReadonly={isReadonly}
           isArtifactVisible={isArtifactVisible}
@@ -135,14 +136,10 @@ export function Chat({
               chatId={id}
               input={input}
               setInput={setInput}
-              handleSubmit={handleSubmit}
               status={status}
               stop={stop}
               attachments={attachments}
               setAttachments={setAttachments}
-              messages={messages}
-              setMessages={setMessages}
-              append={append}
               selectedVisibilityType={visibilityType}
             />
           )}
@@ -158,9 +155,6 @@ export function Chat({
         stop={stop}
         attachments={attachments}
         setAttachments={setAttachments}
-        append={append}
-        messages={messages}
-        setMessages={setMessages}
         reload={reload}
         votes={votes}
         isReadonly={isReadonly}