Skip to main content

strata_sdk/
twap_stream.rs

1use futures_util::{SinkExt, StreamExt};
2use tokio::net::TcpStream;
3use tokio_tungstenite::tungstenite::Message;
4use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
5
6use super::market_stream::platform_websocket_url;
7use super::*;
8
9type PlatformSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
10
11/// One sequenced wallet-scoped TWAP progress stream for one market. The
12/// first event is the snapshot; later events are `twap_update`, heartbeat, or
13/// recovery snapshots on the same stream identity. Any identity or sequence
14/// mismatch fails closed so the caller reconnects and recovers from a fresh
15/// snapshot.
16pub struct TwapStream {
17    market_id: String,
18    wallet_address: String,
19    stream_id: Option<String>,
20    sequence: u64,
21    socket: PlatformSocket,
22}
23
24impl std::fmt::Debug for TwapStream {
25    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        formatter
27            .debug_struct("TwapStream")
28            .field("market_id", &self.market_id)
29            .field("wallet_address", &self.wallet_address)
30            .field("stream_id", &self.stream_id)
31            .field("sequence", &self.sequence)
32            .finish_non_exhaustive()
33    }
34}
35
36impl TwapStream {
37    pub(crate) async fn connect(
38        client: &StrataClient,
39        market_id: &str,
40        wallet_address: &str,
41    ) -> Result<Self, SdkError> {
42        let market_id = validate_platform_market_id(market_id)?;
43        let wallet_address = canonical_public_key(wallet_address, "wallet_address")?;
44        let url = platform_websocket_url(
45            &client.base_url,
46            &format!("v2/markets/{market_id}/account/{wallet_address}/twaps/stream"),
47        )?;
48        let (socket, _) = tokio_tungstenite::connect_async(url.as_str())
49            .await
50            .map_err(|error| SdkError::Stream(error.to_string()))?;
51        Ok(Self {
52            market_id,
53            wallet_address,
54            stream_id: None,
55            sequence: 0,
56            socket,
57        })
58    }
59
60    pub fn market_id(&self) -> &str {
61        &self.market_id
62    }
63
64    pub fn wallet_address(&self) -> &str {
65        &self.wallet_address
66    }
67
68    /// Receive and validate the next event. `Ok(None)` means the peer closed
69    /// cleanly. Callers should reconnect after any error or close.
70    pub async fn next_event(&mut self) -> Result<Option<PlatformTwapEvent>, SdkError> {
71        loop {
72            let Some(frame) = self.socket.next().await else {
73                return Ok(None);
74            };
75            let frame = frame.map_err(|error| SdkError::Stream(error.to_string()))?;
76            match frame {
77                Message::Text(text) => {
78                    let event: PlatformTwapEvent = serde_json::from_str(&text)
79                        .map_err(|error| SdkError::InvalidResponse(error.to_string()))?;
80                    if let Err(error) = self.validate_event(&event) {
81                        let _ = self.socket.close(None).await;
82                        return Err(error);
83                    }
84                    return Ok(Some(event));
85                }
86                Message::Ping(payload) => {
87                    self.socket
88                        .send(Message::Pong(payload))
89                        .await
90                        .map_err(|error| SdkError::Stream(error.to_string()))?;
91                }
92                Message::Pong(_) => {}
93                Message::Close(_) => return Ok(None),
94                _ => {
95                    let _ = self.socket.close(None).await;
96                    return Err(SdkError::InvalidResponse(
97                        "TWAP stream sent a non-text data frame".to_owned(),
98                    ));
99                }
100            }
101        }
102    }
103
104    pub async fn close(&mut self) -> Result<(), SdkError> {
105        self.socket
106            .close(None)
107            .await
108            .map_err(|error| SdkError::Stream(error.to_string()))
109    }
110
111    fn validate_event(&mut self, event: &PlatformTwapEvent) -> Result<(), SdkError> {
112        match event {
113            PlatformTwapEvent::TwapsSnapshot {
114                schema_version,
115                contract_version,
116                market_id,
117                wallet_address,
118                stream_id,
119                sequence,
120                twaps,
121                ..
122            } => {
123                self.validate_identity(
124                    *schema_version,
125                    contract_version,
126                    market_id,
127                    wallet_address,
128                )?;
129                if !valid_handle(stream_id, "twap_stream_") {
130                    return Err(SdkError::InvalidResponse(
131                        "TWAP stream identity is invalid".to_owned(),
132                    ));
133                }
134                let next = validate_response_atoms(sequence, "sequence", false)?;
135                if let Some(current) = &self.stream_id {
136                    if current != stream_id || next <= self.sequence {
137                        return Err(SdkError::InvalidResponse(
138                            "TWAP recovery snapshot did not advance its sequence".to_owned(),
139                        ));
140                    }
141                }
142                validate_twap_rows(twaps)?;
143                self.stream_id = Some(stream_id.clone());
144                self.sequence = next;
145            }
146            PlatformTwapEvent::TwapUpdate {
147                schema_version,
148                contract_version,
149                market_id,
150                wallet_address,
151                stream_id,
152                sequence,
153                previous_sequence,
154                twap,
155                ..
156            } => {
157                self.validate_identity(
158                    *schema_version,
159                    contract_version,
160                    market_id,
161                    wallet_address,
162                )?;
163                self.validate_sequence(stream_id, sequence, previous_sequence)?;
164                validate_twap_rows(std::slice::from_ref(twap))?;
165            }
166            PlatformTwapEvent::Heartbeat {
167                schema_version,
168                contract_version,
169                market_id,
170                wallet_address,
171                stream_id,
172                sequence,
173                previous_sequence,
174                ..
175            } => {
176                self.validate_identity(
177                    *schema_version,
178                    contract_version,
179                    market_id,
180                    wallet_address,
181                )?;
182                self.validate_sequence(stream_id, sequence, previous_sequence)?;
183            }
184        }
185        Ok(())
186    }
187
188    fn validate_identity(
189        &self,
190        schema_version: u16,
191        contract_version: &str,
192        market_id: &str,
193        wallet_address: &str,
194    ) -> Result<(), SdkError> {
195        validate_platform_market_response(
196            schema_version,
197            contract_version,
198            market_id,
199            &self.market_id,
200        )?;
201        if wallet_address != self.wallet_address {
202            return Err(SdkError::InvalidResponse(
203                "TWAP stream wallet does not match the request".to_owned(),
204            ));
205        }
206        Ok(())
207    }
208
209    fn validate_sequence(
210        &mut self,
211        stream_id: &str,
212        sequence: &str,
213        previous_sequence: &str,
214    ) -> Result<(), SdkError> {
215        let Some(current) = &self.stream_id else {
216            return Err(SdkError::InvalidResponse(
217                "TWAP event arrived without its snapshot".to_owned(),
218            ));
219        };
220        let next = validate_response_atoms(sequence, "sequence", false)?;
221        let previous = validate_response_atoms(previous_sequence, "previous_sequence", false)?;
222        if stream_id != current || previous != self.sequence || next != previous.saturating_add(1) {
223            return Err(SdkError::InvalidResponse(
224                "TWAP stream sequence gap detected".to_owned(),
225            ));
226        }
227        self.sequence = next;
228        Ok(())
229    }
230}
231
232/// Every TWAP row must carry opaque identities and consistent progress; the
233/// same rule the HTTP read applies, plus a per-frame uniqueness check.
234pub(crate) fn validate_twap_rows(twaps: &[PlatformTwap]) -> Result<(), SdkError> {
235    if twaps.len() > 2_000 {
236        return Err(SdkError::InvalidResponse(
237            "TWAP stream rows exceed the bounded size".to_owned(),
238        ));
239    }
240    let mut ids = HashSet::new();
241    for twap in twaps {
242        if !valid_handle(&twap.twap_id, "twap_")
243            || !ids.insert(twap.twap_id.as_str())
244            || twap.slices_executed > twap.slices_total
245            || twap.fills.len() > usize::from(twap.slices_total)
246        {
247            return Err(SdkError::InvalidResponse(
248                "TWAP stream contains an invalid schedule".to_owned(),
249            ));
250        }
251        let total = validate_response_atoms(&twap.total_size_atoms, "total_size_atoms", false)?;
252        let executed =
253            validate_response_atoms(&twap.executed_size_atoms, "executed_size_atoms", true)?;
254        validate_response_atoms(&twap.limit_price_atoms, "limit_price_atoms", false)?;
255        validate_response_atoms(
256            &twap.gross_quote_executed_atoms,
257            "gross_quote_executed_atoms",
258            true,
259        )?;
260        if executed > total {
261            return Err(SdkError::InvalidResponse(
262                "TWAP executed size exceeds its schedule".to_owned(),
263            ));
264        }
265        for fill in &twap.fills {
266            if !valid_handle(&fill.fill_id, "twap_fill_") {
267                return Err(SdkError::InvalidResponse(
268                    "TWAP fill identity is invalid".to_owned(),
269                ));
270            }
271            validate_response_atoms(&fill.size_atoms, "size_atoms", false)?;
272            validate_response_atoms(&fill.price_atoms, "price_atoms", false)?;
273        }
274    }
275    Ok(())
276}