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.
@@ -591,31 +612,52 @@ impl Client {
         Ok(())
     }
 
-    /// Shared missing-key repair step for both sync paths: request the given keys and,
-    /// only if a fresh request actually went out (the per-key dedup didn't suppress it),
-    /// wait briefly for the primary to re-share. Returns true iff the caller should
-    /// refetch (a request was sent and we waited); false means nothing was requested
-    /// (empty or deduped), so the caller proceeds without stalling.
-    /// Request the missing decode keys, wait briefly for the re-share, then VERIFY they
-    /// actually landed. Returns true only when every requested key is now stored (the
-    /// caller may process); false means the share didn't arrive in time and the caller
-    /// must NOT process -- doing so would abort with KeyNotFound -- and should skip the
-    /// collection so it re-syncs on a later cycle. Empty input returns true (nothing to
-    /// wait for). Waits even when the per-key dedup suppressed the send: a deduped
-    /// 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 {
+    /// Request the missing decode keys, wait up to `timeout` for the re-share, then
+    /// VERIFY they actually landed. Returns true only when every requested key is now
+    /// stored (the caller may process); false means the share didn't arrive in time and
+    /// the caller must NOT process -- doing so would abort with KeyNotFound -- and should
+    /// skip the collection so it re-syncs on a later cycle. Empty input returns true
+    /// (nothing to wait for). Waits even when the per-key dedup suppressed the send: a
+    /// deduped 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>>, 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;
+        let deadline = wacore::time::Instant::now() + timeout;
+        let mut requested = false;
+        loop {
+            // Register the listener before the store check: the notifier is not sticky,
+            // so a key-share persisted in the check→listen gap would be lost.
+            let listener = self.initial_keys_synced_notifier.listen();
+            if self.all_sync_keys_present(&missing).await {
+                return true;
+            }
+            // Send the explicit re-request once, after confirming the keys are still
+            // missing.
+            if !requested {
+                self.request_missing_keys_with_dedup(missing.clone()).await;
+                requested = true;
+                debug!(target: "Client/AppState", "Requested {count} missing app-state key(s); waiting up to {timeout:?} for the re-share");
+            }
+            let remaining = deadline.saturating_duration_since(wacore::time::Instant::now());
+            if remaining.is_zero() {
+                return false;
+            }
+            // A single share may cover only part of `missing`, and the notifier is global
+            // (an unrelated share wakes us too), so keep re-checking on each wake until
+            // every key lands or the deadline passes — don't give up on the first wake
+            // while budget remains.
+            let _ = rt_timeout(&*self.runtime, remaining, listener).await;
+        }
+    }
+
+    /// True iff every given app-state sync-key id is already stored.
+    async fn all_sync_keys_present(&self, ids: &[Vec<u8>]) -> bool {
         let backend = self.persistence_manager.backend();
-        for id in &missing {
+        for id in ids {
             if backend.get_sync_key(id).await.ok().flatten().is_none() {
                 return false;
             }

src/client/node_io.rs

@@ -914,33 +914,23 @@ 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);
+                // Explicit "critical sync completed" signal for the watchdog. A push_name
+                // check is not a reliable proxy: a business account gets push_name set
+                // from business_name at pairing (src/pair.rs) while still needing the
+                // full sync, so the watchdog would wrongly stand down on a failed sync.
+                let critical_sync_done =
+                    std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
                 let timeout_client = client_clone.clone();
                 let timeout_generation = task_generation;
                 let timeout_rt = client_clone.runtime.clone();
+                let timeout_done = critical_sync_done.clone();
                 let critical_sync_timeout_handle = timeout_rt.spawn(Box::pin(async move {
                     timeout_client.runtime.sleep(Duration::from_secs(CRITICAL_SYNC_TIMEOUT_SECS)).await;
                     // Check generation — if connection was replaced, this timeout is stale
@@ -949,38 +939,68 @@ impl Client {
                     {
                         return;
                     }
-                    // Matches WhatsApp Web's $16(): check if SettingPushName was synced.
-                    // If push_name is still empty after 180s, critical sync failed.
-                    let push_name = timeout_client.get_push_name();
-                    if push_name.is_empty() {
+                    if timeout_done.load(Ordering::SeqCst) {
+                        debug!(
+                            target: "Client/AppState",
+                            "Critical sync timeout fired but critical sync already completed"
+                        );
+                    } else {
                         warn!(
                             target: "Client/AppState",
-                            "Critical app state sync timed out after {CRITICAL_SYNC_TIMEOUT_SECS}s \
-                             (push_name not synced). Reconnecting to retry."
+                            "Critical app state sync did not complete within {CRITICAL_SYNC_TIMEOUT_SECS}s. \
+                             Reconnecting to retry."
                         );
                         // WhatsApp Web does socketLogout here which clears device identity.
                         // We reconnect instead — preserving credentials and keeping the
                         // run loop active so auto-reconnect can retry the sync.
                         timeout_client.reconnect_immediately().await;
-                    } else {
-                        debug!(
-                            target: "Client/AppState",
-                            "Critical sync timeout fired but push_name was already synced"
-                        );
                     }
                 }));
 
+                // 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(()) => {
-                        // Critical sync completed — cancel the timeout timer
+                        // Critical sync completed — signal the watchdog, then cancel it.
+                        critical_sync_done.store(true, Ordering::SeqCst);
                         critical_sync_timeout_handle.abort();
 
                         check_generation!();
@@ -999,9 +1019,12 @@ impl Client {
                     }
                     Err(e) => {
                         client_clone.log_sync_error("critical app state sync", &e);
-                        // Don't abort the timeout or dispatch Connected — the sync failed,
-                        // so the timeout watchdog should remain active to force a reconnect
-                        // if needed. Return early to avoid emitting a spurious Connected event.
+                        // The sync failed — the watchdog must stay alive to force a reconnect.
+                        // detach() so this early return doesn't abort it on drop (AbortHandle
+                        // aborts the task when dropped); without this the watchdog would be
+                        // cancelled exactly when the deadline-bound wait fails, and no
+                        // reconnect would happen.
+                        critical_sync_timeout_handle.detach();
                         return;
                     }
                 }
@@ -1016,11 +1039,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!(