fix(appstate): bound pairing key-share wait by the 180s critical deadline

src/client/app_state.rs

@@ -123,10 +123,17 @@ impl Client {
     /// Sync multiple collections in a single IQ request, re-fetching those with `has_more_patches`.
     /// Matches WA Web's `serverSync()` outer loop (`3JJWKHeu5-P.js:54278-54305`).
     /// Max 5 iterations (WA Web's `C=5` constant).
+    ///
+    /// `key_wait_deadline` bounds how long a missing app-state decode key may be
+    /// awaited. The initial critical bootstrap passes the shared 180s critical-sync
+    /// deadline so the explicit `AppStateSyncKeyRequest` fallback can recover a
+    /// late/never-auto-shared key on the same connection; other callers pass `None`
+    /// for the fixed short default.
     #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.appstate.sync_batched", level = "debug", skip_all, fields(count = collections.len()), err(Debug)))]
     pub(crate) async fn sync_collections_batched(
         &self,
         collections: Vec<WAPatchName>,
+        key_wait_deadline: Option<wacore::time::Instant>,
     ) -> anyhow::Result<()> {
         if collections.is_empty() {
             return Ok(());
@@ -153,7 +160,9 @@ impl Client {
         // Track all collections for cleanup
         let all_collections: Vec<WAPatchName> = pending.clone();
 
-        let result = self.sync_collections_batched_inner(pending).await;
+        let result = self
+            .sync_collections_batched_inner(pending, key_wait_deadline)
+            .await;
 
         // Always clean up in-flight set
         {
@@ -169,6 +178,7 @@ impl Client {
     async fn sync_collections_batched_inner(
         &self,
         mut pending: Vec<WAPatchName>,
+        key_wait_deadline: Option<wacore::time::Instant>,
     ) -> anyhow::Result<()> {
         use wacore::appstate::patch_decode::CollectionSyncError;
         const MAX_ITERATIONS: usize = 5;
@@ -300,7 +310,14 @@ impl Client {
                     missing_all.extend(m);
                 }
             }
-            if !missing_all.is_empty() && !self.request_keys_and_wait(missing_all).await {
+            // Bound the key wait by the critical-sync deadline when one was given
+            // (initial bootstrap), so a late/never-auto-shared key still recovers via
+            // the explicit request on this connection; otherwise a fixed short wait.
+            let key_wait = match key_wait_deadline {
+                Some(deadline) => deadline.saturating_duration_since(wacore::time::Instant::now()),
+                None => Duration::from_secs(10),
+            };
+            if !missing_all.is_empty() && !self.request_keys_and_wait(missing_all, key_wait).await {
                 // The re-shared key didn't land in time. Report failure rather than a
                 // false success: the initial critical-sync path treats Ok as permission
                 // to cancel its retry watchdog and dispatch Connected, which would leave
@@ -547,7 +564,11 @@ impl Client {
                 .missing_key_ids_after_inline(&mut pl, &download)
                 .await
                 .unwrap_or_default();
-            if !missing.is_empty() && !self.request_keys_and_wait(missing).await {
+            if !missing.is_empty()
+                && !self
+                    .request_keys_and_wait(missing, Duration::from_secs(10))
+                    .await
+            {
                 // Report failure (not a partial success) so the caller retries instead of
                 // treating the collection as synced; it re-syncs once the share lands.
                 // Pages already decoded this run have their version persisted.
@@ -605,15 +626,15 @@ impl Client {
     /// request means an earlier one is still in flight, so the key may yet land here,
     /// and a re-verify that fails can't be masked by treating "request sent" as success
     /// or by a wake from an unrelated key share.
-    async fn request_keys_and_wait(&self, missing: Vec<Vec<u8>>) -> bool {
+    async fn request_keys_and_wait(&self, missing: Vec<Vec<u8>>, timeout: Duration) -> bool {
         if missing.is_empty() {
             return true;
         }
         let count = missing.len();
         let listener = self.initial_keys_synced_notifier.listen();
         self.request_missing_keys_with_dedup(missing.clone()).await;
-        debug!(target: "Client/AppState", "Requested {count} missing app-state key(s); waiting up to 10s for the re-share");
-        let _ = rt_timeout(&*self.runtime, Duration::from_secs(10), listener).await;
+        debug!(target: "Client/AppState", "Requested {count} missing app-state key(s); waiting up to {timeout:?} for the re-share");
+        let _ = rt_timeout(&*self.runtime, timeout, listener).await;
         let backend = self.persistence_manager.backend();
         for id in &missing {
             if backend.get_sync_key(id).await.ok().flatten().is_none() {

src/client/node_io.rs

@@ -914,30 +914,13 @@ impl Client {
                     "Starting Initial App State Sync (flag_set={flag_set}, needs_pushname={needs_pushname_from_sync})"
                 );
 
-                if !client_clone
-                    .initial_app_state_keys_received
-                    .load(Ordering::Relaxed)
-                {
-                    debug!(
-                        target: "Client/AppState",
-                        "Waiting up to 5s for app state keys..."
-                    );
-                    let _ = rt_timeout(
-                        &*client_clone.runtime,
-                        Duration::from_secs(5),
-                        client_clone.initial_keys_synced_notifier.listen(),
-                    )
-                    .await;
-
-                    // Check if connection was replaced while waiting
-                    check_generation!();
-                }
-
-                // Start the critical sync timeout timer matching WhatsApp Web's
-                // WAWebSyncBootstrap.$15 (setSyncDCriticalDataSyncTimeout).
-                // WhatsApp Web uses 180s and calls socketLogout(SyncdTimeout) if
-                // the critical data hasn't synced by then.
+                // Single deadline for the whole critical path (key-share grace + batched
+                // IQ + missing-key fallback). Matches WhatsApp Web's WAWebSyncBootstrap
+                // 180s critical-data deadline. Armed before the wait so every step below
+                // is bounded by the same clock.
                 const CRITICAL_SYNC_TIMEOUT_SECS: u64 = 180;
+                let critical_deadline = wacore::time::Instant::now()
+                    + Duration::from_secs(CRITICAL_SYNC_TIMEOUT_SECS);
                 let timeout_client = client_clone.clone();
                 let timeout_generation = task_generation;
                 let timeout_rt = client_clone.runtime.clone();
@@ -970,13 +953,45 @@ impl Client {
                     }
                 }));
 
+                // Brief grace for the auto-shared key that the primary sends at pairing
+                // (the WA Web primary path). The listener is registered before the flag
+                // check because the notifier is not sticky — a key-share landing in the
+                // load→listen gap would otherwise be missed. This wait is only an
+                // optimization to avoid a redundant explicit key request in the common
+                // fast case; if the key is late (heavy history sync) or never
+                // auto-shared, the batched sync below falls back to an explicit
+                // AppStateSyncKeyRequest bounded by `critical_deadline`, so correctness
+                // does not depend on this grace.
+                const KEY_SHARE_GRACE_SECS: u64 = 10;
+                let key_share_listener = client_clone.initial_keys_synced_notifier.listen();
+                if !client_clone
+                    .initial_app_state_keys_received
+                    .load(Ordering::Relaxed)
+                {
+                    debug!(
+                        target: "Client/AppState",
+                        "Waiting up to {KEY_SHARE_GRACE_SECS}s for the auto-shared app state key..."
+                    );
+                    let _ = rt_timeout(
+                        &*client_clone.runtime,
+                        Duration::from_secs(KEY_SHARE_GRACE_SECS),
+                        key_share_listener,
+                    )
+                    .await;
+
+                    // Check if connection was replaced while waiting
+                    check_generation!();
+                }
+
                 // Await critical collections via batched IQ before dispatching Connected.
+                // The deadline lets the missing-key fallback recover a late/never-shared
+                // key on this connection instead of stalling to the watchdog.
                 check_generation!();
                 match client_clone
-                    .sync_collections_batched(vec![
-                        WAPatchName::CriticalBlock,
-                        WAPatchName::CriticalUnblockLow,
-                    ])
+                    .sync_collections_batched(
+                        vec![WAPatchName::CriticalBlock, WAPatchName::CriticalUnblockLow],
+                        Some(critical_deadline),
+                    )
                     .await
                 {
                     Ok(()) => {
@@ -1016,11 +1031,14 @@ impl Client {
                     }
 
                     if let Err(e) = sync_client
-                        .sync_collections_batched(vec![
-                            WAPatchName::RegularLow,
-                            WAPatchName::RegularHigh,
-                            WAPatchName::Regular,
-                        ])
+                        .sync_collections_batched(
+                            vec![
+                                WAPatchName::RegularLow,
+                                WAPatchName::RegularHigh,
+                                WAPatchName::Regular,
+                            ],
+                            None,
+                        )
                         .await
                     {
                         sync_client.log_sync_error("non-critical app state sync", &e);

src/handlers/ib.rs

@@ -93,13 +93,16 @@ async fn handle_ib_impl(client: Arc<Client>, node: &wacore_binary::NodeRef<'_>)
                         if needs_resync && !client_clone.is_shutting_down() {
                             info!("syncd_app_state dirty -- re-syncing all app state collections");
                             if let Err(e) = client_clone
-                                .sync_collections_batched(vec![
-                                    WAPatchName::CriticalBlock,
-                                    WAPatchName::CriticalUnblockLow,
-                                    WAPatchName::RegularLow,
-                                    WAPatchName::RegularHigh,
-                                    WAPatchName::Regular,
-                                ])
+                                .sync_collections_batched(
+                                    vec![
+                                        WAPatchName::CriticalBlock,
+                                        WAPatchName::CriticalUnblockLow,
+                                        WAPatchName::RegularLow,
+                                        WAPatchName::RegularHigh,
+                                        WAPatchName::Regular,
+                                    ],
+                                    None,
+                                )
                                 .await
                                 && !client_clone.is_shutting_down()
                             {

src/handlers/notification/groups.rs

@@ -91,7 +91,7 @@ pub(crate) fn handle_server_sync_notification(
                         log::debug!(target: "Client/AppState", "server_sync task cancelled: connection generation changed during version check");
                         return;
                     }
-                    if let Err(e) = client_clone.sync_collections_batched(to_sync).await
+                    if let Err(e) = client_clone.sync_collections_batched(to_sync, None).await
                         && !client_clone.is_shutting_down()
                     {
                         warn!(

wacore/appstate/src/processor.rs

@@ -590,6 +590,77 @@ mod tests {
         assert!(matches!(err, AppStateError::SnapshotMACMismatch));
     }
 
+    /// Deterministic reproduction of the fresh-pairing race that PR #972 works
+    /// around. The critical `critical_unblock_low` snapshot (the account's saved
+    /// contacts + push name) can arrive before the encrypted app-state key-share
+    /// has been processed, when a heavy history sync saturates the stream at
+    /// pairing time. The SAME snapshot fails to decode with `KeyNotFound` while
+    /// the key is still in flight, and decodes cleanly the instant the key lands
+    /// — proving the failure is purely a key-ORDERING race, not a bad snapshot.
+    ///
+    /// Mirrors the field symptom: `critical_unblock_low v3: N records` failing
+    /// with "didn't find app state key" (`AppStateProcessor::get_app_state_key`
+    /// -> `backend.get_sync_key` returning `None` -> this `get_keys` closure
+    /// returning `KeyNotFound`).
+    #[test]
+    fn critical_snapshot_fails_key_not_found_until_key_share_lands() {
+        let master_key = [7u8; 32];
+        let keys = expand_app_state_keys(&master_key);
+        let key_id = b"appstate-sync-key-1".to_vec();
+
+        // A critical_unblock_low-style snapshot carrying a contact record.
+        let record = create_encrypted_record(
+            wa::syncd_mutation::SyncdOperation::SET,
+            &[1u8; 32],
+            &keys,
+            &key_id,
+            1_700_000_000,
+        );
+        let snapshot = wa::SyncdSnapshot {
+            version: buffa::MessageField::some(wa::SyncdVersion { version: Some(3) }),
+            records: vec![record],
+            key_id: buffa::MessageField::some(wa::KeyId {
+                id: Some(key_id.clone()),
+            }),
+            ..Default::default()
+        };
+
+        // Leg 1 — key-share NOT yet processed: the decode fails with KeyNotFound,
+        // exactly the "didn't find app state key" the paired companion hits.
+        let key_missing = |_: &[u8]| -> Result<Arc<ExpandedAppStateKeys>, AppStateError> {
+            Err(AppStateError::KeyNotFound)
+        };
+        let mut state = HashState::default();
+        let err = process_snapshot(
+            &snapshot,
+            &mut state,
+            key_missing,
+            false,
+            "critical_unblock_low",
+        )
+        .expect_err("must fail while the key-share is still in flight");
+        assert!(
+            matches!(err, AppStateError::KeyNotFound),
+            "expected KeyNotFound (the 'didn't find app state key' failure), got {err:?}"
+        );
+
+        // Leg 2 — key-share lands: the SAME snapshot decodes cleanly. The failure
+        // was ordering, not the snapshot — so the fix is about ensuring the key is
+        // present (event-driven), never about the snapshot or a longer fixed wait.
+        let key_present = |_: &[u8]| Ok(Arc::new(keys.clone()));
+        let mut state2 = HashState::default();
+        let result = process_snapshot(
+            &snapshot,
+            &mut state2,
+            key_present,
+            false,
+            "critical_unblock_low",
+        )
+        .expect("the same snapshot must decode once the key is present");
+        assert_eq!(result.state.version, 3);
+        assert_eq!(result.mutations.len(), 1, "the contact record must apply");
+    }
+
     #[test]
     fn test_process_patch_basic() {
         let master_key = [7u8; 32];