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#[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::SessionsCreated { session_ids }),
97 });
98 }
99 }
100 Ok(session)
101 }
102
103 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 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 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 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 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 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 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 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 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}