[WIP] Fix DM retry handler to delete sessions for all known devices

[Copilot]

Thanks for asking me to work on this. I will get started on it and keep this PR's description up to date as I form a plan and make progress.

---

This section details on the original issue you should resolve

<issue_title>fix: DM retry handler should delete sessions for all known devices, not just device 0</issue_title>
<issue_description>## Bug

The DM retry handler in handle_retry_receipt only deletes the Signal session for the bare/device-0 JID of the requester. When the resend goes through send_message_impl → per-device DM fanout, companion devices (device 33) still have stale sessions with consumed prekey IDs. The recipient fails to decrypt with InvalidPreKeyId.

Root Cause

Two pieces of code that are now inconsistent:

1. send_message_impl DM fanout (send.rs:1071–1121)

Since the per-device DM fanout change, the send path encrypts for all known recipient devices individually:

// send.rs:1091-1098
let mut all_dm_jids = match recipient_cached {
    Some(mut devices) => {
        devices.retain(|j| !j.is_hosted());
        devices  // [B_lid, B_lid:33, ...]
    },
    None => vec![recipient_bare],
};
// ...
self.ensure_e2e_sessions(&all_dm_jids).await?;

ensure_e2e_sessions checks has_session per device. If a session exists, it's reused as-is.

2. handle_retry_receipt session deletion (retry.rs:388–469)

The DM branch resolves the requester to a bare LID and deletes only that session:

// retry.rs:123-129
let participant_jid = if is_group_or_status {
    nr.attrs().optional_jid("participant").unwrap_or(...)
} else {
    receipt.source.sender.clone()  // B_phone@s.whatsapp.net (no device info)
};

// retry.rs:200
let resolved_jid = self.resolve_encryption_jid(&participant_jid).await;
// → B_lid@lid (device 0)

// retry.rs:393, 467
let signal_address = resolved_jid.to_protocol_address();
// → "100000024691356@lid" (device 0 only!)
self.signal_cache.delete_session(&signal_address).await;

Only 100000024691356@lid (device 0) is deleted. The session for 100000024691356:33@lid (device 33) survives.

The interaction

When the resend goes through send_message_impl:

Device	has_session?	What happens
B_lid (device 0)	false (deleted)	Fresh prekey fetch → new PreKeyMessage ✅
B_lid:33 (device 33)	true (stale!)	Reuses existing session → stale PreKeyMessage with consumed prekey ❌

The recipient's companion device (33) receives a PreKeyMessage referencing a one-time prekey that was already consumed in the initial message exchange → InvalidPreKeyId → retry fails silently.

Reproduction Logs

From a test with two clients (A sends to B, B deletes session, retry flow):

# Phase 1: Initial message works fine
[DEBUG] decrypted PreKey message from 100000012345678@lid.0 with one-time prekey ...
[DEBUG] Successfully decrypted message from 559980000001@s.whatsapp.net

# Phase 2: B deletes session, A sends new message
[WARN] Decryption failed for pkmsg from 559980000001@s.whatsapp.net due to InvalidPreKeyId
[WARN] Sending retry receipt jlucaso1/whatsapp-rust#1 for message 3EB0C9AFF6... from 559980000001@s.whatsapp.net

# A receives retry receipt, processes it:
[WARN] Failed to process key bundle from retry receipt: <keys> child missing
[INFO] Deleted session for 100000024691356@lid.0 due to retry receipt
                          ^^^^^^^^^^^^^^^^^^^^^^^^ only device 0!
[INFO] Resending message 3EB0C9AFF6... to 559980000002@s.whatsapp.net

# A resends — fetches fresh prekeys for device 0, but REUSES stale session for device 33:
[DEBUG] ensure_e2e_sessions: has_session(100000024691356@lid) = false → fetch prekeys
[DEBUG] ensure_e2e_sessions: has_session(100000024691356:33@lid) = true → skip (stale!)

