xmtp_mls/subscriptions/

mod.rs

use futures::{Stream, StreamExt};
use prost::Message;
use std::sync::Arc;
use tokio::sync::oneshot;
use xmtp_common::snippet::Snippet;
use xmtp_proto::backend_v1::ServerEnvelope;

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

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

const DEFAULT_STREAM_CONSENT_STATES: [ConsentState; 2] =
    [ConsentState::Allowed, ConsentState::Unknown];

// 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 through the shared receiver.
pub mod barrier;
pub mod catch_up;
pub(crate) mod connection_state;
pub mod incoming;
#[doc(hidden)]
pub mod internal;
pub mod local_delivery;
pub mod message_reader;
pub(crate) mod policy;
pub(crate) mod recovery;
mod stream_all;
pub 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, DeletedBy},
};
use internal::InternalEvent;
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,
};
use xmtp_events::{EventFilter, Subscription};

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
    }
}

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 Subscription<InternalEvent> {
    #[instrument(level = "trace", skip_all)]
    fn stream_consent_updates(self) -> impl Stream<Item = Result<Vec<StoredConsentRecord>>> {
        futures::stream::unfold(self, |subscription| async move {
            loop {
                let item = subscription.next().await?;
                if let Some(xmtp_events::ClientEvent::Lagged(lagged)) = &item.client {
                    tracing::warn!(discarded = lagged.discarded, "legacy consent stream lagged");
                }
                if let Some(InternalEvent::PreferencesChanged { updates, .. }) = item.internal {
                    let records = updates
                        .into_iter()
                        .filter_map(|update| match update {
                            PreferenceUpdate::Consent(record) => Some(record),
                            _ => None,
                        })
                        .collect();
                    return Some((Ok(records), subscription));
                }
            }
        })
    }

    #[instrument(level = "trace", skip_all)]
    fn stream_preference_updates(self) -> impl Stream<Item = Result<Vec<PreferenceUpdate>>> {
        futures::stream::unfold(self, |subscription| async move {
            loop {
                let item = subscription.next().await?;
                if let Some(xmtp_events::ClientEvent::Lagged(lagged)) = &item.client {
                    tracing::warn!(
                        discarded = lagged.discarded,
                        "legacy preference stream lagged"
                    );
                }
                if let Some(InternalEvent::PreferencesChanged { updates, .. }) = item.internal {
                    return Some((Ok(updates), subscription));
                }
            }
        })
    }

    #[instrument(level = "trace", skip_all)]
    fn stream_message_deletions(self) -> impl Stream<Item = Result<DecodedMessage>> {
        futures::stream::unfold(self, |subscription| async move {
            loop {
                let item = subscription.next().await?;
                if let Some(xmtp_events::ClientEvent::Lagged(lagged)) = &item.client {
                    tracing::warn!(
                        discarded = lagged.discarded,
                        "legacy deletion stream lagged"
                    );
                }
                let decoded = match item.internal {
                    // An item can wait in the broadcast buffer, for example
                    // behind a slow callback, until after its message expired.
                    // Check the expiry again at delivery.
                    Some(InternalEvent::MessagesDeleted(deleted)) => {
                        let now = xmtp_common::time::now_ns();
                        deletion_items(deleted, |m| m.expire_at_ns.is_some_and(|at| at <= now))
                    }
                    Some(InternalEvent::MessagesExpired(deleted)) => {
                        deletion_items(deleted, |_| true)
                    }
                    _ => continue,
                };
                return Some((futures::stream::iter(decoded), subscription));
            }
        })
        .flatten()
    }
}

/// Deletion stream items for `deleted`. A row for which `is_expired` holds
/// is the deleted-message placeholder, with no body and no fallback. Other
/// rows keep their body; content that does not decode stays in the body as
/// undecodable content, so every row produces an item.
// implements: META-051
fn deletion_items(
    deleted: internal::DeletedMessages,
    is_expired: impl Fn(&StoredGroupMessage) -> bool,
) -> Vec<Result<DecodedMessage>> {
    let internal::DeletedMessages {
        messages,
        deleted_by_inbox_id,
    } = deleted;
    messages
        .into_iter()
        .map(|message| {
            if !is_expired(&message) {
                return Ok(DecodedMessage::from(message));
            }
            // With no actor, `Sender`; see `DecodedMessage::expired`.
            let deleted_by = deleted_by_inbox_id
                .as_deref()
                .map_or(DeletedBy::Sender, |deleter| {
                    DeletedBy::new(deleter, &message.sender_inbox_id)
                });
            Ok(DecodedMessage::expired(message, deleted_by))
        })
        .collect()
}

#[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 connection is blocked by a fatal configuration failure.
    ///
    /// Either this database is bound to a different deployment or the
    /// deployment now requires a newer client. 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(|error| {
                self.context
                    .server_configuration()
                    .blocked_or(GroupError::from(error))
            })?;
        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 local event subscription, 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 event 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().as_slice().snippet(),
            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().as_slice().snippet(),
            conversation_type = ?conversation_type,
            "stream all messages"
        );

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

    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 {
                // Cancellation can carry a blocked connection cause,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.context.events().subscribe(
                    EventFilter::default().with_internal(|event| {
                        matches!(event, InternalEvent::PreferencesChanged { .. })
                    }),
                    Some(1024),
                );
                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 {
                // Cancellation can carry a blocked connection cause,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.context.events().subscribe(
                    EventFilter::default().with_internal(|event| {
                        matches!(event, InternalEvent::PreferencesChanged { .. })
                    }),
                    Some(1024),
                );
                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 {
                // Cancellation can carry a blocked connection cause,
                // and this stream closes with it rather than silently.
                let cancel = watchdog::StreamCancel::new(&client.context);
                let receiver = client.context.events().subscribe(
                    EventFilter::default().with_internal(InternalEvent::is_message_deletion),
                    Some(1024),
                );
                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().as_slice().snippet(),
            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");
    }
}