1use 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#[derive(Default)]
36struct KnownConversations {
37 welcomes: HashMap<GroupId, Option<i64>>,
39 dms: HashMap<String, GroupId>,
41}
42
43impl KnownConversations {
44 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 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
70pub 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 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 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 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 #[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 #[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 #[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 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 let mut groups = Vec::new();
385 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 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 #[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 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 let dm2 = client2
543 .find_or_create_dm(client1.inbox_id().to_string(), None)
544 .await
545 .unwrap();
546
547 assert_ne!(dm1.group_id, dm2.group_id);
549
550 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}