basal_client/
batch_resources.rs

1//! Resource facades and opt-in Batch workflow helpers. Existing flat methods stay available.
2
3use std::{
4    collections::BTreeMap,
5    fmt,
6    hash::{BuildHasher, Hasher},
7    io::Write,
8    time::Duration,
9};
10
11use http::{
12    header::{HeaderMap, HeaderValue, ACCEPT},
13    Method,
14};
15
16use crate::{
17    Batch, BatchExpiresAfter, BatchInputLine, BatchListOptions, BatchPage, BatchResultLine, Client, CreateBatchRequest,
18    DeletedFile, Error, FileObject, FileUploadOptions,
19};
20
21/// A simplified create request. The SDK serializes and uploads these records before creating the batch.
22/// JSONL is buffered in memory, limited to 50,000 records and 200 MB; the upload adds a multipart copy.
23pub struct CreateBatchFromRequests<'a> {
24    /// Native requests and unique caller identifiers.
25    pub requests: &'a [BatchInputLine],
26    /// Optional owner metadata.
27    pub metadata: Option<BTreeMap<String, String>>,
28    /// Optional output and error file retention.
29    pub output_expires_after: Option<BatchExpiresAfter>,
30}
31
32impl<'a> CreateBatchFromRequests<'a> {
33    /// Prepare native `/v1/basal` records with the server's 24-hour expiry window.
34    pub fn new(requests: &'a [BatchInputLine]) -> Self {
35        Self { requests, metadata: None, output_expires_after: None }
36    }
37}
38
39/// Mutually exclusive input paths for [`Batches::create`]. Prefer the `From` conversions at the call site.
40pub enum BatchCreateInput<'a> {
41    /// Reference a previously uploaded file.
42    File(&'a CreateBatchRequest),
43    /// Serialize/upload native requests, then create a batch.
44    Requests(&'a CreateBatchFromRequests<'a>),
45}
46
47impl<'a> From<&'a CreateBatchRequest> for BatchCreateInput<'a> {
48    fn from(value: &'a CreateBatchRequest) -> Self {
49        Self::File(value)
50    }
51}
52
53impl<'a> From<&'a CreateBatchFromRequests<'a>> for BatchCreateInput<'a> {
54    fn from(value: &'a CreateBatchFromRequests<'a>) -> Self {
55        Self::Requests(value)
56    }
57}
58
59/// Explicit polling settings. The deadline covers both retrieval requests and sleeps.
60#[derive(Clone, Copy, Debug)]
61pub struct BatchWaitOptions {
62    /// Whole client-side wait budget (default ten minutes); timeout never cancels the server job.
63    pub timeout: Duration,
64    /// Delay between successful nonterminal observations (default five seconds).
65    pub poll_interval: Duration,
66}
67
68impl Default for BatchWaitOptions {
69    fn default() -> Self {
70        Self { timeout: Duration::from_secs(600), poll_interval: Duration::from_secs(5) }
71    }
72}
73
74/// Per-file buffering allowance for the results cursor.
75#[derive(Clone, Copy, Debug, Default)]
76pub struct BatchResultsOptions {
77    /// Maximum bytes in each downloaded file, defaulting to the client's bound (8 MiB by default).
78    pub max_bytes: Option<usize>,
79}
80
81/// Typed failures of the opt-in workflow helpers. Debug/Display omit documents, bodies and keys.
82pub enum BatchOperationError {
83    /// Ordinary request/transport failure, with its original details explicitly available.
84    Request(Error),
85    /// Upload succeeded but creation did not return a usable batch. Creation may have succeeded remotely.
86    /// Resume using this file id and the same key; never automatically delete the uploaded input.
87    Submission {
88        /// Uploaded file, usable with `CreateBatchRequest::new`.
89        input_file_id: String,
90        /// Reuse this key together with the same metadata and output retention.
91        idempotency_key: String,
92        /// Original request failure.
93        source: Error,
94    },
95    /// The entire wait budget elapsed; no server cancellation was requested.
96    WaitTimeout {
97        /// Durable batch identifier.
98        batch_id: String,
99        /// Last successfully observed state, when any retrieval finished.
100        last_batch: Option<Box<Batch>>,
101    },
102    /// Results require a terminal batch. Inspect this state or explicitly call `wait`.
103    NotReady {
104        /// Current batch, including counters and validation errors.
105        batch: Box<Batch>,
106    },
107    /// One downloaded JSONL record did not match the result contract.
108    ResultLine {
109        /// File containing the malformed record.
110        file_id: String,
111        /// One-based physical line number.
112        line: usize,
113        /// Parser or envelope validation failure.
114        source: Error,
115    },
116}
117
118impl From<Error> for BatchOperationError {
119    fn from(value: Error) -> Self {
120        Self::Request(value)
121    }
122}
123
124impl fmt::Display for BatchOperationError {
125    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126        match self {
127            Self::Request(error) => fmt::Display::fmt(error, f),
128            Self::Submission { .. } => {
129                f.write_str("Basal batch creation failed after upload; recovery context is available")
130            }
131            Self::WaitTimeout { .. } => f.write_str("Basal batch wait timed out; server work continues"),
132            Self::NotReady { .. } => f.write_str("Basal batch results are not ready"),
133            Self::ResultLine { line, .. } => write!(f, "invalid Basal batch result at line {line}"),
134        }
135    }
136}
137
138impl fmt::Debug for BatchOperationError {
139    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140        fmt::Display::fmt(self, f)
141    }
142}
143
144impl std::error::Error for BatchOperationError {
145    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
146        match self {
147            Self::Request(source) | Self::Submission { source, .. } | Self::ResultLine { source, .. } => Some(source),
148            _ => None,
149        }
150    }
151}
152
153/// Files resource. Borrows the pooled client; never opens local paths or remote URLs.
154#[derive(Clone, Copy)]
155pub struct Files<'a> {
156    client: &'a Client,
157}
158
159/// Batches resource. Workflow polling happens only on an explicit `wait` call.
160#[derive(Clone, Copy)]
161pub struct Batches<'a> {
162    client: &'a Client,
163}
164
165impl Client {
166    /// Files resource with `create`, `retrieve`, `list`, `delete` and `content` operations.
167    pub fn files(&self) -> Files<'_> {
168        Files { client: self }
169    }
170
171    /// Batches resource with both create paths, retrieval and opt-in workflow helpers.
172    pub fn batches(&self) -> Batches<'_> {
173        Batches { client: self }
174    }
175}
176
177impl Files<'_> {
178    /// Upload JSONL bytes with purpose `batch`. Identical to the existing flat upload method.
179    pub async fn create(&self, filename: &str, data: &[u8], options: FileUploadOptions) -> Result<FileObject, Error> {
180        self.client.upload_file(filename, data, options).await
181    }
182    /// Read file metadata.
183    pub async fn retrieve(&self, id: &str) -> Result<FileObject, Error> {
184        self.client.get_file(id).await
185    }
186    /// Read one page; never automatically paginates.
187    pub async fn list(
188        &self,
189        options: BatchListOptions<'_>,
190        purpose: Option<&str>,
191    ) -> Result<BatchPage<FileObject>, Error> {
192        self.client.list_files(options, purpose).await
193    }
194    /// Request deletion.
195    pub async fn delete(&self, id: &str) -> Result<DeletedFile, Error> {
196        self.client.delete_file(id).await
197    }
198    /// Download bytes within the client's existing timeout and byte cap.
199    pub async fn content(&self, id: &str) -> Result<Vec<u8>, Error> {
200        self.client.download_file(id).await
201    }
202    /// Override the successful download byte cap for this call only. HTTP errors keep the client-wide cap.
203    pub async fn content_with_limit(&self, id: &str, max_bytes: usize) -> Result<Vec<u8>, Error> {
204        super::batches::id_component(id, "file_")?;
205        if max_bytes == 0 {
206            return Err(Error::InvalidRequest("download max_bytes must be positive"));
207        }
208        let mut headers = HeaderMap::new();
209        headers.insert(ACCEPT, HeaderValue::from_static("application/octet-stream"));
210        self.client
211            .request_bytes_bounded(Method::GET, &format!("v1/files/{id}/content"), (None, None), headers, max_bytes)
212            .await
213    }
214}
215
216impl<'a> Batches<'a> {
217    /// Create from an uploaded file or from native requests (via `CreateBatchFromRequests`).
218    /// The requests path uploads buffered JSONL and creates the batch, without polling or retries.
219    /// Post-upload failure exposes a recoverable file id/key in `BatchOperationError::Submission`.
220    pub async fn create<'r>(
221        &self,
222        input: impl Into<BatchCreateInput<'r>>,
223        idempotency_key: Option<&str>,
224    ) -> Result<Batch, BatchOperationError> {
225        match input.into() {
226            BatchCreateInput::File(request) => {
227                self.client.create_batch(request, idempotency_key).await.map_err(Into::into)
228            }
229            BatchCreateInput::Requests(request) => {
230                validate_key(idempotency_key)?;
231                validate_common(request)?;
232                let data = encode_requests(request.requests)?;
233                let key = idempotency_key.map(str::to_owned).unwrap_or_else(new_key);
234                let file = self.client.upload_file("batch-input.jsonl", &data, FileUploadOptions::default()).await?;
235                let mut create = CreateBatchRequest::new(&file.id);
236                create.metadata = request.metadata.clone();
237                create.output_expires_after = request.output_expires_after.clone();
238                self.client.create_batch(&create, Some(&key)).await.map_err(|source| BatchOperationError::Submission {
239                    input_file_id: file.id,
240                    idempotency_key: key,
241                    source,
242                })
243            }
244        }
245    }
246    /// Retrieve the latest state without polling.
247    pub async fn retrieve(&self, id: &str) -> Result<Batch, Error> {
248        self.client.get_batch(id).await
249    }
250    /// Retrieve one page without automatic pagination.
251    pub async fn list(&self, options: BatchListOptions<'_>) -> Result<BatchPage<Batch>, Error> {
252        self.client.list_batches(options).await
253    }
254    /// Explicitly request server cancellation; partial results may remain.
255    pub async fn cancel(&self, id: &str) -> Result<Batch, Error> {
256        self.client.cancel_batch(id).await
257    }
258
259    /// Wait until completed/failed/expired/cancelled. HTTP errors propagate without retry.
260    /// Timeout or dropping this future stops client-side polling, never the server batch.
261    pub async fn wait(&self, id: &str, options: BatchWaitOptions) -> Result<Batch, BatchOperationError> {
262        super::batches::id_component(id, "batch_")?;
263        if options.timeout.is_zero()
264            || options.poll_interval.is_zero()
265            || tokio::time::Instant::now().checked_add(options.timeout).is_none()
266            || tokio::time::Instant::now().checked_add(options.poll_interval).is_none()
267        {
268            return Err(Error::InvalidRequest("wait durations must be positive and representable").into());
269        }
270        let mut last_batch = None;
271        let polling = async {
272            loop {
273                let batch = self.retrieve(id).await?;
274                if terminal(&batch.status) {
275                    return Ok(batch);
276                }
277                last_batch = Some(Box::new(batch));
278                tokio::time::sleep(options.poll_interval).await;
279            }
280        };
281        match tokio::time::timeout(options.timeout, polling).await {
282            Ok(result) => result,
283            Err(_) => Err(BatchOperationError::WaitTimeout { batch_id: id.to_owned(), last_batch }),
284        }
285    }
286
287    /// Read terminal output and error files through a bounded, per-file buffered async cursor.
288    /// This never waits for completion. Inspect `cursor.batch()` for batch-level validation errors.
289    pub async fn results(&self, id: &str) -> Result<BatchResults<'a>, BatchOperationError> {
290        self.results_with_options(id, BatchResultsOptions::default()).await
291    }
292
293    /// Results cursor with an explicit per-file byte allowance.
294    pub async fn results_with_options(
295        &self,
296        id: &str,
297        options: BatchResultsOptions,
298    ) -> Result<BatchResults<'a>, BatchOperationError> {
299        if options.max_bytes == Some(0) {
300            return Err(Error::InvalidRequest("download max_bytes must be positive").into());
301        }
302        let batch = self.retrieve(id).await?;
303        if !terminal(&batch.status) {
304            return Err(BatchOperationError::NotReady { batch: Box::new(batch) });
305        }
306        let mut files = Vec::new();
307        for file in [&batch.output_file_id, &batch.error_file_id].into_iter().flatten() {
308            if !files.contains(file) {
309                files.push(file.clone());
310            }
311        }
312        Ok(BatchResults {
313            client: self.client,
314            batch,
315            files,
316            next_file: 0,
317            current: Vec::new(),
318            offset: 0,
319            line: 0,
320            max_bytes: options.max_bytes,
321            finished: false,
322        })
323    }
324}
325
326/// Result cursor. Buffers one file at a time; parsing does not create a second collection of all records.
327/// No background task or connection is retained between `next` calls. Dropping an in-flight `next`
328/// cancels that request; discard the cursor after cancellation rather than resuming uncertain I/O.
329pub struct BatchResults<'a> {
330    client: &'a Client,
331    batch: Batch,
332    files: Vec<String>,
333    next_file: usize,
334    current: Vec<u8>,
335    offset: usize,
336    line: usize,
337    max_bytes: Option<usize>,
338    finished: bool,
339}
340
341impl BatchResults<'_> {
342    /// Final batch, including counters and any batch-level validation errors.
343    pub fn batch(&self) -> &Batch {
344        &self.batch
345    }
346
347    /// Read the next record, output file first then error file. Ends permanently on a read/parse error.
348    /// Record order is not input order; correlate by `custom_id`.
349    pub async fn next(&mut self) -> Result<Option<BatchResultLine>, BatchOperationError> {
350        if self.finished {
351            return Ok(None);
352        }
353        loop {
354            if self.offset < self.current.len() {
355                let rest = &self.current[self.offset..];
356                let length = rest.iter().position(|b| *b == b'\n').unwrap_or(rest.len());
357                let record = parse_result(&rest[..length]);
358                self.offset += length + usize::from(length < rest.len());
359                self.line += 1;
360                return record.map(Some).map_err(|source| {
361                    self.finished = true;
362                    self.current = Vec::new();
363                    BatchOperationError::ResultLine {
364                        file_id: self.files[self.next_file - 1].clone(),
365                        line: self.line,
366                        source,
367                    }
368                });
369            }
370            // Release the preceding file's allocation before buffering the next download.
371            self.current = Vec::new();
372            if self.next_file == self.files.len() {
373                self.finished = true;
374                return Ok(None);
375            }
376            let id = &self.files[self.next_file];
377            // Mark the cursor exhausted before awaiting I/O. Cancellation cannot retry a consumed response.
378            self.finished = true;
379            let bytes = match self.max_bytes {
380                Some(limit) => self.client.files().content_with_limit(id, limit).await?,
381                None => self.client.files().content(id).await?,
382            };
383            self.finished = false;
384            self.current = bytes;
385            self.offset = 0;
386            self.line = 0;
387            self.next_file += 1;
388        }
389    }
390}
391
392fn terminal(status: &str) -> bool {
393    matches!(status, "completed" | "failed" | "expired" | "cancelled")
394}
395
396fn validate_key(key: Option<&str>) -> Result<(), Error> {
397    if key.is_some_and(|k| k.is_empty() || k.len() > 128 || !k.bytes().all(|b| (0x21..=0x7e).contains(&b))) {
398        return Err(Error::InvalidRequest("invalid idempotency key"));
399    }
400    Ok(())
401}
402
403fn validate_common(request: &CreateBatchFromRequests<'_>) -> Result<(), Error> {
404    if request.metadata.as_ref().is_some_and(|m| {
405        m.len() > 16 || m.iter().any(|(k, v)| k.is_empty() || k.chars().count() > 64 || v.chars().count() > 512)
406    }) {
407        return Err(Error::InvalidRequest("invalid batch metadata"));
408    }
409    if request
410        .output_expires_after
411        .as_ref()
412        .is_some_and(|e| e.anchor != "created_at" || !(3600..=2_592_000).contains(&e.seconds))
413    {
414        return Err(Error::InvalidRequest("invalid output retention"));
415    }
416    Ok(())
417}
418
419fn new_key() -> String {
420    let random = std::collections::hash_map::RandomState::new();
421    let mut a = random.build_hasher();
422    a.write_u8(0);
423    let mut b = random.build_hasher();
424    b.write_u8(1);
425    format!("basal-sdk-{:016x}{:016x}", a.finish(), b.finish())
426}
427
428struct JsonlBuffer {
429    bytes: Vec<u8>,
430    limit_hit: bool,
431}
432
433impl Write for JsonlBuffer {
434    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
435        if bytes.len() > 200_000_000usize.saturating_sub(self.bytes.len()) {
436            self.limit_hit = true;
437            return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "batch input exceeds 200 MB"));
438        }
439        self.bytes.extend_from_slice(bytes);
440        Ok(bytes.len())
441    }
442    fn flush(&mut self) -> std::io::Result<()> {
443        Ok(())
444    }
445}
446
447fn encode_requests(requests: &[BatchInputLine]) -> Result<Vec<u8>, Error> {
448    if requests.is_empty() || requests.len() > 50_000 {
449        return Err(Error::InvalidRequest("batch requires 1-50000 requests"));
450    }
451    let mut ids = std::collections::HashSet::new();
452    let mut out = JsonlBuffer { bytes: Vec::new(), limit_hit: false };
453    for record in requests {
454        if record.custom_id.is_empty() || record.custom_id.len() > 256 || !ids.insert(record.custom_id.as_str()) {
455            return Err(Error::InvalidRequest("batch custom_id must be unique and contain 1-256 UTF-8 bytes"));
456        }
457        if record.method != "POST"
458            || record.url != "/v1/basal"
459            || record.body.extra.get("stream").and_then(serde_json::Value::as_bool) == Some(true)
460        {
461            return Err(Error::InvalidRequest("batch records require nonstreaming POST /v1/basal"));
462        }
463        let encoded = serde_json::to_writer(&mut out, record);
464        if out.limit_hit {
465            return Err(Error::InvalidRequest("batch input exceeds 200 MB"));
466        }
467        encoded.map_err(Error::Encode)?;
468        out.write_all(b"\n").map_err(|_| Error::InvalidRequest("batch input exceeds 200 MB"))?;
469    }
470    Ok(out.bytes)
471}
472
473fn parse_result(line: &[u8]) -> Result<BatchResultLine, Error> {
474    let value: serde_json::Value = serde_json::from_slice(line).map_err(Error::Protocol)?;
475    if value.get("response").is_none() || value.get("error").is_none() {
476        return Err(Error::InvalidResponse("batch result requires response and error fields"));
477    }
478    let record: BatchResultLine = serde_json::from_value(value).map_err(Error::Protocol)?;
479    if record.id.is_empty()
480        || record.custom_id.is_empty()
481        || record.response.is_some() == record.error.is_some()
482        || record
483            .response
484            .as_ref()
485            .is_some_and(|r| !(100..=599).contains(&r.status_code) || r.request_id.is_empty() || !r.body.is_object())
486        || record.error.as_ref().is_some_and(|e| e.code.is_empty())
487    {
488        return Err(Error::InvalidResponse("invalid batch result envelope"));
489    }
490    Ok(record)
491}
basal_client/batch_resources.rs

basal_client/
batch_resources.rs
