Channel chat: Add edit message (#9035)

**Summary**:
- Removed reply message from message_menu
- Made render_popover_buttons a bit more reusable
- Fixed issue that you can't close the reply/edit preview when you are
not focusing the message editor
- Notify only the new people that were mentioned inside the edited
message

**Follow up**
- Fix that we update the notification message for the people that we
mentioned already
- Fix that we remove the notification when a message gets deleted.
  - Fix last acknowledge message id is in correct now

**Todo**:
- [x] Add tests
- [x] Change new added bindings to the `Editor::Cancel` event.

Release Notes:

- Added editing of chat messages
([#6707](https://github.com/zed-industries/zed/issues/6707)).

<img width="239" alt="Screenshot 2024-03-09 at 11 55 23"
src="https://github.com/zed-industries/zed/assets/62463826/b0949f0d-0f8b-43e1-ac20-4c6d40ac41e1">
<img width="240" alt="Screenshot 2024-03-13 at 13 34 23"
src="https://github.com/zed-industries/zed/assets/62463826/d0636da2-c5aa-4fed-858e-4bebe5695ba7">

---------

Co-authored-by: Bennet Bo Fenner <53836821+bennetbo@users.noreply.github.com>
Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>

crates/collab/migrations.sqlite/20221109000000_test_schema.sql

@@ -219,6 +219,7 @@ CREATE TABLE IF NOT EXISTS "channel_messages" (
     "sender_id" INTEGER NOT NULL REFERENCES users (id),
     "body" TEXT NOT NULL,
     "sent_at" TIMESTAMP,
+    "edited_at" TIMESTAMP,
     "nonce" BLOB NOT NULL,
     "reply_to_message_id" INTEGER DEFAULT NULL
 );

crates/collab/migrations/20240221151017_add_edited_at_field_to_channel_message.sql

@@ -0,0 +1 @@
+ALTER TABLE channel_messages ADD edited_at TIMESTAMP DEFAULT NULL;

crates/collab/src/db.rs

@@ -458,6 +458,14 @@ pub struct CreatedChannelMessage {
     pub notifications: NotificationBatch,
 }
 
+pub struct UpdatedChannelMessage {
+    pub message_id: MessageId,
+    pub participant_connection_ids: Vec<ConnectionId>,
+    pub notifications: NotificationBatch,
+    pub reply_to_message_id: Option<MessageId>,
+    pub timestamp: PrimitiveDateTime,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
 pub struct Invite {
     pub email_address: String,

crates/collab/src/db/queries/messages.rs

@@ -162,6 +162,9 @@ impl Database {
                         lower_half: nonce.1,
                     }),
                     reply_to_message_id: row.reply_to_message_id.map(|id| id.to_proto()),
+                    edited_at: row
+                        .edited_at
+                        .map(|t| t.assume_utc().unix_timestamp() as u64),
                 }
             })
             .collect::<Vec<_>>();
@@ -199,6 +202,31 @@ impl Database {
         Ok(messages)
     }
 
+    fn format_mentions_to_entities(
+        &self,
+        message_id: MessageId,
+        body: &str,
+        mentions: &[proto::ChatMention],
+    ) -> Result<Vec<tables::channel_message_mention::ActiveModel>> {
+        Ok(mentions
+            .iter()
+            .filter_map(|mention| {
+                let range = mention.range.as_ref()?;
+                if !body.is_char_boundary(range.start as usize)
+                    || !body.is_char_boundary(range.end as usize)
+                {
+                    return None;
+                }
+                Some(channel_message_mention::ActiveModel {
+                    message_id: ActiveValue::Set(message_id),
+                    start_offset: ActiveValue::Set(range.start as i32),
+                    end_offset: ActiveValue::Set(range.end as i32),
+                    user_id: ActiveValue::Set(UserId::from_proto(mention.user_id)),
+                })
+            })
+            .collect::<Vec<_>>())
+    }
+
     /// Creates a new channel message.
     #[allow(clippy::too_many_arguments)]
     pub async fn create_channel_message(
@@ -249,6 +277,7 @@ impl Database {
                 nonce: ActiveValue::Set(Uuid::from_u128(nonce)),
                 id: ActiveValue::NotSet,
                 reply_to_message_id: ActiveValue::Set(reply_to_message_id),
+                edited_at: ActiveValue::NotSet,
             })
             .on_conflict(
                 OnConflict::columns([
@@ -270,23 +299,7 @@ impl Database {
                     let mentioned_user_ids =
                         mentions.iter().map(|m| m.user_id).collect::<HashSet<_>>();
 
-                    let mentions = mentions
-                        .iter()
-                        .filter_map(|mention| {
-                            let range = mention.range.as_ref()?;
-                            if !body.is_char_boundary(range.start as usize)
-                                || !body.is_char_boundary(range.end as usize)
-                            {
-                                return None;
-                            }
-                            Some(channel_message_mention::ActiveModel {
-                                message_id: ActiveValue::Set(message_id),
-                                start_offset: ActiveValue::Set(range.start as i32),
-                                end_offset: ActiveValue::Set(range.end as i32),
-                                user_id: ActiveValue::Set(UserId::from_proto(mention.user_id)),
-                            })
-                        })
-                        .collect::<Vec<_>>();
+                    let mentions = self.format_mentions_to_entities(message_id, body, mentions)?;
                     if !mentions.is_empty() {
                         channel_message_mention::Entity::insert_many(mentions)
                             .exec(&*tx)
@@ -522,4 +535,131 @@ impl Database {
         })
         .await
     }
+
+    /// Updates the channel message with the given ID, body and timestamp(edited_at).
+    pub async fn update_channel_message(
+        &self,
+        channel_id: ChannelId,
+        message_id: MessageId,
+        user_id: UserId,
+        body: &str,
+        mentions: &[proto::ChatMention],
+        edited_at: OffsetDateTime,
+    ) -> Result<UpdatedChannelMessage> {
+        self.transaction(|tx| async move {
+            let channel = self.get_channel_internal(channel_id, &tx).await?;
+            self.check_user_is_channel_participant(&channel, user_id, &tx)
+                .await?;
+
+            let mut rows = channel_chat_participant::Entity::find()
+                .filter(channel_chat_participant::Column::ChannelId.eq(channel_id))
+                .stream(&*tx)
+                .await?;
+
+            let mut is_participant = false;
+            let mut participant_connection_ids = Vec::new();
+            let mut participant_user_ids = Vec::new();
+            while let Some(row) = rows.next().await {
+                let row = row?;
+                if row.user_id == user_id {
+                    is_participant = true;
+                }
+                participant_user_ids.push(row.user_id);
+                participant_connection_ids.push(row.connection());
+            }
+            drop(rows);
+
+            if !is_participant {
+                Err(anyhow!("not a chat participant"))?;
+            }
+
+            let channel_message = channel_message::Entity::find_by_id(message_id)
+                .filter(channel_message::Column::SenderId.eq(user_id))
+                .one(&*tx)
+                .await?;
+
+            let Some(channel_message) = channel_message else {
+                Err(anyhow!("Channel message not found"))?
+            };
+
+            let edited_at = edited_at.to_offset(time::UtcOffset::UTC);
+            let edited_at = time::PrimitiveDateTime::new(edited_at.date(), edited_at.time());
+
+            let updated_message = channel_message::ActiveModel {
+                body: ActiveValue::Set(body.to_string()),
+                edited_at: ActiveValue::Set(Some(edited_at)),
+                reply_to_message_id: ActiveValue::Unchanged(channel_message.reply_to_message_id),
+                id: ActiveValue::Unchanged(message_id),
+                channel_id: ActiveValue::Unchanged(channel_id),
+                sender_id: ActiveValue::Unchanged(user_id),
+                sent_at: ActiveValue::Unchanged(channel_message.sent_at),
+                nonce: ActiveValue::Unchanged(channel_message.nonce),
+            };
+
+            let result = channel_message::Entity::update_many()
+                .set(updated_message)
+                .filter(channel_message::Column::Id.eq(message_id))
+                .filter(channel_message::Column::SenderId.eq(user_id))
+                .exec(&*tx)
+                .await?;
+            if result.rows_affected == 0 {
+                return Err(anyhow!(
+                    "Attempted to edit a message (id: {message_id}) which does not exist anymore."
+                ))?;
+            }
+
+            // we have to fetch the old mentions,
+            // so we don't send a notification when the message has been edited that you are mentioned in
+            let old_mentions = channel_message_mention::Entity::find()
+                .filter(channel_message_mention::Column::MessageId.eq(message_id))
+                .all(&*tx)
+                .await?;
+
+            // remove all existing mentions
+            channel_message_mention::Entity::delete_many()
+                .filter(channel_message_mention::Column::MessageId.eq(message_id))
+                .exec(&*tx)
+                .await?;
+
+            let new_mentions = self.format_mentions_to_entities(message_id, body, mentions)?;
+            if !new_mentions.is_empty() {
+                // insert new mentions
+                channel_message_mention::Entity::insert_many(new_mentions)
+                    .exec(&*tx)
+                    .await?;
+            }
+
+            let mut mentioned_user_ids = mentions.iter().map(|m| m.user_id).collect::<HashSet<_>>();
+            // Filter out users that were mentioned before
+            for mention in old_mentions {
+                mentioned_user_ids.remove(&mention.user_id.to_proto());
+            }
+
+            let mut notifications = Vec::new();
+            for mentioned_user in mentioned_user_ids {
+                notifications.extend(
+                    self.create_notification(
+                        UserId::from_proto(mentioned_user),
+                        rpc::Notification::ChannelMessageMention {
+                            message_id: message_id.to_proto(),
+                            sender_id: user_id.to_proto(),
+                            channel_id: channel_id.to_proto(),
+                        },
+                        false,
+                        &tx,
+                    )
+                    .await?,
+                );
+            }
+
+            Ok(UpdatedChannelMessage {
+                message_id,
+                participant_connection_ids,
+                notifications,
+                reply_to_message_id: channel_message.reply_to_message_id,
+                timestamp: channel_message.sent_at,
+            })
+        })
+        .await
+    }
 }

crates/collab/src/db/tables/channel_message.rs

@@ -11,6 +11,7 @@ pub struct Model {
     pub sender_id: UserId,
     pub body: String,
     pub sent_at: PrimitiveDateTime,
+    pub edited_at: Option<PrimitiveDateTime>,
     pub nonce: Uuid,
     pub reply_to_message_id: Option<MessageId>,
 }

crates/collab/src/rpc.rs

@@ -6,7 +6,7 @@ use crate::{
         self, BufferId, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage,
         Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, Project,
         ProjectId, RemoveChannelMemberResult, ReplicaId, RespondToChannelInvite, RoomId, ServerId,
-        User, UserId,
+        UpdatedChannelMessage, User, UserId,
     },
     executor::Executor,
     AppState, Error, RateLimit, RateLimiter, Result,
@@ -283,6 +283,7 @@ impl Server {
             .add_message_handler(leave_channel_chat)
             .add_request_handler(send_channel_message)
             .add_request_handler(remove_channel_message)
+            .add_request_handler(update_channel_message)
             .add_request_handler(get_channel_messages)
             .add_request_handler(get_channel_messages_by_id)
             .add_request_handler(get_notifications)
@@ -3191,6 +3192,7 @@ async fn send_channel_message(
             },
         )
         .await?;
+
     let message = proto::ChannelMessage {
         sender_id: session.user_id.to_proto(),
         id: message_id.to_proto(),
@@ -3199,6 +3201,7 @@ async fn send_channel_message(
         timestamp: timestamp.unix_timestamp() as u64,
         nonce: Some(nonce),
         reply_to_message_id: request.reply_to_message_id,
+        edited_at: None,
     };
     broadcast(
         Some(session.connection_id),
@@ -3261,6 +3264,71 @@ async fn remove_channel_message(
     Ok(())
 }
 
+async fn update_channel_message(
+    request: proto::UpdateChannelMessage,
+    response: Response<proto::UpdateChannelMessage>,
+    session: Session,
+) -> Result<()> {
+    let channel_id = ChannelId::from_proto(request.channel_id);
+    let message_id = MessageId::from_proto(request.message_id);
+    let updated_at = OffsetDateTime::now_utc();
+    let UpdatedChannelMessage {
+        message_id,
+        participant_connection_ids,
+        notifications,
+        reply_to_message_id,
+        timestamp,
+    } = session
+        .db()
+        .await
+        .update_channel_message(
+            channel_id,
+            message_id,
+            session.user_id,
+            request.body.as_str(),
+            &request.mentions,
+            updated_at,
+        )
+        .await?;
+
+    let nonce = request
+        .nonce
+        .clone()
+        .ok_or_else(|| anyhow!("nonce can't be blank"))?;
+
+    let message = proto::ChannelMessage {
+        sender_id: session.user_id.to_proto(),
+        id: message_id.to_proto(),
+        body: request.body.clone(),
+        mentions: request.mentions.clone(),
+        timestamp: timestamp.assume_utc().unix_timestamp() as u64,
+        nonce: Some(nonce),
+        reply_to_message_id: reply_to_message_id.map(|id| id.to_proto()),
+        edited_at: Some(updated_at.unix_timestamp() as u64),
+    };
+
+    response.send(proto::Ack {})?;
+
+    let pool = &*session.connection_pool().await;
+    broadcast(
+        Some(session.connection_id),
+        participant_connection_ids,
+        |connection| {
+            session.peer.send(
+                connection,
+                proto::ChannelMessageUpdate {
+                    channel_id: channel_id.to_proto(),
+                    message: Some(message.clone()),
+                },
+            )
+        },
+    );
+
+    send_notifications(pool, &session.peer, notifications);
+
+    Ok(())
+}
+
 /// Mark a channel message as read
 async fn acknowledge_channel_message(
     request: proto::AckChannelMessage,

crates/collab/src/tests/channel_message_tests.rs

@@ -466,3 +466,136 @@ async fn test_chat_replies(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext)
         )
     });
 }
+
+#[gpui::test]
+async fn test_chat_editing(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext) {
+    let mut server = TestServer::start(cx_a.executor()).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let client_b = server.create_client(cx_b, "user_b").await;
+
+    let channel_id = server
+        .make_channel(
+            "the-channel",
+            None,
+            (&client_a, cx_a),
+            &mut [(&client_b, cx_b)],
+        )
+        .await;
+
+    // Client A sends a message, client B should see that there is a new message.
+    let channel_chat_a = client_a
+        .channel_store()
+        .update(cx_a, |store, cx| store.open_channel_chat(channel_id, cx))
+        .await
+        .unwrap();
+
+    let channel_chat_b = client_b
+        .channel_store()
+        .update(cx_b, |store, cx| store.open_channel_chat(channel_id, cx))
+        .await
+        .unwrap();
+
+    let msg_id = channel_chat_a
+        .update(cx_a, |c, cx| {
+            c.send_message(
+                MessageParams {
+                    text: "Initial message".into(),
+                    reply_to_message_id: None,
+                    mentions: Vec::new(),
+                },
+                cx,
+            )
+            .unwrap()
+        })
+        .await
+        .unwrap();
+
+    cx_a.run_until_parked();
+
+    channel_chat_a
+        .update(cx_a, |c, cx| {
+            c.update_message(
+                msg_id,
+                MessageParams {
+                    text: "Updated body".into(),
+                    reply_to_message_id: None,
+                    mentions: Vec::new(),
+                },
+                cx,
+            )
+            .unwrap()
+        })
+        .await
+        .unwrap();
+
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+
+    channel_chat_a.update(cx_a, |channel_chat, _| {
+        let update_message = channel_chat.find_loaded_message(msg_id).unwrap();
+
+        assert_eq!(update_message.body, "Updated body");
+        assert_eq!(update_message.mentions, Vec::new());
+    });
+    channel_chat_b.update(cx_b, |channel_chat, _| {
+        let update_message = channel_chat.find_loaded_message(msg_id).unwrap();
+
+        assert_eq!(update_message.body, "Updated body");
+        assert_eq!(update_message.mentions, Vec::new());
+    });
+
+    // test mentions are updated correctly
+
+    client_b.notification_store().read_with(cx_b, |store, _| {
+        assert_eq!(store.notification_count(), 1);
+        let entry = store.notification_at(0).unwrap();
+        assert!(matches!(
+            entry.notification,
+            Notification::ChannelInvitation { .. }
+        ),);
+    });
+
+    channel_chat_a
+        .update(cx_a, |c, cx| {
+            c.update_message(
+                msg_id,
+                MessageParams {
+                    text: "Updated body including a mention for @user_b".into(),
+                    reply_to_message_id: None,
+                    mentions: vec![(37..45, client_b.id())],
+                },
+                cx,
+            )
+            .unwrap()
+        })
+        .await
+        .unwrap();
+
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+
+    channel_chat_a.update(cx_a, |channel_chat, _| {
+        assert_eq!(
+            channel_chat.find_loaded_message(msg_id).unwrap().body,
+            "Updated body including a mention for @user_b",
+        )
+    });
+    channel_chat_b.update(cx_b, |channel_chat, _| {
+        assert_eq!(
+            channel_chat.find_loaded_message(msg_id).unwrap().body,
+            "Updated body including a mention for @user_b",
+        )
+    });
+    client_b.notification_store().read_with(cx_b, |store, _| {
+        assert_eq!(store.notification_count(), 2);
+        let entry = store.notification_at(0).unwrap();
+        assert_eq!(
+            entry.notification,
+            Notification::ChannelMessageMention {
+                message_id: msg_id,
+                sender_id: client_a.id(),
+                channel_id: channel_id.0,
+            }
+        );
+    });
+}