1use crate::{
2 bundle::{
3 jito::{config::JitoConfig, constant::JITO_MOCK_BLOCK_ENGINE_URL, error::JitoError},
4 BundleError,
5 },
6 sanitize_error,
7};
8use reqwest::Client;
9use serde::{Deserialize, Serialize};
10use serde_json::{json, Value};
11
12pub enum JitoBundleClient {
13 Live(JitoClient),
14 Mock(JitoMockClient),
15}
16
17impl JitoBundleClient {
18 pub fn new(config: &JitoConfig) -> Self {
19 if config.block_engine_url == JITO_MOCK_BLOCK_ENGINE_URL {
20 Self::Mock(JitoMockClient::new())
21 } else {
22 Self::Live(JitoClient::new(config))
23 }
24 }
25
26 pub async fn send_bundle(
27 &self,
28 encoded_transactions: &[String],
29 ) -> Result<String, BundleError> {
30 match self {
31 Self::Live(client) => client.send_bundle(encoded_transactions).await,
32 Self::Mock(client) => client.send_bundle(encoded_transactions).await,
33 }
34 }
35
36 pub async fn get_bundle_statuses(
37 &self,
38 bundle_uuids: Vec<String>,
39 ) -> Result<Value, BundleError> {
40 match self {
41 Self::Live(client) => client.get_bundle_statuses(bundle_uuids).await,
42 Self::Mock(client) => client.get_bundle_statuses(bundle_uuids).await,
43 }
44 }
45}
46
47pub struct JitoClient {
48 block_engine_url: String,
49 client: Client,
50}
51
52#[derive(Debug, Serialize, Deserialize)]
53struct JsonRpcRequest {
54 jsonrpc: String,
55 id: u64,
56 method: String,
57 params: Value,
58}
59
60#[derive(Debug, Deserialize)]
61struct JsonRpcResponse {
62 result: Option<Value>,
63 error: Option<JsonRpcError>,
64}
65
66#[derive(Debug, Deserialize)]
67struct JsonRpcError {
68 message: String,
69}
70
71impl JitoClient {
72 pub fn new(config: &JitoConfig) -> Self {
73 Self { block_engine_url: config.block_engine_url.clone(), client: Client::new() }
74 }
75
76 async fn send_request(&self, method: &str, params: Value) -> Result<Value, JitoError> {
77 let url = format!("{}/api/v1/bundles", self.block_engine_url);
78
79 let request = JsonRpcRequest {
80 jsonrpc: "2.0".to_string(),
81 id: 1,
82 method: method.to_string(),
83 params,
84 };
85
86 let response = self
87 .client
88 .post(&url)
89 .header("Content-Type", "application/json")
90 .json(&request)
91 .send()
92 .await
93 .map_err(|e| {
94 JitoError::ApiError(format!("HTTP request failed: {}", sanitize_error!(e)))
95 })?;
96
97 let status = response.status();
98 if !status.is_success() {
99 return Err(JitoError::ApiError(format!("HTTP error: {}", status)));
100 }
101
102 let rpc_response: JsonRpcResponse = response.json().await.map_err(|e| {
103 JitoError::ApiError(format!("Failed to parse response: {}", sanitize_error!(e)))
104 })?;
105
106 if let Some(error) = rpc_response.error {
107 return Err(JitoError::ApiError(error.message));
108 }
109
110 rpc_response
111 .result
112 .ok_or_else(|| JitoError::ApiError("Empty response from Jito".to_string()))
113 }
114
115 pub async fn send_bundle(
116 &self,
117 encoded_transactions: &[String],
118 ) -> Result<String, BundleError> {
119 let params = json!([encoded_transactions, {"encoding": "base64"}]);
120 let result = self.send_request("sendBundle", params).await?;
121
122 result.as_str().map(|s| s.to_string()).ok_or_else(|| {
123 JitoError::ApiError("Invalid bundle UUID in response".to_string()).into()
124 })
125 }
126
127 pub async fn get_bundle_statuses(
128 &self,
129 bundle_uuids: Vec<String>,
130 ) -> Result<Value, BundleError> {
131 let params = json!([bundle_uuids]);
132 Ok(self.send_request("getBundleStatuses", params).await?)
133 }
134}
135
136pub struct JitoMockClient;
137
138impl JitoMockClient {
139 pub fn new() -> Self {
140 Self
141 }
142
143 pub async fn send_bundle(
144 &self,
145 _encoded_transactions: &[String],
146 ) -> Result<String, BundleError> {
147 let random_id: u64 = rand::random();
148 Ok(format!("mock-bundle-{random_id}"))
149 }
150
151 pub async fn get_bundle_statuses(
152 &self,
153 bundle_uuids: Vec<String>,
154 ) -> Result<Value, BundleError> {
155 let mock_statuses: Vec<Value> = bundle_uuids
156 .iter()
157 .map(|uuid| {
158 json!({
159 "bundle_id": uuid,
160 "status": "Landed",
161 "landed_slot": 12345678
162 })
163 })
164 .collect();
165 Ok(json!({ "value": mock_statuses }))
166 }
167}
168
169impl Default for JitoMockClient {
170 fn default() -> Self {
171 Self::new()
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178 use crate::tests::transaction_mock::create_mock_encoded_transaction;
179 use mockito::{Matcher, Server};
180
181 #[tokio::test]
182 async fn test_send_bundle_success() {
183 let mut server = Server::new_async().await;
184
185 let mock = server
186 .mock("POST", "/api/v1/bundles")
187 .match_header("content-type", "application/json")
188 .match_body(Matcher::PartialJson(json!({"method": "sendBundle"})))
189 .with_status(200)
190 .with_header("content-type", "application/json")
191 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":"bundle-uuid-12345"}"#)
192 .create();
193
194 let config = JitoConfig { block_engine_url: server.url() };
195 let client = JitoClient::new(&config);
196
197 let tx = create_mock_encoded_transaction();
198
199 let result = client.send_bundle(&[tx]).await;
200 mock.assert();
201 assert!(result.is_ok());
202 assert_eq!(result.unwrap(), "bundle-uuid-12345");
203 }
204
205 #[tokio::test]
206 async fn test_send_bundle_http_error() {
207 let mut server = Server::new_async().await;
208
209 let mock = server
210 .mock("POST", "/api/v1/bundles")
211 .with_status(500)
212 .with_body("Internal Server Error")
213 .create();
214
215 let config = JitoConfig { block_engine_url: server.url() };
216 let client = JitoClient::new(&config);
217
218 let tx = create_mock_encoded_transaction();
219
220 let result = client.send_bundle(&[tx]).await;
221 mock.assert();
222 assert!(result.is_err());
223 }
224
225 #[tokio::test]
226 async fn test_send_bundle_rpc_error() {
227 let mut server = Server::new_async().await;
228
229 let mock = server
230 .mock("POST", "/api/v1/bundles")
231 .with_status(200)
232 .with_header("content-type", "application/json")
233 .with_body(r#"{"jsonrpc":"2.0","id":1,"error":{"message":"Bundle rejected"}}"#)
234 .create();
235
236 let config = JitoConfig { block_engine_url: server.url() };
237 let client = JitoClient::new(&config);
238
239 let tx = create_mock_encoded_transaction();
240
241 let result = client.send_bundle(&[tx]).await;
242 mock.assert();
243 assert!(result.is_err());
244 }
245
246 #[tokio::test]
247 async fn test_get_bundle_statuses_success() {
248 let mut server = Server::new_async().await;
249
250 let mock = server
251 .mock("POST", "/api/v1/bundles")
252 .match_body(Matcher::PartialJson(json!({"method": "getBundleStatuses"})))
253 .with_status(200)
254 .with_header("content-type", "application/json")
255 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[{"bundle_id":"uuid1","status":"Landed"}]}}"#)
256 .create();
257
258 let config = JitoConfig { block_engine_url: server.url() };
259 let client = JitoClient::new(&config);
260
261 let result = client.get_bundle_statuses(vec!["uuid1".to_string()]).await;
262 mock.assert();
263 assert!(result.is_ok());
264 }
265
266 #[tokio::test]
267 async fn test_send_bundle_invalid_uuid_response() {
268 let mut server = Server::new_async().await;
269
270 let mock = server
271 .mock("POST", "/api/v1/bundles")
272 .with_status(200)
273 .with_header("content-type", "application/json")
274 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":12345}"#)
275 .create();
276
277 let config = JitoConfig { block_engine_url: server.url() };
278 let client = JitoClient::new(&config);
279
280 let tx = create_mock_encoded_transaction();
281
282 let result = client.send_bundle(&[tx]).await;
283 mock.assert();
284 assert!(result.is_err());
285 }
286
287 #[tokio::test]
288 async fn test_send_bundle_empty_result() {
289 let mut server = Server::new_async().await;
290
291 let mock = server
292 .mock("POST", "/api/v1/bundles")
293 .with_status(200)
294 .with_header("content-type", "application/json")
295 .with_body(r#"{"jsonrpc":"2.0","id":1}"#)
296 .create();
297
298 let config = JitoConfig { block_engine_url: server.url() };
299 let client = JitoClient::new(&config);
300
301 let tx = create_mock_encoded_transaction();
302
303 let result = client.send_bundle(&[tx]).await;
304 mock.assert();
305 assert!(result.is_err());
306 }
307
308 #[tokio::test]
309 async fn test_send_bundle_multiple_transactions() {
310 let mut server = Server::new_async().await;
311
312 let mock = server
313 .mock("POST", "/api/v1/bundles")
314 .with_status(200)
315 .with_header("content-type", "application/json")
316 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":"multi-tx-bundle-uuid"}"#)
317 .create();
318
319 let config = JitoConfig { block_engine_url: server.url() };
320 let client = JitoClient::new(&config);
321
322 let txs = vec![
323 create_mock_encoded_transaction(),
324 create_mock_encoded_transaction(),
325 create_mock_encoded_transaction(),
326 ];
327
328 let result = client.send_bundle(&txs).await;
329 mock.assert();
330 assert!(result.is_ok());
331 assert_eq!(result.unwrap(), "multi-tx-bundle-uuid");
332 }
333
334 #[tokio::test]
335 async fn test_send_bundle_malformed_json_response() {
336 let mut server = Server::new_async().await;
337
338 let mock = server
339 .mock("POST", "/api/v1/bundles")
340 .with_status(200)
341 .with_header("content-type", "application/json")
342 .with_body(r#"not valid json"#)
343 .create();
344
345 let config = JitoConfig { block_engine_url: server.url() };
346 let client = JitoClient::new(&config);
347
348 let tx = create_mock_encoded_transaction();
349
350 let result = client.send_bundle(&[tx]).await;
351 mock.assert();
352 assert!(result.is_err());
353 }
354
355 #[tokio::test]
356 async fn test_get_bundle_statuses_empty_uuids() {
357 let mut server = Server::new_async().await;
358
359 let mock = server
360 .mock("POST", "/api/v1/bundles")
361 .match_body(Matcher::PartialJson(
362 json!({"method": "getBundleStatuses", "params": [[]]}),
363 ))
364 .with_status(200)
365 .with_header("content-type", "application/json")
366 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[]}}"#)
367 .create();
368
369 let config = JitoConfig { block_engine_url: server.url() };
370 let client = JitoClient::new(&config);
371
372 let result = client.get_bundle_statuses(vec![]).await;
373 mock.assert();
374 assert!(result.is_ok());
375 }
376
377 #[tokio::test]
378 async fn test_get_bundle_statuses_multiple_uuids() {
379 let mut server = Server::new_async().await;
380
381 let mock = server
382 .mock("POST", "/api/v1/bundles")
383 .match_body(Matcher::PartialJson(json!({
384 "method": "getBundleStatuses",
385 "params": [["uuid1", "uuid2", "uuid3"]]
386 })))
387 .with_status(200)
388 .with_header("content-type", "application/json")
389 .with_body(r#"{"jsonrpc":"2.0","id":1,"result":{"value":[{"bundle_id":"uuid1","status":"Landed"},{"bundle_id":"uuid2","status":"Pending"}]}}"#)
390 .create();
391
392 let config = JitoConfig { block_engine_url: server.url() };
393 let client = JitoClient::new(&config);
394
395 let result = client
396 .get_bundle_statuses(vec![
397 "uuid1".to_string(),
398 "uuid2".to_string(),
399 "uuid3".to_string(),
400 ])
401 .await;
402 mock.assert();
403 assert!(result.is_ok());
404 }
405
406 #[tokio::test]
407 async fn test_jito_bundle_client_dispatches_to_mock() {
408 let config = JitoConfig { block_engine_url: JITO_MOCK_BLOCK_ENGINE_URL.to_string() };
409 let client = JitoBundleClient::new(&config);
410
411 assert!(matches!(client, JitoBundleClient::Mock(_)));
412
413 let tx = create_mock_encoded_transaction();
414 let result = client.send_bundle(&[tx]).await;
415
416 assert!(result.is_ok());
417 let uuid = result.unwrap();
418 assert!(uuid.starts_with("mock-bundle-"), "Expected mock UUID prefix");
419 }
420
421 #[tokio::test]
422 async fn test_jito_bundle_client_dispatches_to_real() {
423 let config = JitoConfig { block_engine_url: "https://example.com".to_string() };
424 let client = JitoBundleClient::new(&config);
425
426 assert!(matches!(client, JitoBundleClient::Live(_)));
427 }
428
429 #[tokio::test]
430 async fn test_mock_client_get_bundle_statuses() {
431 let config = JitoConfig { block_engine_url: JITO_MOCK_BLOCK_ENGINE_URL.to_string() };
432 let client = JitoBundleClient::new(&config);
433
434 let result = client.get_bundle_statuses(vec!["test-uuid".to_string()]).await;
435
436 assert!(result.is_ok());
437 let statuses = result.unwrap();
438 assert!(statuses["value"].is_array());
439 assert_eq!(statuses["value"][0]["status"], "Landed");
440 }
441}