From 3fc5713e6138be0f1b9a46ee25a60b7d6c5736be Mon Sep 17 00:00:00 2001 From: Orual Date: Fri, 17 Oct 2025 00:42:14 -0400 Subject: [PATCH] raw and loosely typed data decode options --- .../jacquard-common/src/xrpc/subscription.rs | 161 ++++++++++++++---- 1 file changed, 132 insertions(+), 29 deletions(-) diff --git a/crates/jacquard-common/src/xrpc/subscription.rs b/crates/jacquard-common/src/xrpc/subscription.rs index 66b46b4c0..0085cd2a7 100644 --- a/crates/jacquard-common/src/xrpc/subscription.rs +++ b/crates/jacquard-common/src/xrpc/subscription.rs @@ -3,6 +3,7 @@ //! This module defines traits and types for typed WebSocket subscriptions, //! mirroring the request/response pattern used for HTTP XRPC endpoints. +use n0_future::stream::Boxed; use serde::{Deserialize, Serialize}; use std::error::Error; use std::future::Future; @@ -10,8 +11,8 @@ use std::marker::PhantomData; use url::Url; use crate::stream::StreamError; -use crate::websocket::{WebSocketClient, WebSocketConnection}; -use crate::{CowStr, IntoStatic}; +use crate::websocket::{WebSocketClient, WebSocketConnection, WsSink, WsStream}; +use crate::{CowStr, Data, IntoStatic, RawData, WsMessage}; /// Encoding format for subscription messages #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -173,8 +174,8 @@ impl SubscriptionStream { pub fn into_stream( self, ) -> ( - crate::websocket::WsSink, - n0_future::stream::Boxed, StreamError>>, + WsSink, + Boxed, StreamError>>, ) where for<'a> StreamMessage<'a, S>: IntoStatic>, @@ -183,13 +184,122 @@ impl SubscriptionStream { let (tx, rx) = self.connection.split(); - let stream: n0_future::stream::Boxed<_> = match S::ENCODING { - MessageEncoding::Json => { - Box::pin(rx.into_inner().filter_map(|msg| decode_json_msg::(msg))) - } - MessageEncoding::DagCbor => { - Box::pin(rx.into_inner().filter_map(|msg| decode_cbor_msg::(msg))) - } + let stream = match S::ENCODING { + MessageEncoding::Json => rx + .into_inner() + .filter_map(|msg| decode_json_msg::(msg)) + .boxed(), + MessageEncoding::DagCbor => rx + .into_inner() + .filter_map(|msg| decode_cbor_msg::(msg)) + .boxed(), + }; + + (tx, stream) + } + + /// Converts the subscription into a stream of raw atproto data. + pub fn into_raw_data_stream(self) -> (WsSink, Boxed, StreamError>>) { + use n0_future::StreamExt as _; + + let (tx, rx) = self.connection.split(); + + fn parse_msg<'a>(bytes: &'a [u8]) -> Result, serde_json::Error> { + serde_json::from_slice(bytes) + } + fn parse_cbor<'a>( + bytes: &'a [u8], + ) -> Result, serde_ipld_dagcbor::DecodeError> + { + serde_ipld_dagcbor::from_slice(bytes) + } + + let stream = match S::ENCODING { + MessageEncoding::Json => rx + .into_inner() + .filter_map(|msg_result| match msg_result { + Ok(WsMessage::Text(text)) => Some( + parse_msg(text.as_ref()) + .map(|v| v.into_static()) + .map_err(StreamError::decode), + ), + Ok(WsMessage::Binary(bytes)) => Some( + parse_msg(&bytes) + .map(|v| v.into_static()) + .map_err(StreamError::decode), + ), + Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())), + Err(e) => Some(Err(e)), + }) + .boxed(), + MessageEncoding::DagCbor => rx + .into_inner() + .filter_map(|msg_result| match msg_result { + Ok(WsMessage::Binary(bytes)) => Some( + parse_cbor(&bytes) + .map(|v| v.into_static()) + .map_err(|e| StreamError::decode(crate::error::DecodeError::from(e))), + ), + Ok(WsMessage::Text(_)) => Some(Err(StreamError::wrong_message_format( + "expected binary frame for CBOR, got text", + ))), + Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())), + Err(e) => Some(Err(e)), + }) + .boxed(), + }; + + (tx, stream) + } + + /// Converts the subscription into a stream of loosely-typed atproto data. + pub fn into_data_stream(self) -> (WsSink, Boxed, StreamError>>) { + use n0_future::StreamExt as _; + + let (tx, rx) = self.connection.split(); + + fn parse_msg<'a>(bytes: &'a [u8]) -> Result, serde_json::Error> { + serde_json::from_slice(bytes) + } + fn parse_cbor<'a>( + bytes: &'a [u8], + ) -> Result, serde_ipld_dagcbor::DecodeError> { + serde_ipld_dagcbor::from_slice(bytes) + } + + let stream = match S::ENCODING { + MessageEncoding::Json => rx + .into_inner() + .filter_map(|msg_result| match msg_result { + Ok(WsMessage::Text(text)) => Some( + parse_msg(text.as_ref()) + .map(|v| v.into_static()) + .map_err(StreamError::decode), + ), + Ok(WsMessage::Binary(bytes)) => Some( + parse_msg(&bytes) + .map(|v| v.into_static()) + .map_err(StreamError::decode), + ), + Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())), + Err(e) => Some(Err(e)), + }) + .boxed(), + MessageEncoding::DagCbor => rx + .into_inner() + .filter_map(|msg_result| match msg_result { + Ok(WsMessage::Binary(bytes)) => Some( + parse_cbor(&bytes) + .map(|v| v.into_static()) + .map_err(|e| StreamError::decode(crate::error::DecodeError::from(e))), + ), + Ok(WsMessage::Text(_)) => Some(Err(StreamError::wrong_message_format( + "expected binary frame for CBOR, got text", + ))), + Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())), + Err(e) => Some(Err(e)), + }) + .boxed(), }; (tx, stream) @@ -205,35 +315,28 @@ impl SubscriptionStream { /// Replaces the internal WebSocket stream with one copy and returns a typed decoded /// stream. Both streams receive all messages. Useful for observing raw messages /// while also processing typed messages. - pub fn tee( - &mut self, - ) -> n0_future::stream::Boxed, StreamError>> + pub fn tee(&mut self) -> Boxed, StreamError>> where for<'a> StreamMessage<'a, S>: IntoStatic>, { use n0_future::StreamExt as _; let rx = self.connection.receiver_mut(); - let (raw_rx, typed_rx_source) = std::mem::replace( - rx, - crate::websocket::WsStream::new(futures::stream::empty()), - ) - .tee(); + let (raw_rx, typed_rx_source) = + std::mem::replace(rx, WsStream::new(n0_future::stream::empty())).tee(); // Put the raw stream back *rx = raw_rx; match S::ENCODING { - MessageEncoding::Json => Box::pin( - typed_rx_source - .into_inner() - .filter_map(|msg| decode_json_msg::(msg)), - ), - MessageEncoding::DagCbor => Box::pin( - typed_rx_source - .into_inner() - .filter_map(|msg| decode_cbor_msg::(msg)), - ), + MessageEncoding::Json => typed_rx_source + .into_inner() + .filter_map(|msg| decode_json_msg::(msg)) + .boxed(), + MessageEncoding::DagCbor => typed_rx_source + .into_inner() + .filter_map(|msg| decode_cbor_msg::(msg)) + .boxed(), } } } -- 2.51.2