Merge pull request #15 from jtclarkjr/fix/unsend-db-messages

jtclarkjr · web-flow · commit 6c35e824eefb · 2025-11-08T12:34:47.000+09:00
fix: unsend initial loaded messages not removed from ui
diff --git a/components/realtime-chat.tsx b/components/realtime-chat.tsx
@@ -49,7 +49,9 @@ export const RealtimeChat = ({
     loading,
     queueStatus,
     clearFailedMessages,
-    onMessageUpdate
+    onMessageUpdate,
+    markMessageAsDeleted,
+    deletedMessageIds
   } = useRealtimeChat({
     roomId,
     username,
@@ -61,7 +63,8 @@ export const RealtimeChat = ({
   const { unsendMessage, isUnsending } = useUnsendMessage({
     userId,
     roomId,
-    onMessageUpdate
+    onMessageUpdate,
+    markMessageAsDeleted
   })
 
   const {
@@ -100,7 +103,8 @@ export const RealtimeChat = ({
     initialMessages,
     realtimeMessages,
     streamingMessages,
-    userId
+    userId,
+    deletedMessageIds
   })
 
   // Smart auto-scroll that only scrolls when appropriate
diff --git a/hooks/chat/use-realtime-chat.tsx b/hooks/chat/use-realtime-chat.tsx
@@ -209,6 +209,42 @@ export function useRealtimeChat({
     )
   }, [missedMessages, optimisticMessages, deletedMessageIds])
 
+  // Callback to update confirmed messages and track deleted messages
+  const handleMessageUpdate = useCallback(
+    (updater: (messages: ChatMessage[]) => ChatMessage[]) => {
+      setConfirmedMessages((current) => {
+        const updated = updater(current)
+        
+        // Check if any messages were marked as deleted and add them to deletedMessageIds
+        updated.forEach((msg) => {
+          if (msg.isDeleted) {
+            setDeletedMessageIds((prev) => new Set(prev).add(msg.id))
+          }
+        })
+        
+        return updated
+      })
+    },
+    []
+  )
+
+  // Direct method to mark a message as deleted (for unsend)
+  const markMessageAsDeleted = useCallback((messageId: string) => {
+    setDeletedMessageIds((prev) => new Set(prev).add(messageId))
+    // Also update confirmed messages if the message exists there
+    setConfirmedMessages((current) =>
+      current.map((msg) =>
+        msg.id === messageId
+          ? {
+              ...msg,
+              isDeleted: true,
+              content: 'This message was deleted'
+            }
+          : msg
+      )
+    )
+  }, [])
+
   return {
     messages: allMessages,
     sendMessage,
@@ -217,6 +253,8 @@ export function useRealtimeChat({
     loading: missedMessagesLoading,
     queueStatus,
     clearFailedMessages,
-    onMessageUpdate: setConfirmedMessages
+    onMessageUpdate: handleMessageUpdate,
+    markMessageAsDeleted,
+    deletedMessageIds
   }
 }
diff --git a/hooks/messages/use-message-merging.tsx b/hooks/messages/use-message-merging.tsx
@@ -8,13 +8,15 @@ interface UseMessageMergingProps {
   realtimeMessages: ChatMessage[]
   streamingMessages: ChatMessage[]
   userId: string
+  deletedMessageIds?: Set<string>
 }
 
 export function useMessageMerging({
   initialMessages,
   realtimeMessages,
   streamingMessages,
-  userId
+  userId,
+  deletedMessageIds = new Set()
 }: UseMessageMergingProps) {
   const allMessages = useMemo(() => {
     const mergedMessages = [...initialMessages, ...realtimeMessages]
@@ -39,6 +41,12 @@ export function useMessageMerging({
     const uniqueMessages = mergedMessages.filter((message, index, self) => {
       if (!message) return false
       if (!message.id) return false
+      
+      // Filter out deleted messages (check both isDeleted flag and deletedMessageIds set)
+      if (message.isDeleted || deletedMessageIds.has(message.id)) {
+        return false
+      }
+      
       // Filter out messages without content or invalid structure
       const isStreamingMessage = streamingMessages.some((sm) => sm === message)
       if (
@@ -90,7 +98,7 @@ export function useMessageMerging({
     })
 
     return sortedMessages
-  }, [initialMessages, realtimeMessages, streamingMessages, userId])
+  }, [initialMessages, realtimeMessages, streamingMessages, userId, deletedMessageIds])
 
   return allMessages
 }
diff --git a/hooks/messages/use-unsend-message.tsx b/hooks/messages/use-unsend-message.tsx
@@ -8,6 +8,7 @@ interface UseUnsendMessageProps {
   userId: string
   roomId: string
   onMessageUpdate: (updater: (messages: ChatMessage[]) => ChatMessage[]) => void
+  markMessageAsDeleted: (messageId: string) => void
 }
 
 interface UseUnsendMessageReturn {
@@ -18,7 +19,8 @@ interface UseUnsendMessageReturn {
 export function useUnsendMessage({
   userId,
   roomId,
-  onMessageUpdate
+  onMessageUpdate,
+  markMessageAsDeleted
 }: UseUnsendMessageProps): UseUnsendMessageReturn {
   const [unsendingMessages, setUnsendingMessages] = useState<Set<string>>(
     new Set()
@@ -41,20 +43,8 @@ export function useUnsendMessage({
         })
 
         if (result.success) {
-          // Update the message in the local state to show it as deleted
-          onMessageUpdate((current) =>
-            current.map((msg) =>
-              msg.id === messageId
-                ? {
-                    ...msg,
-                    isDeleted: true,
-                    deletedAt: result.message.deletedAt,
-                    deletedBy: result.message.deletedBy,
-                    content: 'This message was deleted'
-                  }
-                : msg
-            )
-          )
+          // Mark message as deleted - this handles both confirmed and missed messages
+          markMessageAsDeleted(messageId)
           return true
         } else {
           console.error('Failed to unsend message:', result.error)
@@ -71,7 +61,7 @@ export function useUnsendMessage({
         })
       }
     },
-    [userId, roomId, onMessageUpdate, unsendingMessages, unsendMessageMutation]
+    [userId, roomId, onMessageUpdate, markMessageAsDeleted, unsendingMessages, unsendMessageMutation]
   )
 
   const isUnsending = useCallback(
