fix: WA Web-compliant DM retry handler (chat/requester separation + alternate key lookup) - #550
Conversation
The retry handler conflated the raw `from` JID for both message lookup and session management. When the server includes a device suffix in `from` (e.g. `:33`), the message lookup silently fails because messages are stored under bare JIDs. Mirrors WA Web's `getActualChatInfo` (WAWebHandleRetryRequest): - `chat`: bare JID (no device) for message lookup and stanza addressing - `requester`: device-specific JID for session management Also adds alternate PN/LID key lookup in `take_recent_message`, matching WA Web's `getAlternateMsgKey` (WAWebLidMigrationUtils). This handles PN-to-LID mapping changes between send time and retry time.
📝 WalkthroughWalkthrough
Changes
Sequence Diagram(s)sequenceDiagram
participant Client as "Client"
participant Resolver as "Retry Resolver"
participant Cache as "L1 Cache"
participant Backend as "DB/Backend"
Client->>Resolver: resolve_retry_chat_info(receipt)
Resolver-->>Client: RetryChatInfo { chat, requester, original_from, is_bot }
Client->>Cache: try_take_by_key(primary: ChatMessageId)
alt primary cache hit
Cache-->>Client: Message
Client->>Backend: background take_sent_message(primary)
Client-->>Client: return (Message, None)
else primary cache miss
Client->>Backend: take_sent_message(primary)
alt primary backend hit
Backend-->>Client: Message
Client-->>Client: return (Message, None)
else primary backend miss
Client->>Client: construct alternate key (PN↔LID)
Client->>Cache: try_take_by_key(alternate)
alt alternate hit
Cache-->>Client: Message
Client->>Backend: background take_sent_message(alternate)
Client-->>Client: return (Message, Some(alt_chat))
else
Client-->>Client: return None
end
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@claude review this nitpicky |
|
Claude finished @jlucaso1's task in 2m 25s —— View job Nitpicky Code Review
Overall this is a solid, well-motivated change. The separation of chat vs. requester JIDs is correct and the test coverage is thorough. Here's the nitpicky pass you asked for:
|
There was a problem hiding this comment.
Pull request overview
Improves retry-receipt handling by correctly separating “chat” vs “requester” JIDs (especially for DM retries with device-suffixed from JIDs), and enhances recent-message lookup to tolerate PN↔LID mapping changes between send and retry time.
Changes:
- Added
RetryChatInfo+resolve_retry_chat_info()and updated retry handling to use bare chat JIDs for lookup/addressing and device-scoped requester JIDs for session ops. - Extended
take_recent_message()with PN↔LID alternate-key fallback viaalternate_message_key()and refactored lookup intotry_take_by_key(). - Added/updated tests covering DM/group/status retry chat resolution and PN→LID fallback lookup behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/retry.rs |
Introduces RetryChatInfo resolution and updates retry flow to use separated chat/requester JIDs; adds targeted tests. |
src/client/sender_keys.rs |
Refactors recent-message retrieval to support alternate PN/LID key fallback and reduces unnecessary cloning in hot paths. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| self.add_recent_message(&info.chat, &message_id, &original_msg) | ||
| .await; | ||
|
|
||
| // Resolved JID for session operations; keep original for stanza addressing |
There was a problem hiding this comment.
The inline comment says "keep original for stanza addressing", but the code below uses info.chat (bare/non-AD) when building both DM and group retry stanzas. This is a bit misleading/confusing now that chat/requester are separated; please update the comment to reflect that resolved_jid is for session/encryption operations, while info.chat is used for stanza addressing and message lookup.
| // Resolved JID for session operations; keep original for stanza addressing | |
| // Resolve the requester's JID for session/encryption operations; use | |
| // info.chat for stanza addressing and recent-message lookup. |
Benchmark Results1 regression(s) detected (>2% threshold):
58 unchanged benchmark(s)
|
- Rewrite DM bare-JID path: to_non_ad() first, single clone in bare case instead of two - Use info.requester in duplicate-retry debug log for consistency - Fix stale make_test_receipt docstring - Simplify DM test nodes (attrs unused in DM branch) - Hoist ChatMessageId import to file level in sender_keys.rs - Add LID DM with device suffix test
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b8f1f5ca3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let chat = from.to_non_ad(); | ||
| let requester = if from.device() == 0 && from.agent == 0 { |
There was a problem hiding this comment.
Preserve AD fields when deriving DM retry chat JID
In resolve_retry_chat_info, the DM branch always normalizes from with to_non_ad(), which clears agent/device for every non-group chat. This regresses retries for non-PN/LID direct chats (for example @bot, @hosted, or other AD JIDs): add_recent_message stores under the original chat JID via make_chat_message_id, but take_recent_message now looks up a stripped JID, so the original message cannot be found and the retry is skipped. Before this change, these fields were preserved because lookup used receipt.source.chat directly.
Useful? React with 👍 / 👎.
When take_recent_message finds the message via the alternate key (e.g., stored under PN but retry arrives as LID), the Signal session lives in the stored message's namespace. Normalize the requester to match and use it directly as the encryption target, skipping resolve_encryption_jid which would map it back to the wrong namespace. Mirrors WA Web's getActualChatInfo: d.isLid() ? toLid(e.from) : toPn(e.from)
- Extract shared swap_pn_lid_namespace() helper to client/lid_pn.rs, used by both alternate_message_key and normalize_requester_namespace - Add normalize_requester_after_alternate_key_hit test verifying both LID→PN and PN→LID namespace swaps with device preservation
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/retry.rs (1)
232-251:⚠️ Potential issue | 🟠 MajorKeep the receipt requester separate from the normalized session target.
After Line 226 mutates
info.requester, thehas_device()check and peer detection here run against the swapped PN/LID user instead of the actual requester from the receipt. That can reject legitimate retries onceadd_lid_pn_mapping()has migrated device-registry entries to the LID namespace, because the retry came from LID but this code now looks up the old PN user. Store the normalized JID in a separate variable for Signal/session work and leaveinfo.requesteras the on-the-wire requester.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/retry.rs` around lines 232 - 251, The code currently mutates info.requester earlier (via add_lid_pn_mapping) so subsequent checks like has_device(&sender_user, ...) and the peer detection logic using device_snapshot and is_peer run against the normalized/ swapped JID instead of the original on-the-wire requester; to fix, preserve the original requester by copying it to a new variable (e.g., let original_requester = info.requester.clone() or let normalized_target = ... depending on intent) and use original_requester for has_device and peer detection while using a separate normalized variable for Signal/session operations; update references in this block (info.requester, sender_user, has_device, device_snapshot, is_peer) to use the correct variable so lookups remain against the true incoming requester and session targeting uses the normalized JID.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/retry.rs`:
- Around line 225-230: The alternate-namespace resolved_jid is being overwritten
because process_retry_key_bundle() re-runs resolve_encryption_jid() instead of
using the already-final resolved_jid computed in the caller; update the call
site around resolved_jid (where normalize_requester_namespace() and
resolve_encryption_jid() are used) to pass the final resolved_jid into
process_retry_key_bundle(), and modify process_retry_key_bundle() (and any
helpers it calls) to accept an optional/pre-resolved JID (or a skip-resolve
flag) and use that value directly instead of calling
resolve_encryption_jid(&info.requester) when the pre-resolved JID is supplied;
reference symbols: resolved_jid, normalize_requester_namespace,
resolve_encryption_jid, process_retry_key_bundle, and info.requester.
---
Outside diff comments:
In `@src/retry.rs`:
- Around line 232-251: The code currently mutates info.requester earlier (via
add_lid_pn_mapping) so subsequent checks like has_device(&sender_user, ...) and
the peer detection logic using device_snapshot and is_peer run against the
normalized/ swapped JID instead of the original on-the-wire requester; to fix,
preserve the original requester by copying it to a new variable (e.g., let
original_requester = info.requester.clone() or let normalized_target = ...
depending on intent) and use original_requester for has_device and peer
detection while using a separate normalized variable for Signal/session
operations; update references in this block (info.requester, sender_user,
has_device, device_snapshot, is_peer) to use the correct variable so lookups
remain against the true incoming requester and session targeting uses the
normalized JID.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b85fa0b1-4b94-47f2-9275-128b375c54d5
📒 Files selected for processing (3)
src/client/lid_pn.rssrc/client/sender_keys.rssrc/retry.rs
- Remove normalize_requester_namespace (did a second cache lookup for the same PN/LID mapping already fetched by alternate_message_key) - Return alt_chat Jid from take_recent_message instead of bool, so the caller builds the normalized requester directly from it — zero extra cache fetches - Remove redundant to_non_ad() in alternate_message_key (input key.chat is already bare from make_chat_message_id) - Update test to exercise swap_pn_lid_namespace directly (the shared helper both paths use)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c08f35376e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| agent: requester.agent, | ||
| integrator: requester.integrator, | ||
| }; | ||
| info.requester.clone() |
There was a problem hiding this comment.
Re-establish retry session in the same namespace
When take_recent_message hits the alternate PN/LID key path, this branch returns resolved_jid in the alternate namespace, but the DM flow later calls ensure_e2e_sessions(&resolved_jid) which normalizes PN→LID (resolve_lid_mappings in src/client/sessions.rs). Because delete_dm_retry_session_target has already deleted the session for resolved_jid, we can end up recreating a session under LID while prepare_dm_retry_stanza still encrypts with the original PN resolved_jid, causing retry encryption to fail in the exact PN/LID migration scenario this change adds.
Useful? React with 👍 / 👎.
When resolve_encryption_jid changed the namespace (PN→LID), the original input `to` IS the alternate key. Use it directly instead of round-tripping through swap_pn_lid_namespace (which just recovers the original). Removes alternate_message_key method entirely — logic is now inline. Cache fetches in alternate-hit path: 2 → 1 (common case), 2 (LID input).
- alternate_key_lookup_pn_input_server_changed: PN input where resolve_encryption_jid changes namespace — exercises the server != server optimization that skips the cache round-trip - no_alternate_without_mapping: no PN/LID mapping → no alternate tried, primary hit returns None alt_chat - alternate_key_both_miss: mapping exists but no message stored under either key → returns None
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 40509439db
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let resolved_jid = if let Some(alt_chat) = alt_chat | ||
| && !is_group_or_status |
There was a problem hiding this comment.
Preserve alt namespace through retry key-bundle processing
When an alternate PN/LID key hit occurs, this branch intentionally keeps resolved_jid in the alternate namespace, but process_retry_key_bundle later re-runs resolve_encryption_jid on that JID and establishes the refreshed session under the primary namespace instead. In PN↔LID migration retries that include <keys>, the flow can then delete/use the alternate resolved_jid while the rebuilt session exists under the other namespace, so prepare_dm_retry_stanza may encrypt against a missing/stale session and the resend fails. Fresh evidence: this mismatch is already introduced before ensure_e2e_sessions runs because process_retry_key_bundle always re-resolves its input.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/retry.rs`:
- Around line 224-238: In process_retry_key_bundle, stop re-resolving the
requester JID via resolve_encryption_jid and instead pass the already-determined
requester_jid directly into to_protocol_address so the prekey is stored under
the correct Signal address; update the code in process_retry_key_bundle to
remove the resolve_encryption_jid call/variable and use requester_jid when
calling to_protocol_address (preserving the PN namespace when alt_chat produced
the requester_jid).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
process_retry_key_bundle called resolve_encryption_jid on the requester JID, undoing the alternate-path namespace normalization (PN back to LID). The caller already passes the correctly-resolved JID — use it directly via to_protocol_address().
Gap 1 - ensure_e2e_sessions re-resolution: Extract ensure_sessions_inner, add ensure_e2e_sessions_resolved that skips resolve_lid_mappings. DM retries use the resolved variant since the caller already normalized the JID. Matches WA Web where ensureE2ESessions uses input JIDs as-is (E2ESessionsJob.js:56-84). Gap 2 - recipient attribute for peer/bot retries: resolve_retry_chat_info now takes own_pn/own_lid and reads the `recipient` node attr. For peer devices (from=our JID) and bots (from.is_bot()), uses recipient as the chat target. Matches WA Web getTargetChat (RetryRequest.js:339-371). Gap 3 - bot JID namespace normalization exception: Bots skip the alternate PN/LID namespace normalization. Matches WA Web getActualChatInfo: `e.from.isBot() ? (p = e.from)` (RetryRequest.js:311). Gap 4 - stanza `to` uses original_from: Added original_from field to RetryChatInfo. DM retry stanza now uses the raw `from` JID for the `to` attribute, matching WA Web variable `m` (RetryRequest.js:30,97).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 46404129c0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let requester = if from.device() == 0 && from.agent == 0 { | ||
| chat.clone() |
There was a problem hiding this comment.
Keep requester bound to raw from JID for recipient retries
In resolve_retry_chat_info, the bare-from fallback assigns requester = chat.clone(). That is incorrect when chat was rewritten from recipient (peer/device-sync or bot retries), because requester then points at the conversation partner instead of the retrying sender. In that case the downstream has_device/session/key-bundle path and prepare_dm_retry_stanza participant JID are built for the wrong user, so bare-from retries with a recipient attribute can be skipped or resent to an invalid target.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/client/sessions.rs`:
- Line 168: The inner redundant import "use wacore::types::jid::JidExt;" should
be removed — JidExt is already imported at module scope, so delete the duplicate
use statement (the inner use of JidExt) to avoid redundancy and keep imports
clean; verify no other local shadowing depends on that inner import.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 2c32d42d-34be-47ce-9df4-65ccede767b8
📒 Files selected for processing (2)
src/client/sessions.rssrc/retry.rs
|
|
||
| /// 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; |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Redundant inner import.
JidExt is already imported at module level (line 7). This inner use can be removed.
♻️ Suggested fix
/// 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;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| use wacore::types::jid::JidExt; | |
| /// Core session-check + prekey-fetch logic shared by both entry points. | |
| async fn ensure_sessions_inner(&self, jids: Vec<Jid>) -> Result<()> { | |
| let device_store = self.persistence_manager.get_device_arc().await; |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/client/sessions.rs` at line 168, The inner redundant import "use
wacore::types::jid::JidExt;" should be removed — JidExt is already imported at
module scope, so delete the duplicate use statement (the inner use of JidExt) to
avoid redundancy and keep imports clean; verify no other local shadowing depends
on that inner import.
Summary
getActualChatInfo+getTargetChat(WAWebHandleRetryRequest). The rawfromJID is split into: bare chat (message lookup), device-specific requester (session ops), and original_from (stanzatoattribute).take_recent_message, matching WA Web'sgetAlternateMsgKey(WAWebLidMigrationUtils). Handles PN-to-LID mapping changes between send time and retry time with zero redundant cache fetches in the common case.d.isLid() ? toLid(e.from) : toPn(e.from)+createDeviceWidFromUserAndDevice. Skipsresolve_encryption_jidto avoid undoing the normalization.recipientnode attribute, matching WA Web'sgetTargetChat(RetryRequest.js:339-371).ensure_e2e_sessions_resolvedthat skipsresolve_lid_mappings, matching WA Web whereensureE2ESessionsuses input JIDs as-is (E2ESessionsJob.js:56-84).process_retry_key_bundlethat undid alternate-path normalization.Problem
The DM retry handler used the raw
fromJID for both message lookup and session management. When the server includes a device suffix (e.g.:33), message lookup silently fails becauseChatMessageIdequality includes the device field but messages are stored under bare JIDs. When the server sends barefrom, session operations could target the wrong device.WA Web compliance achieved
getTargetChat— bot/peer/recipient routingresolve_retry_chat_infowithown_pn,own_lid,recipientattrgetActualChatInfo— chat/requester splitRetryChatInfo { chat, requester, original_from, is_bot }getAlternateMsgKey— PN/LID fallbacktake_recent_messagewith server-comparison optimizationd.isLid() ? toLid(e.from) : toPn(e.from)alt_chatnamespace normalization inlinee.from.isBot() ? (p = e.from)!info.is_botguard on normalizationm = e.fromfor stanzatooriginal_fromfieldensureE2ESessions([g])as-isensure_e2e_sessions_resolved(no re-resolution)processKeyBundlewith normalized JIDresolve_encryption_jidinprocess_retry_key_bundleFiles changed
src/retry.rsRetryChatInfo+resolve_retry_chat_infowith peer/bot/recipient handling. Namespace normalization. All call sites updated.src/client/sender_keys.rstake_recent_messagereturns(msg, Option<Jid>).try_take_by_keyextracted. Alternate key with server-comparison optimization.src/client/lid_pn.rsswap_pn_lid_namespaceshared helper.src/client/sessions.rsensure_e2e_sessions_resolved+ extractedensure_sessions_inner.Test plan
resolve_retry_chat_info_dm_with_device— PN DM, device stripped for chat, preserved for requesterresolve_retry_chat_info_lid_dm_with_device— LID DM, same behaviorresolve_retry_chat_info_dm_bare— bare from, chat == requesterresolve_retry_chat_info_group— group JID preserved, participant extractedresolve_retry_chat_info_status_broadcast— participant from attrresolve_retry_chat_info_status_broadcast_no_participant— fallback to senderresolve_retry_chat_info_peer_device_with_recipient— chat = recipientresolve_retry_chat_info_peer_device_without_recipient— fallback warningresolve_retry_chat_info_peer_via_lid— peer detected via LID matchresolve_retry_chat_info_bot_with_recipient— bot + recipient → chat = recipientresolve_retry_chat_info_bot_without_recipient— fallback to fromresolve_retry_chat_info_preserves_original_from— device suffix keptdm_retry_message_lookup_uses_bare_jid— store/take/re-add round-tripalternate_key_lookup_pn_to_lid— LID input, message under PN, swap fallbackalternate_key_lookup_pn_input_server_changed— PN input, server-comparison optimizationno_alternate_without_mapping— no mapping → primary onlyalternate_key_both_miss— both keys miss → Noneswap_pn_lid_namespace_preserves_device— LID↔PN swap + device + group→None