fix: WA Web-compliant DM retry handler (chat/requester separation + alternate key lookup)

src/client/lid_pn.rs

@@ -165,6 +165,32 @@ impl Client {
         }
     }
 
+    /// Swap a JID's namespace between PN and LID, preserving device/agent/integrator.
+    /// Returns `None` if no mapping exists or the JID is neither PN nor LID.
+    pub(crate) async fn swap_pn_lid_namespace(&self, jid: &Jid) -> Option<Jid> {
+        if jid.is_lid() {
+            let pn_user = self.lid_pn_cache.get_phone_number(&jid.user).await?;
+            Some(Jid {
+                user: pn_user.into(),
+                server: wacore_binary::Server::Pn,
+                device: jid.device,
+                agent: jid.agent,
+                integrator: jid.integrator,
+            })
+        } else if jid.is_pn() {
+            let lid_user = self.lid_pn_cache.get_current_lid(&jid.user).await?;
+            Some(Jid {
+                user: lid_user.into(),
+                server: wacore_binary::Server::Lid,
+                device: jid.device,
+                agent: jid.agent,
+                integrator: jid.integrator,
+            })
+        } else {
+            None
+        }
+    }
+
     /// Migrate Signal sessions and identity keys from PN to LID address.
     ///
     /// All reads/writes go through `signal_cache` to avoid reading stale data

src/client/sender_keys.rs

@@ -1,6 +1,7 @@
 //! Sender key tracking and message cache methods for Client.
 
 use anyhow::Result;
+use wacore::types::message::ChatMessageId;
 use wacore_binary::Jid;
 use waproto::whatsapp as wa;
 
@@ -63,38 +64,77 @@ impl Client {
     }
 
     /// Take a sent message for retry handling. Checks L1 cache first (if enabled),
