From 64d1aad2bbde4bea9ab3003e9a7d08c1b53585db Mon Sep 17 00:00:00 2001
From: Jeremy <git@jrmyphlmn.com>
Date: Thu, 5 Dec 2024 15:29:33 +0300
Subject: [PATCH] feat: edit and resubmit messages (#592)

---
 app/(chat)/actions.ts           |  13 ++++
 app/(chat)/api/chat/route.ts    |   9 ++-
 components/block-messages.tsx   |  12 +++-
 components/block.tsx            |   6 ++
 components/chat.tsx             |   5 ++
 components/message-editor.tsx   | 113 ++++++++++++++++++++++++++++++++
 components/message.tsx          |  95 +++++++++++++++++++++------
 components/messages.tsx         |  12 +++-
 components/multimodal-input.tsx |   2 +-
 components/use-block-stream.tsx |  18 ++++-
 hooks/use-user-message-id.ts    |  10 +++
 lib/db/queries.ts               |  32 ++++++++-
 12 files changed, 302 insertions(+), 25 deletions(-)
 create mode 100644 components/message-editor.tsx
 create mode 100644 hooks/use-user-message-id.ts

diff --git a/app/(chat)/actions.ts b/app/(chat)/actions.ts
index 0d4b2c9..da4c5ea 100644
--- a/app/(chat)/actions.ts
+++ b/app/(chat)/actions.ts
@@ -4,6 +4,10 @@ import { type CoreUserMessage, generateText } from 'ai';
 import { cookies } from 'next/headers';
 
 import { customModel } from '@/lib/ai';
+import {
+  deleteMessagesByChatIdAfterTimestamp,
+  getMessageById,
+} from '@/lib/db/queries';
 
 export async function saveModelId(model: string) {
   const cookieStore = await cookies();
@@ -27,3 +31,12 @@ export async function generateTitleFromUserMessage({
 
   return title;
 }
+
+export async function deleteTrailingMessages({ id }: { id: string }) {
+  const [message] = await getMessageById({ id });
+
+  await deleteMessagesByChatIdAfterTimestamp({
+    chatId: message.chatId,
+    timestamp: message.createdAt,
+  });
+}
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
index b5efc43..8452f3b 100644
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@@ -81,14 +81,21 @@ export async function POST(request: Request) {
     await saveChat({ id, userId: session.user.id, title });
   }
 
+  const userMessageId = generateUUID();
+
   await saveMessages({
     messages: [
-      { ...userMessage, id: generateUUID(), createdAt: new Date(), chatId: id },
+      { ...userMessage, id: userMessageId, createdAt: new Date(), chatId: id },
     ],
   });
 
   const streamingData = new StreamData();
 
+  streamingData.append({
+    type: 'user-message-id',
+    content: userMessageId,
+  });
+
   const result = streamText({
     model: customModel(model.apiIdentifier),
     system: systemPrompt,
diff --git a/components/block-messages.tsx b/components/block-messages.tsx
index 47f263f..2d83068 100644
--- a/components/block-messages.tsx
+++ b/components/block-messages.tsx
@@ -3,7 +3,7 @@ import { UIBlock } from './block';
 import { PreviewMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Vote } from '@/lib/db/schema';
-import { Message } from 'ai';
+import { ChatRequestOptions, Message } from 'ai';
 
 interface BlockMessagesProps {
   chatId: string;
@@ -12,6 +12,12 @@ interface BlockMessagesProps {
   isLoading: boolean;
   votes: Array<Vote> | undefined;
   messages: Array<Message>;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }
 
 function PureBlockMessages({
@@ -21,6 +27,8 @@ function PureBlockMessages({
   isLoading,
   votes,
   messages,
+  setMessages,
+  reload,
 }: BlockMessagesProps) {
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>();
@@ -43,6 +51,8 @@ function PureBlockMessages({
               ? votes.find((vote) => vote.messageId === message.id)
               : undefined
           }
+          setMessages={setMessages}
+          reload={reload}
         />
       ))}
 
diff --git a/components/block.tsx b/components/block.tsx
index 7f8cd75..d426cf9 100644
--- a/components/block.tsx
+++ b/components/block.tsx
@@ -58,6 +58,7 @@ function PureBlock({
   setBlock,
   messages,
   setMessages,
+  reload,
   votes,
 }: {
   chatId: string;
@@ -82,6 +83,9 @@ function PureBlock({
     },
     chatRequestOptions?: ChatRequestOptions,
   ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }) {
   const {
     data: documents,
@@ -286,6 +290,8 @@ function PureBlock({
               setBlock={setBlock}
               votes={votes}
               messages={messages}
+              setMessages={setMessages}
+              reload={reload}
             />
 
             <form className="flex flex-row gap-2 relative items-end w-full px-4 pb-4">
diff --git a/components/chat.tsx b/components/chat.tsx
index ad65857..0eee5a8 100644
--- a/components/chat.tsx
+++ b/components/chat.tsx
@@ -36,8 +36,10 @@ export function Chat({
     append,
     isLoading,
     stop,
+    reload,
     data: streamingData,
   } = useChat({
+    id,
     body: { id, modelId: selectedModelId },
     initialMessages,
     onFinish: () => {
@@ -81,6 +83,8 @@ export function Chat({
           isLoading={isLoading}
           votes={votes}
           messages={messages}
+          setMessages={setMessages}
+          reload={reload}
         />
 
         <form className="flex mx-auto px-4 bg-background pb-4 md:pb-6 gap-2 w-full md:max-w-3xl">
@@ -116,6 +120,7 @@ export function Chat({
             setBlock={setBlock}
             messages={messages}
             setMessages={setMessages}
+            reload={reload}
             votes={votes}
           />
         )}
diff --git a/components/message-editor.tsx b/components/message-editor.tsx
new file mode 100644
index 0000000..35fbcf9
--- /dev/null
+++ b/components/message-editor.tsx
@@ -0,0 +1,113 @@
+'use client';
+
+import { ChatRequestOptions, Message } from 'ai';
+import { Button } from './ui/button';
+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;
+  setMode: Dispatch<SetStateAction<'view' | 'edit'>>;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
+};
+
+export function MessageEditor({
+  message,
+  setMode,
+  setMessages,
+  reload,
+}: MessageEditorProps) {
+  const { userMessageIdFromServer } = useUserMessageId();
+  const [isSubmitting, setIsSubmitting] = useState<boolean>(false);
+
+  const [draftContent, setDraftContent] = useState<string>(message.content);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  useEffect(() => {
+    if (textareaRef.current) {
+      adjustHeight();
+    }
+  }, []);
+
+  const adjustHeight = () => {
+    if (textareaRef.current) {
+      textareaRef.current.style.height = 'auto';
+      textareaRef.current.style.height = `${textareaRef.current.scrollHeight + 2}px`;
+    }
+  };
+
+  const handleInput = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
+    setDraftContent(event.target.value);
+    adjustHeight();
+  };
+
+  return (
+    <div className="flex flex-col gap-2 w-full">
+      <Textarea
+        ref={textareaRef}
+        className="bg-transparent outline-none overflow-hidden resize-none !text-base rounded-xl w-full"
+        value={draftContent}
+        onChange={handleInput}
+      />
+
+      <div className="flex flex-row gap-2 justify-end">
+        <Button
+          variant="outline"
+          className="h-fit py-2 px-3"
+          onClick={() => {
+            setMode('view');
+          }}
+        >
+          Cancel
+        </Button>
+        <Button
+          variant="default"
+          className="h-fit py-2 px-3"
+          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,
+            });
+
+            setMessages((messages) => {
+              const index = messages.findIndex((m) => m.id === message.id);
+
+              if (index !== -1) {
+                const updatedMessage = {
+                  ...message,
+                  content: draftContent,
+                };
+
+                return [...messages.slice(0, index), updatedMessage];
+              }
+
+              return messages;
+            });
+
+            setMode('view');
+            reload();
+          }}
+        >
+          {isSubmitting ? 'Sending...' : 'Send'}
+        </Button>
+      </div>
+    </div>
+  );
+}
diff --git a/components/message.tsx b/components/message.tsx
index 7b6ad09..0a5658c 100644
--- a/components/message.tsx
+++ b/components/message.tsx
@@ -1,20 +1,24 @@
 'use client';
 
-import type { Message } from 'ai';
+import type { ChatRequestOptions, Message } from 'ai';
 import cx from 'classnames';
 import { motion } from 'framer-motion';
-import { memo, type Dispatch, type SetStateAction } from 'react';
+import { memo, useState, type Dispatch, type SetStateAction } from 'react';
 
 import type { Vote } from '@/lib/db/schema';
 
 import type { UIBlock } from './block';
 import { DocumentToolCall, DocumentToolResult } from './document';
-import { SparklesIcon } from './icons';
+import { PencilEditIcon, SparklesIcon } from './icons';
 import { Markdown } from './markdown';
 import { MessageActions } from './message-actions';
 import { PreviewAttachment } from './preview-attachment';
 import { Weather } from './weather';
 import equal from 'fast-deep-equal';
+import { cn } from '@/lib/utils';
+import { Button } from './ui/button';
+import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
+import { MessageEditor } from './message-editor';
 
 const PurePreviewMessage = ({
   chatId,
@@ -23,6 +27,8 @@ const PurePreviewMessage = ({
   setBlock,
   vote,
   isLoading,
+  setMessages,
+  reload,
 }: {
   chatId: string;
   message: Message;
@@ -30,7 +36,15 @@ const PurePreviewMessage = ({
   setBlock: Dispatch<SetStateAction<UIBlock>>;
   vote: Vote | undefined;
   isLoading: boolean;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }) => {
+  const [mode, setMode] = useState<'view' | 'edit'>('view');
+
   return (
     <motion.div
       className="w-full mx-auto max-w-3xl px-4 group/message"
@@ -39,8 +53,12 @@ const PurePreviewMessage = ({
       data-role={message.role}
     >
       <div
-        className={cx(
-          'group-data-[role=user]/message:bg-primary group-data-[role=user]/message:text-primary-foreground flex gap-4 group-data-[role=user]/message:px-3 w-full group-data-[role=user]/message:w-fit group-data-[role=user]/message:ml-auto group-data-[role=user]/message:max-w-2xl group-data-[role=user]/message:py-2 rounded-xl',
+        className={cn(
+          'flex gap-4 w-full group-data-[role=user]/message:ml-auto group-data-[role=user]/message:max-w-2xl',
+          {
+            'w-full': mode === 'edit',
+            'group-data-[role=user]/message:w-fit': mode !== 'edit',
+          },
         )}
       >
         {message.role === 'assistant' && (
@@ -50,9 +68,58 @@ const PurePreviewMessage = ({
         )}
 
         <div className="flex flex-col gap-2 w-full">
-          {message.content && (
-            <div className="flex flex-col gap-4">
-              <Markdown>{message.content as string}</Markdown>
+          {message.experimental_attachments && (
+            <div className="flex flex-row justify-end gap-2">
+              {message.experimental_attachments.map((attachment) => (
+                <PreviewAttachment
+                  key={attachment.url}
+                  attachment={attachment}
+                />
+              ))}
+            </div>
+          )}
+
+          {message.content && mode === 'view' && (
+            <div className="flex flex-row gap-2 items-start">
+              {message.role === 'user' && (
+                <Tooltip>
+                  <TooltipTrigger asChild>
+                    <Button
+                      variant="ghost"
+                      className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
+                      onClick={() => {
+                        setMode('edit');
+                      }}
+                    >
+                      <PencilEditIcon />
+                    </Button>
+                  </TooltipTrigger>
+                  <TooltipContent>Edit message</TooltipContent>
+                </Tooltip>
+              )}
+
+              <div
+                className={cn('flex flex-col gap-4', {
+                  'bg-primary text-primary-foreground px-3 py-2 rounded-xl':
+                    message.role === 'user',
+                })}
+              >
+                <Markdown>{message.content as string}</Markdown>
+              </div>
+            </div>
+          )}
+
+          {message.content && mode === 'edit' && (
+            <div className="flex flex-row gap-2 items-start">
+              <div className="size-8" />
+
+              <MessageEditor
+                key={message.id}
+                message={message}
+                setMode={setMode}
+                setMessages={setMessages}
+                reload={reload}
+              />
             </div>
           )}
 
@@ -129,17 +196,6 @@ const PurePreviewMessage = ({
             </div>
           )}
 
-          {message.experimental_attachments && (
-            <div className="flex flex-row gap-2">
-              {message.experimental_attachments.map((attachment) => (
-                <PreviewAttachment
-                  key={attachment.url}
-                  attachment={attachment}
-                />
-              ))}
-            </div>
-          )}
-
           <MessageActions
             key={`action-${message.id}`}
             chatId={chatId}
@@ -158,6 +214,7 @@ export const PreviewMessage = memo(
   (prevProps, nextProps) => {
     if (prevProps.isLoading !== nextProps.isLoading) return false;
     if (prevProps.isLoading && nextProps.isLoading) return false;
+    if (prevProps.message.content && nextProps.message.content) return false;
     if (!equal(prevProps.vote, nextProps.vote)) return false;
     return true;
   },
diff --git a/components/messages.tsx b/components/messages.tsx
index 06fbda4..14bba95 100644
--- a/components/messages.tsx
+++ b/components/messages.tsx
@@ -1,4 +1,4 @@
-import { Message } from 'ai';
+import { ChatRequestOptions, Message } from 'ai';
 import { PreviewMessage, ThinkingMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Overview } from './overview';
@@ -13,6 +13,12 @@ interface MessagesProps {
   isLoading: boolean;
   votes: Array<Vote> | undefined;
   messages: Array<Message>;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }
 
 function PureMessages({
@@ -22,6 +28,8 @@ function PureMessages({
   isLoading,
   votes,
   messages,
+  setMessages,
+  reload,
 }: MessagesProps) {
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>();
@@ -46,6 +54,8 @@ function PureMessages({
               ? votes.find((vote) => vote.messageId === message.id)
               : undefined
           }
+          setMessages={setMessages}
+          reload={reload}
         />
       ))}
 
diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx
index 690cd0b..f9b8942 100644
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@@ -228,7 +228,7 @@ function PureMultimodalInput({
         value={input}
         onChange={handleInput}
         className={cx(
-          'min-h-[24px] max-h-[calc(75dvh)] overflow-hidden resize-none rounded-xl text-base bg-muted',
+          'min-h-[24px] max-h-[calc(75dvh)] overflow-hidden resize-none rounded-xl !text-base bg-muted',
           className,
         )}
         rows={3}
diff --git a/components/use-block-stream.tsx b/components/use-block-stream.tsx
index 46d669f..1386f4c 100644
--- a/components/use-block-stream.tsx
+++ b/components/use-block-stream.tsx
@@ -5,9 +5,18 @@ import { useSWRConfig } from 'swr';
 import type { Suggestion } from '@/lib/db/schema';
 
 import type { UIBlock } from './block';
+import { useUserMessageId } from '@/hooks/use-user-message-id';
 
 type StreamingDelta = {
-  type: 'text-delta' | 'title' | 'id' | 'suggestion' | 'clear' | 'finish';
+  type:
+    | 'text-delta'
+    | 'title'
+    | 'id'
+    | 'suggestion'
+    | 'clear'
+    | 'finish'
+    | 'user-message-id';
+
   content: string | Suggestion;
 };
 
@@ -23,6 +32,8 @@ export function useBlockStream({
     Array<Suggestion>
   >([]);
 
+  const { setUserMessageIdFromServer } = useUserMessageId();
+
   useEffect(() => {
     if (optimisticSuggestions && optimisticSuggestions.length > 0) {
       const [optimisticSuggestion] = optimisticSuggestions;
@@ -37,6 +48,11 @@ export function useBlockStream({
 
     const delta = mostRecentDelta as StreamingDelta;
 
+    if (delta.type === 'user-message-id') {
+      setUserMessageIdFromServer(delta.content as string);
+      return;
+    }
+
     setBlock((draftBlock) => {
       switch (delta.type) {
         case 'id':
diff --git a/hooks/use-user-message-id.ts b/hooks/use-user-message-id.ts
new file mode 100644
index 0000000..63a3262
--- /dev/null
+++ b/hooks/use-user-message-id.ts
@@ -0,0 +1,10 @@
+'use client';
+
+import useSWR from 'swr';
+
+export function useUserMessageId() {
+  const { data: userMessageIdFromServer, mutate: setUserMessageIdFromServer } =
+    useSWR('userMessageIdFromServer', null);
+
+  return { userMessageIdFromServer, setUserMessageIdFromServer };
+}
diff --git a/lib/db/queries.ts b/lib/db/queries.ts
index d78fdf8..2bd400d 100644
--- a/lib/db/queries.ts
+++ b/lib/db/queries.ts
@@ -1,7 +1,7 @@
 import 'server-only';
 
 import { genSaltSync, hashSync } from 'bcrypt-ts';
-import { and, asc, desc, eq, gt } from 'drizzle-orm';
+import { and, asc, desc, eq, gt, gte } from 'drizzle-orm';
 import { drizzle } from 'drizzle-orm/postgres-js';
 import postgres from 'postgres';
 
@@ -279,3 +279,33 @@ export async function getSuggestionsByDocumentId({
     throw error;
   }
 }
+
+export async function getMessageById({ id }: { id: string }) {
+  try {
+    return await db.select().from(message).where(eq(message.id, id));
+  } catch (error) {
+    console.error('Failed to get message by id from database');
+    throw error;
+  }
+}
+
+export async function deleteMessagesByChatIdAfterTimestamp({
+  chatId,
+  timestamp,
+}: {
+  chatId: string;
+  timestamp: Date;
+}) {
+  try {
+    return await db
+      .delete(message)
+      .where(
+        and(eq(message.chatId, chatId), gte(message.createdAt, timestamp)),
+      );
+  } catch (error) {
+    console.error(
+      'Failed to delete messages by id after timestamp from database',
+    );
+    throw error;
+  }
+}