fix(prekeys): act on the device a rejected fetch names

src/client/context_impl.rs

@@ -18,13 +18,14 @@ impl SendContextResolver for Client {
         &self,
         jids: &[Jid],
     ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
-        self.fetch_pre_keys(jids, None).await
+        // The fan-out has its own batch-level handling; it wants the bundles only.
+        self.fetch_pre_keys(jids, None).await.map(|o| o.bundles)
     }
 
     async fn fetch_prekeys_for_identity_check(
         &self,
         jids: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
         self.fetch_pre_keys(jids, Some(PreKeyFetchReason::Identity))
             .await
             .map_err(|e| {

src/client/sessions.rs

@@ -323,9 +323,13 @@ impl Client {
 /// `wacore::request::ServerErrorCode` to cross the crate boundary. A downcast
 /// to either one alone silently answers `false` for the other, and the failure
 /// mode of that is the send failing exactly as it did before.
+/// The `<error code>` the server attaches to a device it no longer knows.
+const UNREGISTERED_DEVICE_CODE: u16 = 406;
+
 fn is_device_unregistered(err: &anyhow::Error) -> bool {
     use crate::error::ErrorChainExt;
-    err.server_rejection().is_some_and(|r| r.code == 406)
+    err.server_rejection()
+        .is_some_and(|r| r.code == UNREGISTERED_DEVICE_CODE)
 }
 
 /// The distinct users named by `jids`, in first-seen order.
@@ -457,6 +461,30 @@ impl Client {
             Err(e) => return Err(e),
         };
 
+        // The server named these individually, which is the per-device signal a
+        // batch-wide failure cannot give: refresh exactly their device lists and
+        // leave the rest of the batch alone. The send continues, because the
+        // devices that did come back with a bundle are unaffected and skipping
+        // them would deliver to fewer devices for a reason that only concerns
+        // the named ones.
+        if !prekey_bundles.rejected.is_empty() {
+            let rejected: Vec<Jid> = prekey_bundles
+                .rejected
+                .iter()
+                .filter(|device| device.code == UNREGISTERED_DEVICE_CODE)
+                .map(|device| device.jid.clone())
+                .collect();
+            if !rejected.is_empty() {
+                log::debug!(
+                    "prekey fetch rejected {} of {} device(s) as unregistered; \
+                     refreshing their device lists",
+                    rejected.len(),
+                    jids.len()
+                );
+                self.invalidate_device_caches_for(&rejected).await;
+            }
+        }
+
         let mut adapter = self.signal_adapter().await;
         let mut rng = rand::make_rng::<StdRng>();
 
@@ -465,7 +493,10 @@ impl Client {
         let mut failed_count = 0;
 
         for jid in jids {
-            if let Some(bundle) = prekey_bundles.get(&jid.normalize_for_prekey_bundle()) {
+            if let Some(bundle) = prekey_bundles
+                .bundles
+                .get(&jid.normalize_for_prekey_bundle())
+            {
                 match self
                     .install_prekey_bundle_cached(jid, bundle, &mut adapter, &mut rng)
                     .await

src/prekeys.rs

@@ -13,7 +13,7 @@ use std::sync::atomic::Ordering;
 use wacore::iq::prekeys::{
     DigestKeyBundleSpec, PreKeyCountSpec, PreKeyFetchReason, PreKeyFetchSpec, PreKeyUploadSpec,
 };
-use wacore::libsignal::protocol::{KeyPair, PreKeyBundle, PublicKey};
+use wacore::libsignal::protocol::{KeyPair, PublicKey};
 use wacore::libsignal::store::record_helpers::encode_pre_key_record_to;
 use wacore::store::commands::DeviceCommand;
 use wacore_binary::Jid;
@@ -147,7 +147,7 @@ impl Client {
         &self,
         jids: &[Jid],
         reason: Option<PreKeyFetchReason>,
-    ) -> Result<std::collections::HashMap<Jid, PreKeyBundle>, anyhow::Error> {
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
         let spec = match reason {
             Some(r) => PreKeyFetchSpec::with_reason(jids.to_vec(), r),
             None => PreKeyFetchSpec::new(jids.to_vec()),
@@ -161,13 +161,13 @@ impl Client {
         // identity as the fallback in validateADVwithIdentityKey).
         let spec = spec.with_account_identities(self.collect_account_identities(jids).await);
 
-        let bundles = self.execute(spec).await?;
+        let outcome = self.execute(spec).await?;
 
-        for jid in bundles.keys() {
+        for jid in outcome.bundles.keys() {
             log::debug!("Successfully parsed pre-key bundle for {}", jid.observe());
         }
 
-        Ok(bundles)
+        Ok(outcome)
     }
 
     /// Load, for each companion JID, its account (device 0) identity key from the

wacore/benches/send_receive_benchmark.rs

@@ -441,8 +441,8 @@ impl SendContextResolver for MockResolver {
     async fn fetch_prekeys_for_identity_check(
         &self,
         _: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
-        Ok(HashMap::new())
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
+        Ok(wacore::prekeys::PreKeyFetchOutcome::default())
     }
     async fn resolve_group_info(
         &self,

wacore/src/client/context.rs

@@ -220,10 +220,13 @@ pub trait SendContextResolver: crate::sync_marker::MaybeSendSync {
         jids: &[Jid],
     ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error>;
 
+    /// Returns the bundles alongside the devices the server rejected by name,
+    /// so a per-device rejection is not flattened into "no bundle" before the
+    /// fan-out can tell the two apart.
     async fn fetch_prekeys_for_identity_check(
         &self,
         jids: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error>;
+    ) -> Result<crate::prekeys::PreKeyFetchOutcome, anyhow::Error>;
 
     async fn resolve_group_info(&self, jid: &Jid) -> Result<Arc<GroupInfo>, anyhow::Error>;
 

wacore/src/iq/prekeys.rs

@@ -144,7 +144,7 @@ impl PreKeyFetchSpec {
 }
 
 impl IqSpec for PreKeyFetchSpec {
-    type Response = std::collections::HashMap<Jid, PreKeyBundle>;
+    type Response = crate::prekeys::PreKeyFetchOutcome;
 
     fn build_iq(&self) -> InfoQuery<'static> {
         let content = PreKeyUtils::build_fetch_prekeys_request(