Optimistic update when user send a message

Author: BrunoQuaresma

chat/src/components/ChatInterface.tsx

@@ -11,9 +11,15 @@ import { Alert, AlertTitle, AlertDescription } from "./ui/alert";
 import { ModeToggle } from "./mode-toggle";
 
 interface Message {
+  id: number;
   role: string;
   content: string;
-  id: number;
+}
+
+// Draft messages are used to optmistically update the UI
+// before the server responds.
+interface DraftMessage extends Omit<Message, "id"> {
+  id?: number;
 }
 
 interface MessageUpdateEvent {
@@ -28,7 +34,7 @@ interface StatusChangeEvent {
 }
 
 export default function ChatInterface() {
-  const [messages, setMessages] = useState<Message[]>([]);
+  const [messages, setMessages] = useState<(Message | DraftMessage)[]>([]);
   const [loading, setLoading] = useState<boolean>(false);
   const [serverStatus, setServerStatus] = useState<string>("unknown");
   const searchParams = useSearchParams();
@@ -87,12 +93,18 @@ export default function ChatInterface() {
         const data: MessageUpdateEvent = JSON.parse(event.data);
 
         setMessages((prevMessages) => {
+          // Clean up draft messages
+          const updatedMessages = [...prevMessages].filter(
+            (m) => m.id !== undefined
+          );
+
           // Check if message with this ID already exists
-          const existingIndex = prevMessages.findIndex((m) => m.id === data.id);
+          const existingIndex = updatedMessages.findIndex(
+            (m) => m.id === data.id
+          );
 
           if (existingIndex !== -1) {
             // Update existing message
-            const updatedMessages = [...prevMessages];
             updatedMessages[existingIndex] = {
               role: data.role,
               content: data.message,
@@ -102,7 +114,7 @@ export default function ChatInterface() {
           } else {
             // Add new message
             return [
-              ...prevMessages,
+              ...updatedMessages,
               {
                 role: data.role,
                 content: data.message,
@@ -164,6 +176,10 @@ export default function ChatInterface() {
 
     // For raw messages, don't set loading state as it's usually fast
     if (type === "user") {
+      setMessages((prevMessages) => [
+        ...prevMessages,
+        { role: "user", content },
+      ]);
       setLoading(true);
     }
 
@@ -263,7 +279,7 @@ export default function ChatInterface() {
           </div>
         )}
 
-        <MessageList messages={messages} loading={loading} />
+        <MessageList messages={messages} />
         <MessageInput onSendMessage={sendMessage} disabled={loading} />
       </main>
     </div>

chat/src/components/MessageList.tsx

@@ -6,15 +6,17 @@ interface Message {
   id: number;
 }
 
+// Draft messages are used to optmistically update the UI
+// before the server responds.
+interface DraftMessage extends Omit<Message, "id"> {
+  id?: number;
+}
+
 interface MessageListProps {
-  messages: Message[];
-  loading?: boolean;
+  messages: (Message | DraftMessage)[];
 }
 
-export default function MessageList({
-  messages,
-  loading = false,
-}: MessageListProps) {
+export default function MessageList({ messages }: MessageListProps) {
   // If no messages, show a placeholder
   if (messages.length === 0) {
     return (
@@ -73,15 +75,15 @@ export default function MessageList({
       <div className="p-4 flex flex-col gap-4 max-w-4xl mx-auto">
         {messages.map((message) => (
           <div
-            key={message.id}
+            key={message.id ?? "draft"}
             className={`${message.role === "user" ? "text-right" : ""}`}
           >
             <div
               className={`inline-block rounded-lg ${
                 message.role === "user"
                   ? "bg-accent-foreground rounded-lg max-w-[90%] px-4 py-3 text-accent"
                   : "max-w-[80ch]"
-              }`}
+              } ${message.id === undefined ? "animate-pulse" : ""}`}
             >
               <div
                 className={`whitespace-pre-wrap break-words text-left text-sm ${
@@ -97,13 +99,6 @@ export default function MessageList({
             </div>
           </div>
         ))}
-
-        {/* Loading indicator for message being sent */}
-        {loading && (
-          <div className="w-fit self-end">
-            <LoadingDots />
-          </div>
-        )}
       </div>
     </div>
   );