xmtp_mls/subscriptions/

mod.rs

use futures::{Stream, StreamExt};
use prost::Message;
use std::sync::Arc;
use tokio::sync::{broadcast, oneshot};
use tokio_stream::wrappers::BroadcastStream;
use xmtp_proto::backend_v1::ServerEnvelope;
use xmtp_proto::types::GroupId;

use tracing::instrument;
use xmtp_db::prelude::*;
use xmtp_proto::api_client::XmtpMlsStreams;

use stream_all::StreamAllMessages;
use stream_conversations::StreamConversations;

// Live backend tests require native full-duplex HTTP/2.
#[cfg(all(test, not(target_arch = "wasm32")))]
mod bidi_tests;
// Randomized delivery fuzz over the live node (same gating as `bidi_tests`).
#[cfg(all(test, not(target_arch = "wasm32")))]
mod bidi_fuzz_tests;
// One-shot bounded catch-up over the bidi wire (native-only, like the
// connection it rides).
pub mod barrier;
#[cfg(not(target_arch = "wasm32"))]
pub mod catch_up;
pub mod incoming;
pub mod local_delivery;
pub mod message_reader;
pub(crate) mod policy;
mod stream_all;
mod stream_conversations;
pub mod stream_failure;
pub mod stream_messages;
// Live integration tests for the router (v3 wire; same gating rationale as
// `bidi_tests` above).
#[cfg(all(test, not(target_arch = "wasm32")))]
mod delivery_integration_tests;
// Native callback adapters over the shared backend transport.
#[cfg(not(target_arch = "wasm32"))]
pub mod router_callbacks;
#[cfg(all(test, not(target_arch = "wasm32")))]
mod router_callbacks_tests;
pub(crate) mod watchdog;

use crate::messages::enrichment::EnrichMessageError;
#[cfg(any(test, feature = "test-utils"))]
use crate::subscriptions::stream_messages::stream_stats::{StreamStatsWrapper, StreamWithStats};

use crate::worker::device_sync::preference_sync::PreferenceUpdate;
use crate::{
    Client,
    context::XmtpSharedContext,
    groups::{GroupError, MlsGroup, mls_sync::GroupMessageProcessingError},
    messages::decoded_message::DecodedMessage,
};
use thiserror::Error;
use xmtp_common::{ErrorCode, MaybeSend, RetryableError, StreamHandle, retryable};
use xmtp_db::{
    NotFound, StorageError,
    consent_record::{ConsentState, StoredConsentRecord},
    group::ConversationType,
    group_message::StoredGroupMessage,
};

pub(crate) type Result<T> = std::result::Result<T, SubscribeError>;

#[derive(Debug, Error)]
pub enum LocalEventError {
    #[error("Unable to send event: {0}")]
    Send(String),
}

impl RetryableError for LocalEventError {
    fn is_retryable(&self) -> bool {
        true
    }
}

/// Events local to this client
/// are broadcast across all senders/receivers of streams
#[derive(Debug, Clone)]
pub enum LocalEvents {
    // a new group was created
    NewGroup(GroupId),
    /// A committed local message can be read. This is a hint, not a delivery event.
    MessagesStored,
    PreferencesChanged(Vec<PreferenceUpdate>),
    // a message was deleted (contains the decoded message that was deleted)
    MsgsDeleted(Vec<StoredGroupMessage>),
}

#[derive(Clone)]
pub enum SyncWorkerEvent {
    NewSyncGroupFromWelcome(Vec<u8>),
    NewSyncGroupMsg,
    // The sync worker will auto-sync these with other devices.
    SyncPreferences(Vec<PreferenceUpdate>),
    CycleHMAC,
    Tick,
}

impl std::fmt::Debug for SyncWorkerEvent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NewSyncGroupFromWelcome(arg0) => f
                .debug_tuple("NewSyncGroupFromWelcome")
                .field(&hex::encode(arg0))
                .finish(),
            Self::NewSyncGroupMsg => write!(f, "NewSyncGroupMsg"),
            Self::SyncPreferences(arg0) => f.debug_tuple("SyncPreferences").field(arg0).finish(),
            Self::CycleHMAC => write!(f, "CycleHMAC"),
            Self::Tick => write!(f, "Tick"),
        }
    }
}

