(feat): add message editing

(qol): update todos

packages/chat/src/components/message.tsx

@@ -27,6 +27,10 @@ import MessageContextMenu from "./messageContextMenu";
 import Media from "./media";
 import { useStorage } from "@tensamin/storage/context";
 import { useMTP } from "@tensamin/mtp";
+import Input from "@tensamin/markdown/input";
+import { useChat } from "../context";
+import { encryptChatText } from "@tensamin/crypto/chatSecret";
+import { log, toast } from "@tensamin/shared/log";
 
 function MessageComponent({
   grouped,
@@ -138,6 +142,62 @@ function MessageComponent({
     }
   }, [message.Content]);
 
+  // Message editing
+  const { chatSecret, editMessageContent, userId } = useChat();
+  const [editing, setEditing] = useState(false);
+  const [editDraft, setEditDraft] = useState(message.Content);
+  useEffect(() => {
+    if (!editing) {
+      setEditDraft(message.Content);
+    }
+  }, [editing, message.Content]);
+  const editMessage = useCallback(
+    async (newContent: string) => {
+      if (!chatSecret) return;
+
+      const encryptedContent = await encryptChatText(
+        chatSecret,
+        newContent,
+      ).catch((err) => {
+        toast("error", "Failed to encrypt edit", String(err));
+      });
+
+      if (!encryptedContent) return;
+
+      const previousContent = message.Content;
+      editMessageContent(message.SendTime, newContent);
+
+      try {
+        const response = await send("MessageEdit", {
+          ChatPartnerId: userId,
+          Content: encryptedContent,
+          SendTime: message.SendTime,
+        });
+
+        if (response.type.startsWith("Error")) {
+          throw new Error(response.type);
+        }
+      } catch (err) {
+        editMessageContent(
+          message.SendTime,
+          previousContent,
+          message.Edited ?? false,
+        );
+        log(1, "chat", "red", "Failed to edit message", err);
+        toast("error", "Failed to edit message", String(err));
+      }
+    },
+    [
+      chatSecret,
+      editMessageContent,
+      userId,
+      message.Content,
+      message.Edited,
+      message.SendTime,
+      send,
+    ],
+  );
+
   return (
     <div
       // pt-3 is to get a gap between messages
@@ -147,6 +207,8 @@ function MessageComponent({
         <MessageContextMenu
           content={message.Content}
           messageId={message.SendTime}
+          onSetEditing={setEditing}
+          senderId={message.SenderId}
         >
           <div
             className={cn(
@@ -221,7 +283,16 @@ function MessageComponent({
                     </div>
                   </div>
                 )}
-                {isValidURL ? (
+                {editing ? (
+                  <Input
+                    setValue={setEditDraft}
+                    value={editDraft}
+                    onSubmit={() => {
+                      editMessage(editDraft);
+                      setEditing(false);
+                    }}
+                  />
+                ) : isValidURL ? (
                   <Media link={message.Content} />
                 ) : (
                   <Text value={message.Content} />

packages/chat/src/components/messageContextMenu.tsx

@@ -17,9 +17,10 @@ import {
   useIsMobile,
 } from "@tensamin/ui";
 import { Pin, Clipboard, Pen, Reply, Forward, Trash } from "lucide-react";
-import { useMemo, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
 import type { ReactElement, ReactNode } from "react";
 import { useChat } from "../context";
+import { useStorage } from "@tensamin/storage/context";
 
 async function copyText(text: string) {
   await navigator.clipboard.writeText(text);
@@ -136,6 +137,8 @@ function MessageMenuContent({
   messageId,
   onAddReaction,
   showReactionItems = true,
+  onSetEditing,
+  senderId,
 }: {
   components: MenuComponents;
   content: string;
@@ -143,12 +146,20 @@ function MessageMenuContent({
   messageId: number;
   onAddReaction?: () => void | Promise<void>;
   showReactionItems?: boolean;
+  onSetEditing: (value: boolean) => void;
+  senderId: number;
 }) {
   const { Content, Group, Item, Separator, Sub, SubContent, SubTrigger } =
     components;
 
+  const { load } = useStorage();
   const { setReplyTo } = useChat();
 
+  const [ownId, setOwnId] = useState(0);
+  useEffect(() => {
+    load("user_id").then(setOwnId);
+  }, [load]);
+
   return (
     <Content>
       <Group>
@@ -172,9 +183,16 @@ function MessageMenuContent({
       </Group>
       <Separator />
       <Group>
-        <Item disabled className="flex justify-between">
-          <p>Edit Message</p> <Pen />
-        </Item>
+        {senderId === ownId && (
+          <Item
+            className="flex justify-between"
+            onClick={() => {
+              onSetEditing(true);
+            }}
+          >
+            <p>Edit Message</p> <Pen />
+          </Item>
+        )}
         <Item
           className="flex justify-between"
           onClick={() => {
@@ -214,10 +232,14 @@ export default function MessageContextMenu({
   children,
   content,
   messageId,
+  onSetEditing,
+  senderId,
 }: {
   children: ReactElement;
   content: string;
   messageId: number;
+  onSetEditing: (value: boolean) => void;
+  senderId: number;
 }) {
   const isMobile = useIsMobile();
   const devEnabled = useMemo(
@@ -257,6 +279,8 @@ export default function MessageContextMenu({
             messageId={messageId}
             onAddReaction={() => setReactionDrawerOpen(true)}
             showReactionItems={false}
+            onSetEditing={onSetEditing}
+            senderId={senderId}
           />
         </Drawer>
         <Drawer open={reactionDrawerOpen} onOpenChange={setReactionDrawerOpen}>
@@ -282,6 +306,8 @@ export default function MessageContextMenu({
         content={content}
         devEnabled={devEnabled}
         messageId={messageId}
+        onSetEditing={onSetEditing}
+        senderId={senderId}
       />
     </ContextMenu>
   );

packages/chat/src/context.tsx

@@ -86,6 +86,39 @@ function updateMessageStateBySendTime<
   };
 }
 
+function updateMessageContentBySendTime<
+  T extends { SendTime: number; Content: string; Edited?: boolean },
+>(
+  messages: T[],
+  sendTime: number,
+  content: string,
+  edited = true,
+): { next: T[]; updated: boolean } {
+  let updated = false;
+
+  const next = messages.map((item) => {
+    if (item.SendTime !== sendTime) {
+      return item;
+    }
+
+    if (item.Content === content && item.Edited === edited) {
+      return item;
+    }
+
+    updated = true;
+    return {
+      ...item,
+      Content: content,
+      Edited: edited,
+    };
+  });
+
+  return {
+    next,
+    updated,
+  };
+}
+
 function assertProtocolSuccess(type: string, response: { type: string }) {
   if (response.type.startsWith("Error")) {
     throw new Error(`${type} failed: ${response.type}`);
@@ -413,9 +446,112 @@ export default function Provider({ children }: { children: ReactNode }) {
     setLiveMessagesState([]);
   }, []);
 
+  const editMessageContent = useCallback(
+    (sendTime: number, content: string, edited = true) => {
+      setLiveMessagesState((prev) => {
+        const { next, updated } = updateMessageContentBySendTime(
+          prev,
+          sendTime,
+          content,
+          edited,
+        );
+        return updated ? next : prev;
+      });
+
+      const queryKey = [
+        "chat-messages",
+        String(userIdValue),
+        currentChatSecret !== null,
+      ] as const;
+      queryClient.setQueryData<InfiniteData<RawMessages>>(
+        queryKey,
+        (current) => {
+          if (!current) {
+            return current;
+          }
+
+          let updated = false;
+
+          const pages = current.pages.map((page) => {
+            const nextPage = updateMessageContentBySendTime(
+              page,
+              sendTime,
+              content,
+              edited,
+            );
+
+            if (nextPage.updated) {
+              updated = true;
+            }
+
+            return nextPage.next;
+          });
+
+          if (!updated) {
+            return current;
+          }
+
+          return {
+            ...current,
+            pages,
+          };
+        },
+      );
+    },
+    [currentChatSecret, userIdValue],
+  );
+
   // Get live updates for message states
   useEffect(() => {
     return subscribePush((message) => {
+      if (message.type === "MessageEditLive") {
+        if (!currentChatSecret) return;
+
+        const rawData = message.data as {
+          ChatPartnerId: unknown;
+          SendTime: unknown;
+          Content: string;
+        };
+
+        const chatPartnerId = Number(rawData.ChatPartnerId);
+        const sendTime = Number(rawData.SendTime);
+
+        if (!Number.isFinite(chatPartnerId) || !Number.isFinite(sendTime)) {
+          log(
+            3,
+            "chat",
+            "yellow",
+            "Cancel message edit update due to invalid data",
+          );
+          return;
+        }
+
+        if (chatPartnerId !== userIdValue) {
+          log(
+            3,
+            "chat",
+            "yellow",
+            "Cancel message edit update due to user ID mismatch",
+            {
+              expected: userIdValue,
+              received: chatPartnerId,
+            },
+          );
+          return;
+        }
+
+        void decryptChatText(currentChatSecret, rawData.Content)
+          .then((content) => {
+            editMessageContent(sendTime, content);
+          })
+          .catch((err) => {
+            log(1, "chat", "red", "Failed to decrypt message edit", err, {
+              SendTime: sendTime,
+            });
+          });
+        return;
+      }
+
       if (message.type !== "MessageState") return;
 
       const rawData = message.data as {
@@ -505,7 +641,14 @@ export default function Provider({ children }: { children: ReactNode }) {
         },
       );
     });
-  }, [addLiveMessage, currentChatSecret, send, subscribePush, userIdValue]);
+  }, [
+    addLiveMessage,
+    currentChatSecret,
+    editMessageContent,
+    send,
+    subscribePush,
+    userIdValue,
+  ]);
 
   // Replys
   const [replyTo, setReplyTo] = useState<number | undefined>(undefined);
@@ -518,6 +661,7 @@ export default function Provider({ children }: { children: ReactNode }) {
           getChatSecret,
           liveMessages: () => liveMessagesState,
           addLiveMessage,
+          editMessageContent,
           clearLiveMessages,
           chatSecret: currentChatSecret,
           userId: userIdValue,
@@ -541,6 +685,11 @@ type contextType = {
   addLiveMessage: (message: RawMessage) => {
     setFailed: (failed: boolean) => void;
   };
+  editMessageContent: (
+    sendTime: number,
+    content: string,
+    edited?: boolean,
+  ) => void;
   clearLiveMessages: () => void;
   chatSecret: Uint8Array | null;
   userId: number;