Skip to main content

simulator_client/
client.rs

1use std::time::Duration;
2
3use bon::Builder;
4use chrono::{DateTime, Utc};
5use simulator_api::{
6    AvailableRange, BacktestResponse, BundleBuildRequest, BundleBuildStatusResponse,
7    SnapshotSlotResponse, usage::UsageReport,
8};
9use tokio_tungstenite::{
10    connect_async,
11    tungstenite::{client::IntoClientRequest, http::HeaderValue},
12};
13
14use crate::{
15    BacktestClientError, BacktestClientResult, BacktestSession, CreateSession,
16    session::CreateRequestResult,
17    urls::{backtest_ws_url, http_base_from_ws_url},
18};
19
20const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
21
22/// Backtest WebSocket client configured with a base URL and API key.
23///
24/// Build with [`BacktestClient::builder`], which supports optional timeouts and
25/// raw message logging.
26#[derive(Debug, Clone, Builder)]
27#[builder(on(String, into))]
28pub struct BacktestClient {
29    /// WebSocket endpoint. Accepts a bare hostname (normalized to
30    /// `wss://{host}/backtest`) or an explicit `ws(s)://` URL, whose
31    /// `/backtest` path is appended only if missing.
32    #[builder(with = |url: impl Into<String>| backtest_ws_url(&url.into()))]
33    url: String,
34    /// API key injected as the `X-API-Key` header.
35    #[builder(default)]
36    api_key: String,
37
38    /// Timeout for the initial connect handshake.
39    #[builder(default = DEFAULT_CONNECT_TIMEOUT)]
40    connect_timeout: Duration,
41
42    /// Default timeout for request/response operations.
43    request_timeout: Option<Duration>,
44
45    /// Log raw inbound responses at debug level.
46    #[builder(default)]
47    log_raw: bool,
48}
49
50impl BacktestClient {
51    async fn connect(&self) -> BacktestClientResult<BacktestSession> {
52        let mut request = self.url.clone().into_client_request().map_err(|source| {
53            BacktestClientError::BuildRequest {
54                url: self.url.clone(),
55                source: Box::new(source),
56            }
57        })?;
58
59        request
60            .headers_mut()
61            .insert("X-API-Key", HeaderValue::from_str(&self.api_key)?);
62
63        let (stream, _) = tokio::time::timeout(self.connect_timeout, connect_async(request))
64            .await
65            .map_err(|_| BacktestClientError::Timeout {
66                action: "connecting",
67                duration: self.connect_timeout,
68            })?
69            .map_err(|source| BacktestClientError::Connect {
70                url: self.url.clone(),
71                source: Box::new(source),
72            })?;
73        Ok(BacktestSession::new(
74            stream,
75            self.request_timeout,
76            self.log_raw,
77        ))
78    }
79
80    /// Create a backtest session by connecting and sending a `CreateBacktestSession` request.
81    pub async fn create_session(
82        &self,
83        create: CreateSession,
84    ) -> BacktestClientResult<BacktestSession> {
85        let request = create.into_request()?;
86        let rpc_base_url = http_base_from_ws_url(&self.url);
87        let mut session = self.connect().await?;
88        match session
89            .create_with_request(request, rpc_base_url, None)
90            .await?
91        {
92            CreateRequestResult::Single { .. } => {}
93            CreateRequestResult::Parallel { session_ids, .. } => {
94                return Err(BacktestClientError::UnexpectedResponse {
95                    context: "creating single session",
96                    response: Box::new(BacktestResponse::SessionsCreatedV2 {
97                        control_session_id: String::new(),
98                        session_ids,
99                        task_ids: Vec::new(),
100                        start_slots: Vec::new(),
101                        end_slots: Vec::new(),
102                    }),
103                });
104            }
105        }
106        Ok(session)
107    }
108
109    /// Create one or many sessions and return the resulting session IDs.
110    ///
111    /// When `CreateSession.parallel` is true, this may return many IDs.
112    pub async fn create_sessions(
113        &self,
114        create: CreateSession,
115    ) -> BacktestClientResult<Vec<String>> {
116        self.create_sessions_with_progress(create, |_| {}).await
117    }
118
119    /// Create one or many sessions and stream each successfully created session ID.
120    ///
121    /// `on_session_created` is called for each streamed `SessionCreated` response.
122    pub async fn create_sessions_with_progress(
123        &self,
124        create: CreateSession,
125        mut on_session_created: impl FnMut(String) + Send,
126    ) -> BacktestClientResult<Vec<String>> {
127        let request = create.into_request()?;
128        let rpc_base_url = http_base_from_ws_url(&self.url);
129        let mut session = self.connect().await?;
130        match session
131            .create_with_request(request, rpc_base_url, Some(&mut on_session_created))
132            .await?
133        {
134            CreateRequestResult::Single { session_id, .. } => Ok(vec![session_id]),
135            CreateRequestResult::Parallel { session_ids, .. } => Ok(session_ids),
136        }
137    }
138
139    /// Like [`Self::create_sessions`], but additionally returns the opaque
140    /// per-session `task_id` reported by the server (if any).
141    pub async fn create_sessions_with_task_ids(
142        &self,
143        create: CreateSession,
144    ) -> BacktestClientResult<Vec<(String, Option<String>)>> {
145        let request = create.into_request()?;
146        let rpc_base_url = http_base_from_ws_url(&self.url);
147        let mut session = self.connect().await?;
148        match session
149            .create_with_request(request, rpc_base_url, None)
150            .await?
151        {
152            CreateRequestResult::Single {
153                session_id,
154                task_id,
155            } => Ok(vec![(session_id, task_id)]),
156            CreateRequestResult::Parallel {
157                session_ids,
158                task_ids,
159            } => Ok(session_ids.into_iter().zip(task_ids).collect()),
160        }
161    }
162
163    /// Fetch the available slot ranges from the server's `/available-ranges` endpoint.
164    pub async fn available_ranges(&self) -> BacktestClientResult<Vec<AvailableRange>> {
165        let base = http_base_from_ws_url(&self.url);
166        let url = format!("{base}/available-ranges");
167        // Forwarded (when set) so the endpoint adds this key's own ranges,
168        // keeping the client's split in sync with the server's. Omitted when
169        // unset so a keyless client gets global ranges, not an empty-key 401.
170        let mut request = reqwest::Client::new().get(&url);
171        if !self.api_key.is_empty() {
172            request = request.header("X-API-Key", &self.api_key);
173        }
174        let ranges = request
175            .send()
176            .await
177            .map_err(|e| BacktestClientError::Http {
178                url: url.clone(),
179                source: Box::new(e),
180            })?
181            .error_for_status()
182            .map_err(|e| BacktestClientError::Http {
183                url: url.clone(),
184                source: Box::new(e),
185            })?
186            .json::<Vec<AvailableRange>>()
187            .await
188            .map_err(|e| BacktestClientError::Http {
189                url: url.clone(),
190                source: Box::new(e),
191            })?;
192        Ok(ranges)
193    }
194
195    /// Fetch per-API-key usage metrics from the server's `/usage` endpoint.
196    ///
197    /// `since` and `until` narrow the reporting window; both default to the
198    /// server's configured window when omitted.
199    pub async fn usage(
200        &self,
201        since: Option<DateTime<Utc>>,
202        until: Option<DateTime<Utc>>,
203    ) -> BacktestClientResult<UsageReport> {
204        let base = http_base_from_ws_url(&self.url);
205        let url = format!("{base}/usage");
206        let mut req = reqwest::Client::new()
207            .get(&url)
208            .header("X-API-Key", &self.api_key);
209        if let Some(t) = since {
210            req = req.query(&[(
211                "since",
212                t.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
213            )]);
214        }
215        if let Some(t) = until {
216            req = req.query(&[(
217                "until",
218                t.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
219            )]);
220        }
221        let resp = req.send().await.map_err(|e| BacktestClientError::Http {
222            url: url.clone(),
223            source: Box::new(e),
224        })?;
225        let status = resp.status();
226        if !status.is_success() {
227            let body = resp.text().await.unwrap_or_default();
228            return Err(BacktestClientError::HttpStatus { url, status, body });
229        }
230        resp.json::<UsageReport>()
231            .await
232            .map_err(|e| BacktestClientError::Http {
233                url: url.clone(),
234                source: Box::new(e),
235            })
236    }
237
238    /// Most-recent available snapshot slot (`GET /snapshot-slot`).
239    pub async fn get_latest_snapshot_slot(&self) -> BacktestClientResult<SnapshotSlotResponse> {
240        let base = http_base_from_ws_url(&self.url);
241        let url = format!("{base}/snapshot-slot");
242        let resp = reqwest::Client::new()
243            .get(&url)
244            .header("X-API-Key", &self.api_key)
245            .send()
246            .await
247            .map_err(|e| BacktestClientError::Http {
248                url: url.clone(),
249                source: Box::new(e),
250            })?;
251        let status = resp.status();
252        if !status.is_success() {
253            let body = resp.text().await.unwrap_or_default();
254            return Err(BacktestClientError::HttpStatus { url, status, body });
255        }
256        resp.json::<SnapshotSlotResponse>()
257            .await
258            .map_err(|e| BacktestClientError::Http {
259                url: url.clone(),
260                source: Box::new(e),
261            })
262    }
263
264    /// Issue a bundle build (`POST /build-bundle`); the build runs async, so
265    /// poll [`Self::get_bundle_status`] for completion.
266    pub async fn build_bundle(
267        &self,
268        request: &BundleBuildRequest,
269    ) -> BacktestClientResult<BundleBuildStatusResponse> {
270        let base = http_base_from_ws_url(&self.url);
271        let url = format!("{base}/build-bundle");
272        let resp = reqwest::Client::new()
273            .post(&url)
274            .header("X-API-Key", &self.api_key)
275            .json(request)
276            .send()
277            .await
278            .map_err(|e| BacktestClientError::Http {
279                url: url.clone(),
280                source: Box::new(e),
281            })?;
282        let status = resp.status();
283        if !status.is_success() {
284            let body = resp.text().await.unwrap_or_default();
285            return Err(BacktestClientError::HttpStatus { url, status, body });
286        }
287        resp.json::<BundleBuildStatusResponse>()
288            .await
289            .map_err(|e| BacktestClientError::Http {
290                url: url.clone(),
291                source: Box::new(e),
292            })
293    }
294
295    /// Fetch a build's status (`GET /build-bundle/{request_id}`).
296    pub async fn get_bundle_status(
297        &self,
298        request_id: &str,
299    ) -> BacktestClientResult<BundleBuildStatusResponse> {
300        let base = http_base_from_ws_url(&self.url);
301        let url = format!("{base}/build-bundle/{request_id}");
302        let resp = reqwest::Client::new()
303            .get(&url)
304            .header("X-API-Key", &self.api_key)
305            .send()
306            .await
307            .map_err(|e| BacktestClientError::Http {
308                url: url.clone(),
309                source: Box::new(e),
310            })?;
311        let status = resp.status();
312        if !status.is_success() {
313            let body = resp.text().await.unwrap_or_default();
314            return Err(BacktestClientError::HttpStatus { url, status, body });
315        }
316        resp.json::<BundleBuildStatusResponse>()
317            .await
318            .map_err(|e| BacktestClientError::Http {
319                url: url.clone(),
320                source: Box::new(e),
321            })
322    }
323
324    /// Attach to an existing backtest session over the control websocket.
325    pub async fn attach_session(
326        &self,
327        session_id: impl Into<String>,
328        last_sequence: Option<u64>,
329    ) -> BacktestClientResult<BacktestSession> {
330        let rpc_base_url = http_base_from_ws_url(&self.url);
331        let mut session = self.connect().await?;
332        session
333            .attach(session_id.into(), last_sequence, rpc_base_url)
334            .await?;
335        Ok(session)
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn builder_normalizes_bare_hosts_and_keeps_full_urls() {
345        let from_host = BacktestClient::builder()
346            .url("simulator.termina.technology")
347            .build();
348        assert_eq!(from_host.url, "wss://simulator.termina.technology/backtest");
349
350        let from_full_url = BacktestClient::builder()
351            .url("ws://localhost:8900/backtest".to_string())
352            .build();
353        assert_eq!(from_full_url.url, "ws://localhost:8900/backtest");
354    }
355}