refactor: use generateId() for id creation (#736)

app/(chat)/actions.ts

@@ -1,6 +1,6 @@
 'use server';
 
-import { type CoreUserMessage, generateText } from 'ai';
+import { type CoreUserMessage, generateText, Message } from 'ai';
 import { cookies } from 'next/headers';
 
 import { customModel } from '@/lib/ai';
@@ -19,7 +19,7 @@ export async function saveModelId(model: string) {
 export async function generateTitleFromUserMessage({
   message,
 }: {
-  message: CoreUserMessage;
+  message: Message;
 }) {
   const { text: title } = await generateText({
     model: customModel('gpt-4o-mini'),

app/(chat)/api/chat/route.ts

@@ -1,6 +1,5 @@
 import {
   type Message,
-  convertToCoreMessages,
   createDataStreamResponse,
   smoothStream,
   streamText,
@@ -65,8 +64,7 @@ export async function POST(request: Request) {
     return new Response('Model not found', { status: 404 });
   }
 
-  const coreMessages = convertToCoreMessages(messages);
-  const userMessage = getMostRecentUserMessage(coreMessages);
+  const userMessage = getMostRecentUserMessage(messages);
 
   if (!userMessage) {
     return new Response('No user message found', { status: 400 });
@@ -79,28 +77,20 @@ export async function POST(request: Request) {
     await saveChat({ id, userId: session.user.id, title });
   }
 
-  const userMessageId = generateUUID();
-
   await saveMessages({
-    messages: [
-      { ...userMessage, id: userMessageId, createdAt: new Date(), chatId: id },
-    ],
+    messages: [{ ...userMessage, createdAt: new Date(), chatId: id }],
   });
 
   return createDataStreamResponse({
     execute: (dataStream) => {
-      dataStream.writeData({
-        type: 'user-message-id',
-        content: userMessageId,
-      });
-
       const result = streamText({
         model: customModel(model.apiIdentifier),
         system: systemPrompt,
-        messages: coreMessages,
+        messages,
         maxSteps: 5,
         experimental_activeTools: allTools,
         experimental_transform: smoothStream({ chunking: 'word' }),
+        experimental_generateMessageId: generateUUID,
         tools: {
           getWeather,
           createDocument: createDocument({ session, dataStream, model }),
@@ -120,16 +110,8 @@ export async function POST(request: Request) {
               await saveMessages({
                 messages: responseMessagesWithoutIncompleteToolCalls.map(
                   (message) => {
-                    const messageId = generateUUID();
-
-                    if (message.role === 'assistant') {
-                      dataStream.writeMessageAnnotation({
-                        messageIdFromServer: messageId,
-                      });
-                    }
-
                     return {
-                      id: messageId,
+                      id: message.id,
                       chatId: id,
                       role: message.role,
                       content: message.content,

components/chat.tsx

@@ -7,7 +7,7 @@ import useSWR, { useSWRConfig } from 'swr';
 
 import { ChatHeader } from '@/components/chat-header';
 import type { Vote } from '@/lib/db/schema';
-import { fetcher } from '@/lib/utils';
+import { fetcher, generateUUID } from '@/lib/utils';
 
 import { Block } from './block';
 import { MultimodalInput } from './multimodal-input';
@@ -45,6 +45,8 @@ export function Chat({
     body: { id, modelId: selectedModelId },
     initialMessages,
     experimental_throttle: 100,
+    sendExtraMessageFields: true,
+    generateId: generateUUID,
     onFinish: () => {
       mutate('/api/history');
     },

components/data-stream-handler.tsx

@@ -5,7 +5,6 @@ import { useEffect, useRef } from 'react';
 import { blockDefinitions, BlockKind } from './block';
 import { Suggestion } from '@/lib/db/schema';
 import { initialBlockData, useBlock } from '@/hooks/use-block';
-import { useUserMessageId } from '@/hooks/use-user-message-id';
 
 export type DataStreamDelta = {
   type:
@@ -17,14 +16,12 @@ export type DataStreamDelta = {
     | 'suggestion'
     | 'clear'
     | 'finish'
-    | 'user-message-id'
     | 'kind';
   content: string | Suggestion;
 };
 
 export function DataStreamHandler({ id }: { id: string }) {
   const { data: dataStream } = useChat({ id });
-  const { setUserMessageIdFromServer } = useUserMessageId();
   const { block, setBlock, setMetadata } = useBlock();
   const lastProcessedIndex = useRef(-1);
 
@@ -35,11 +32,6 @@ export function DataStreamHandler({ id }: { id: string }) {
     lastProcessedIndex.current = dataStream.length - 1;
 
     (newDeltas as DataStreamDelta[]).forEach((delta: DataStreamDelta) => {
-      if (delta.type === 'user-message-id') {
-        setUserMessageIdFromServer(delta.content as string);
-        return;
-      }
-
       const blockDefinition = blockDefinitions.find(
         (blockDefinition) => blockDefinition.kind === block.kind,
       );
@@ -97,7 +89,7 @@ export function DataStreamHandler({ id }: { id: string }) {
         }
       });
     });
-  }, [dataStream, setBlock, setUserMessageIdFromServer, setMetadata, block]);
+  }, [dataStream, setBlock, setMetadata, block]);
 
   return null;
 }

components/message-actions.tsx

@@ -4,7 +4,6 @@ import { useSWRConfig } from 'swr';
 import { useCopyToClipboard } from 'usehooks-ts';
 
 import type { Vote } from '@/lib/db/schema';
-import { getMessageIdFromAnnotations } from '@/lib/utils';
 
 import { CopyIcon, ThumbDownIcon, ThumbUpIcon } from './icons';
 import { Button } from './ui/button';
@@ -62,13 +61,11 @@ export function PureMessageActions({
               disabled={vote?.isUpvoted}
               variant="outline"
               onClick={async () => {
-                const messageId = getMessageIdFromAnnotations(message);
-
                 const upvote = fetch('/api/vote', {
                   method: 'PATCH',
                   body: JSON.stringify({
                     chatId,
-                    messageId,
+                    messageId: message.id,
                     type: 'up',
                   }),
                 });
@@ -116,13 +113,11 @@ export function PureMessageActions({
               variant="outline"
               disabled={vote && !vote.isUpvoted}
               onClick={async () => {
-                const messageId = getMessageIdFromAnnotations(message);
-
                 const downvote = fetch('/api/vote', {
                   method: 'PATCH',
                   body: JSON.stringify({
                     chatId,
-                    messageId,
+                    messageId: message.id,
                     type: 'down',
                   }),
                 });

components/message-editor.tsx

@@ -6,7 +6,6 @@ import { Dispatch, SetStateAction, useEffect, useRef, useState } from 'react';
 import { Textarea } from './ui/textarea';
 import { deleteTrailingMessages } from '@/app/(chat)/actions';
 import { toast } from 'sonner';
-import { useUserMessageId } from '@/hooks/use-user-message-id';
 
 export type MessageEditorProps = {
   message: Message;
@@ -25,7 +24,6 @@ export function MessageEditor({
   setMessages,
   reload,
 }: MessageEditorProps) {
-  const { userMessageIdFromServer } = useUserMessageId();
   const [isSubmitting, setIsSubmitting] = useState<boolean>(false);
 
   const [draftContent, setDraftContent] = useState<string>(message.content);
@@ -74,16 +72,9 @@ export function MessageEditor({
           disabled={isSubmitting}
           onClick={async () => {
             setIsSubmitting(true);
-            const messageId = userMessageIdFromServer ?? message.id;
-
-            if (!messageId) {
-              toast.error('Something went wrong, please try again!');
-              setIsSubmitting(false);
-              return;
-            }
 
             await deleteTrailingMessages({
-              id: messageId,
+              id: message.id,
             });
 
             setMessages((messages) => {

hooks/use-user-message-id.ts

@@ -1,10 +0,0 @@
-'use client';
-
-import useSWR from 'swr';
-
-export function useUserMessageId() {
-  const { data: userMessageIdFromServer, mutate: setUserMessageIdFromServer } =
-    useSWR('userMessageIdFromServer', null);
-
-  return { userMessageIdFromServer, setUserMessageIdFromServer };
-}

lib/utils.ts

@@ -126,9 +126,12 @@ export function convertToUIMessages(
   }, []);
 }
 
+type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage;
+type ResponseMessage = ResponseMessageWithoutId & { id: string };
+
 export function sanitizeResponseMessages(
-  messages: Array<CoreToolMessage | CoreAssistantMessage>,
-): Array<CoreToolMessage | CoreAssistantMessage> {
+  messages: Array<ResponseMessage>,
+): Array<ResponseMessage> {
   const toolResultIds: Array<string> = [];
 
   for (const message of messages) {
@@ -198,7 +201,7 @@ export function sanitizeUIMessages(messages: Array<Message>): Array<Message> {
   );
 }
 
-export function getMostRecentUserMessage(messages: Array<CoreMessage>) {
+export function getMostRecentUserMessage(messages: Array<Message>) {
   const userMessages = messages.filter((message) => message.role === 'user');
   return userMessages.at(-1);
 }
@@ -212,13 +215,3 @@ export function getDocumentTimestampByIndex(
 
   return documents[index].createdAt;
 }
-
-export function getMessageIdFromAnnotations(message: Message) {
-  if (!message.annotations) return message.id;
-
-  const [annotation] = message.annotations;
-  if (!annotation) return message.id;
-
-  // @ts-expect-error messageIdFromServer is not defined in MessageAnnotation
-  return annotation.messageIdFromServer;
-}