# A builds enc for both devices:
[DEBUG] Building PreKeyWhisperMessage for: 100000024691356@lid.0 with preKeyId: 13  ← fresh ✅
[DEBUG] Building PreKeyWhisperMessage for: 100000024691356:33@lid.0 with preKeyId: 12  ← consumed! ❌

# B's companion (device 33) receives the stale enc:
[INFO] processing PreKey message from 100000012345678@lid.0 with one-time prekey 12
[ERROR] Message from 100000012345678@lid.0 failed to decrypt: No current session
[WARN] Decryption failed for pkmsg due to InvalidPreKeyId
                                          ^^^^^^^^^^^^^^^ prekey 12 was already consumed!

The retry loop continues until MAX_DECRYPT_RETRIES is hit, then gives up silently.

Before the per-device DM fanout

This didn't manifest before because the old DM send path only encrypted for the bare recipient JID (device 0):

// OLD send.rs DM fanout:
let mut all_dm_jids = Vec::with_capacity(1 + own_devices.len());
all_dm_jids.push(recipient_bare);  // Only device 0!
all_dm_jids.extend(own_devices);

The server handled multi-device fanout. Since A only established a session with device 0 (and the retry handler correctly deleted it), the resend always used fresh prekeys.

Proposed Fix

In the DM branch of handle_retry_receipt, delete sessions for all known devices of the requester, not just the bare JID:

// retry.rs DM branch (currently lines 388-469)
} else {
    // Delete sessions for ALL known devices of the requester.
    // The per-device DM fanout (send.rs) encrypts for each device individually,
    // so stale sessions on any device cause InvalidPreKeyId on resend.
    let requester_bare = resolved_jid.to_non_ad();
    let known_devices = self.get_devices_from_registry(&requester_bare).await;

    let device_jids: Vec<Jid> = match known_devices {
        Some(devices) => devices,
        None => vec![resolved_jid.clone()],
    };

    for device_jid in &device_jids {
        let signal_address = device_jid.to_protocol_address();

        // Base key collision tracking (existing logic, applied per-device)
        // ... (keep existing base key check logic but iterate over each device)

        self.signal_cache.delete_session(&signal_address).await;
        info!("Deleted session for {signal_address} due to retry receipt");
    }
    self.flush_signal_cache().await?;
}

This ensures that when send_message_impl calls ensure_e2e_sessions for the resend, ALL of the requester's devices will have has_session = false, triggering fresh prekey fetches and new sessions.

Alternative: Minimal fix

If iterating all devices feels too broad, a minimal fix would be to just skip the has_session optimization during retries — but send_message_impl doesn't currently know it's being called from a retry context (the force_key_distribution flag is for sender keys, not session establishment).

Test plan

• Existing test_retry_flow_after_session_deletion e2e test (currently fails, would pass with this fix)
• Unit test: verify handle_retry_receipt DM branch deletes sessions for all known devices
• Unit test: verify the resend after retry uses fresh PreKeyMessages for all devices</issue_description>

Comments on the Issue (you are @copilot in this section)

@jlucaso1 ## Proposed E2E Test

This test validates the fix by exercising the exact failure path: delete session → retry → resend must succeed for all recipient devices.

It follows the existing patterns in tests/e2e/tests/ (uses TestClient, send_and_expect_text, text_msg, event waiting, in-memory SQLite isolation).

// tests/e2e/tests/retry_dm_multidevice.rs

//! E2E test: DM retry after session deletion must recover across all devices.
//!
//! Validates that when a recipient deletes their Signal session and the retry
//! flow triggers, the sender re-establishes sessions for ALL known recipient
//! devices — not just device 0. This is a regression test for the interaction
//! between per-device DM fanout (send.rs) and the retry handler (retry.rs).
//!
//! Flow:
//! 1. A sends DM to B → establishes Signal sessions for all B's devices
//! 2. B replies to A → confirms bidirectional session
//! 3. B deletes Signal session with A
//! 4. A sends DM to B → B can't decrypt → retry receipt → A resends
//! 5. Verify B receives the retried message (proves all device sessions refreshed)
//! 6. A sends another DM → B receives it (proves session fully recovered)

