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#[derive(Debug, Clone, Builder)]
27#[builder(on(String, into))]
28pub struct BacktestClient {
29 #[builder(with = |url: impl Into<String>| backtest_ws_url(&url.into()))]
33 url: String,
34 #[builder(default)]
36 api_key: String,
37
38 #[builder(default = DEFAULT_CONNECT_TIMEOUT)]
40 connect_timeout: Duration,
41
42 request_timeout: Option<Duration>,
44
45 #[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 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 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 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 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 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 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 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 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 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 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 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}