impl LocalEvents {
    fn consent_filter(self) -> Option<Vec<StoredConsentRecord>> {
        match self {
            Self::PreferencesChanged(updates) => {
                let updates = updates
                    .into_iter()
                    .filter_map(|pu| match pu {
                        PreferenceUpdate::Consent(cr) => Some(cr),
                        _ => None,
                    })
                    .collect();
                Some(updates)
            }

            _ => None,
        }
    }

    fn preference_filter(self) -> Option<Vec<PreferenceUpdate>> {
        match self {
            Self::PreferencesChanged(updates) => Some(updates),
            _ => None,
        }
    }

    fn message_deletion_filter(self) -> Option<Vec<StoredGroupMessage>> {
        match self {
            Self::MsgsDeleted(msgs) => Some(msgs),
            _ => None,
        }
    }
}

pub(crate) trait StreamMessages {
    fn stream_consent_updates(self) -> impl Stream<Item = Result<Vec<StoredConsentRecord>>>;
    fn stream_preference_updates(self) -> impl Stream<Item = Result<Vec<PreferenceUpdate>>>;
    fn stream_message_deletions(self) -> impl Stream<Item = Result<DecodedMessage>>;
}

impl StreamMessages for broadcast::Receiver<LocalEvents> {
    #[instrument(level = "trace", skip_all)]
    fn stream_consent_updates(self) -> impl Stream<Item = Result<Vec<StoredConsentRecord>>> {
        BroadcastStream::new(self).filter_map(|event| async {
            xmtp_common::optify!(event, "Missed message due to event queue lag")
                .and_then(LocalEvents::consent_filter)
                .map(Result::Ok)
        })
    }

    #[instrument(level = "trace", skip_all)]
    fn stream_preference_updates(self) -> impl Stream<Item = Result<Vec<PreferenceUpdate>>> {
        BroadcastStream::new(self).filter_map(|event| async {
            xmtp_common::optify!(event, "Missed message due to event queue lag")
                .and_then(LocalEvents::preference_filter)
                .map(Result::Ok)
        })
    }

    #[instrument(level = "trace", skip_all)]
    fn stream_message_deletions(self) -> impl Stream<Item = Result<DecodedMessage>> {
        BroadcastStream::new(self)
            .filter_map(|event| async {
                xmtp_common::optify!(event, "Missed message due to event queue lag")
                    .and_then(LocalEvents::message_deletion_filter)
                    .map(futures::stream::iter)
            })
            .flatten()
            // let caller handle any potential decode failures
            // this should be rare since the message already in db
            .map(|m| DecodedMessage::try_from(m).map_err(Into::into))
    }
}

