Skip to main content

xmtp_mls/subscriptions/
stream_conversations.rs

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