From f962f0d242b5be1c966935896707484e94b31e65 Mon Sep 17 00:00:00 2001
From: Ben Lopata <benlopata@gmail.com>
Date: Fri, 18 Oct 2024 14:33:31 -0500
Subject: [PATCH] Fix message reaction display.

---
 www/app/page.tsx              | 81 +++++++++++++++++++++--------------
 www/components/messagebox.tsx | 20 ++++++---
 www/utils/api.ts              |  2 +
 3 files changed, 64 insertions(+), 39 deletions(-)

diff --git a/www/app/page.tsx b/www/app/page.tsx
index 3edbf85..9555194 100644
--- a/www/app/page.tsx
+++ b/www/app/page.tsx
@@ -35,7 +35,11 @@ const URL = process.env.NEXT_PUBLIC_API_URL;
 export default function Home() {
   const [userId, setUserId] = useState<string>();
 
-  const [isThoughtsOpen, setIsThoughtsOpen] = useState<boolean>(false);
+  const [isThoughtsOpenState, setIsThoughtsOpenState] =
+    useState<boolean>(false);
+  const [openThoughtMessageId, setOpenThoughtMessageId] = useState<
+    string | null
+  >(null);
   const [isSidebarOpen, setIsSidebarOpen] = useState<boolean>(false);
 
   const [thought, setThought] = useState<string>('');
@@ -57,6 +61,14 @@ export default function Home() {
 
   const [isSubscribed, setIsSubscribed] = useState(false);
 
+  const setIsThoughtsOpen = (
+    isOpen: boolean,
+    messageId: string | null = null
+  ) => {
+    setIsThoughtsOpenState(isOpen);
+    setOpenThoughtMessageId(isOpen ? messageId : null);
+  };
+
   useEffect(() => {
     (async () => {
       const {
@@ -112,17 +124,17 @@ export default function Home() {
     return api.getConversations();
   };
 
-  const {
-    data: conversations,
-    mutate: mutateConversations,
-    error,
-  } = useSWR(userId, conversationsFetcher, {
-    onSuccess: (conversations) => {
-      setConversationId(conversations[0].conversationId);
-      setCanSend(true);
-    },
-    revalidateOnFocus: false,
-  });
+  const { data: conversations, mutate: mutateConversations } = useSWR(
+    userId,
+    conversationsFetcher,
+    {
+      onSuccess: (conversations) => {
+        setConversationId(conversations[0].conversationId);
+        setCanSend(true);
+      },
+      revalidateOnFocus: false,
+    }
+  );
 
   const messagesFetcher = async (conversationId: string) => {
     if (!userId) return Promise.resolve([]);
@@ -136,7 +148,6 @@ export default function Home() {
     data: messages,
     mutate: mutateMessages,
     isLoading: messagesLoading,
-    error: _,
   } = useSWR(conversationId, messagesFetcher, { revalidateOnFocus: false });
 
   const handleReactionAdded = async (messageId: string, reaction: Reaction) => {
@@ -148,21 +159,24 @@ export default function Home() {
       await api.addOrRemoveReaction(conversationId, messageId, reaction);
 
       // Optimistically update the local data
-      mutateMessages((currentMessages) => {
-        if (!currentMessages) return currentMessages;
-        return currentMessages.map((msg) => {
-          if (msg.id === messageId) {
-            return {
-              ...msg,
-              metadata: {
-                ...msg.metadata,
-                reaction,
-              },
-            };
-          }
-          return msg;
-        });
-      }, true);
+      mutateMessages(
+        (currentMessages) => {
+          if (!currentMessages) return currentMessages;
+          return currentMessages.map((msg) => {
+            if (msg.id === messageId) {
+              return {
+                ...msg,
+                metadata: {
+                  ...msg.metadata,
+                  reaction,
+                },
+              };
+            }
+            return msg;
+          });
+        },
+        { revalidate: false }
+      );
     } catch (error) {
       console.error('Failed to update reaction:', error);
     }
@@ -248,7 +262,7 @@ export default function Home() {
 
         mutateMessages(
           [
-            ...newMessages?.slice(0, -1)!,
+            ...(newMessages?.slice(0, -1) || []),
             {
               text: currentModelOutput,
               isUser: false,
@@ -336,7 +350,10 @@ export default function Home() {
               loading={messagesLoading}
               conversationId={conversationId}
               setThought={setThought}
-              setIsThoughtsOpen={setIsThoughtsOpen}
+              isThoughtOpen={openThoughtMessageId === message.id}
+              setIsThoughtsOpen={(isOpen) =>
+                setIsThoughtsOpen(isOpen, message.id)
+              }
               onReactionAdded={handleReactionAdded}
             />
           )) || (
@@ -403,8 +420,8 @@ export default function Home() {
       </div>
       <Thoughts
         thought={thought}
-        setIsThoughtsOpen={setIsThoughtsOpen}
-        isThoughtsOpen={isThoughtsOpen}
+        setIsThoughtsOpen={(isOpen: boolean) => setIsThoughtsOpen(isOpen, null)}
+        isThoughtsOpen={isThoughtsOpenState}
       />
     </main>
   );
diff --git a/www/components/messagebox.tsx b/www/components/messagebox.tsx
index c2bdd7e..da6b00e 100644
--- a/www/components/messagebox.tsx
+++ b/www/components/messagebox.tsx
@@ -17,7 +17,7 @@ interface MessageBoxProps {
   conversationId?: string;
   message: Message;
   loading?: boolean;
-  isThoughtsOpen?: boolean;
+  isThoughtOpen?: boolean;
   setIsThoughtsOpen: (isOpen: boolean) => void;
   setThought: (thought: string) => void;
   onReactionAdded: (messageId: string, reaction: Reaction) => Promise<void>;
@@ -29,6 +29,7 @@ export default function MessageBox({
   URL,
   message,
   loading = false,
+  isThoughtOpen,
   setIsThoughtsOpen,
   conversationId,
   onReactionAdded,
@@ -49,10 +50,7 @@ export default function MessageBox({
 
     try {
       const reactionToSend = reaction === newReaction ? null : newReaction;
-      await onReactionAdded(
-        messageId,
-        reactionToSend as Exclude<Reaction, null>
-      );
+      await onReactionAdded(messageId, reactionToSend as Reaction);
     } catch (err) {
       console.error(err);
       setError('Failed to update reaction.');
@@ -63,7 +61,11 @@ export default function MessageBox({
 
   const handleFetchThought = async () => {
     if (!messageId || !conversationId || !userId || !URL) return;
-
+    if (isThoughtOpen) {
+      // If thought is already open, close it
+      setIsThoughtsOpen(false);
+      return;
+    }
     setIsThoughtLoading(true);
     setError(null);
 
@@ -140,7 +142,11 @@ export default function MessageBox({
               </div>
             </button>
             <button
-              className="p-2 rounded-full bg-gray-200 dark:bg-gray-700"
+              className={`p-2 rounded-full ${
+                isThoughtOpen
+                  ? 'bg-neon-green text-gray-800'
+                  : 'bg-gray-200 dark:bg-gray-700'
+              }`}
               onClick={handleFetchThought}
               disabled={isThoughtLoading}
             >
diff --git a/www/utils/api.ts b/www/utils/api.ts
index de920cb..88fb306 100644
--- a/www/utils/api.ts
+++ b/www/utils/api.ts
@@ -165,9 +165,11 @@ export class API {
       if (!rawMessages) return [];
       const messages: Message[] = rawMessages.map((rawMessage: any) => {
         return {
+          ...rawMessage,
           text: rawMessage.content,
           isUser: rawMessage.isUser,
           id: rawMessage.id,
+          metadata: rawMessage.metadata,
         };
       });