use e2e_tests::{send_and_expect_text, text_msg, TestClient};
use log::info;
use wacore::types::events::Event;

#[tokio::test]
async fn test_dm_retry_recovers_all_device_sessions() -> anyhow::Result<()> {
    let _ = env_logger::builder().is_test(true).try_init();

    // --- Setup: two clients ---
    let mut client_a = TestClient::connect("e2e_retry_dm_a").await?;
    let mut client_b = TestClient::connect("e2e_retry_dm_b").await?;

    let jid_a = client_a.jid().await;
    let jid_b = client_b.jid().await;
    assert_ne!(jid_a, jid_b, "Clients must have different phone numbers");

    // === Phase 1: Establish bidirectional Signal session ===
    info!("Phase 1: Establishing bidirectional session");

    send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "phase1-a2b", 30).await?;
    send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "phase1-b2a", 30).await?;

    info!("Phase 1 complete: bidirectional session established");

    // === Phase 2: Break session and verify retry recovery ===
    info!("Phase 2: Deleting B's session with A and triggering retry");

    // Delete B's Signal session (and identity) for A — simulates session corruption/reinstall
    client_b
        .client
        .signal()
        .delete_sessions(std::slice::from_ref(&jid_a))
        .await?;

    // Brief pause for session flush
    tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

    // A sends message — B can't decrypt → retry receipt → A re-establishes session → B decrypts
    // This is the critical assertion: the retry must succeed even with per-device DM fanout.
    // Before the fix, A only refreshed device 0's session, leaving device 33 stale.
    send_and_expect_text(
        &client_a.client,
        &mut client_b,
        &jid_b,
        "phase2-retry-msg",
        30, // generous timeout: retry involves multiple roundtrips
    )
    .await?;

    info!("Phase 2 complete: retry recovery succeeded");

    // === Phase 3: Verify session fully recovered ===
    info!("Phase 3: Verifying session recovery with bidirectional messages");

    send_and_expect_text(
        &client_a.client,
        &mut client_b,
        &jid_b,
        "phase3-post-retry",
        10,
    )
    .await?;
    send_and_expect_text(
        &client_b.client,
        &mut client_a,
        &jid_a,
        "phase3-reply",
        10,
    )
    .await?;

    info!("Phase 3 complete: bidirectional messaging works after retry");

    // === Phase 4: Verify B experienced a decryption failure (proves session was actually broken) ===
    // Check that at least one UndecryptableMessage event was emitted during Phase 2.
    // This ensures we're actually testing the retry path, not a no-op.
    let had_decrypt_failure = client_b
        .wait_for_event(0, |e| matches!(e, Event::UndecryptableMessage(_)))
        .await
        .is_ok();

    // Note: The UndecryptableMessage event may have already been consumed by wait_for_text.
    // If the retry was fast enough, B might have received the retried message before we
    // check for failures. The real proof is that Phase 2's send_and_expect_text succeeded
    // despite the deleted session — that can only happen via the retry flow.
    info!(
        "Decryption failure observed: {} (expected: true, but fast retries may consume it)",
        had_decrypt_failure
    );

    // === Cleanup ===
    client_a.disconnect().await;
    client_b.disconnect().await;

    Ok(())
}

What it validates

Step	What's tested	Without fix	With fix
Phase 1	Baseline session establishment	✅ passes	✅ passes
Phase 2	Retry after session deletion	❌ times out (device 33 gets stale prekey)	�� all device sessions refreshed
Phase 3	Post-retry bidirectional messaging	❌ never reached	✅ session fully recovered

Running it

# Start bartender mock server first
cd bartender/mock-server && cargo run

# Run the test
cd whatsapp-rust && cargo test -p e2e-tests --test retry_dm_multidevice -- --nocapture

The test uses send_and_expect_text which internally calls send_message + wait_for_text with a 30s timeout. The retry flow involves: B decrypt failure → B sends retry receipt → server forwards to A → A re-encrypts → B decrypts. This typically completes in 2-5 seconds when working correctly, and times out at 30s when the bug manifests.</comment_new>