rustbox_sdk/
lib.rs

1//! Official Rust client for [Rustbox](https://rustbox.sh).
2//!
3//! ```no_run
4//! use rustbox_sdk::{Rustbox, SubmitRequest};
5//! # async fn run() -> Result<(), rustbox_sdk::RustboxError> {
6//! let client = Rustbox::new(&std::env::var("RUSTBOX_API_KEY").unwrap())?;
7//! let res = client.run(&SubmitRequest {
8//!     language: "python".into(),
9//!     code: "print('hi')".into(),
10//!     ..Default::default()
11//! }).await?;
12//! println!("{}", res["result"]["verdict"]);
13//! # Ok(()) }
14//! ```
15
16use std::time::Duration;
17
18use serde::Serialize;
19use thiserror::Error;
20use tokio::time::sleep;
21
22/// SDK version. Sent in `User-Agent`.
23pub const VERSION: &str = "0.4.0";
24
25/// Production endpoint.
26pub const DEFAULT_BASE_URL: &str = "https://api.rustbox.sh";
27
28const DEFAULT_TIMEOUT: Duration = Duration::from_secs(65);
29const DEFAULT_MAX_RETRIES: u32 = 2;
30
31/// Execution profile.
32///
33/// - `Profile::Judge` (default): short evaluation runs.
34/// - `Profile::Agent`: longer jobs with egress proxy + per-key byte
35///   budgets. Requires a non-trial API key.
36#[derive(Serialize, Debug, Clone, Copy, PartialEq, Eq)]
37#[serde(rename_all = "lowercase")]
38pub enum Profile {
39    Judge,
40    Agent,
41}
42
43#[derive(Serialize, Debug, Clone, Default)]
44pub struct SubmitRequest {
45    pub language: String,
46    pub code: String,
47    pub stdin: String,
48    /// Optional profile override. None falls back to server-side default ("judge").
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub profile: Option<Profile>,
51}
52
53/// Errors returned by the Rustbox client.
54#[derive(Debug, Error)]
55pub enum RustboxError {
56    #[error("api_key required")]
57    MissingApiKey,
58    #[error("invalid base_url")]
59    InvalidBaseUrl,
60    #[error("invalid or missing API key (HTTP {0})")]
61    Auth(u16),
62    #[error("rate limit exceeded (HTTP 429)")]
63    RateLimit,
64    #[error("server error (HTTP {0})")]
65    Server(u16),
66    #[error("API error (HTTP {status}): {body}")]
67    Api { status: u16, body: String },
68    #[error("request timed out")]
69    Timeout,
70    #[error(transparent)]
71    Transport(#[from] reqwest::Error),
72    #[error("response decode failed: {0}")]
73    Decode(String),
74}
75
76/// Optional submit-only knobs that don't belong in the request body.
77#[derive(Debug, Clone, Default)]
78pub struct SubmitOptions {
79    /// `Idempotency-Key` header value. Safe to retry POST /api/submit when set.
80    pub idempotency_key: Option<String>,
81}
82
83pub struct Rustbox {
84    api_key: String,
85    base_url: String,
86    client: reqwest::Client,
87    max_retries: u32,
88}
89
90impl Rustbox {
91    /// Construct a Rustbox client. `api_key` is required (must be non-empty).
92    /// Base URL defaults to `DEFAULT_BASE_URL`; override with `with_base_url`.
93    pub fn new(api_key: &str) -> Result<Self, RustboxError> {
94        if api_key.is_empty() {
95            return Err(RustboxError::MissingApiKey);
96        }
97        let client = reqwest::Client::builder()
98            .timeout(DEFAULT_TIMEOUT)
99            .build()
100            .map_err(RustboxError::Transport)?;
101        Ok(Self {
102            api_key: api_key.to_string(),
103            base_url: DEFAULT_BASE_URL.to_string(),
104            client,
105            max_retries: DEFAULT_MAX_RETRIES,
106        })
107    }
108
109    /// Override the API base URL. Use for staging.
110    /// Trailing slashes are trimmed.
111    pub fn with_base_url(mut self, base_url: &str) -> Result<Self, RustboxError> {
112        if base_url.is_empty() {
113            return Err(RustboxError::InvalidBaseUrl);
114        }
115        self.base_url = base_url.trim_end_matches('/').to_string();
116        Ok(self)
117    }
118
119    /// Override the per-request timeout. Set `Duration::ZERO` to disable.
120    pub fn with_timeout(mut self, timeout: Duration) -> Result<Self, RustboxError> {
121        let mut builder = reqwest::Client::builder();
122        if !timeout.is_zero() {
123            builder = builder.timeout(timeout);
124        }
125        self.client = builder.build().map_err(RustboxError::Transport)?;
126        Ok(self)
127    }
128
129    /// Override the retry budget on transient (5xx, network) failures.
130    pub fn with_max_retries(mut self, n: u32) -> Self {
131        self.max_retries = n;
132        self
133    }
134
135    pub fn base_url(&self) -> &str {
136        &self.base_url
137    }
138
139    fn backoff_delay(&self, attempt: u32) -> Duration {
140        Duration::from_millis((100u64 * (1u64 << attempt.min(8))).min(5_000))
141    }
142
143    async fn send_with_retry(
144        &self,
145        build: impl Fn() -> reqwest::RequestBuilder,
146    ) -> Result<reqwest::Response, RustboxError> {
147        let mut last_err: Option<RustboxError> = None;
148        for attempt in 0..=self.max_retries {
149            let req = build()
150                .header("X-API-Key", &self.api_key)
151                .header("User-Agent", format!("rustbox-sdk-rust/{VERSION}"));
152            match req.send().await {
153                Ok(resp) => {
154                    if resp.status().as_u16() >= 500 && attempt < self.max_retries {
155                        sleep(self.backoff_delay(attempt)).await;
156                        continue;
157                    }
158                    return Ok(resp);
159                }
160                Err(e) => {
161                    let is_timeout = e.is_timeout();
162                    last_err = Some(if is_timeout {
163                        RustboxError::Timeout
164                    } else {
165                        RustboxError::Transport(e)
166                    });
167                    if attempt >= self.max_retries {
168                        return Err(last_err.unwrap());
169                    }
170                    sleep(self.backoff_delay(attempt)).await;
171                }
172            }
173        }
174        Err(last_err.unwrap_or(RustboxError::Decode("retry exhausted".into())))
175    }
176
177    async fn handle(&self, resp: reqwest::Response) -> Result<serde_json::Value, RustboxError> {
178        let status = resp.status();
179        let code = status.as_u16();
180        if status.is_success() || code == 408 {
181            return resp
182                .json()
183                .await
184                .map_err(|e| RustboxError::Decode(e.to_string()));
185        }
186        match code {
187            401 | 403 => Err(RustboxError::Auth(code)),
188            429 => Err(RustboxError::RateLimit),
189            500..=599 => Err(RustboxError::Server(code)),
190            _ => {
191                let body = resp.text().await.unwrap_or_default();
192                Err(RustboxError::Api { status: code, body })
193            }
194        }
195    }
196
197    pub async fn submit(
198        &self,
199        req: &SubmitRequest,
200        wait: bool,
201        opts: SubmitOptions,
202    ) -> Result<serde_json::Value, RustboxError> {
203        let url = format!("{}/api/submit?wait={}", self.base_url, wait);
204        let body = serde_json::to_vec(req).map_err(|e| RustboxError::Decode(e.to_string()))?;
205
206        let resp = self
207            .send_with_retry(|| {
208                let mut rb = self
209                    .client
210                    .post(&url)
211                    .header("Content-Type", "application/json")
212                    .body(body.clone());
213                if let Some(ref key) = opts.idempotency_key {
214                    rb = rb.header("Idempotency-Key", key);
215                }
216                rb
217            })
218            .await?;
219        self.handle(resp).await
220    }
221
222    pub async fn get_result(&self, id: &str) -> Result<serde_json::Value, RustboxError> {
223        let url = format!("{}/api/result/{}", self.base_url, id);
224        let resp = self.send_with_retry(|| self.client.get(&url)).await?;
225        self.handle(resp).await
226    }
227
228    pub async fn get_languages(&self) -> Result<Vec<String>, RustboxError> {
229        let url = format!("{}/api/languages", self.base_url);
230        let resp = self.send_with_retry(|| self.client.get(&url)).await?;
231        let val = self.handle(resp).await?;
232        serde_json::from_value(val).map_err(|e| RustboxError::Decode(e.to_string()))
233    }
234
235    pub async fn get_health(&self) -> Result<serde_json::Value, RustboxError> {
236        let url = format!("{}/api/health", self.base_url);
237        let resp = self.send_with_retry(|| self.client.get(&url)).await?;
238        self.handle(resp).await
239    }
240
241    pub async fn get_ready(&self) -> Result<serde_json::Value, RustboxError> {
242        let url = format!("{}/api/health/ready", self.base_url);
243        let resp = self.send_with_retry(|| self.client.get(&url)).await?;
244        self.handle(resp).await
245    }
246
247    /// Submit + wait (sync) + auto-poll fallback. Auto-generates an
248    /// Idempotency-Key so the underlying POST is safe to retry.
249    pub async fn run(&self, req: &SubmitRequest) -> Result<serde_json::Value, RustboxError> {
250        let opts = SubmitOptions {
251            idempotency_key: Some(idempotency_id()),
252        };
253        let mut res = self.submit(req, true, opts).await?;
254        if res
255            .pointer("/result/verdict")
256            .and_then(|v| v.as_str())
257            .is_some()
258        {
259            return Ok(res);
260        }
261
262        let id = match res.get("id").and_then(|v| v.as_str()) {
263            Some(i) => i.to_string(),
264            None => return Ok(res),
265        };
266
267        for i in 0..45 {
268            let delay_ms = (40.0 * (1.5_f64).powi(i)).min(600.0) as u64;
269            sleep(Duration::from_millis(delay_ms)).await;
270
271            res = self.get_result(&id).await?;
272            if res
273                .pointer("/result/verdict")
274                .and_then(|v| v.as_str())
275                .is_some()
276            {
277                return Ok(res);
278            }
279        }
280        Ok(res)
281    }
282}
283
284fn idempotency_id() -> String {
285    uuid::Uuid::new_v4().to_string()
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291    use wiremock::matchers::{method, path};
292    use wiremock::{Mock, MockServer, ResponseTemplate};
293
294    fn req() -> SubmitRequest {
295        SubmitRequest {
296            language: "python".into(),
297            code: "print(1)".into(),
298            stdin: "".into(),
299            profile: None,
300        }
301    }
302
303    #[tokio::test]
304    async fn new_should_default_base_url_to_production() {
305        let client = Rustbox::new("k").unwrap();
306        assert_eq!(client.base_url(), DEFAULT_BASE_URL);
307    }
308
309    #[tokio::test]
310    async fn new_should_return_err_when_api_key_empty() {
311        let r = Rustbox::new("");
312        assert!(matches!(r, Err(RustboxError::MissingApiKey)));
313    }
314
315    #[tokio::test]
316    async fn with_base_url_should_override_default_and_trim_slash() {
317        let client = Rustbox::new("k")
318            .unwrap()
319            .with_base_url("https://custom.example.com/")
320            .unwrap();
321        assert_eq!(client.base_url(), "https://custom.example.com");
322    }
323
324    #[tokio::test]
325    async fn with_base_url_should_return_err_when_empty() {
326        let r = Rustbox::new("k").unwrap().with_base_url("");
327        assert!(matches!(r, Err(RustboxError::InvalidBaseUrl)));
328    }
329
330    #[tokio::test]
331    async fn run_should_return_verdict_on_first_response_when_complete() {
332        let mock_server = MockServer::start().await;
333        Mock::given(method("POST"))
334            .and(path("/api/submit"))
335            .respond_with(
336                ResponseTemplate::new(200)
337                    .set_body_json(serde_json::json!({"id": "1", "result": {"verdict": "AC"}})),
338            )
339            .mount(&mock_server)
340            .await;
341
342        let client = Rustbox::new("test")
343            .unwrap()
344            .with_base_url(&mock_server.uri())
345            .unwrap();
346        let res = client.run(&req()).await.unwrap();
347        assert_eq!(
348            res.pointer("/result/verdict").unwrap().as_str().unwrap(),
349            "AC"
350        );
351    }
352
353    #[tokio::test]
354    async fn run_should_poll_until_verdict_when_initial_returns_408() {
355        let mock_server = MockServer::start().await;
356        Mock::given(method("POST"))
357            .and(path("/api/submit"))
358            .respond_with(ResponseTemplate::new(408).set_body_json(serde_json::json!({"id": "1"})))
359            .mount(&mock_server)
360            .await;
361
362        Mock::given(method("GET"))
363            .and(path("/api/result/1"))
364            .respond_with(
365                ResponseTemplate::new(200)
366                    .set_body_json(serde_json::json!({"id": "1", "result": {"verdict": "TLE"}})),
367            )
368            .mount(&mock_server)
369            .await;
370
371        let client = Rustbox::new("test")
372            .unwrap()
373            .with_base_url(&mock_server.uri())
374            .unwrap();
375        let res = client.run(&req()).await.unwrap();
376        assert_eq!(
377            res.pointer("/result/verdict").unwrap().as_str().unwrap(),
378            "TLE"
379        );
380    }
381
382    #[tokio::test]
383    async fn submit_should_return_auth_err_on_401() {
384        let mock_server = MockServer::start().await;
385        Mock::given(method("POST"))
386            .and(path("/api/submit"))
387            .respond_with(ResponseTemplate::new(401))
388            .mount(&mock_server)
389            .await;
390
391        let client = Rustbox::new("test")
392            .unwrap()
393            .with_base_url(&mock_server.uri())
394            .unwrap();
395        let err = client
396            .submit(&req(), false, SubmitOptions::default())
397            .await
398            .unwrap_err();
399        assert!(matches!(err, RustboxError::Auth(401)));
400    }
401
402    #[tokio::test]
403    async fn submit_should_return_rate_limit_on_429() {
404        let mock_server = MockServer::start().await;
405        Mock::given(method("POST"))
406            .and(path("/api/submit"))
407            .respond_with(ResponseTemplate::new(429))
408            .mount(&mock_server)
409            .await;
410
411        let client = Rustbox::new("test")
412            .unwrap()
413            .with_base_url(&mock_server.uri())
414            .unwrap();
415        let err = client
416            .submit(&req(), false, SubmitOptions::default())
417            .await
418            .unwrap_err();
419        assert!(matches!(err, RustboxError::RateLimit));
420    }
421
422    #[tokio::test]
423    async fn submit_should_return_server_err_on_503_after_retries() {
424        let mock_server = MockServer::start().await;
425        Mock::given(method("POST"))
426            .and(path("/api/submit"))
427            .respond_with(ResponseTemplate::new(503))
428            .mount(&mock_server)
429            .await;
430
431        let client = Rustbox::new("test")
432            .unwrap()
433            .with_base_url(&mock_server.uri())
434            .unwrap()
435            .with_max_retries(1);
436        let err = client
437            .submit(&req(), false, SubmitOptions::default())
438            .await
439            .unwrap_err();
440        assert!(matches!(err, RustboxError::Server(503)));
441    }
442
443    #[tokio::test]
444    async fn submit_should_send_user_agent_header() {
445        let mock_server = MockServer::start().await;
446        Mock::given(method("POST"))
447            .and(path("/api/submit"))
448            .and(wiremock::matchers::header_regex(
449                "user-agent",
450                r"^rustbox-sdk-rust/",
451            ))
452            .respond_with(
453                ResponseTemplate::new(200)
454                    .set_body_json(serde_json::json!({"id": "1", "result": {"verdict": "AC"}})),
455            )
456            .mount(&mock_server)
457            .await;
458
459        let client = Rustbox::new("test")
460            .unwrap()
461            .with_base_url(&mock_server.uri())
462            .unwrap();
463        let res = client
464            .submit(&req(), false, SubmitOptions::default())
465            .await
466            .unwrap();
467        assert_eq!(
468            res.pointer("/result/verdict").unwrap().as_str().unwrap(),
469            "AC"
470        );
471    }
472}
rustbox_sdk/lib.rs

rustbox_sdk/
lib.rs
