xmtp_mls/groups/
subscriptions.rs

1use super::MlsGroup;
2use crate::{
3    context::XmtpSharedContext,
4    subscriptions::{
5        Result, stream_messages::StreamGroupMessages, watchdog::spawn_watchdog_stream,
6    },
7};
8use futures::Stream;
9use prost::Message;
10use xmtp_proto::backend_v1::ServerEnvelope;
11
12use xmtp_common::MaybeSend;
13use xmtp_common::StreamHandle;
14use xmtp_db::group_message::StoredGroupMessage;
15use xmtp_proto::api_client::XmtpMlsStreams;
16use xmtp_proto::types::GroupId;
17
18impl<Context> MlsGroup<Context>
19where
20    Context: XmtpSharedContext + 'static,
21{
22    /// Use a push envelope only as a target for ordered receipt and processing.
23    pub async fn process_streamed_group_message(
24        &self,
25        envelope_bytes: Vec<u8>,
26    ) -> Result<Vec<StoredGroupMessage>> {
27        use xmtp_db::prelude::*;
28        let wire = ServerEnvelope::decode(envelope_bytes.as_slice())?;
29        let meta = wire
30            .meta
31            .as_ref()
32            .ok_or(xmtp_api::ApiError::InvalidResponse("group metadata"))?;
33        let (topic, cursor, _) = xmtp_api_backend::envelope::metadata(
34            meta,
35            xmtp_proto::types::TopicKind::GroupMessagesV1,
36        )?;
37        if topic != xmtp_proto::types::Topic::new_group_message(self.group_id) {
38            return Err(xmtp_api::ApiError::InvalidResponse("group topic").into());
39        }
40        crate::subscriptions::barrier::wait_through(&self.context, [(topic, cursor)].into(), None)
41            .await
42            .map_err(super::GroupError::from)?;
43        Ok(self
44            .context
45            .db()
46            .get_group_message_by_cursor(self.group_id, cursor)?
47            .into_iter()
48            .collect())
49    }
50
51    #[tracing::instrument(err, skip_all, fields(operation = "stream.stream_group_messages"))]
52    pub async fn stream<'a>(
53        &'a self,
54    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + use<'a, Context>>
55    where
56        Context::ApiClient: XmtpMlsStreams + 'a,
57    {
58        StreamGroupMessages::new(&self.context, vec![self.group_id]).await
59    }
60
61    /// create a stream that is not attached to any lifetime
62    #[tracing::instrument(
63        err,
64        skip_all,
65        fields(operation = "stream.stream_group_messages_owned")
66    )]
67    pub async fn stream_owned(
68        &self,
69    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + 'static>
70    where
71        Context: 'static,
72        Context::ApiClient: XmtpMlsStreams + 'static,
73        Context::Db: 'static,
74    {
75        StreamGroupMessages::new_owned(self.context.clone(), vec![self.group_id]).await
76    }
77
78    pub fn stream_with_callback(
79        context: Context,
80        group_id: GroupId,
81        callback: impl FnMut(Result<StoredGroupMessage>) + MaybeSend + 'static,
82        on_close: impl FnOnce() + MaybeSend + 'static,
83    ) -> impl StreamHandle<StreamOutput = Result<()>>
84    where
85        Context: 'static,
86        Context::ApiClient: XmtpMlsStreams + 'static,
87    {
88        stream_messages_with_callback(
89            context.clone(),
90            vec![group_id].into_iter(),
91            callback,
92            on_close,
93        )
94    }
95}
96
97/// Deliver stored messages for these groups and share ordered network receipt.
98pub(crate) fn stream_messages_with_callback<Context>(
99    context: Context,
100    active_conversations: impl Iterator<Item = GroupId> + MaybeSend + 'static,
101    callback: impl FnMut(Result<StoredGroupMessage>) + MaybeSend + 'static,
102    on_close: impl FnOnce() + MaybeSend + 'static,
103) -> impl StreamHandle<StreamOutput = Result<()>>
104where
105    Context: XmtpSharedContext + 'static,
106    Context::ApiClient: XmtpMlsStreams + 'static,
107    Context::Db: 'static,
108{
109    let cancel = crate::subscriptions::watchdog::StreamCancel::new(&context);
110    let groups: Vec<GroupId> = active_conversations.collect();
111    // Reopening reads saved D. A dropped, unacknowledged item remains available.
112    spawn_watchdog_stream(
113        cancel,
114        "stream_messages",
115        move || {
116            let context = context.clone();
117            let groups = groups.clone();
118            async move { StreamGroupMessages::new_owned(context, groups).await }
119        },
120        callback,
121        on_close,
122    )
123}
124
125#[cfg(test)]
126pub(crate) mod tests {
127    use crate::context::XmtpSharedContext;
128    use crate::groups::send_message_opts::SendMessageOpts;
129    use futures::StreamExt;
130    use std::sync::Arc;
131
132    use crate::builder::ClientBuilder;
133    use prost::Message as ProstMessage;
134    use std::time::Duration;
135    use xmtp_cryptography::utils::generate_local_wallet;
136    use xmtp_db::group_message::GroupMessageKind;
137
138    #[xmtp_common::timeout(Duration::from_secs(10))]
139    #[rstest::rstest]
140    #[xmtp_common::test(flavor = "current_thread")]
141    async fn test_subscribe_messages() {
142        let amal = ClientBuilder::new_test_client(&generate_local_wallet()).await;
143        let bola = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
144
145        let amal_group = amal.create_group(None, None).unwrap();
146        // Add bola
147        amal_group.add_members(&[bola.inbox_id()]).await.unwrap();
148
149        // Get bola's version of the same group
150        let bola_groups = bola.sync_welcomes().await.unwrap();
151        let bola_group = bola_groups.first().unwrap();
152        bola_group.receive().await.unwrap();
153        let retained = bola_group.find_messages(&Default::default()).unwrap();
154
155        let stream = bola_group.stream().await.unwrap();
156        futures::pin_mut!(stream);
157        for expected in retained {
158            assert_eq!(stream.next().await.unwrap().unwrap(), expected);
159        }
160
161        amal_group
162            .send_message("hello".as_bytes(), SendMessageOpts::default())
163            .await
164            .unwrap();
165        let first_val = stream.next().await.unwrap().unwrap();
166        assert_eq!(first_val.decrypted_message_bytes, "hello".as_bytes());
167
168        amal_group
169            .send_message("goodbye".as_bytes(), SendMessageOpts::default())
170            .await
171            .unwrap();
172        let second_val = stream.next().await.unwrap().unwrap();
173        assert_eq!(second_val.decrypted_message_bytes, "goodbye".as_bytes());
174    }
175
176    // TODO: THIS TESTS ALSO LOSES MESSAGES
177    #[xmtp_common::timeout(Duration::from_secs(10))]
178    #[rstest::rstest]
179    #[xmtp_common::test(flavor = "multi_thread")]
180    #[cfg_attr(target_arch = "wasm32", ignore)]
181    async fn test_subscribe_multiple() {
182        let amal = Arc::new(ClientBuilder::new_test_client_vanilla(&generate_local_wallet()).await);
183        let group = amal.create_group(None, None).unwrap();
184
185        let stream = group.stream().await.unwrap();
186        futures::pin_mut!(stream);
187
188        for i in 0..10 {
189            group
190                .send_message(
191                    format!("hello {}", i).as_bytes(),
192                    SendMessageOpts::default(),
193                )
194                .await
195                .unwrap();
196        }
197
198        // Limit the stream so that it closes after 10 messages
199        let limited_stream = stream.take(10);
200        let values = limited_stream.collect::<Vec<_>>().await;
201        assert_eq!(values.len(), 10);
202        for value in values {
203            assert!(
204                value
205                    .unwrap()
206                    .decrypted_message_bytes
207                    .starts_with("hello".as_bytes())
208            );
209        }
210    }
211
212    #[xmtp_common::timeout(Duration::from_secs(5))]
213    #[rstest::rstest]
214    #[xmtp_common::test]
215    async fn test_subscribe_membership_changes() {
216        let amal = Arc::new(ClientBuilder::new_test_client(&generate_local_wallet()).await);
217        let bola = ClientBuilder::new_test_client(&generate_local_wallet()).await;
218
219        let amal_group = amal.create_group(None, None).unwrap();
220
221        let stream = amal_group.stream().await.unwrap();
222        futures::pin_mut!(stream);
223
224        amal_group.add_members(&[bola.inbox_id()]).await.unwrap();
225
226        let first_val = stream.next().await.unwrap().unwrap();
227        assert_eq!(first_val.kind, GroupMessageKind::MembershipChange);
228
229        amal_group
230            .send_message("hello".as_bytes(), SendMessageOpts::default())
231            .await
232            .unwrap();
233        let second_val = stream.next().await.unwrap().unwrap();
234        assert_eq!(second_val.decrypted_message_bytes, "hello".as_bytes());
235    }
236
237    #[xmtp_common::test(unwrap_try = true)]
238    async fn test_process_streamed_group_message() {
239        crate::tester!(alix);
240        crate::tester!(bo);
241        let group = alix.create_group(None, None)?;
242        group.add_members(&[bo.inbox_id()]).await?;
243        let bo_groups = bo.sync_welcomes().await?;
244        let bo_group = bo_groups.first().unwrap();
245        group
246            .send_message(b"test message", SendMessageOpts::default())
247            .await?;
248        let envelopes = alix
249            .context
250            .api()
251            .query_all(
252                std::collections::HashMap::from([(
253                    xmtp_proto::types::Topic::new_group_message(group.group_id),
254                    xmtp_proto::types::Cursor(0),
255                )]),
256                alix.context.api().limits().max_query_limit as u32,
257            )
258            .await?;
259        let envelope = envelopes.last().unwrap();
260        let result = bo_group
261            .process_streamed_group_message(envelope.encode_to_vec())
262            .await;
263        assert!(
264            result.is_ok(),
265            "Backend processing must succeed: {result:?}"
266        );
267        let messages = result?;
268        assert!(!messages.is_empty());
269        assert_eq!(messages.len(), 1);
270        assert_eq!(messages[0].decrypted_message_bytes, b"test message");
271    }
272}
xmtp_mls/groups/subscriptions.rs

xmtp_mls/groups/
subscriptions.rs
