fix(client): preserve recipient in <ack>, soften unknown stream:error

src/client.rs

@@ -3222,8 +3222,11 @@ impl Client {
     }
 
     pub(crate) async fn handle_stream_error(&self, node: &wacore_binary::NodeRef<'_>) {
-        self.is_logged_in.store(false, Ordering::Relaxed);
-
+        // is_logged_in handling: opt-in branches (515/516/401/409/conflict) clear it
+        // in the disconnect block below; 429/503 clear it inline because the server
+        // explicitly rejected the session and outgoing sends should bail fast; the
+        // unknown/code-less catch-all keeps it true so is_fully_ready()-gated work
+        // (notably prekey uploads) survives ack-shaped routing errors.
         let mut attrs = node.attrs();
         let code_cow = attrs.optional_string("code");
         let code = code_cow.as_deref().unwrap_or("");
@@ -3302,17 +3305,29 @@ impl Client {
                     should_disconnect = true;
                 }
                 "429" => {
+                    // Server signalled rate-limit on this session: mark logged-out so
+                    // outgoing sends bail fast instead of being interpreted as abuse
+                    // while we wait for the (likely-imminent) reconnect.
                     warn!(
                         "Got 429 stream error (rate limited). Will auto-reconnect with extended backoff."
                     );
+                    self.is_logged_in.store(false, Ordering::Relaxed);
                     self.auto_reconnect_errors.fetch_add(5, Ordering::Relaxed);
                 }
                 "503" => {
+                    // Server is going down/restarting: mark logged-out so sends fail
+                    // fast against the soon-to-die socket. Auto-reconnect handles recovery.
                     info!("Got 503 service unavailable, will auto-reconnect.");
+                    self.is_logged_in.store(false, Ordering::Relaxed);
                 }
                 _ => {
-                    error!("Unknown stream error: {}", DisplayableNodeRef(node));
-                    self.expected_disconnect.store(true, Ordering::Relaxed);
+                    // Server wraps per-stanza routing failures in <stream:error> without a
+                    // code (e.g. <ack/>): treat as informational so we don't trigger reconnect
+                    // storms. is_logged_in stays true on purpose — whatsmeow clears it eagerly,
+                    // but here is_fully_ready() gates prekey uploads and we want them to keep
+                    // working while the socket is still alive. Severity is warn!, not error!,
+                    // because the connection is intentionally preserved.
+                    warn!("Unknown stream error: {}", DisplayableNodeRef(node));
                     self.core.event_bus.dispatch(Event::StreamError(
                         crate::types::events::StreamError {
                             code: code.to_string(),
@@ -3323,8 +3338,11 @@ impl Client {
             }
         }
 
-        // Single transport lock acquisition for all branches that need disconnect.
+        // Single is_logged_in clear + transport disconnect for every opt-in branch
+        // (515/516/401/409 and conflict). 429/503/unknown fall through so the
+        // socket layer notices a real teardown without us forcing one.
         if should_disconnect {
+            self.is_logged_in.store(false, Ordering::Relaxed);
             let transport_opt = self.transport.lock().await.clone();
             if let Some(transport) = transport_opt {
                 self.runtime
@@ -3333,10 +3351,9 @@ impl Client {
                     }))
                     .detach();
             }
+            info!("Notifying connection shutdown from stream error handler");
+            self.notify_connection_shutdown();
         }
-
-        info!("Notifying connection shutdown from stream error handler");
-        self.notify_connection_shutdown();
     }
 
     pub(crate) async fn handle_connect_failure(&self, node: &wacore_binary::NodeRef<'_>) {
@@ -3912,6 +3929,10 @@ fn encode_ack_bytes(
         let from_str = from_val.as_str();
         p_str.as_ref() != from_str.as_ref()
     });
+    // Server expects `recipient` echoed back so it can route the ack to the
+    // origin companion/device (hosted-companion, peer, LID-routed stanzas).
+    // Dropping it makes the server close the stream with `<stream:error><ack/>`.
+    let recipient_val = node.get_attr("recipient");
     let tag = node.tag.as_ref();
 
     let typ_val = if tag != "message" && !is_encrypt_identity_notification(node) {
@@ -3922,16 +3943,18 @@ fn encode_ack_bytes(
 
     let include_from = tag == "message" && own_device_pn.is_some();
 
-    // Count attrs: class + id + to + optional(from, participant, type)
+    // Count attrs: class + id + to + optional(from, participant, recipient, type)
     let attr_count = 3
         + usize::from(include_from)
         + usize::from(participant_val.is_some())
+        + usize::from(recipient_val.is_some())
         + usize::from(typ_val.is_some());
 
     struct AckNode<'a> {
         id: &'a wacore_binary::node::ValueRef<'a>,
         from: &'a wacore_binary::node::ValueRef<'a>,
         participant: Option<&'a wacore_binary::node::ValueRef<'a>>,
+        recipient: Option<&'a wacore_binary::node::ValueRef<'a>>,
         typ: Option<&'a wacore_binary::node::ValueRef<'a>>,
         own_pn: Option<&'a Jid>,
         tag_str: &'a str,
@@ -3966,6 +3989,10 @@ fn encode_ack_bytes(
                 enc.write_string("participant")?;
                 p.encode_value(enc)?;
             }
+            if let Some(r) = self.recipient {
+                enc.write_string("recipient")?;
+                r.encode_value(enc)?;
+            }
             if let Some(t) = self.typ {
                 enc.write_string("type")?;
                 t.encode_value(enc)?;
@@ -3984,6 +4011,7 @@ fn encode_ack_bytes(
         id: id_val,
         from: from_val,
         participant: participant_val,
+        recipient: recipient_val,
         typ: typ_val,
         own_pn: if include_from { own_device_pn } else { None },
         tag_str: tag,
@@ -4007,13 +4035,14 @@ fn build_ack_node(node: &wacore_binary::NodeRef<'_>, own_device_pn: Option<&Jid>
         .get_attr("participant")
         .filter(|p| p.as_str().as_ref() != from_ref.as_str().as_ref())
         .map(|v| v.to_node_value());
+    let recipient = node.get_attr("recipient").map(|v| v.to_node_value());
     let tag = node.tag.as_ref();
     let typ = if tag != "message" && !is_encrypt_identity_notification(node) {
         node.get_attr("type").map(|v| v.to_node_value())
     } else {
         None
     };
-    let mut attrs = Attrs::with_capacity(6);
+    let mut attrs = Attrs::with_capacity(7);
     attrs.insert("class", NodeValue::from(tag));
     attrs.insert("id", id);
     attrs.insert("to", from);
@@ -4025,6 +4054,9 @@ fn build_ack_node(node: &wacore_binary::NodeRef<'_>, own_device_pn: Option<&Jid>
     if let Some(p) = participant {
         attrs.insert("participant", p);
     }
+    if let Some(r) = recipient {
+        attrs.insert("recipient", r);
+    }
     if let Some(t) = typ {
         attrs.insert("type", t);
     }
@@ -5873,6 +5905,135 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_build_ack_node_for_message_with_recipient_preserves_recipient() {
+        // Peer / hosted-companion / LID-routed messages carry `recipient`.
+        // The server uses it to route the ack back to the origin device;
+        // without it the stream is torn down with <stream:error><ack/></stream:error>.
+        let incoming = NodeBuilder::new("message")
+            .attr("from", "166361967902821@lid")
+            .attr("id", "2A32F960553696093D99")
+            .attr("type", "text")
+            .attr("recipient", "146991363395800@lid")
+            .build();
+        let own_device_pn: Jid = "155500012345:48@s.whatsapp.net"
+            .parse()
+            .expect("own device PN JID should parse");
+
+        let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn))
+            .expect("message ack should be buildable");
+
+        assert!(ack.attrs.get("class").is_some_and(|v| v == "message"));
+        assert!(
+            ack.attrs
+                .get("recipient")
+                .is_some_and(|v| v == "146991363395800@lid"),
+            "message ACK must echo the incoming `recipient` attribute"
+        );
+    }
+
+    #[test]
+    fn test_build_ack_node_for_receipt_with_recipient_preserves_recipient() {
+        // Receipt acks must also echo `recipient` when the incoming carries it.
+        let incoming = NodeBuilder::new("receipt")
+            .attr("from", "120363098765432100@g.us")
+            .attr("id", "RCPT-WITH-RECIPIENT")
+            .attr("type", "read")
+            .attr("recipient", "242395589390497@lid")
+            .build();
+        let own_device_pn: Jid = "155500012345:48@s.whatsapp.net"
+            .parse()
+            .expect("own device PN JID should parse");
+
+        let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn))
+            .expect("receipt ack should be buildable");
+
+        assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt"));
+        assert!(
+            ack.attrs
+                .get("recipient")
+                .is_some_and(|v| v == "242395589390497@lid"),
+            "receipt ACK must echo the incoming `recipient` attribute"
+        );
+    }
+
+    #[test]
+    fn test_build_ack_node_for_message_without_recipient_omits_recipient() {
+        // Regression guard: never synthesise a `recipient` field if the
+        // incoming stanza did not carry one — server would reject the ack.
+        let incoming = NodeBuilder::new("message")
+            .attr("from", "120363161500776365@g.us")
+            .attr("id", "A5791A5392EF60E3FB06")
+            .attr("type", "text")
+            .attr("participant", "181531758878822@lid")
+            .build();
+        let own_device_pn: Jid = "155500012345:48@s.whatsapp.net"
+            .parse()
+            .expect("own device PN JID should parse");
+
+        let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn))
+            .expect("message ack should be buildable");
+
+        assert!(
+            !ack.attrs.contains_key("recipient"),
+            "ACK must NOT add `recipient` when the incoming stanza has none"
+        );
+    }
+
+    #[test]
+    fn test_encode_ack_bytes_roundtrip_recipient() {
+        // Exercises the real wire encoder (`encode_ack_bytes`), not just the
+        // `build_ack_node` test mirror: serialize, decode the bytes back, and
+        // assert the parsed ACK echoes `recipient` when present and omits it
+        // when absent. Guards against the two builders silently diverging.
+        let own_device_pn: Jid = "155500012345:48@s.whatsapp.net"
+            .parse()
+            .expect("own device PN JID should parse");
+
+        let with_recipient = NodeBuilder::new("message")
+            .attr("from", "166361967902821@lid")
+            .attr("id", "2A32F960553696093D99")
+            .attr("type", "text")
+            .attr("recipient", "146991363395800@lid")
+            .build();
+        let buf = encode_ack_bytes(&with_recipient.as_node_ref(), Some(&own_device_pn))
+            .expect("encode_ack_bytes should not error")
+            .expect("encode_ack_bytes should produce bytes");
+        // The Encoder prepends a leading format byte (see `marshal`); the
+        // decoder wants raw protocol bytes — same handling as `node_to_owned_ref`.
+        let decoded =
+            wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode");
+        assert_eq!(decoded.tag, "ack");
+        assert!(
+            decoded
+                .get_attr("class")
+                .is_some_and(|v| v.as_str() == "message"),
+            "decoded ack must have class=message"
+        );
+        assert!(
+            decoded
+                .get_attr("recipient")
+                .is_some_and(|v| v.as_str() == "146991363395800@lid"),
+            "encode_ack_bytes must echo `recipient` onto the wire"
+        );
+
+        let without_recipient = NodeBuilder::new("message")
+            .attr("from", "120363161500776365@g.us")
+            .attr("id", "A5791A5392EF60E3FB06")
+            .attr("type", "text")
+            .attr("participant", "181531758878822@lid")
+            .build();
+        let buf = encode_ack_bytes(&without_recipient.as_node_ref(), Some(&own_device_pn))
+            .expect("encode_ack_bytes should not error")
+            .expect("encode_ack_bytes should produce bytes");
+        let decoded =
+            wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode");
+        assert!(
+            decoded.get_attr("recipient").is_none(),
+            "encode_ack_bytes must not synthesise `recipient` when absent"
+        );
+    }
+
     /// Smoke test: server ping with xmlns but no id attribute is handled.
     #[tokio::test]
     async fn test_handle_iq_ping_without_id() {
@@ -5976,13 +6137,22 @@ mod tests {
     #[tokio::test]
     async fn test_stream_error_429_keeps_reconnect_with_backoff() {
         let client = create_offline_sync_test_client().await;
+        client.is_logged_in.store(true, Ordering::Relaxed);
         let before = client.auto_reconnect_errors.load(Ordering::Relaxed);
         let node = NodeBuilder::new("stream:error").attr("code", "429").build();
         client.handle_stream_error(&node.as_node_ref()).await;
         assert!(
             client.enable_auto_reconnect.load(Ordering::Relaxed),
             "429 should keep auto-reconnect enabled"
         );
+        assert!(
+            !client.is_logged_in.load(Ordering::Relaxed),
+            "429 must clear is_logged_in so sends bail before the server flags abuse"
+        );
+        assert!(
+            !client.expected_disconnect.load(Ordering::Relaxed),
+            "429 must not mark the disconnect as expected (auto-reconnect path)"
+        );
         let after = client.auto_reconnect_errors.load(Ordering::Relaxed);
         assert_eq!(
             after,
@@ -5994,12 +6164,72 @@ mod tests {
     #[tokio::test]
     async fn test_stream_error_503_keeps_reconnect() {
         let client = create_offline_sync_test_client().await;
+        client.is_logged_in.store(true, Ordering::Relaxed);
         let node = NodeBuilder::new("stream:error").attr("code", "503").build();
         client.handle_stream_error(&node.as_node_ref()).await;
         assert!(
             client.enable_auto_reconnect.load(Ordering::Relaxed),
             "503 should keep auto-reconnect enabled"
         );
+        assert!(
+            !client.is_logged_in.load(Ordering::Relaxed),
+            "503 must clear is_logged_in so sends bail against the dying socket"
+        );
+        assert!(
+            !client.expected_disconnect.load(Ordering::Relaxed),
+            "503 must not mark the disconnect as expected (auto-reconnect path)"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_stream_error_unknown_keeps_connection_alive() {
+        // Unknown stream:error (no `code` attribute) must mirror whatsmeow's
+        // default branch: log + dispatch event, but NOT mark this as an
+        // expected disconnect. Setting that flag silently swallows the next
+        // real disconnect and races the read loop into shutdown.
+        let client = create_offline_sync_test_client().await;
+        // Simulate an authenticated session before the stream error arrives.
+        client.is_logged_in.store(true, Ordering::Relaxed);
+        let node = NodeBuilder::new("stream:error").build();
+        client.handle_stream_error(&node.as_node_ref()).await;
+        assert!(
+            client.is_logged_in.load(Ordering::Relaxed),
+            "unknown stream:error must NOT log the client out"
+        );
+        assert!(
+            !client.expected_disconnect.load(Ordering::Relaxed),
+            "unknown stream:error must not mark the disconnect as expected"
+        );
+        assert!(
+            client.enable_auto_reconnect.load(Ordering::Relaxed),
+            "unknown stream:error must keep auto-reconnect enabled"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_stream_error_ack_shaped_does_not_force_shutdown() {
+        // Server wraps per-stanza routing failures in `<stream:error><ack/>`
+        // with no `code` attribute. Treat as informational, not as a fatal
+        // stream teardown.
+        let client = create_offline_sync_test_client().await;
+        client.is_logged_in.store(true, Ordering::Relaxed);
+        let ack_child = NodeBuilder::new("ack")
+            .attr("class", "message")
+            .attr("type", "text")
+            .attr("id", "2A32F960553696093D99")
+            .build();
+        let node = NodeBuilder::new("stream:error")
+            .children([ack_child])
+            .build();
+        client.handle_stream_error(&node.as_node_ref()).await;
+        assert!(
+            client.is_logged_in.load(Ordering::Relaxed),
+            "ack-shaped stream:error must NOT log the client out"
+        );
+        assert!(
+            !client.expected_disconnect.load(Ordering::Relaxed),
+            "ack-shaped stream:error must not mark the disconnect as expected"
+        );
     }
 
     #[tokio::test]