Skip to main content

trybench_sdk/
platform.rs

1//! Explicit headless management of Bench, independent of tracing and local tests.
2use serde_json::{json, Value};
3use std::{collections::HashMap, fmt, time::Duration};
4
5const CATALOG: &str = include_str!("operations.json");
6#[derive(Debug)]
7pub struct PlatformError {
8    pub status: u16,
9    pub code: String,
10    pub message: String,
11    pub reference: Option<String>,
12    pub details: Value,
13}
14impl PlatformError {
15    fn local(code: &str, message: &str) -> Self {
16        Self {
17            status: 0,
18            code: code.into(),
19            message: message.into(),
20            reference: None,
21            details: json!({}),
22        }
23    }
24    pub fn retryable(&self) -> bool {
25        self.status == 429 || self.status >= 500
26    }
27}
28impl fmt::Display for PlatformError {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        f.write_str(&self.message)
31    }
32}
33impl std::error::Error for PlatformError {}
34
35pub struct PlatformFile {
36    pub name: String,
37    pub content: Vec<u8>,
38}
39#[derive(Default)]
40pub struct PlatformRequest {
41    pub path: HashMap<String, String>,
42    pub query: HashMap<String, String>,
43    pub body: Option<Value>,
44    pub form: serde_json::Map<String, Value>,
45    pub files: Vec<PlatformFile>,
46}
47
48pub struct BenchPlatform {
49    endpoint: String,
50    credential: String,
51    http: reqwest::Client,
52    operations: HashMap<String, Value>,
53}
54impl BenchPlatform {
55    /// Supply a Bench API key or short-lived Bench OAuth token. An empty
56    /// credential permits public login/catalog operations only.
57    pub fn new(credential: impl Into<String>, endpoint: &str) -> Result<Self, PlatformError> {
58        let url = reqwest::Url::parse(endpoint)
59            .map_err(|_| PlatformError::local("invalid_endpoint", "Invalid API origin."))?;
60        if !url.username().is_empty()
61            || url.password().is_some()
62            || url.query().is_some()
63            || url.fragment().is_some()
64            || url.path() != "/"
65            || !(url.scheme() == "https"
66                || (url.scheme() == "http"
67                    && matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "[::1]"))))
68        {
69            return Err(PlatformError::local(
70                "invalid_endpoint",
71                "Use an HTTPS API origin or loopback HTTP origin.",
72            ));
73        }
74        let catalog: Value = serde_json::from_str(CATALOG)
75            .map_err(|_| PlatformError::local("catalog_error", "Invalid operation catalog."))?;
76        let operations = catalog["operations"]
77            .as_array()
78            .ok_or_else(|| PlatformError::local("catalog_error", "Missing operations."))?
79            .iter()
80            .map(|op| {
81                (
82                    op["id"].as_str().unwrap_or_default().to_string(),
83                    op.clone(),
84                )
85            })
86            .collect();
87        let http = reqwest::Client::builder()
88            .redirect(reqwest::redirect::Policy::none())
89            .timeout(Duration::from_secs(120))
90            .build()
91            .map_err(|_| PlatformError::local("client_error", "Cannot create HTTP client."))?;
92        Ok(Self {
93            endpoint: endpoint.trim_end_matches('/').into(),
94            credential: credential.into(),
95            http,
96            operations,
97        })
98    }
99    /// Replace an expired OAuth token; never put credentials in source or logs.
100    pub fn set_credential(&mut self, credential: impl Into<String>) {
101        self.credential = credential.into();
102    }
103    pub fn operations(&self) -> Value {
104        serde_json::from_str(CATALOG).expect("validated catalog")
105    }
106    /// One operation, one request. No retries or implicit evaluation spending.
107    /// NDJSON results are returned as an event array; stream errors are errors.
108    pub async fn call(
109        &self,
110        operation: &str,
111        input: PlatformRequest,
112    ) -> Result<Value, PlatformError> {
113        let op = self
114            .operations
115            .get(operation)
116            .ok_or_else(|| PlatformError::local("unknown_operation", "Unknown Bench operation."))?;
117        let mut route = op["path"].as_str().unwrap_or_default().to_string();
118        for key in op["path_parameters"]
119            .as_object()
120            .into_iter()
121            .flat_map(|v| v.keys())
122        {
123            let value = input
124                .path
125                .get(key)
126                .ok_or_else(|| PlatformError::local("invalid_path", "Missing path parameter."))?;
127            if value.is_empty()
128                || [".", ".."].contains(&value.as_str())
129                || value
130                    .chars()
131                    .any(|c| c.is_whitespace() || "/?#\\%".contains(c))
132            {
133                return Err(PlatformError::local(
134                    "invalid_path",
135                    "Invalid path parameter.",
136                ));
137            }
138            let mut segment = reqwest::Url::parse("https://path.invalid/").expect("static URL");
139            segment.path_segments_mut().expect("base URL").push(value);
140            route = route.replace(
141                &format!("{{{key}}}"),
142                segment.path().trim_start_matches('/'),
143            );
144        }
145        let mut url = reqwest::Url::parse(&(self.endpoint.clone() + &route))
146            .map_err(|_| PlatformError::local("invalid_path", "Invalid path."))?;
147        if !input.query.is_empty() {
148            url.query_pairs_mut().extend_pairs(input.query.iter());
149        }
150        let method =
151            reqwest::Method::from_bytes(op["method"].as_str().unwrap_or_default().as_bytes())
152                .map_err(|_| PlatformError::local("catalog_error", "Invalid operation method."))?;
153        let mut request = self
154            .http
155            .request(method, url)
156            .header("Accept", "application/json, application/x-ndjson");
157        if op["auth"] != "public" {
158            if self.credential.trim().is_empty() {
159                return Err(PlatformError {
160                    status: 401,
161                    code: "unauthorized".into(),
162                    message: "Provide a Bench credential or authenticate through MCP OAuth.".into(),
163                    reference: None,
164                    details: json!({}),
165                });
166            }
167            request = request.bearer_auth(&self.credential);
168        }
169        if op["multipart"].as_bool().unwrap_or(false) {
170            if input.body.is_some() {
171                return Err(PlatformError::local(
172                    "invalid_input",
173                    "Use form and files for multipart operations.",
174                ));
175            }
176            let boundary = format!("bench-{}", uuid::Uuid::new_v4());
177            let mut data = Vec::new();
178            for (key, value) in input.form {
179                if key.contains(['\r', '\n', '"']) {
180                    return Err(PlatformError::local(
181                        "invalid_input",
182                        "Invalid form field name.",
183                    ));
184                }
185                let text = value
186                    .as_str()
187                    .map(str::to_owned)
188                    .unwrap_or_else(|| value.to_string());
189                data.extend_from_slice(format!("--{boundary}\r\nContent-Disposition: form-data; name=\"{key}\"\r\n\r\n{text}\r\n").as_bytes());
190            }
191            let mut total = 0;
192            let limit = op["max_file_bytes"].as_u64().unwrap_or(4 * 1024 * 1024) as usize;
193            let field = op["file_field"].as_str().unwrap_or("file");
194            for file in input.files {
195                if file.name.is_empty() || file.name.contains(['\r', '\n', '/', '\\', '"']) {
196                    return Err(PlatformError::local(
197                        "invalid_input",
198                        "Use a file name without directories or line breaks.",
199                    ));
200                }
201                total += file.content.len();
202                if total > limit {
203                    return Err(PlatformError::local(
204                        "upload_too_large",
205                        "Upload exceeds operation file-size limit.",
206                    ));
207                }
208                data.extend_from_slice(format!("--{boundary}\r\nContent-Disposition: form-data; name=\"{field}\"; filename=\"{}\"\r\nContent-Type: application/octet-stream\r\n\r\n", file.name).as_bytes());
209                data.extend_from_slice(&file.content);
210                data.extend_from_slice(b"\r\n");
211            }
212            data.extend_from_slice(format!("--{boundary}--\r\n").as_bytes());
213            request = request
214                .header(
215                    "Content-Type",
216                    format!("multipart/form-data; boundary={boundary}"),
217                )
218                .body(data);
219        } else {
220            if !input.files.is_empty() || !input.form.is_empty() {
221                return Err(PlatformError::local(
222                    "invalid_input",
223                    "Operation does not accept multipart uploads.",
224                ));
225            }
226            if let Some(body) = input.body {
227                request = request.json(&body);
228            }
229        }
230        let mut response = request.send().await.map_err(|_| {
231            PlatformError::local(
232                "network_error",
233                "Could not reach Bench; check connection and timeout.",
234            )
235        })?;
236        let status = response.status().as_u16();
237        let ndjson = response
238            .headers()
239            .get("Content-Type")
240            .and_then(|v| v.to_str().ok())
241            .unwrap_or("")
242            .contains("ndjson");
243        let mut data = Vec::new();
244        while let Some(chunk) = response
245            .chunk()
246            .await
247            .map_err(|_| PlatformError::local("network_error", "Could not read Bench response."))?
248        {
249            if data.len() + chunk.len() > 16 * 1024 * 1024 {
250                return Err(PlatformError::local(
251                    "response_too_large",
252                    "Response exceeds 16 MB; use pagination.",
253                ));
254            }
255            data.extend_from_slice(&chunk);
256        }
257        if !(200..300).contains(&status) {
258            let value: Value = serde_json::from_slice(&data).unwrap_or(json!({}));
259            return Err(PlatformError {
260                status,
261                code: value["error"]["code"]
262                    .as_str()
263                    .map(str::to_owned)
264                    .unwrap_or_else(|| format!("http_{status}")),
265                message: value["error"]["message"]
266                    .as_str()
267                    .unwrap_or("Bench request failed.")
268                    .into(),
269                reference: value["error"]["reference"].as_str().map(str::to_owned),
270                details: value["error"].clone(),
271            });
272        }
273        if status == 204 || data.is_empty() {
274            return Ok(Value::Null);
275        }
276        if ndjson {
277            let mut events = Vec::new();
278            for line in data
279                .split(|c| *c == b'\n')
280                .filter(|line| !line.iter().all(u8::is_ascii_whitespace))
281            {
282                let event: Value = serde_json::from_slice(line).map_err(|_| {
283                    PlatformError::local("invalid_response", "Invalid Bench stream event.")
284                })?;
285                if event["type"] == "error" {
286                    return Err(PlatformError {
287                        status,
288                        code: event["code"].as_str().unwrap_or("stream_error").into(),
289                        message: event["message"]
290                            .as_str()
291                            .or_else(|| event["error"].as_str())
292                            .unwrap_or("Operation failed.")
293                            .into(),
294                        reference: None,
295                        details: json!({}),
296                    });
297                }
298                events.push(event);
299            }
300            return Ok(Value::Array(events));
301        }
302        serde_json::from_slice(&data)
303            .map_err(|_| PlatformError::local("invalid_response", "Invalid Bench JSON response."))
304    }
305}