fix(xmtp_mls,xmtp_proto): durable AddMissingInstallations reconciliation on the TaskRunner

crates/xmtp_db/src/encrypted_store/tasks.rs

@@ -550,7 +550,9 @@ pub(crate) mod tests {
     /// constants without a row-migration story.
     #[xmtp_common::test]
     fn data_hash_encoding_is_pinned() {
-        use xmtp_proto::xmtp::mls::database::{KpDeletion, KpRotation, PullInDeadline};
+        use xmtp_proto::xmtp::mls::database::{
+            AddMissingInstallations, KpDeletion, KpRotation, PullInDeadline,
+        };
         let rotation = TaskProto {
             task: Some(TaskKind::KpRotation(KpRotation {})),
         };
@@ -579,6 +581,19 @@ pub(crate) mod tests {
             hex::encode(data_hash_for(&pull_in)),
             "16b424873a34096e5157ab9f0a31e80dff1d23ebd8b1aab2b948a4732abfc849"
         );
+        // AddMissingInstallations (bytes group_id): outer tag for oneof field 7
+        // is (7<<3)|2 = 0x3a, LEN 0x22 (34 = 2-byte inner header + 32 bytes);
+        // inner field 1 tag 0x0a, LEN 0x20.
+        let add = TaskProto {
+            task: Some(TaskKind::AddMissingInstallations(AddMissingInstallations {
+                group_id: vec![0x22; 32],
+            })),
+        };
+        assert_eq!(add.encode_to_vec()[..4], [0x3a, 0x22, 0x0a, 0x20]);
+        assert_eq!(
+            hex::encode(data_hash_for(&add)),
+            "2180ffa08703d0dcc600a070da9d162da6821e29176af24e7a9955af2cf43764"
+        );
         // Determinism under repetition and a decode round-trip.
         let bytes = pull_in.encode_to_vec();
         for _ in 0..100 {

crates/xmtp_mls/src/groups/intents/queue.rs

@@ -1,11 +1,6 @@
-use futures::{StreamExt, TryFutureExt, stream};
-use std::collections::HashSet;
-use std::future::Future;
-
 use crate::groups::intents::GROUP_KEY_ROTATION_INTERVAL_NS;
 use crate::groups::{GroupError, MlsGroup, XmtpSharedContext};
 use derive_builder::Builder;
-use itertools::Itertools;
 use xmtp_db::TransactionOutcome::Continue;
 use xmtp_db::{
     DbQuery, TransactionOutcome,
@@ -41,71 +36,6 @@ impl QueueIntentBuilder {
             })
             .map(TransactionOutcome::into_continued)
     }
-
-    /// Queue this intent for each group
-    /// Accepts a hashset to ensure each group is unique
-    /// Accepts a closure that returns the intent data for a group
-    pub async fn queue_for_each<'a, C, F, Fut, E>(
-        &mut self,
-        groups: HashSet<MlsGroup<C>>,
-        data: F,
-    ) -> Result<Vec<StoredGroupIntent>, GroupError>
-    where
-        C: XmtpSharedContext + 'a,
-        F: Fn(MlsGroup<C>) -> Fut,
-        Fut: Future<Output = Result<Vec<u8>, E>>,
-        GroupError: From<E>,
-        E: std::fmt::Debug + std::error::Error,
-    {
-        if groups.is_empty() {
-            return Ok(vec![]);
-        }
-
-        let groups: Vec<MlsGroup<C>> = Vec::from_iter(groups);
-        let context: C = groups.first().expect("checked for empty").context.clone();
-
-        // get the intent data for each group
-        let (groups, errors): (Vec<_>, Vec<_>) = stream::iter(groups)
-            .map(|group| data(group.clone()).map_ok(move |d| (group, d)))
-            .buffered(10)
-            .collect::<Vec<Result<_, _>>>()
-            .await
-            .into_iter()
-            .partition_result();
-        let mut errors = errors.into_iter().map(GroupError::from).collect::<Vec<_>>();
-
-        let (intents, errs): (Vec<StoredGroupIntent>, Vec<_>) = context
-            .mls_storage()
-            .transaction(|conn| {
-                let (intents, errs): (Vec<StoredGroupIntent>, Vec<_>) = groups
-                    .into_iter()
-                    .map(|(group, data)| {
-                        let storage = conn.key_store();
-
-                        // nesting a transaction uses SQLite Savepoints
-                        // https://sqlite.org/lang_savepoint.html
-                        // Each savepoint here always commits or errors, so
-                        // unwrap the committed value directly.
-                        storage
-                            .savepoint(|conn| {
-                                let intent = self.clone().data(data).build()?;
-                                let storage = conn.key_store();
-                                let db = storage.db();
-                                intent.queue_with_conn(&db, &group).map(Continue)
-                            })
-                            .map(TransactionOutcome::into_continued)
-                    })
-                    .partition_result();
-                Ok::<_, GroupError>(Continue((intents, errs)))
-            })?
-            .into_continued();
-        errors.extend(errs);
-
-        for error in errors {
-            tracing::warn!("failed to queue intent {error}");
-        }
-        Ok(intents)
-    }
 }
 
 impl QueueIntent {
@@ -259,67 +189,3 @@ impl QueueIntent {
         Ok(())
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use std::{iter, sync::Arc};
-    use tokio::sync::Mutex;
-    use xmtp_common::Generate;
-    use xmtp_db::group::{GroupMembershipState, StoredGroup};
-    use xmtp_proto::types::GroupId;
-
-    use crate::test::mock::{NewMockContext, context};
-
-    use super::*;
-    use rstest::*;
-
-    fn group<C: XmtpSharedContext>(context: &C, id: Option<GroupId>) -> MlsGroup<&C> {
-        let id = id.unwrap_or_else(GroupId::generate);
-        StoredGroup::builder()
-            .id(id)
-            .created_at_ns(1)
-            .membership_state(GroupMembershipState::Allowed)
-            .added_by_inbox_id("bob")
-            .should_publish_commit_log(false)
-            .build()
-            .unwrap()
-            .store(&context.mls_storage().db())
-            .unwrap();
-
-        MlsGroup {
-            group_id: id,
-            dm_id: None,
-            created_at_ns: 1,
-            context,
-            mls_commit_lock: Arc::new(Default::default()),
-            mutex: Arc::new(Mutex::new(())),
-            conversation_type: xmtp_db::group::ConversationType::Dm,
-        }
-    }
-
-    #[rstest]
-    #[xmtp_common::test]
-    async fn can_queue_intent_for_each_group(context: NewMockContext) {
-        let groups = iter::repeat_with(|| group(&context, None)).take(10);
-        let intents = QueueIntent::update_group_membership()
-            .data(vec![0, 1, 2])
-            .queue_for_each(groups.collect(), async |_| Ok::<_, GroupError>(vec![]))
-            .await
-            .unwrap();
-        assert_eq!(intents.len(), 10);
-    }
-
-    #[rstest]
-    #[xmtp_common::test]
-    async fn only_queues_for_unique(context: NewMockContext) {
-        let groups = iter::repeat_n(group(&context, Some(GroupId::ZERO)), 10);
-
-        let stored = QueueIntent::update_group_membership()
-            .data(vec![0, 1, 2])
-            .queue_for_each(groups.collect(), async |_g| Ok::<_, GroupError>(vec![]))
-            .await
-            .unwrap();
-
-        assert_eq!(stored.len(), 1);
-    }
-}

crates/xmtp_mls/src/test/client_test_utils.rs

@@ -73,14 +73,20 @@ where
         let other_sync = other.device_sync_client();
         let mut sync_group = this_sync.get_sync_group().await?;
         let mut other_sync_group = other_sync.get_sync_group().await?;
-        for i in 0..10 {
+        for _ in 0..10 {
             sync_group = this_sync.get_sync_group().await?;
             other_sync_group = other_sync.get_sync_group().await?;
 
             sync_group.sync().await?;
             other_sync_group.sync().await?;
 
-            if sync_group.group_id == other_sync_group.group_id {
+            // Converge on epoch too, not just the group id: a commit published
+            // right after one side joins (e.g. the post-join KeyUpdate that
+            // rides the first HMAC-cycle message) is in flight exactly when
+            // the ids first match.
+            if sync_group.group_id == other_sync_group.group_id
+                && sync_group.epoch().await? == other_sync_group.epoch().await?
+            {
                 break;
             }
         }

crates/xmtp_mls/src/worker/device_sync/mod.rs

@@ -2,20 +2,16 @@ use crate::{
     client::ClientError,
     context::XmtpSharedContext,
     groups::{
-        GroupError, MlsGroup, PreconfiguredPolicies, intents::QueueIntent, send_message_opts,
-        summary::SyncSummary, welcome_sync::WelcomeService,
+        GroupError, MlsGroup, PreconfiguredPolicies, send_message_opts, summary::SyncSummary,
+        welcome_sync::WelcomeService,
     },
     mls_store::{MlsStore, MlsStoreError},
     subscriptions::{SubscribeError, SyncWorkerEvent},
     worker::{NeedsDbReconnect, metrics::WorkerMetrics},
 };
-use futures::{StreamExt, TryStreamExt, stream};
 use owo_colors::OwoColorize;
 use prost::Message;
-use std::{
-    collections::{HashMap, HashSet},
-    sync::Arc,
-};
+use std::{collections::HashMap, sync::Arc};
 use thiserror::Error;
 use tokio::sync::broadcast::error::RecvError;
 use tracing::instrument;
@@ -24,21 +20,28 @@ use xmtp_archive::{ArchiveError, BackupMetadata};
 use xmtp_common::ErrorCode;
 use xmtp_common::{NS_IN_DAY, RetryableError, time::now_ns};
 use xmtp_content_types::encoded_content_to_bytes;
+use xmtp_db::tasks::NewTask;
 use xmtp_db::{
     NotFound, StorageError, consent_record::ConsentState, group::GroupQueryArgs,
     group_message::StoredGroupMessage,
 };
 use xmtp_db::{XmtpDb, group::ConversationType, prelude::*};
 use xmtp_id::{InboxIdRef, associations::DeserializationError};
 use xmtp_mls_common::group::GroupMetadataOptions;
-use xmtp_proto::types::InstallationId;
+use xmtp_proto::types::{GroupId, InstallationId};
 use xmtp_proto::xmtp::{
     device_sync::content::{
         DeviceSyncContent as DeviceSyncContentProto, device_sync_content::Content as ContentProto,
     },
-    mls::message_contents::{
-        ContentTypeId, EncodedContent, PlaintextEnvelope,
-        plaintext_envelope::{Content, V1},
+    mls::{
+        database::{
+            AddMissingInstallations as AddMissingInstallationsProto, Task as TaskProto,
+            task::Task as TaskKindProto,
+        },
+        message_contents::{
+            ContentTypeId, EncodedContent, PlaintextEnvelope,
+            plaintext_envelope::{Content, V1},
+        },
     },
 };
 
@@ -364,7 +367,17 @@ where
                     hex::encode(self.context.installation_id()),
                     hex::encode(sync_group.group_id)
                 );
-                sync_group.add_missing_installations().await?;
+                if let Err(inline_err) = sync_group.add_missing_installations().await {
+                    // The group row is already persisted, so this add is never
+                    // re-attempted (later calls take the `Some` branch) — arm
+                    // the durable reconcile task so the TaskRunner heals it,
+                    // then surface the original error. Armed only on failure:
+                    // enqueue-first duplicated the reconcile (and its identity
+                    // fetch) on every sync-group creation.
+                    self.schedule_add_missing_installations_task(sync_group.group_id)
+                        .map_err(Box::new)?;
+                    return Err(inline_err);
+                }
                 sync_group.sync_with_conn().await?;
 
                 self.metrics.increment_metric(SyncMetric::SyncGroupCreated);
@@ -378,36 +391,47 @@ where
 
     /// This should be triggered when a new sync group appears,
     /// indicating the presence of a new installation.
+    ///
+    /// Schedules one durable AddMissingInstallations task per eligible group
+    /// on the TaskRunner (deduped by payload hash) so the membership add is
+    /// retried with backoff instead of being lost if a single inline attempt
+    /// fails (e.g. identity propagation lag → MissingSequenceId).
     #[cfg_attr(
         any(test, feature = "test-utils"),
         tracing::instrument(level = "info", skip_all)
     )]
-    pub async fn add_new_installation_to_groups(&self) -> Result<(), DeviceSyncError> {
+    pub fn schedule_add_installations_to_groups(&self) -> Result<usize, DeviceSyncError> {
         let groups = self.mls_store.find_groups(GroupQueryArgs {
             last_activity_after_ns: Some(now_ns() - NS_IN_DAY * 90),
             consent_states: Some(vec![ConsentState::Allowed, ConsentState::Unknown]),
             ..Default::default()
         })?;
 
-        let groups = HashSet::from_iter(groups);
-        let intents = QueueIntent::update_group_membership()
-            .queue_for_each(groups, move |group| async move {
-                let intent = group.get_membership_update_intent(&[], &[]).await?;
-                let intent: Vec<u8> = intent.into();
-                Ok::<_, GroupError>(intent)
-            })
-            .await?;
-
-        let context = &self.context;
-        stream::iter(intents)
-            .map(Ok::<_, GroupError>)
-            .try_for_each_concurrent(10, |intent| async move {
-                let (group, _) = MlsGroup::new_cached(context, &intent.group_id)?;
-                group.sync_until_intent_resolved(intent.id).await?;
-                Ok(())
-            })
-            .await?;
+        for group in &groups {
+            self.schedule_add_missing_installations_task(group.group_id)?;
+        }
+        Ok(groups.len())
+    }
 
+    /// Durably enqueue one AddMissingInstallations task for `group_id`
+    /// (deduped by payload hash against any pending row for the same group)
+    /// and wake the TaskRunner.
+    pub(crate) fn schedule_add_missing_installations_task(
+        &self,
+        group_id: GroupId,
+    ) -> Result<(), DeviceSyncError> {
+        let task = NewTask::builder()
+            .originating_message_sequence_id(0)
+            .originating_message_originator_id(0)
+            .build(TaskProto {
+                task: Some(TaskKindProto::AddMissingInstallations(
+                    AddMissingInstallationsProto {
+                        group_id: group_id.to_vec(),
+                    },
+                )),
+            })?;
+        self.context.db().create_or_ignore_task(task)?;
+        self.context.task_channels().wake();
         Ok(())
     }
 }