diff --git a/native/chat/text-message-tooltip-modal.react.js b/native/chat/text-message-tooltip-modal.react.js
--- a/native/chat/text-message-tooltip-modal.react.js
+++ b/native/chat/text-message-tooltip-modal.react.js
@@ -60,8 +60,8 @@
       inputState,
       'inputState should be set in TextMessageTooltipModal.onPressEdit',
     );
-    inputState.setEditedMessageID(messageInfo.id);
-  }, [inputState, messageInfo.id]);
+    inputState.setEditedMessage(messageInfo);
+  }, [inputState, messageInfo]);
   const renderEditIcon = React.useCallback(
     style => <SWMansionIcon name="edit-1" style={style} size={16} />,
     [],
diff --git a/native/input/input-state-container.react.js b/native/input/input-state-container.react.js
--- a/native/input/input-state-container.react.js
+++ b/native/input/input-state-container.react.js
@@ -60,6 +60,7 @@
   type RawMultimediaMessageInfo,
   type SendMessageResult,
   type SendMessagePayload,
+  type MessageInfo,
 } from 'lib/types/message-types.js';
 import type { RawImagesMessageInfo } from 'lib/types/messages/images.js';
 import type {
@@ -159,7 +160,7 @@
   state: State = {
     pendingUploads: {},
     editState: {
-      editedMessageID: null,
+      editedMessage: null,
     },
   };
   sendCallbacks: Array<() => void> = [];
@@ -401,7 +402,7 @@
       reportURIDisplayed: this.reportURIDisplayed,
       setPendingThreadUpdateHandler: this.setPendingThreadUpdateHandler,
       editState,
-      setEditedMessageID: this.setEditedMessageID,
+      setEditedMessage: this.setEditedMessage,
     }),
   );
 
@@ -1199,9 +1200,9 @@
     this.replyCallbacks.push(callbackReply);
   };
 
-  setEditedMessageID = (messageID: ?string) => {
+  setEditedMessage = (editedMessage: ?MessageInfo) => {
     this.setState({
-      editState: { editedMessageID: messageID },
+      editState: { editedMessage },
     });
   };
 
diff --git a/native/input/input-state.js b/native/input/input-state.js
--- a/native/input/input-state.js
+++ b/native/input/input-state.js
@@ -3,6 +3,7 @@
 import * as React from 'react';
 
 import type { NativeMediaSelection } from 'lib/types/media-types.js';
+import type { MessageInfo } from 'lib/types/message-types.js';
 import type { RawTextMessageInfo } from 'lib/types/messages/text.js';
 import type { ThreadInfo } from 'lib/types/thread-types.js';
 
@@ -23,7 +24,7 @@
 };
 
 export type EditState = {
-  +editedMessageID: ?string,
+  +editedMessage: ?MessageInfo,
 };
 
 export type InputState = {
@@ -55,7 +56,10 @@
     pendingThreadUpdateHandler: ?(ThreadInfo) => mixed,
   ) => void,
   +editState: EditState,
-  +setEditedMessageID: (editedMessageID: ?string) => void,
+  +setEditedMessage: (
+    editedMessage: ?MessageInfo,
+    callback?: () => void,
+  ) => void,
 };
 
 const InputStateContext: React.Context<?InputState> =