Skip to main content

strata_sdk/
execution_stream.rs

1use std::collections::HashSet;
2use std::time::Duration;
3
4use futures_util::{SinkExt, StreamExt};
5use tokio::net::TcpStream;
6use tokio_tungstenite::tungstenite::Message;
7use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
8
9use super::market_stream::platform_websocket_url;
10use super::*;
11
12const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10);
13pub const MAX_WATCHED_EXECUTIONS: usize = 64;
14
15type PlatformSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
16
17/// Sequenced state for the executions an agent prepared in one market. The
18/// stream opens with a `watch` frame for the requested handles and begins with
19/// a snapshot; later `watch` calls add handles on the same stream. Any market
20/// or sequence mismatch fails closed so the caller reconnects.
21pub struct ExecutionStream {
22    market_id: String,
23    stream_id: String,
24    sequence: u64,
25    watched: HashSet<String>,
26    initial_snapshot: Option<PlatformExecutionEvent>,
27    socket: PlatformSocket,
28}
29
30impl std::fmt::Debug for ExecutionStream {
31    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        formatter
33            .debug_struct("ExecutionStream")
34            .field("market_id", &self.market_id)
35            .field("stream_id", &self.stream_id)
36            .field("sequence", &self.sequence)
37            .field("watched", &self.watched.len())
38            .finish_non_exhaustive()
39    }
40}
41
42impl ExecutionStream {
43    pub(crate) async fn connect(
44        client: &StrataClient,
45        market_id: &str,
46        execution_ids: &[String],
47    ) -> Result<Self, SdkError> {
48        let market_id = validate_platform_market_id(market_id)?;
49        let ids = checked_execution_ids(execution_ids)?;
50        let url = platform_websocket_url(
51            &client.base_url,
52            &format!("v2/markets/{market_id}/executions/stream"),
53        )?;
54        let (mut socket, _) = tokio_tungstenite::connect_async(url.as_str())
55            .await
56            .map_err(|error| SdkError::Stream(error.to_string()))?;
57        send_watch(&mut socket, &ids).await?;
58        let frame = tokio::time::timeout(SNAPSHOT_TIMEOUT, socket.next())
59            .await
60            .map_err(|_| SdkError::Stream("execution snapshot timed out".to_owned()))?
61            .ok_or_else(|| SdkError::Stream("socket closed before execution snapshot".to_owned()))?
62            .map_err(|error| SdkError::Stream(error.to_string()))?;
63        let Message::Text(text) = frame else {
64            return Err(SdkError::InvalidResponse(
65                "expected a text execution snapshot".to_owned(),
66            ));
67        };
68        let snapshot: PlatformExecutionEvent = serde_json::from_str(&text)
69            .map_err(|error| SdkError::InvalidResponse(error.to_string()))?;
70        let (stream_id, sequence) = match &snapshot {
71            PlatformExecutionEvent::ExecutionsSnapshot {
72                schema_version,
73                contract_version,
74                market_id: response_market,
75                stream_id,
76                sequence,
77                executions,
78                unknown_execution_ids,
79                ..
80            } => {
81                validate_platform_market_response(
82                    *schema_version,
83                    contract_version,
84                    response_market,
85                    &market_id,
86                )?;
87                if !valid_handle(stream_id, "execution_stream_") {
88                    return Err(SdkError::InvalidResponse(
89                        "execution stream identity is invalid".to_owned(),
90                    ));
91                }
92                validate_execution_rows(executions, unknown_execution_ids, &market_id)?;
93                (
94                    stream_id.clone(),
95                    validate_response_atoms(sequence, "sequence", false)?,
96                )
97            }
98            _ => {
99                return Err(SdkError::InvalidResponse(
100                    "execution stream did not begin with a snapshot".to_owned(),
101                ))
102            }
103        };
104        Ok(Self {
105            market_id,
106            stream_id,
107            sequence,
108            watched: ids.into_iter().collect(),
109            initial_snapshot: Some(snapshot),
110            socket,
111        })
112    }
113
114    pub fn market_id(&self) -> &str {
115        &self.market_id
116    }
117
118    /// Watch additional handles; the server answers with an update or unknown
119    /// event per new handle on the same sequence.
120    pub async fn watch(&mut self, execution_ids: &[String]) -> Result<(), SdkError> {
121        let fresh: Vec<String> = checked_execution_ids(execution_ids)?
122            .into_iter()
123            .filter(|id| !self.watched.contains(id))
124            .collect();
125        if fresh.is_empty() {
126            return Ok(());
127        }
128        if self.watched.len().saturating_add(fresh.len()) > MAX_WATCHED_EXECUTIONS {
129            return Err(SdkError::InvalidRequest(format!(
130                "at most {MAX_WATCHED_EXECUTIONS} executions can be watched per stream"
131            )));
132        }
133        send_watch(&mut self.socket, &fresh).await?;
134        self.watched.extend(fresh);
135        Ok(())
136    }
137
138    /// The first call returns the snapshot received during connection. Later
139    /// calls return sequenced update, expired, unknown, heartbeat, or recovery
140    /// snapshot events.
141    pub async fn next_event(&mut self) -> Result<Option<PlatformExecutionEvent>, SdkError> {
142        if let Some(snapshot) = self.initial_snapshot.take() {
143            return Ok(Some(snapshot));
144        }
145        loop {
146            let Some(frame) = self.socket.next().await else {
147                return Ok(None);
148            };
149            let frame = frame.map_err(|error| SdkError::Stream(error.to_string()))?;
150            match frame {
151                Message::Text(text) => {
152                    let event: PlatformExecutionEvent = serde_json::from_str(&text)
153                        .map_err(|error| SdkError::InvalidResponse(error.to_string()))?;
154                    if let Err(error) = self.validate_event(&event) {
155                        let _ = self.socket.close(None).await;
156                        return Err(error);
157                    }
158                    return Ok(Some(event));
159                }
160                Message::Ping(payload) => {
161                    self.socket
162                        .send(Message::Pong(payload))
163                        .await
164                        .map_err(|error| SdkError::Stream(error.to_string()))?;
165                }
166                Message::Pong(_) => {}
167                Message::Close(_) => return Ok(None),
168                _ => {
169                    let _ = self.socket.close(None).await;
170                    return Err(SdkError::InvalidResponse(
171                        "execution stream sent a non-text data frame".to_owned(),
172                    ));
173                }
174            }
175        }
176    }
177
178    pub async fn close(&mut self) -> Result<(), SdkError> {
179        self.socket
180            .close(None)
181            .await
182            .map_err(|error| SdkError::Stream(error.to_string()))
183    }
184
185    fn validate_event(&mut self, event: &PlatformExecutionEvent) -> Result<(), SdkError> {
186        match event {
187            PlatformExecutionEvent::ExecutionsSnapshot {
188                schema_version,
189                contract_version,
190                market_id,
191                stream_id,
192                sequence,
193                executions,
194                unknown_execution_ids,
195                ..
196            } => {
197                validate_platform_market_response(
198                    *schema_version,
199                    contract_version,
200                    market_id,
201                    &self.market_id,
202                )?;
203                let next = validate_response_atoms(sequence, "sequence", false)?;
204                if stream_id != &self.stream_id || next <= self.sequence {
205                    return Err(SdkError::InvalidResponse(
206                        "execution recovery snapshot did not advance its sequence".to_owned(),
207                    ));
208                }
209                validate_execution_rows(executions, unknown_execution_ids, &self.market_id)?;
210                self.sequence = next;
211            }
212            PlatformExecutionEvent::ExecutionUpdate {
213                schema_version,
214                contract_version,
215                market_id,
216                stream_id,
217                sequence,
218                previous_sequence,
219                execution,
220                ..
221            } => {
222                validate_platform_market_response(
223                    *schema_version,
224                    contract_version,
225                    market_id,
226                    &self.market_id,
227                )?;
228                self.validate_sequence(stream_id, sequence, previous_sequence)?;
229                validate_execution_rows(std::slice::from_ref(execution), &[], &self.market_id)?;
230            }
231            PlatformExecutionEvent::ExecutionExpired {
232                schema_version,
233                contract_version,
234                market_id,
235                stream_id,
236                sequence,
237                previous_sequence,
238                execution_id,
239                ..
240            }
241            | PlatformExecutionEvent::ExecutionUnknown {
242                schema_version,
243                contract_version,
244                market_id,
245                stream_id,
246                sequence,
247                previous_sequence,
248                execution_id,
249                ..
250            } => {
251                validate_platform_market_response(
252                    *schema_version,
253                    contract_version,
254                    market_id,
255                    &self.market_id,
256                )?;
257                self.validate_sequence(stream_id, sequence, previous_sequence)?;
258                if !valid_handle(execution_id, "se_") {
259                    return Err(SdkError::InvalidResponse(
260                        "execution stream handle is invalid".to_owned(),
261                    ));
262                }
263            }
264            PlatformExecutionEvent::Heartbeat {
265                schema_version,
266                contract_version,
267                market_id,
268                stream_id,
269                sequence,
270                previous_sequence,
271                ..
272            } => {
273                validate_platform_market_response(
274                    *schema_version,
275                    contract_version,
276                    market_id,
277                    &self.market_id,
278                )?;
279                self.validate_sequence(stream_id, sequence, previous_sequence)?;
280            }
281        }
282        Ok(())
283    }
284
285    fn validate_sequence(
286        &mut self,
287        stream_id: &str,
288        sequence: &str,
289        previous_sequence: &str,
290    ) -> Result<(), SdkError> {
291        let next = validate_response_atoms(sequence, "sequence", false)?;
292        let previous = validate_response_atoms(previous_sequence, "previous_sequence", false)?;
293        if stream_id != self.stream_id
294            || previous != self.sequence
295            || next != previous.saturating_add(1)
296        {
297            return Err(SdkError::InvalidResponse(
298                "execution stream sequence gap detected".to_owned(),
299            ));
300        }
301        self.sequence = next;
302        Ok(())
303    }
304}
305
306async fn send_watch(socket: &mut PlatformSocket, ids: &[String]) -> Result<(), SdkError> {
307    socket
308        .send(Message::Text(
309            serde_json::to_string(&PlatformExecutionCommand::Watch {
310                execution_ids: ids.to_vec(),
311            })
312            .map_err(|error| SdkError::InvalidRequest(error.to_string()))?
313            .into(),
314        ))
315        .await
316        .map_err(|error| SdkError::Stream(error.to_string()))
317}
318
319fn checked_execution_ids(execution_ids: &[String]) -> Result<Vec<String>, SdkError> {
320    if execution_ids.is_empty() {
321        return Err(SdkError::InvalidRequest(
322            "at least one execution_id is required".to_owned(),
323        ));
324    }
325    let mut seen = HashSet::new();
326    let mut ids = Vec::with_capacity(execution_ids.len());
327    for id in execution_ids {
328        let id = id.trim().to_owned();
329        if !valid_handle(&id, "se_") {
330            return Err(SdkError::InvalidRequest(
331                "execution_id must be an opaque Strata execution handle".to_owned(),
332            ));
333        }
334        if seen.insert(id.clone()) {
335            ids.push(id);
336        }
337    }
338    if ids.len() > MAX_WATCHED_EXECUTIONS {
339        return Err(SdkError::InvalidRequest(format!(
340            "at most {MAX_WATCHED_EXECUTIONS} executions can be watched per stream"
341        )));
342    }
343    Ok(ids)
344}
345
346fn validate_execution_rows(
347    executions: &[PlatformExecutionRow],
348    unknown_execution_ids: &[String],
349    market_id: &str,
350) -> Result<(), SdkError> {
351    if executions.len() > MAX_WATCHED_EXECUTIONS
352        || unknown_execution_ids.len() > MAX_WATCHED_EXECUTIONS
353    {
354        return Err(SdkError::InvalidResponse(
355            "execution stream rows exceed the bounded size".to_owned(),
356        ));
357    }
358    let mut ids = HashSet::new();
359    for row in executions {
360        let confirmed = row.status == PlatformExecutionState::Confirmed;
361        if !valid_handle(&row.execution_id, "se_")
362            || !ids.insert(row.execution_id.as_str())
363            || row.market_id != market_id
364            || confirmed != row.signature.is_some()
365            || confirmed != (row.settlement == PlatformSettlementState::Confirmed)
366        {
367            return Err(SdkError::InvalidResponse(
368                "execution stream contains an inconsistent execution".to_owned(),
369            ));
370        }
371    }
372    for id in unknown_execution_ids {
373        if !valid_handle(id, "se_") || !ids.insert(id.as_str()) {
374            return Err(SdkError::InvalidResponse(
375                "execution stream unknown handles are invalid".to_owned(),
376            ));
377        }
378    }
379    Ok(())
380}