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    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::SessionsCreated { session_ids }),
97                });
98            }
99        }
100        Ok(session)
101    }
102
103    /// Create one or many sessions and return the resulting session IDs.
104    ///
105    /// When `CreateSession.parallel` is true, this may return many IDs.
106    pub async fn create_sessions(
107        &self,
108        create: CreateSession,
109    ) -> BacktestClientResult<Vec<String>> {
110        self.create_sessions_with_progress(create, |_| {}).await
111    }
112
113    /// Create one or many sessions and stream each successfully created session ID.
114    ///
115    /// `on_session_created` is called for each streamed `SessionCreated` response.
116    pub async fn create_sessions_with_progress(
117        &self,
118        create: CreateSession,
119        mut on_session_created: impl FnMut(String) + Send,
120    ) -> BacktestClientResult<Vec<String>> {
121        let request = create.into_request()?;
122        let rpc_base_url = http_base_from_ws_url(&self.url);
123        let mut session = self.connect().await?;
124        match session
125            .create_with_request(request, rpc_base_url, Some(&mut on_session_created))
126            .await?
127        {
128            CreateRequestResult::Single { session_id, .. } => Ok(vec![session_id]),
129            CreateRequestResult::Parallel { session_ids, .. } => Ok(session_ids),
130        }
131    }
132
133    /// Like [`Self::create_sessions`], but additionally returns the opaque
134    /// per-session `task_id` reported by the server (if any).
135    pub async fn create_sessions_with_task_ids(
136        &self,
137        create: CreateSession,
138    ) -> BacktestClientResult<Vec<(String, Option<String>)>> {
139        let request = create.into_request()?;
140        let rpc_base_url = http_base_from_ws_url(&self.url);
141        let mut session = self.connect().await?;
142        match session
143            .create_with_request(request, rpc_base_url, None)
144            .await?
145        {
146            CreateRequestResult::Single {
147                session_id,
148                task_id,
149            } => Ok(vec![(session_id, task_id)]),
150            CreateRequestResult::Parallel {
151                session_ids,
152                task_ids,
153            } => Ok(session_ids.into_iter().zip(task_ids).collect()),
154        }
155    }
156
157    /// Fetch the available slot ranges from the server's `/available-ranges` endpoint.
158    pub async fn available_ranges(&self) -> BacktestClientResult<Vec<AvailableRange>> {
159        let base = http_base_from_ws_url(&self.url);
160        let url = format!("{base}/available-ranges");
161        // Forwarded (when set) so the endpoint adds this key's own ranges,
162        // keeping the client's split in sync with the server's. Omitted when
163        // unset so a keyless client gets global ranges, not an empty-key 401.
164        let mut request = reqwest::Client::new().get(&url);
165        if !self.api_key.is_empty() {
166            request = request.header("X-API-Key", &self.api_key);
167        }
168        let ranges = request
169            .send()
170            .await
171            .map_err(|e| BacktestClientError::Http {
172                url: url.clone(),
173                source: Box::new(e),
174            })?
175            .error_for_status()
176            .map_err(|e| BacktestClientError::Http {
177                url: url.clone(),
178                source: Box::new(e),
179            })?
180            .json::<Vec<AvailableRange>>()
181            .await
182            .map_err(|e| BacktestClientError::Http {
183                url: url.clone(),
184                source: Box::new(e),
185            })?;
186        Ok(ranges)
187    }
188
189    /// Fetch per-API-key usage metrics from the server's `/usage` endpoint.
190    ///
191    /// `since` and `until` narrow the reporting window; both default to the
192    /// server's configured window when omitted.
193    pub async fn usage(
194        &self,
195        since: Option<DateTime<Utc>>,
196        until: Option<DateTime<Utc>>,
197    ) -> BacktestClientResult<UsageReport> {
198        let base = http_base_from_ws_url(&self.url);
199        let url = format!("{base}/usage");
200        let mut req = reqwest::Client::new()
201            .get(&url)
202            .header("X-API-Key", &self.api_key);
203        if let Some(t) = since {
204            req = req.query(&[(
205                "since",
206                t.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
207            )]);
208        }
209        if let Some(t) = until {
210            req = req.query(&[(
211                "until",
212                t.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
213            )]);
214        }
215        let resp = req.send().await.map_err(|e| BacktestClientError::Http {
216            url: url.clone(),
217            source: Box::new(e),
218        })?;
219        let status = resp.status();
220        if !status.is_success() {
221            let body = resp.text().await.unwrap_or_default();
222            return Err(BacktestClientError::HttpStatus { url, status, body });
223        }
224        resp.json::<UsageReport>()
225            .await
226            .map_err(|e| BacktestClientError::Http {
227                url: url.clone(),
228                source: Box::new(e),
229            })
230    }
231
232    /// Issue a bundle build (`POST /build-bundle`); the build runs async, so
233    /// poll [`Self::get_bundle_status`] for completion.
234    pub async fn build_bundle(
235        &self,
236        request: &BundleBuildRequest,
237    ) -> BacktestClientResult<BundleBuildStatusResponse> {
238        let base = http_base_from_ws_url(&self.url);
239        let url = format!("{base}/build-bundle");
240        let resp = reqwest::Client::new()
241            .post(&url)
242            .header("X-API-Key", &self.api_key)
243            .json(request)
244            .send()
245            .await
246            .map_err(|e| BacktestClientError::Http {
247                url: url.clone(),
248                source: Box::new(e),
249            })?;
250        let status = resp.status();
251        if !status.is_success() {
252            let body = resp.text().await.unwrap_or_default();
253            return Err(BacktestClientError::HttpStatus { url, status, body });
254        }
255        resp.json::<BundleBuildStatusResponse>()
256            .await
257            .map_err(|e| BacktestClientError::Http {
258                url: url.clone(),
259                source: Box::new(e),
260            })
261    }
262
263    /// Fetch a build's status (`GET /build-bundle/{request_id}`).
264    pub async fn get_bundle_status(
265        &self,
266        request_id: &str,
267    ) -> BacktestClientResult<BundleBuildStatusResponse> {
268        let base = http_base_from_ws_url(&self.url);
269        let url = format!("{base}/build-bundle/{request_id}");
270        let resp = reqwest::Client::new()
271            .get(&url)
272            .header("X-API-Key", &self.api_key)
273            .send()
274            .await
275            .map_err(|e| BacktestClientError::Http {
276                url: url.clone(),
277                source: Box::new(e),
278            })?;
279        let status = resp.status();
280        if !status.is_success() {
281            let body = resp.text().await.unwrap_or_default();
282            return Err(BacktestClientError::HttpStatus { url, status, body });
283        }
284        resp.json::<BundleBuildStatusResponse>()
285            .await
286            .map_err(|e| BacktestClientError::Http {
287                url: url.clone(),
288                source: Box::new(e),
289            })
290    }
291
292    /// Attach to an existing backtest session over the control websocket.
293    pub async fn attach_session(
294        &self,
295        session_id: impl Into<String>,
296        last_sequence: Option<u64>,
297    ) -> BacktestClientResult<BacktestSession> {
298        let rpc_base_url = http_base_from_ws_url(&self.url);
299        let mut session = self.connect().await?;
300        session
301            .attach(session_id.into(), last_sequence, rpc_base_url)
302            .await?;
303        Ok(session)
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn builder_normalizes_bare_hosts_and_keeps_full_urls() {
313        let from_host = BacktestClient::builder()
314            .url("simulator.termina.technology")
315            .build();
316        assert_eq!(from_host.url, "wss://simulator.termina.technology/backtest");
317
318        let from_full_url = BacktestClient::builder()
319            .url("ws://localhost:8900/backtest".to_string())
320            .build();
321        assert_eq!(from_full_url.url, "ws://localhost:8900/backtest");
322    }
323}