#[derive(thiserror::Error, Debug, ErrorCode)]
pub enum SubscribeError {
    /// Local message delivery failed. May be retryable for a storage failure.
    #[error(transparent)]
    #[error_code(inherit)]
    LocalDelivery(#[from] local_delivery::LocalDeliveryError),
    /// The shared native transport failed. May be retryable.
    #[cfg(not(target_arch = "wasm32"))]
    #[error(transparent)]
    Transport(#[from] xmtp_api_backend::TransportError),
    /// Group error.
    ///
    /// Group operation failed during subscription. May be retryable.
    #[error(transparent)]
    Group(#[from] Box<GroupError>),
    /// Not found.
    ///
    /// Subscribed resource not found. Retryable.
    #[error(transparent)]
    NotFound(#[from] NotFound),
    /// Group message not found.
    ///
    /// Expected message missing from database. Retryable.
    // TODO: Add this to `NotFound`
    #[error("group message expected in database but is missing")]
    GroupMessageNotFound,
    /// Receive group error.
    ///
    /// Processing streamed group message failed. May be retryable.
    #[error("processing group message in stream: {0}")]
    ReceiveGroup(#[from] Box<GroupMessageProcessingError>),
    /// Storage error.
    ///
    /// Database operation failed. May be retryable.
    #[error(transparent)]
    Storage(#[from] StorageError),
    /// Decode error.
    ///
    /// Protobuf decoding failed. Not retryable.
    #[error(transparent)]
    Decode(#[from] prost::DecodeError),
    /// API client error.
    ///
    /// Network request failed. Retryable.
    #[error(transparent)]
    ApiClient(#[from] xmtp_api::ApiError),
    /// Boxed error.
    ///
    /// Wrapped dynamic error. May be retryable.
    #[error("{0}")]
    BoxError(Box<dyn RetryableError>),
    /// Database connection error.
    ///
    /// Database connection failed. Retryable.
    #[error(transparent)]
    Db(#[from] xmtp_db::ConnectionError),
    /// Conversion error.
    ///
    /// Proto conversion failed. Not retryable.
    #[error(transparent)]
    Conversion(#[from] xmtp_proto::ConversionError),
    /// Envelope error.
    ///
    /// Invalid backend envelope. Not retryable.
    #[error(transparent)]
    Envelope(#[from] xmtp_api_backend::envelope::EnvelopeError),
    /// Enriched Message Error.
    #[error("error occured during subscription {0}")]
    Enriched(#[from] EnrichMessageError),
    /// The client latched a fatal configuration failure.
    ///
    /// Either this database is bound to a different deployment (CFG-051) or the
    /// deployment now requires a newer client (CFG-061). Every open stream is
    /// closed with it and every later call fails with it. Not retryable.
    #[error(transparent)]
    #[error_code(inherit)]
    Configuration(Box<crate::client::ClientError>),
    /// Stream liveness watchdog tripped.
    ///
    /// No activity arrived within the idle timeout, so the stream was terminated to force
    /// a reconnect (resuming from the persisted cursor). Retryable.
    #[error("stream went stale: no activity within the idle timeout")]
    StreamStale,
}

impl SubscribeError {
    pub fn dyn_err(other: impl RetryableError + 'static) -> Self {
        SubscribeError::BoxError(Box::new(other) as _)
    }
}

impl From<GroupError> for SubscribeError {
    fn from(value: GroupError) -> Self {
        SubscribeError::Group(Box::new(value))
    }
}

impl From<GroupMessageProcessingError> for SubscribeError {
    fn from(value: GroupMessageProcessingError) -> Self {
        SubscribeError::ReceiveGroup(Box::new(value))
    }
}

impl RetryableError for SubscribeError {
    fn is_retryable(&self) -> bool {
        use SubscribeError::*;
        match self {
            #[cfg(not(target_arch = "wasm32"))]
            Transport(error) => error.is_retryable(),
            LocalDelivery(e) => retryable!(e),
            Group(e) => retryable!(e),
            GroupMessageNotFound => true,
            ReceiveGroup(e) => retryable!(e),
            Storage(e) => retryable!(e),
            Decode(_) => false,
            NotFound(e) => retryable!(e),
            ApiClient(e) => retryable!(e),
            BoxError(e) => retryable!(e),
            Db(c) => retryable!(c),
            Conversion(c) => retryable!(c),
            Envelope(c) => retryable!(c),
            Enriched(c) => retryable!(c),
            Configuration(_) => false,
            StreamStale => true,
        }
    }
}

impl crate::worker::NeedsDbReconnect for SubscribeError {
    /// Forwards a dropped-pool signal so the device-sync worker stops on
    /// disconnect. `BoxError` wraps an opaque error we can't introspect → `false`.
    fn needs_db_reconnect(&self) -> bool {
        use SubscribeError::*;
        match self {
            Group(e) => e.needs_db_reconnect(),
            Storage(e) => e.db_needs_connection(),
            LocalDelivery(local_delivery::LocalDeliveryError::Storage(e)) => {
                e.db_needs_connection()
            }
            LocalDelivery(_) => false,
            Db(c) => c.db_needs_connection(),
            #[cfg(not(target_arch = "wasm32"))]
            Transport(_) => false,
            GroupMessageNotFound | ReceiveGroup(_) | Decode(_) | NotFound(_) | ApiClient(_)
            | BoxError(_) | Conversion(_) | Envelope(_) | Enriched(_) | Configuration(_)
            | StreamStale => false,
        }
    }
}

impl<Context> Client<Context>
where
    Context: XmtpSharedContext + 'static,
{
    /// Use push metadata as a target. Fetch its complete prefix before processing.
    pub async fn process_streamed_welcome_message(
        &self,
        envelope_bytes: Vec<u8>,
    ) -> Result<Vec<MlsGroup<Context>>> {
        let wire = ServerEnvelope::decode(envelope_bytes.as_slice())?;
        let meta = wire
            .meta
            .as_ref()
            .ok_or(xmtp_api::ApiError::InvalidResponse("welcome metadata"))?;
        let (topic, cursor, _) = xmtp_api_backend::envelope::metadata(
            meta,
            xmtp_proto::types::TopicKind::WelcomeMessagesV1,
        )?;
        if topic != xmtp_proto::types::Topic::new_welcome_message(self.context.installation_id()) {
            return Err(xmtp_api::ApiError::InvalidResponse("welcome installation").into());
        }
        barrier::wait_through(&self.context, [(topic, cursor)].into(), None)
            .await
            .map_err(GroupError::from)?;
        let Some(group) = self.context.db().find_group_by_sequence_id(cursor)? else {
            return Ok(Vec::new());
        };
        Ok(vec![MlsGroup::new(
            self.context.clone(),
            group.id,
            group.dm_id,
            group.conversation_type,
            group.created_at_ns,
        )])
    }

    #[xmtp_common::span(prefix = "stream")]
    pub async fn stream_conversations(
        &self,
        conversation_type: Option<ConversationType>,
        include_duplicate_dms: bool,
    ) -> Result<impl Stream<Item = Result<MlsGroup<Context>>> + use<'_, Context>>
    where
        Context::ApiClient: XmtpMlsStreams,
    {
        StreamConversations::new(
            &self.context,
            conversation_type,
            include_duplicate_dms,
            None,
        )
        .await
    }

    /// Stream conversations but decouple the lifetime of 'self' from the stream.
    #[xmtp_common::span(prefix = "stream")]
    pub async fn stream_conversations_owned(
        &self,
        conversation_type: Option<ConversationType>,
        include_duplicate_dms: bool,
    ) -> Result<impl Stream<Item = Result<MlsGroup<Context>>> + 'static + use<Context>>
    where
        Context::ApiClient: XmtpMlsStreams,
    {
        StreamConversations::new_owned(
            self.context.clone(),
            conversation_type,
            include_duplicate_dms,
            None,
        )
        .await
    }
}

impl<Context> Client<Context>
where
    Context: XmtpSharedContext + 'static,
    Context::ApiClient: XmtpMlsStreams + 'static,
    Context::MlsStorage: 'static,
{
    pub fn stream_conversations_with_callback(
        client: Arc<Client<Context>>,
        conversation_type: Option<ConversationType>,
        convo_callback: impl FnMut(Result<MlsGroup<Context>>) + MaybeSend + 'static,
        on_close: impl FnOnce() + MaybeSend + 'static,
        include_duplicate_dms: bool,
    ) -> impl StreamHandle<StreamOutput = Result<()>> {
        let cancel = watchdog::StreamCancel::new(&client.context);
        // Re-subscribing recreates the underlying `LocalEvents` broadcast receiver, which
        // has no replay; the watchdog runner establishes the new subscription *before* its
        // reconnect wait, so the new receiver is attached while we pause. Network welcomes
        // are caught up from the persisted cursor, so the only residual gap is a *locally*
        // created group (`LocalEvents::NewGroup`) broadcast in the brief window while the new
        // subscription is being built — bounded, since the caller already holds that group.
        watchdog::spawn_watchdog_stream(
            cancel,
            "stream_conversations",
            move || {
                let client = client.clone();
                async move {
                    client
                        .stream_conversations_owned(conversation_type, include_duplicate_dms)
                        .await
                }
            },
            convo_callback,
            on_close,
        )
    }

    #[xmtp_common::span(prefix = "stream")]
    pub async fn stream_all_messages(
        &self,
        conversation_type: Option<ConversationType>,
        consent_state: Option<Vec<ConsentState>>,
    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + '_> {
        tracing::debug!(
            inbox_id = self.inbox_id(),
            installation_id = %self.context.installation_id(),
            conversation_type = ?conversation_type,
            "stream all messages"
        );

        StreamAllMessages::new(&self.context, conversation_type, consent_state).await
    }

    #[xmtp_common::span(prefix = "stream")]
    pub async fn stream_all_messages_owned(
        &self,
        conversation_type: Option<ConversationType>,
        consent_state: Option<Vec<ConsentState>>,
    ) -> Result<impl Stream<Item = Result<StoredGroupMessage>> + 'static + use<Context>> {
        tracing::debug!(
            inbox_id = self.inbox_id(),
            installation_id = %self.context.installation_id(),
            conversation_type = ?conversation_type,
            "stream all messages"
        );

        StreamAllMessages::new_owned(self.context.clone(), conversation_type, consent_state).await
    }

    pub fn stream_all_messages_with_callback(
        context: Context,
        conversation_type: Option<ConversationType>,
        consent_state: Option<Vec<ConsentState>>,
        callback: impl FnMut(Result<StoredGroupMessage>) + MaybeSend + 'static,
        on_close: impl FnOnce() + MaybeSend + 'static,
    ) -> impl StreamHandle<StreamOutput = Result<()>> {
        let cancel = watchdog::StreamCancel::new(&context);
        watchdog::spawn_watchdog_stream(
            cancel,
            "stream_all_messages",
            move || {
                let context = context.clone();
                let consent_state = consent_state.clone();
                async move {
                    StreamAllMessages::new_owned(context, conversation_type, consent_state).await
                }
            },
            callback,
            on_close,
        )
    }

    pub fn stream_consent_with_callback(
        client: Arc<Client<Context>>,
        mut callback: impl FnMut(Result<Vec<StoredConsentRecord>>) + MaybeSend + 'static,
        on_close: impl FnOnce() + MaybeSend + 'static,
    ) -> impl StreamHandle<StreamOutput = Result<()>> {
        let (tx, rx) = oneshot::channel();

        xmtp_common::spawn(
            Some(rx),
            xmtp_common::bind_task_hub(async move {
                // CFG-051 and CFG-061: cancellation can carry a latched reason,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.local_events.subscribe();
                let stream = receiver.stream_consent_updates();

                futures::pin_mut!(stream);
                let _ = tx.send(());
                let cancelled = loop {
                    tokio::select! {
                        _ = cancel.cancelled() => break true,
                        next = stream.next() => match next {
                            Some(message) => callback(message),
                            None => break false,
                        }
                    }
                };
                tracing::debug!("`stream_consent` stream ended, dropping stream");
                let result = watchdog::close_reason(&cancel, cancelled, &mut callback);
                on_close();
                result
            }),
        )
    }

    pub fn stream_preferences_with_callback(
        client: Arc<Client<Context>>,
        mut callback: impl FnMut(Result<Vec<PreferenceUpdate>>) + MaybeSend + 'static,
        on_close: impl FnOnce() + MaybeSend + 'static,
    ) -> impl StreamHandle<StreamOutput = Result<()>> {
        let (tx, rx) = oneshot::channel();

        xmtp_common::spawn(
            Some(rx),
            xmtp_common::bind_task_hub(async move {
                // CFG-051 and CFG-061: cancellation can carry a latched reason,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.local_events.subscribe();
                let stream = receiver.stream_preference_updates();

                futures::pin_mut!(stream);
                let _ = tx.send(());
                let cancelled = loop {
                    tokio::select! {
                        _ = cancel.cancelled() => break true,
                        next = stream.next() => match next {
                            Some(message) => callback(message),
                            None => break false,
                        }
                    }
                };
                tracing::debug!("`stream_preferences` stream ended, dropping stream");
                let result = watchdog::close_reason(&cancel, cancelled, &mut callback);
                on_close();
                result
            }),
        )
    }

    pub fn stream_message_deletions_with_callback(
        client: Arc<Client<Context>>,
        mut callback: impl FnMut(Result<DecodedMessage>) + MaybeSend + 'static,
        on_close: impl FnOnce() + MaybeSend + 'static,
    ) -> impl StreamHandle<StreamOutput = Result<()>> {
        let (tx, rx) = oneshot::channel();

        xmtp_common::spawn(
            Some(rx),
            xmtp_common::bind_task_hub(async move {
                // CFG-051 and CFG-061: cancellation can carry a latched reason,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.local_events.subscribe();
                let stream = receiver.stream_message_deletions();

                futures::pin_mut!(stream);
                let _ = tx.send(());
                let cancelled = loop {
                    tokio::select! {
                        _ = cancel.cancelled() => break true,
                        next = stream.next() => match next {
                            Some(message) => callback(message),
                            None => break false,
                        }
                    }
                };
                tracing::debug!("`stream_message_deletions` stream ended, dropping stream");
                let result = watchdog::close_reason(&cancel, cancelled, &mut callback);
                on_close();
                result
            }),
        )
    }
}

impl<Context> Client<Context>
where
    Context: XmtpSharedContext + 'static,
    Context::ApiClient: XmtpMlsStreams + 'static,
    Context::MlsStorage: 'static,
    <Context::ApiClient as XmtpMlsStreams>::GroupMessageStream: Unpin,
    <Context::ApiClient as XmtpMlsStreams>::WelcomeMessageStream: Unpin,
{
    #[tracing::instrument(level = "trace", skip_all)]
    #[cfg(any(test, feature = "test-utils"))]
    pub async fn stream_all_messages_owned_with_stats(
        &self,
        conversation_type: Option<ConversationType>,
        consent_state: Option<Vec<ConsentState>>,
    ) -> Result<impl StreamWithStats<Item = Result<StoredGroupMessage>> + 'static> {
        tracing::debug!(
            inbox_id = self.inbox_id(),
            installation_id = %self.context.installation_id(),
            conversation_type = ?conversation_type,
            "stream all messages"
        );

        let stream =
            StreamAllMessages::new_owned(self.context.clone(), conversation_type, consent_state)
                .await?;

        Ok(StreamStatsWrapper::new(stream))
    }
}

#[derive(Debug, Clone, Copy)]
pub enum StreamKind {
    All,
    Conversations,
    Messages,
}

#[cfg(test)]
pub(crate) mod tests {
    use crate::context::XmtpSharedContext;
    use crate::tester;

    /// A macro for asserting that a stream yields a specific decrypted message.
    ///
    /// # Example
    /// ```rust
    /// assert_msg!(stream, b"first");
    /// ```
    #[macro_export]
    macro_rules! assert_msg {
        ($stream:expr, $expected:expr) => {
            let next = $stream
                .next()
                .await
                .unwrap()
                .inspect_err(|e| tracing::error!("{}", e.to_string()))
                .unwrap();

            assert_eq!(
                String::from_utf8_lossy(next.decrypted_message_bytes.as_slice()),
                String::from_utf8_lossy($expected.as_bytes())
            );
        };
    }

    /// A macro for asserting that a stream yields a specific decrypted message.
    ///
    /// # Example
    /// ```rust
    /// assert_msg!(stream, b"first");
    /// ```
    #[macro_export]
    macro_rules! assert_msg_exists {
        ($stream:expr) => {
            assert!(
                !$stream
                    .next()
                    .await
                    .unwrap()
                    .unwrap()
                    .decrypted_message_bytes
                    .is_empty()
            );
        };
    }

    #[xmtp_common::test(flavor = "multi_thread", worker_threads = 5, unwrap_try = true)]
    async fn test_process_streamed_welcome_message() {
        use prost::Message;
        use xmtp_proto::types::{Cursor, Topic};
        tester!(alix);
        tester!(bo);
        let alix_group = alix.create_group(None, None)?;
        alix_group.add_members(&[bo.inbox_id()]).await?;
        let envelopes = alix
            .context
            .api()
            .query_all(
                [(
                    Topic::new_welcome_message(bo.context.installation_id()),
                    Cursor(0),
                )]
                .into(),
                bo.context.api().limits().max_query_limit as u32,
            )
            .await?;
        assert!(!envelopes.is_empty(), "Should have at least one welcome");
        let groups = bo
            .process_streamed_welcome_message(envelopes[0].encode_to_vec())
            .await?;
        assert_eq!(groups.len(), 1, "Should have exactly one group");
    }
}