-    /// then falls back to DB. Matches WA Web's getMessageTable().get() pattern.
-    pub(crate) async fn take_recent_message(&self, to: &Jid, id: &str) -> Option<wa::Message> {
+    /// then falls back to DB. On miss, tries an alternate PN/LID key to handle
+    /// mapping changes between send time and retry time (WAWebLidMigrationUtils
+    /// `getAlternateMsgKey`).
+    /// Returns `(message, alternate_chat)`. When the message was found via the
+    /// alternate PN/LID key, `alternate_chat` contains the namespace that
+    /// matched -- the caller should use it for session operations instead of
+    /// `resolve_encryption_jid` (which would map back to the primary).
+    pub(crate) async fn take_recent_message(
+        &self,
+        to: &Jid,
+        id: &str,
+    ) -> Option<(wa::Message, Option<Jid>)> {
+        let primary_key = self.make_chat_message_id(to, id).await;
+        if let Some(msg) = self.try_take_by_key(&primary_key).await {
+            return Some((msg, None));
+        }
+
+        // Primary miss -- try alternate PN<->LID key.
+        // If resolve_encryption_jid changed the namespace (PN→LID), the
+        // original `to` is already the alternate -- skip the cache lookup.
+        // Otherwise (LID input), swap via cache to try the PN form.
+        let alt_chat = if primary_key.chat.server != to.server {
+            Some(to.clone())
+        } else {
+            self.swap_pn_lid_namespace(&primary_key.chat).await
+        };
+
+        if let Some(alt_chat) = alt_chat {
+            log::debug!(
+                "Primary key miss for {}:{}, trying alternate {}",
+                primary_key.chat,
+                id,
+                alt_chat
+            );
+            let alt_key = ChatMessageId {
+                chat: alt_chat,
+                id: primary_key.id,
+            };
+            if let Some(msg) = self.try_take_by_key(&alt_key).await {
+                return Some((msg, Some(alt_key.chat)));
+            }
+        }
+
+        None
+    }
+
+    /// Look up and consume a message by exact `ChatMessageId` (L1 cache then DB).
+    async fn try_take_by_key(&self, key: &ChatMessageId) -> Option<wa::Message> {
         use prost::Message;
-        let key = self.make_chat_message_id(to, id).await;
         let chat_str = key.chat.to_string();
         let has_l1_cache = self.cache_config.recent_messages.capacity > 0;
 
         // L1 cache check (if capacity > 0)
-        if has_l1_cache && let Some(bytes) = self.recent_messages.remove(&key).await {
+        if has_l1_cache && let Some(bytes) = self.recent_messages.remove(key).await {
             if let Ok(msg) = wa::Message::decode(bytes.as_slice()) {
                 // Cache hit — consume the DB row in the background to avoid orphans.
-                // Note: if the background DB write from add_recent_message hasn't completed
-                // yet, this delete may run first and the write creates an orphan. This is
-                // harmless — periodic cleanup (sent_message_ttl_secs) purges it. The race
-                // window is negligible since retry receipts arrive seconds after send.
                 let backend = self.persistence_manager.backend();
-                let cs = chat_str.clone();
                 let mid = key.id.clone();
                 self.runtime
                     .spawn(Box::pin(async move {
-                        if let Err(e) = backend.take_sent_message(&cs, &mid).await {
-                            log::warn!("Failed to clean up sent message {cs}:{mid}: {e}");
+                        if let Err(e) = backend.take_sent_message(&chat_str, &mid).await {
+                            log::warn!("Failed to clean up sent message {chat_str}:{mid}: {e}");
                         }
                     }))
                     .detach();
                 return Some(msg);
             }
-            // Cache decode failed — fall through to DB
             log::warn!(
                 "Failed to decode cached message for {}:{}, trying DB",
-                to,
-                id
+                key.chat,
+                key.id
             );
         }
 
@@ -108,16 +148,21 @@ impl Client {
             Ok(Some(bytes)) => match wa::Message::decode(bytes.as_slice()) {
                 Ok(msg) => Some(msg),
                 Err(e) => {
-                    log::warn!("Failed to decode DB message for {}:{}: {}", to, id, e);
+                    log::warn!(
+                        "Failed to decode DB message for {}:{}: {}",
+                        key.chat,
+                        key.id,
+                        e
+                    );
                     None
                 }
             },
             Ok(None) => None,
             Err(e) => {
                 log::warn!(
                     "Failed to read sent message from DB for {}:{}: {}",
-                    to,
-                    id,
+                    key.chat,
+                    key.id,
                     e
                 );
                 None

src/client/sessions.rs

@@ -144,21 +144,35 @@ impl Client {
     /// Ensure E2E sessions exist for the given device JIDs.
     /// Waits for offline delivery, resolves LID mappings, then batches prekey fetches.
     pub(crate) async fn ensure_e2e_sessions(&self, device_jids: &[Jid]) -> Result<()> {
-        use wacore::types::jid::JidExt;
-
         if device_jids.is_empty() {
             return Ok(());
         }
-
         self.wait_for_offline_delivery_end().await;
         let resolved_jids = self.resolve_lid_mappings(device_jids).await;
+        self.ensure_sessions_inner(resolved_jids).await
+    }
+
+    /// Like `ensure_e2e_sessions` but skips `resolve_lid_mappings`. Use when the
+    /// caller already resolved JIDs to the correct namespace (e.g., after
+    /// alternate PN/LID key normalization in retry handling).
+    pub(crate) async fn ensure_e2e_sessions_resolved(&self, jids: &[Jid]) -> Result<()> {
+        if jids.is_empty() {
+            return Ok(());
+        }
+        self.wait_for_offline_delivery_end().await;
+        self.ensure_sessions_inner(jids.to_vec()).await
+    }
+
+    /// Core session-check + prekey-fetch logic shared by both entry points.
+    async fn ensure_sessions_inner(&self, jids: Vec<Jid>) -> Result<()> {
+        use wacore::types::jid::JidExt;
 
         let device_store = self.persistence_manager.get_device_arc().await;
-        let mut jids_needing_sessions = Vec::with_capacity(resolved_jids.len());
+        let mut jids_needing_sessions = Vec::with_capacity(jids.len());
 
         {
             let device_guard = device_store.read().await;
-            for jid in resolved_jids {
+            for jid in jids {
                 let signal_addr = jid.to_protocol_address();
                 // Check cache first (includes unflushed sessions), fall back to backend
                 match self