From 8dd73568ce4b8a591ec7a89945d90c86c124960b Mon Sep 17 00:00:00 2001
From: Alois <aloisianer@proton.me>
Date: Tue, 7 Jul 2026 18:13:43 +0200
Subject: [PATCH] (feat): add message editing (qol): update todos

---
 packages/chat/src/components/message.tsx      |  73 ++++++++-
 .../src/components/messageContextMenu.tsx     |  34 +++-
 packages/chat/src/context.tsx                 | 151 +++++++++++++++++-
 packages/chat/todo.md                         |   1 +
 packages/mtp/src/context.tsx                  |   1 +
 packages/shared/src/data.ts                   |  47 ++++++
 todo.md                                       |   1 +
 type-maps.yaml                                |   8 +
 8 files changed, 310 insertions(+), 6 deletions(-)

diff --git a/packages/chat/src/components/message.tsx b/packages/chat/src/components/message.tsx
index 9848fbf..aa050fd 100644
--- a/packages/chat/src/components/message.tsx
+++ b/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} />
diff --git a/packages/chat/src/components/messageContextMenu.tsx b/packages/chat/src/components/messageContextMenu.tsx
index a4be6f7..7d3415e 100644
--- a/packages/chat/src/components/messageContextMenu.tsx
+++ b/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>
   );
diff --git a/packages/chat/src/context.tsx b/packages/chat/src/context.tsx
index a2b55ff..a2ed7b3 100644
--- a/packages/chat/src/context.tsx
+++ b/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;
diff --git a/packages/chat/todo.md b/packages/chat/todo.md
index 4bff343..1bed8c5 100644
--- a/packages/chat/todo.md
+++ b/packages/chat/todo.md
@@ -3,3 +3,4 @@
 - Placeholder image if media fails to load
 - Signature verifications via ed25519 key
 - Confirmation when exiting with text in the input box.
+- Add arrow up hotkey to edit last message
diff --git a/packages/mtp/src/context.tsx b/packages/mtp/src/context.tsx
index 9f96555..a781c77 100644
--- a/packages/mtp/src/context.tsx
+++ b/packages/mtp/src/context.tsx
@@ -200,6 +200,7 @@ export function Provider(props: {
 
     const unsubscribers = [
       "MessageLive",
+      "MessageEditLive",
       "MessageState",
       "CallInvite",
       "ErrorNoIota",
diff --git a/packages/shared/src/data.ts b/packages/shared/src/data.ts
index 4d1fdb4..5bba99a 100644
--- a/packages/shared/src/data.ts
+++ b/packages/shared/src/data.ts
@@ -66,6 +66,11 @@ const callSecretEnvelopeRequest = z.object({
   WrappingScheme: z.string(),
 });
 
+export const Reaction = z.object({
+  Reaction: z.string(),
+  SenderId: z.number(),
+});
+
 export const Message = z.object({
   NotEncrypted: z.boolean().optional(),
   SenderId: z.number(),
@@ -79,6 +84,8 @@ export const Message = z.object({
   MessageState: z
     .enum(["read", "received", "sent", "sending", "awaiting"]) // awaiting for 'internal' use
     .default("received"),
+  Edited: z.boolean().optional(),
+  Reactions: z.array(Reaction).optional(),
 });
 
 export const failedUser = {
@@ -195,6 +202,46 @@ export const mtp = {
       PingIota: z.number(),
     }),
   },
+  MessageEditLive: {
+    request: z.object({}),
+    response: z.object({
+      Content: z.base64(),
+      ChatPartnerId: z.number(),
+      SendTime: z.number(),
+    }),
+  },
+  MessageEdit: {
+    request: z.object({
+      Content: z.base64(),
+      ChatPartnerId: z.number(),
+      SendTime: z.number(),
+    }),
+    response: z.object({}),
+  },
+  MessageReactionAdd: {
+    request: z.object({
+      Reaction: z.string(),
+      SendTime: z.number(),
+      ChatPartnerId: z.number(),
+    }),
+    response: z.object(),
+  },
+  MessageReactionRemove: {
+    request: z.object({
+      Reaction: z.string(),
+      SendTime: z.number(),
+      ChatPartnerId: z.number(),
+    }),
+    response: z.object(),
+  },
+  MessageReactionLive: {
+    request: z.object(),
+    response: z.object({
+      Reaction: z.string(),
+      SendTime: z.number(),
+      ChatPartnerId: z.number(),
+    }),
+  },
   MessageLive: {
     request: z.object({}).optional(),
     response: z.object({
diff --git a/todo.md b/todo.md
index 6927b1a..71b7816 100644
--- a/todo.md
+++ b/todo.md
@@ -1,3 +1,4 @@
 - Move legal to extra onboarding package
 - Add a bunch of tests
 - Add packages/cache/ to handle caching
+- Full accessability
diff --git a/type-maps.yaml b/type-maps.yaml
index 15f49d7..1e9583b 100644
--- a/type-maps.yaml
+++ b/type-maps.yaml
@@ -143,6 +143,11 @@ type_maps:
       GetChatSecret: 140
       ChatSecretResponse: 141
       ChatSecretForward: 142
+      MessageEditLive: 144
+      MessageEdit: 145
+      MessageReactionAdd: 146
+      MessageReactionRemove: 147
+      MessageReactionLive: 148
     DataTypes:
       ErrorType: 32
       ErrorProtocol: 33
@@ -257,3 +262,6 @@ type_maps:
       SenderUserId: 152
       RecipientUserId: 153
       Recipients: 154
+      Edited: 155
+      Reactions: 156
+      Reaction: 157