xmtp_mls/subscriptions/
stream_conversations.rs

1//! Live conversation notifications from committed group discovery.
2
3use super::{
4    Result,
5    incoming::{IncomingCoordinator, IncomingLease, IncomingScope, stream_api::GroupSetGate},
6    internal::InternalEvent,
7    local_delivery::LocalDeliveryError,
8};
9use crate::{context::XmtpSharedContext, groups::MlsGroup};
10use futures::Stream;
11use std::{
12    collections::{HashMap, VecDeque},
13    pin::Pin,
14    sync::{Arc, Weak},
15    task::{Context, Poll},
16};
17use xmtp_common::{BoxDynStream, time::sleep};
18use xmtp_db::{
19    consent_record::ConsentState,
20    group::{ConversationType, GroupQueryArgs, StoredGroup},
21    prelude::*,
22};
23use xmtp_events::EventFilter;
24use xmtp_proto::{
25    api_client::XmtpMlsStreams,
26    types::{GroupId, Topic},
27};
28
29const ALL_CONSENT_STATES: [ConsentState; 3] = [
30    ConsentState::Allowed,
31    ConsentState::Unknown,
32    ConsentState::Denied,
33];
34
35/// Local notification history, not durable receipt or processing progress.
36#[derive(Default)]
37struct KnownConversations {
38    /// A later Welcome for an existing group is a new rejoin notification.
39    welcomes: HashMap<GroupId, Option<i64>>,
40    /// Keep one selected group per participant pair for this stream's lifetime.
41    dms: HashMap<String, GroupId>,
42}
43
44impl KnownConversations {
45    /// Exclude groups already committed when the live stream starts.
46    fn from_groups(groups: Vec<StoredGroup>) -> Self {
47        let mut known = Self::default();
48        for group in groups {
49            known.welcomes.insert(group.id, group.sequence_id);
50            if let Some(dm_id) = group.dm_id {
51                known.dms.entry(dm_id).or_insert(group.id);
52            }
53        }
54        known
55    }
56
57    /// Observe committed discovery without letting DM list ranking change its identity.
58    fn observe(&mut self, group: &StoredGroup, include_duplicate_dms: bool) -> bool {
59        if self.welcomes.insert(group.id, group.sequence_id) == Some(group.sequence_id)
60            || ConversationType::virtual_types().contains(&group.conversation_type)
61        {
62            return false;
63        }
64        if !include_duplicate_dms && let Some(dm_id) = &group.dm_id {
65            return *self.dms.entry(dm_id.clone()).or_insert(group.id) == group.id;
66        }
67        true
68    }
69}
70
71/// The groups the stream selects, or none while the group-set generation
72/// equals the one `gate` recorded at its last successful read. A new group, a
73/// later Welcome and a consent change all change the generation. A message
74/// does not, though it can make another group a stitched DM's winner: when
75/// that DM's groups hold different consent, a consent-filtered stream reports
76/// it at the next change of the group set instead of the next pass.
77// implements: PROC-015
78fn changed_groups<C: XmtpSharedContext>(
79    context: &C,
80    query: &GroupQueryArgs,
81    gate: &mut GroupSetGate,
82) -> std::result::Result<Vec<StoredGroup>, xmtp_db::ConnectionError> {
83    let db = context.db();
84    let generation = db.group_set_generation()?;
85    if !gate.changed(generation) {
86        return Ok(Vec::new());
87    }
88    let groups = if query.consent_states.as_ref().is_some_and(Vec::is_empty) {
89        Vec::new()
90    } else {
91        #[cfg(any(test, feature = "test-hooks"))]
92        test_hooks::count_groups_read();
93        db.find_groups(query)?
94    };
95    gate.record(generation);
96    Ok(groups)
97}
98
99/// Notify one subscriber of committed local creation and Welcome joins.
100pub struct StreamConversations<C: XmtpSharedContext> {
101    inner: BoxDynStream<'static, Result<MlsGroup<C>>>,
102    lease: Weak<IncomingLease>,
103}
104
105impl<C: XmtpSharedContext + 'static> StreamConversations<C> {
106    pub async fn new(
107        context: &C,
108        conversation_type: Option<ConversationType>,
109        include_duplicate_dms: bool,
110        consent_states: Option<Vec<ConsentState>>,
111    ) -> Result<Self>
112    where
113        C::ApiClient: XmtpMlsStreams,
114    {
115        Self::new_owned(
116            context.clone(),
117            conversation_type,
118            include_duplicate_dms,
119            consent_states,
120        )
121        .await
122    }
123
124    /// Capture the local discovery baseline before starting the shared receiver.
125    pub async fn new_owned(
126        context: C,
127        conversation_type: Option<ConversationType>,
128        include_duplicate_dms: bool,
129        consent_states: Option<Vec<ConsentState>>,
130    ) -> Result<Self>
131    where
132        C::ApiClient: XmtpMlsStreams,
133    {
134        context.server_configuration().check()?;
135        let events = context.events().subscribe(
136            EventFilter::default().with_internal(|event| {
137                matches!(event, InternalEvent::GroupJoined { is_sync: false, .. })
138            }),
139            Some(10),
140        );
141        let known = KnownConversations::from_groups(context.db().find_groups(GroupQueryArgs {
142            include_sync_groups: true,
143            include_duplicate_dms: true,
144            consent_states: Some(ALL_CONSENT_STATES.to_vec()),
145            ..Default::default()
146        })?);
147        let coordinator = IncomingCoordinator::for_context(&context);
148        let lease = Arc::new(coordinator.acquire_stream(IncomingScope::Topics(vec![
149            Topic::new_welcome_message(context.installation_id()),
150        ])));
151        let query = GroupQueryArgs {
152            conversation_type,
153            consent_states: Some(consent_states.unwrap_or_else(|| ALL_CONSENT_STATES.to_vec())),
154            include_duplicate_dms,
155            ..Default::default()
156        };
157        let observer = Arc::downgrade(&lease);
158        let stream = futures::stream::unfold(
159            Some((
160                context,
161                events,
162                lease,
163                known,
164                VecDeque::new(),
165                query,
166                GroupSetGate::default(),
167            )),
168            |state| async move {
169                let (context, events, lease, mut known, mut ready, query, mut gate) = state?;
170                loop {
171                    if let Err(error) = context.server_configuration().check() {
172                        lease.close();
173                        return Some((Err(error.into()), None));
174                    }
175                    if context.is_closed() {
176                        return None;
177                    }
178                    if let Err(error) = lease.check_recovery() {
179                        lease.close();
180                        return Some((Err(LocalDeliveryError::from(error).into()), None));
181                    }
182                    if let Some(group) = ready.pop_front() {
183                        return Some((
184                            Ok(group),
185                            Some((context, events, lease, known, ready, query, gate)),
186                        ));
187                    }
188                    #[cfg(any(test, feature = "test-hooks"))]
189                    test_hooks::count_pass();
190                    match changed_groups(&context, &query, &mut gate) {
191                        Ok(groups) => {
192                            for group in groups {
193                                if known.observe(&group, query.include_duplicate_dms) {
194                                    ready.push_back(MlsGroup::new(
195                                        context.clone(),
196                                        group.id,
197                                        group.dm_id,
198                                        group.conversation_type,
199                                        group.created_at_ns,
200                                    ));
201                                }
202                            }
203                        }
204                        Err(error) => {
205                            // The database operation has already used its normal
206                            // retry policy. Only the app can start another stream.
207                            lease.close();
208                            return Some((Err(error.into()), None));
209                        }
210                    }
211                    if !ready.is_empty() {
212                        continue;
213                    }
214                    tokio::select! {
215                        _ = context.cancellation_token().cancelled() => {},
216                        _ = lease.changed() => {},
217                        _ = sleep(context.stream_policy().active_database_poll_interval) => {},
218                        event = events.next() => { event?; },
219                    }
220                }
221            },
222        );
223        Ok(Self {
224            inner: Box::pin(stream),
225            lease: observer,
226        })
227    }
228
229    /// Observe the same connection that supplies this conversation stream.
230    pub fn lease(&self) -> Option<Arc<IncomingLease>> {
231        self.lease.upgrade()
232    }
233}
234
235impl<C: XmtpSharedContext> Stream for StreamConversations<C> {
236    type Item = Result<MlsGroup<C>>;
237    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
238        self.inner.as_mut().poll_next(cx)
239    }
240}
241
242#[cfg(any(test, feature = "test-hooks"))]
243mod test_hooks;
244
245#[cfg(test)]
246mod test {
247    use crate::prelude::*;
248    use crate::utils::ClientTester;
249    use std::sync::Arc;
250
251    use super::*;
252    use crate::builder::ClientBuilder;
253    use crate::groups::send_message_opts::SendMessageOpts;
254    use crate::tester;
255    use crate::utils::fixtures::{alix, bo};
256    use xmtp_db::group::GroupQueryArgs;
257
258    use futures::StreamExt;
259    use std::cell::Cell;
260    use xmtp_cryptography::utils::generate_local_wallet;
261
262    use super::test_hooks::{GROUPS_READS, PASSES};
263
264    /// Idle passes do not reread the groups table, and an invitation that
265    /// changes the group set is still delivered on the next pass.
266    // verifies: CONS-044, PROC-015
267    #[xmtp_common::test(unwrap_try = true)]
268    async fn conversation_stream_rereads_groups_only_after_a_change() {
269        tester!(alix, disable_workers);
270        tester!(bo, disable_workers);
271        let passes = || PASSES.with(Cell::get);
272        let reads = || GROUPS_READS.with(Cell::get);
273        let mut stream = StreamConversations::new(&bo.context, None, false, None).await?;
274        tokio::select! {
275            _ = stream.next() => panic!("the stream yielded with no new conversation"),
276            waited = xmtp_common::wait_for_ge(|| async { passes() }, 3) => waited?,
277        }
278        assert_eq!(reads(), 1, "an idle pass reread the groups table");
279
280        let group = alix.create_group(None, None)?;
281        group.add_members(&[bo.inbox_id()]).await?;
282        assert_eq!(stream.next().await.unwrap()?.group_id, group.group_id);
283        assert_eq!(reads(), 2);
284    }
285
286    // verifies: CONF-075
287    #[xmtp_common::test(unwrap_try = true)]
288    async fn a_blocked_connection_fences_queued_conversations() {
289        use crate::{server_configuration::BlockedConnection, subscriptions::SubscribeError};
290
291        tester!(alix, disable_workers);
292        let mut stream = StreamConversations::new(&alix.context, None, false, None).await?;
293        alix.create_group(None, None)?;
294        alix.create_group(None, None)?;
295        stream.next().await.unwrap()?;
296
297        alix.context.server_configuration().block_connection(
298            BlockedConnection::ClientVersionTooOld {
299                client: "1.0.0".into(),
300                minimum: "9999.0.0".into(),
301            },
302        );
303        alix.context.cancellation_token().cancel();
304        let error = stream.next().await.unwrap().unwrap_err();
305        assert!(matches!(
306            error,
307            SubscribeError::Blocked(BlockedConnection::ClientVersionTooOld { ref minimum, .. })
308                if minimum == "9999.0.0"
309        ));
310        assert!(stream.next().await.is_none());
311    }
312
313    xmtp_common::if_native! {
314    // verifies: PROC-040
315    #[xmtp_common::test(unwrap_try = true)]
316    async fn storage_error_ends_notifications_once_and_releases_the_lease() {
317        use xmtp_db::ConnectionExt;
318        tester!(alix, persistent_db, disable_workers);
319        let coordinator = IncomingCoordinator::for_context(&alix.context);
320        let baseline_owners = Arc::strong_count(&coordinator);
321        let mut stream = StreamConversations::new(&alix.context, None, false, None).await?;
322        assert_eq!(Arc::strong_count(&coordinator), baseline_owners + 1);
323        alix.context.db().disconnect()?;
324        let expected = alix.context.db().find_groups(GroupQueryArgs::default()).unwrap_err();
325        let error = stream.next().await.unwrap().unwrap_err();
326        assert_eq!(error.to_string(), expected.to_string());
327        assert!(matches!(error, super::super::SubscribeError::Db(
328            xmtp_db::ConnectionError::Platform(xmtp_db::PlatformStorageError::PoolNeedsConnection)
329        )), "{error:?}");
330        assert_eq!(Arc::strong_count(&coordinator), baseline_owners);
331        assert!(stream.next().await.is_none());
332        assert!(!alix.context.is_closed());
333
334        alix.context.db().reconnect()?;
335        let mut replacement = StreamConversations::new(&alix.context, None, false, None).await?;
336        let group = alix.create_group(None, None)?;
337        assert_eq!(replacement.next().await.unwrap()?.group_id, group.group_id);
338    }
339    }
340
341    #[xmtp_common::timeout(std::time::Duration::from_secs(10))]
342    #[rstest::rstest]
343    #[case::two_conversations(2)]
344    #[case::five_conversations(5)]
345    #[xmtp_common::test]
346    #[awt]
347    async fn stream_welcomes(
348        #[future] alix: ClientTester,
349        #[future] bo: ClientTester,
350        #[case] group_size: usize,
351    ) {
352        let mut groups = vec![];
353        let mut stream = StreamConversations::new(&bo.context, None, false, None)
354            .await
355            .unwrap();
356        for _ in 0..group_size {
357            let alix_bo_group = alix.create_group(None, None).unwrap();
358            groups.push(alix_bo_group.group_id);
359            alix_bo_group.add_members(&[bo.inbox_id()]).await.unwrap();
360        }
361        while !groups.is_empty() {
362            let bo_received_groups = stream.next().await.unwrap().unwrap();
363            let index = groups
364                .iter()
365                .position(|group_id| bo_received_groups.group_id == *group_id)
366                .expect("group must be found");
367            groups.remove(index);
368        }
369
370        assert!(groups.is_empty(), "Groups must have all been received");
371    }
372
373    #[rstest::rstest]
374    #[xmtp_common::test(unwrap_try = true)]
375    async fn test_sync_groups_are_not_streamed() {
376        tester!(alix, sync_worker);
377        let stream = alix.stream_conversations(None, false).await?;
378        futures::pin_mut!(stream);
379
380        tester!(_alix2, from: alix);
381
382        let result =
383            xmtp_common::time::timeout(std::time::Duration::from_millis(100), stream.next()).await;
384        assert!(result.is_err(), "Sync group should not stream");
385    }
386
387    #[rstest::rstest]
388    #[case(ConversationType::Dm, "Unexpectedly received a Group")]
389    #[case(ConversationType::Group, "Unexpectedly received a DM")]
390    #[xmtp_common::test]
391    //TODO: case 2 consistently fails on timeout only in CI in webassembly
392    // difficult to tell why. not able to repro locally.
393    // CI might have issues with http connection limits
394    #[cfg_attr(target_arch = "wasm32", ignore)]
395    async fn test_dm_stream_filter(
396        #[case] conversation_type: ConversationType,
397        #[case] expected: &str,
398    ) {
399        tester!(alix);
400        tester!(bo);
401        let stream = alix
402            .stream_conversations(Some(conversation_type), false)
403            .await
404            .unwrap();
405        futures::pin_mut!(stream);
406
407        alix.find_or_create_dm(bo.inbox_id().to_string(), None)
408            .await
409            .unwrap();
410
411        let group = alix.create_group(None, None).unwrap();
412        group.add_members(&[bo.inbox_id()]).await.unwrap();
413
414        let group = stream.next().await.unwrap();
415        let metadata = group.unwrap().metadata().await.unwrap();
416
417        assert_eq!(
418            metadata.conversation_type, conversation_type,
419            "{}",
420            expected
421        );
422        // there is only one item on the stream
423        let result =
424            xmtp_common::time::timeout(std::time::Duration::from_millis(100), stream.next()).await;
425        assert!(result.is_err(), "should only be one item in the stream");
426    }
427
428    #[rstest::rstest]
429    #[xmtp_common::test]
430    async fn test_dm_stream_all_conversation_types() {
431        let alix = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
432        let bo = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
433        let davon = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
434        let eri = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
435
436        // Start a stream with all conversations
437        let mut groups = Vec::new();
438        // Wait for 2 seconds for the group creation to be streamed
439        let stream = alix.stream_conversations(None, false).await.unwrap();
440        futures::pin_mut!(stream);
441
442        alix.find_or_create_dm(davon.inbox_id().to_string(), None)
443            .await
444            .unwrap();
445        let group = stream.next().await.unwrap();
446        assert!(group.is_ok());
447        groups.push(group.unwrap());
448
449        let dm = eri
450            .find_or_create_dm(alix.inbox_id().to_string(), None)
451            .await
452            .unwrap();
453        dm.add_members(&[alix.inbox_id()]).await.unwrap();
454        let group = stream.next().await.unwrap();
455        assert!(group.is_ok());
456        groups.push(group.unwrap());
457
458        let group = alix.create_group(None, None).unwrap();
459        group.add_members(&[bo.inbox_id()]).await.unwrap();
460        let group = stream.next().await.unwrap();
461        assert!(group.is_ok());
462        groups.push(group.unwrap());
463
464        assert_eq!(groups.len(), 3);
465    }
466
467    #[xmtp_common::timeout(std::time::Duration::from_secs(10))]
468    #[rstest::rstest]
469    #[xmtp_common::test]
470    async fn test_self_group_creation() {
471        tester!(alix);
472        tester!(bo);
473
474        let stream = alix
475            .stream_conversations(Some(ConversationType::Group), false)
476            .await
477            .unwrap();
478        futures::pin_mut!(stream);
479
480        alix.create_group(None, None).unwrap();
481        let _self_group = stream.next().await.unwrap();
482
483        let group = bo.create_group(None, None).unwrap();
484        group.add_members(&[alix.inbox_id()]).await.unwrap();
485        let _bo_group = stream.next().await.unwrap();
486
487        // Verify syncing welcomes while streaming causes no issues
488        alix.sync_welcomes().await.unwrap();
489        let find_groups_results = alix.find_groups(GroupQueryArgs::default()).unwrap();
490        assert_eq!(2, find_groups_results.len());
491    }
492
493    /// Consent filtering cannot turn a group from the live baseline into a new join.
494    #[rstest::rstest]
495    #[case::unfiltered(None, 2)]
496    #[case::allowed_only(Some(vec![ConsentState::Allowed]), 1)]
497    #[case::empty_filter(Some(Vec::new()), 0)]
498    #[xmtp_common::test(unwrap_try = true)]
499    async fn conversation_consent_filter_preserves_live_baseline(
500        #[case] consent_states: Option<Vec<ConsentState>>,
501        #[case] expected_count: usize,
502    ) {
503        use xmtp_common::time::{Duration, timeout};
504
505        tester!(alix, disable_workers);
506        let old = alix.create_group(None, None).unwrap();
507        old.update_consent_state(ConsentState::Denied).unwrap();
508        let mut stream = StreamConversations::new(
509            &alix.context,
510            Some(ConversationType::Group),
511            false,
512            consent_states,
513        )
514        .await
515        .unwrap();
516
517        old.update_consent_state(ConsentState::Allowed).unwrap();
518        let denied = alix.create_group(None, None).unwrap();
519        denied.update_consent_state(ConsentState::Denied).unwrap();
520        let allowed = alix.create_group(None, None).unwrap();
521        let expected = match expected_count {
522            2 => vec![denied.group_id, allowed.group_id],
523            1 => vec![allowed.group_id],
524            _ => Vec::new(),
525        };
526        for group_id in expected {
527            let observed = timeout(Duration::from_secs(5), stream.next())
528                .await
529                .unwrap()
530                .unwrap()
531                .unwrap();
532            assert_eq!(observed.group_id, group_id);
533        }
534        assert!(
535            timeout(Duration::from_millis(100), stream.next())
536                .await
537                .is_err()
538        );
539    }
540
541    #[xmtp_common::timeout(std::time::Duration::from_secs(5))]
542    #[rstest::rstest]
543    #[xmtp_common::test]
544    async fn test_add_remove_re_add() {
545        tester!(alix);
546        tester!(bo);
547
548        let alix_group = alix
549            .create_group_with_members(&[bo.inbox_id().to_string()], None, None)
550            .await
551            .unwrap();
552
553        alix_group.remove_members(&[bo.inbox_id()]).await.unwrap();
554        bo.sync_welcomes().await.unwrap();
555        let stream = bo
556            .stream_conversations(Some(ConversationType::Group), false)
557            .await
558            .unwrap();
559        futures::pin_mut!(stream);
560        alix_group
561            .add_members(&[bo.inbox_id().to_string()])
562            .await
563            .unwrap();
564
565        let group_result = stream.next().await.unwrap();
566        if let Err(error) = group_result {
567            panic!("Error streaming group: {:?}", error);
568        }
569    }
570
571    #[xmtp_common::timeout(std::time::Duration::from_secs(15))]
572    #[rstest::rstest]
573    #[xmtp_common::test]
574    async fn test_duplicate_dm_not_streamed() {
575        use xmtp_cryptography::utils::generate_local_wallet;
576
577        let client1 = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
578        let client2 = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
579
580        let mut stream = client1.stream_conversations(None, false).await.unwrap();
581
582        // First DM - should stream
583        let dm1 = client1
584            .find_or_create_dm(client2.inbox_id().to_string(), None)
585            .await
586            .unwrap();
587
588        let streamed_dm1 = stream.next().await.unwrap();
589        assert!(streamed_dm1.is_ok());
590        assert_eq!(streamed_dm1.unwrap().group_id, dm1.group_id);
591
592        // Create a second DM with same participants — triggers duplicate logic.
593        // client2 has never synced welcomes, so its database holds no DM for
594        // this pair and it must build a second group rather than find dm1.
595        let dm2 = client2
596            .find_or_create_dm(client1.inbox_id().to_string(), None)
597            .await
598            .unwrap();
599
600        // Make sure it's actually a new group
601        assert_ne!(dm1.group_id, dm2.group_id);
602
603        // It should NOT appear in the stream. Wait longer than the stream's
604        // database poll interval, otherwise the stream is merely being given
605        // too little time to emit and a broken filter would still pass here.
606        let wait = xmtp_configuration::ACTIVE_DATABASE_POLL_INTERVAL * 3;
607        let result = xmtp_common::time::timeout(wait, stream.next()).await;
608        assert!(result.is_err(), "Duplicate DM was unexpectedly streamed");
609    }
610
611    #[xmtp_common::timeout(std::time::Duration::from_secs(120))]
612    #[rstest::rstest]
613    #[case::five_dms(5)]
614    #[case::onehundred_dms(100)]
615    #[xmtp_common::test]
616    #[awt]
617
618    async fn test_many_concurrent_dm_invites(#[future] alix: ClientTester, #[case] dms: usize) {
619        let alix_inbox_id = Arc::new(alix.inbox_id().to_string());
620        let mut clients = vec![];
621        for _ in 0..dms {
622            let client =
623                Arc::new(ClientBuilder::new_test_client_vanilla(&generate_local_wallet()).await);
624            clients.push(client);
625        }
626
627        let stream = alix.stream_all_messages(None, None).await.unwrap();
628        for client in clients.iter().take(dms) {
629            xmtp_common::task::spawn({
630                let id = alix_inbox_id.clone();
631                let c = client.clone();
632                async move {
633                    xmtp_common::time::sleep(std::time::Duration::from_millis(100)).await;
634                    let dm = c.find_or_create_dm(id.as_ref(), None).await?;
635                    dm.send_message(b"hi", SendMessageOpts::default()).await?;
636                    Ok::<_, crate::client::ClientError>(())
637                }
638            });
639        }
640        futures::pin_mut!(stream);
641        for _ in 0..dms {
642            let _welcome = stream.next().await;
643        }
644    }
645}
xmtp_mls/subscriptions/stream_conversations.rs

xmtp_mls/subscriptions/
stream_conversations.rs
