state: serialize per-chat get→mutate→set with ChatStore::update

Concurrent handler tasks (the batch-forward design spawns several per
chat) each snapshotted the same ChatData and last-writer-wins silently
dropped mutations — e.g. a second edit_message record, leaving one
prompt's Forward button dead. All write cycles now run under a per-chat
async lock; read-only callers keep get().

crates/xmedia-bot/src/handlers.rs

@@ -238,9 +238,11 @@ async fn execute_command(
         Command::SetForwardChannel(channel) => {
             let result = match set_forward_channel_handler(bot, message, channel).await {
                 Ok(channel_id) => {
-                    let mut chat_data = CHAT_STORE.get(message.chat.id.0).await;
-                    chat_data.forward_channel_id = Some(channel_id);
-                    CHAT_STORE.set(message.chat.id.0, &chat_data).await;
+                    CHAT_STORE
+                        .update(message.chat.id.0, |data| {
+                            data.forward_channel_id = Some(channel_id);
+                        })
+                        .await;
                     "Add successfully.".to_string()
                 }
                 Err(SetForwardChannelError::EmptyParameter) => {
@@ -264,31 +266,34 @@ async fn execute_command(
         }
         Command::RemoveForwardChannel => {
             let chat_id = message.chat.id.0;
-            let mut chat_data = CHAT_STORE.get(chat_id).await;
-            let text = if chat_data.forward_channel_id.is_some() {
-                chat_data.forward_channel_id = None;
-                CHAT_STORE.set(chat_id, &chat_data).await;
+            let text = CHAT_STORE
+                .update(chat_id, |data| {
+                    if data.forward_channel_id.is_some() {
+                        data.forward_channel_id = None;
                         "Remove successfully.".to_string()
                     } else {
                         "No channel to remove.".to_string()
-            };
+                    }
+                })
+                .await;
             reply(bot.clone(), message.clone(), text).await?;
         }
         Command::EditBeforeForward => {
             let chat_id = message.chat.id.0;
-            let mut chat_data = CHAT_STORE.get(chat_id).await;
-            let text = if chat_data.forward_channel_id.is_none() {
+            let text = CHAT_STORE
+                .update(chat_id, |data| {
+                    if data.forward_channel_id.is_none() {
                         "Please enable forward channel first.".to_string()
-            } else if chat_data.edit_before_forward {
-                chat_data.edit_before_forward = false;
-                chat_data.edit_message.clear();
-                CHAT_STORE.set(chat_id, &chat_data).await;
+                    } else if data.edit_before_forward {
+                        data.edit_before_forward = false;
+                        data.edit_message.clear();
                         "Disable edit before forward.".to_string()
                     } else {
-                chat_data.edit_before_forward = true;
-                CHAT_STORE.set(chat_id, &chat_data).await;
+                        data.edit_before_forward = true;
                         "Enable edit before forward.".to_string()
-            };
+                    }
+                })
+                .await;
             reply(bot.clone(), message.clone(), text).await?;
         }
         Command::SetTemplate(name) => {
@@ -302,11 +307,14 @@ async fn execute_command(
                     } else if name.is_empty() {
                         "Please provide a name for the template.".to_string()
                     } else {
-                        let mut chat_data = CHAT_STORE.get(chat_id).await;
-                        chat_data
-                            .template
-                            .insert(name, html_escape::encode_text(reply_text).into_owned());
-                        CHAT_STORE.set(chat_id, &chat_data).await;
+                        CHAT_STORE
+                            .update(chat_id, |data| {
+                                data.template.insert(
+                                    name,
+                                    html_escape::encode_text(reply_text).into_owned(),
+                                );
+                            })
+                            .await;
                         "Template set.".to_string()
                     }
                 }
@@ -344,9 +352,11 @@ async fn execute_command(
                 .await?;
                 return Ok(());
             }
-            let mut chat_data = CHAT_STORE.get(chat_id).await;
-            chat_data.message_format.insert(site.to_string(), format);
-            CHAT_STORE.set(chat_id, &chat_data).await;
+            CHAT_STORE
+                .update(chat_id, |data| {
+                    data.message_format.insert(site.to_string(), format);
+                })
+                .await;
             reply(bot.clone(), message.clone(), "Format set.").await?;
         }
         Command::ClearCache(arg) => {
@@ -793,7 +803,7 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
     let chat_id = message.chat().id.0;
     let prompt_message_id = message.id().0 as i64;
     let ttl_secs = CONFIG.edit_message_ttl.as_secs() as i64;
-    let mut chat_data = CHAT_STORE.get(chat_id).await;
+    let chat_data = CHAT_STORE.get(chat_id).await;
     let edit = chat_data.edit_message.get(&prompt_message_id).cloned();
     let Some(edit) = edit else {
         log::info!(
@@ -807,8 +817,11 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
     };
     // Lazy expiry: a stale record (past the TTL, not yet swept) is dropped.
     if edit.created_at + ttl_secs <= unix_now() {
-        chat_data.edit_message.remove(&prompt_message_id);
-        CHAT_STORE.set(chat_id, &chat_data).await;
+        CHAT_STORE
+            .update(chat_id, |data| {
+                data.edit_message.remove(&prompt_message_id);
+            })
+            .await;
         bot.answer_callback_query(callback_query_id)
             .text("Expired")
             .await?;
@@ -844,8 +857,11 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
                         let _ = bot
                             .delete_message(ChatId(chat_id), MessageId(prompt_message_id as i32))
                             .await;
-                        chat_data.edit_message.remove(&prompt_message_id);
-                        CHAT_STORE.set(chat_id, &chat_data).await;
+                        CHAT_STORE
+                            .update(chat_id, |data| {
+                                data.edit_message.remove(&prompt_message_id);
+                            })
+                            .await;
                     }
                     Err(send::SendError::Retryable {
                         delay_seconds,
@@ -884,10 +900,13 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
                 .caption(template_html)
                 .parse_mode(ParseMode::Html)
                 .await;
-            if let Some(entry) = chat_data.edit_message.get_mut(&prompt_message_id) {
+            CHAT_STORE
+                .update(chat_id, |data| {
+                    if let Some(entry) = data.edit_message.get_mut(&prompt_message_id) {
                         entry.template = name.to_string();
                     }
-            CHAT_STORE.set(chat_id, &chat_data).await;
+                })
+                .await;
             log::info!("template '{name}' applied to prompt {prompt_message_id}");
         }
         bot.answer_callback_query(callback_query_id).await?;

crates/xmedia-bot/src/send.rs

@@ -1048,8 +1048,7 @@ pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec<i64>) {
     };
 
     if edit_before_forward {
-        let mut chat_data = CHAT_STORE.get(chat_id).await;
-        let keyboard = build_edit_markup(&chat_data.template);
+        let keyboard = build_edit_markup(&CHAT_STORE.get(chat_id).await.template);
         match bot
             .send_message(ChatId(chat_id), "Reply to edit message.")
             .reply_markup(keyboard)
@@ -1064,8 +1063,12 @@ pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec<i64>) {
                     prompt.id.0,
                     message_ids.len()
                 );
-                chat_data.edit_message.insert(
-                    prompt.id.0 as i64,
+                let prompt_id = prompt.id.0 as i64;
+                let source_url = source_url.clone();
+                CHAT_STORE
+                    .update(chat_id, move |data| {
+                        data.edit_message.insert(
+                            prompt_id,
                             EditMessage {
                                 url: source_url,
                                 chat_id,
@@ -1074,7 +1077,8 @@ pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec<i64>) {
                                 created_at: unix_now(),
                             },
                         );
-                CHAT_STORE.set(chat_id, &chat_data).await;
+                    })
+                    .await;
             }
             Err(e) => log::error!("failed to send edit prompt: {e}"),
         }

crates/xmedia-bot/src/state.rs

@@ -6,6 +6,7 @@ use rusqlite::params;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::path::Path;
+use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 #[derive(Serialize, Deserialize, Default, Clone, Debug)]
@@ -34,6 +35,9 @@ pub struct EditMessage {
 pub struct ChatStore {
     /// In-memory cache; the DB is the source of truth on first access.
     cache: Mutex<HashMap<i64, ChatData>>,
+    /// Per-chat async locks serializing get→mutate→set so concurrent handler
+    /// tasks (batch-forwards, callbacks) cannot clobber each other's writes.
+    locks: Mutex<HashMap<i64, Arc<tokio::sync::Mutex<()>>>>,
     db_path: String,
 }
 
@@ -62,6 +66,7 @@ impl ChatStore {
         drop(conn);
         Ok(ChatStore {
             cache: Mutex::new(HashMap::new()),
+            locks: Mutex::new(HashMap::new()),
             db_path: path.to_string(),
         })
     }
@@ -111,6 +116,26 @@ impl ChatStore {
         }
     }
 
+    /// Serializes a get→mutate→set cycle per chat: concurrent handler tasks
+    /// (the batch-forward design spawns several per chat) each snapshot the
+    /// same `ChatData` and last-writer-wins would silently drop mutations,
+    /// e.g. a second `edit_message` record. The per-chat lock makes the
+    /// cycle atomic. Returns the closure's result.
+    pub async fn update<R>(&self, chat_id: i64, f: impl FnOnce(&mut ChatData) -> R) -> R {
+        let lock = {
+            let mut locks = self.locks.lock();
+            locks
+                .entry(chat_id)
+                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
+                .clone()
+        };
+        let _guard = lock.lock().await;
+        let mut data = self.get(chat_id).await;
+        let r = f(&mut data);
+        self.set(chat_id, &data).await;
+        r
+    }
+
     /// Removes edit-before-forward records whose `created_at + ttl` is in the
     /// past. Returns the removed `(chat_id, prompt_message_id)` pairs so the
     /// caller can clear the prompt's buttons.
@@ -152,3 +177,45 @@ impl ChatStore {
         removed
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn concurrent_updates_do_not_lose_edit_records() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = std::sync::Arc::new(
+            ChatStore::open(dir.path().join("s.db").to_str().unwrap()).unwrap(),
+        );
+        let mut handles = Vec::new();
+        for i in 0..4 {
+            let store = Arc::clone(&store);
+            handles.push(tokio::spawn(async move {
+                store
+                    .update(1001, |data| {
+                        data.edit_message.insert(
+                            i,
+                            EditMessage {
+                                url: format!("https://x.com/u/status/{i}"),
+                                chat_id: 1001,
+                                forward_message_ids: vec![i],
+                                template: String::new(),
+                                created_at: 0,
+                            },
+                        );
+                    })
+                    .await;
+            }));
+        }
+        for h in handles {
+            h.await.unwrap();
+        }
+        let data = store.get(1001).await;
+        assert_eq!(
+            data.edit_message.len(),
+            4,
+            "concurrent get→mutate→set must not drop records"
+        );
+    }
+}