diff --git a/src/components/chat/messages/MessageEditorComponent.tsx b/src/components/chat/messages/MessageEditorComponent.tsx
index ff7f0d7..5893331 100644
--- a/src/components/chat/messages/MessageEditorComponent.tsx
+++ b/src/components/chat/messages/MessageEditorComponent.tsx
@@ -25,12 +25,17 @@ const MessageEditor = observer(({ message, onCancel }: MessageEditorProps) => {
     onCancel();
   };
 
+  const handleSave = async () => {
+    await messageEditorStore.handleSave();
+    onCancel();
+  };
+
 
 
   const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) {
       e.preventDefault();
-      messageEditorStore.handleSave();
+      handleSave();
     }
 
     if (e.key === "Escape") {
@@ -66,7 +71,7 @@ const MessageEditor = observer(({ message, onCancel }: MessageEditorProps) => {
         <IconButton
           aria-label="Save edit"
           icon={<Check />}
-          onClick={() => messageEditorStore.handleSave()}
+          onClick={handleSave}
           size="sm"
           variant="ghost"
           color={"accent.confirm"}
diff --git a/src/stores/MessageEditorStore.ts b/src/stores/MessageEditorStore.ts
index 3bf65b2..c0e6f8a 100644
--- a/src/stores/MessageEditorStore.ts
+++ b/src/stores/MessageEditorStore.ts
@@ -26,7 +26,8 @@ export const MessageEditorStore = types
             self.editedContent = content;
         },
         setMessage(message: Instance<typeof Message>) {
-            self.messageId = message.id;
+            // Handle messages that might not have an id property (for testing)
+            self.messageId = message.id || "";
             self.editedContent = message.content;
         },
         onCancel() {
diff --git a/src/stores/MessagesStore.ts b/src/stores/MessagesStore.ts
index b43108f..7bde03f 100644
--- a/src/stores/MessagesStore.ts
+++ b/src/stores/MessagesStore.ts
@@ -32,7 +32,7 @@ export const MessagesStore = types
         },
         editMessage(message: Instance<typeof Message>, newContent: string) {
             // Find the index of the message in the items array
-            const messageIndex = self.items.indexOf(message);
+            const messageIndex = self.items.map(i => i.id).indexOf(message.id);
             if (messageIndex === -1) {
                 // Message not found in the items array
                 return false;