Skip to main content

openrouter_batch/
transport.rs

1//! Wire types and the transport seam. The only async I/O surface of the
2//! batch path; e2e wire fakes sit below the trait, unit fakes above it.
3
4use async_trait::async_trait;
5use serde::{Deserialize, Serialize};
6use std::fmt;
7
8/// The API shape every request in one normalization batch uses.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum BatchEndpoint {
11    ChatCompletions,
12}
13
14impl BatchEndpoint {
15    pub fn as_str(self) -> &'static str {
16        match self {
17            Self::ChatCompletions => "/v1/chat/completions",
18        }
19    }
20}
21
22impl fmt::Display for BatchEndpoint {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.write_str(self.as_str())
25    }
26}
27
28/// One request inside a batch: a deterministic join key plus the endpoint
29/// body. `custom_id` must be unique within the batch.
30#[derive(Debug, Clone, Serialize)]
31pub struct BatchItem {
32    pub custom_id: String,
33    pub body: serde_json::Value,
34}
35
36/// A batch submission.
37///
38/// Field declaration order is wire order: the API stream-parses the body and
39/// rejects it with 400 when `requests` is serialized before `endpoint` and
40/// `model` (quickstart warning; pinned by a serialization test).
41#[derive(Debug, Clone, Serialize)]
42pub struct SubmitBatch {
43    endpoint: &'static str,
44    model: String,
45    requests: Vec<BatchItem>,
46}
47
48/// One synchronous embedding request with array input.
49#[derive(Debug, Clone, Serialize)]
50pub struct EmbeddingRequest {
51    model: String,
52    input: Vec<String>,
53    encoding_format: &'static str,
54}
55
56impl EmbeddingRequest {
57    pub fn new(model: impl Into<String>, input: Vec<String>) -> Self {
58        Self {
59            model: model.into(),
60            input,
61            encoding_format: "float",
62        }
63    }
64}
65
66/// One indexed vector in a synchronous embedding response.
67#[derive(Debug, Clone, Deserialize)]
68pub struct EmbeddingData {
69    pub index: i64,
70    pub embedding: Vec<f32>,
71}
72
73/// The typed fields needed from a synchronous embedding response. Internal
74/// to the client: the public decode is `facade::embedding_vectors` over the
75/// stored raw payload.
76#[derive(Debug, Clone, Deserialize)]
77pub struct EmbeddingResponseBody {
78    pub data: Vec<EmbeddingData>,
79}
80
81/// A successful HTTP response before typed result validation.
82#[derive(Debug, Clone)]
83pub struct RawEmbeddingResponse {
84    pub raw: serde_json::Value,
85    pub provider_request_id: Option<String>,
86}
87
88impl SubmitBatch {
89    pub fn new(
90        endpoint: BatchEndpoint,
91        model: impl Into<String>,
92        requests: Vec<BatchItem>,
93    ) -> Self {
94        Self {
95            endpoint: endpoint.as_str(),
96            model: model.into(),
97            requests,
98        }
99    }
100
101    /// The batch's request items, for tests and callers that must inspect
102    /// the exact wire payload after submission.
103    pub fn requests(&self) -> &[BatchItem] {
104        &self.requests
105    }
106}
107
108/// Batch status lifecycle: `validating -> in_progress -> finalizing ->
109/// completed`, with `failed`, `expired`, and `cancelled` as the other
110/// terminal outcomes.
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
112#[serde(rename_all = "snake_case")]
113pub enum BatchStatus {
114    Validating,
115    InProgress,
116    Finalizing,
117    Completed,
118    Failed,
119    Expired,
120    Cancelling,
121    Cancelled,
122    /// A status added after this client was written. Never terminal: the
123    /// caller keeps polling instead of guessing.
124    #[serde(other)]
125    Unknown,
126}
127
128impl BatchStatus {
129    pub fn is_terminal(self) -> bool {
130        matches!(
131            self,
132            Self::Completed | Self::Failed | Self::Expired | Self::Cancelled
133        )
134    }
135}
136
137/// `request_counts.completed + failed == total` is a completeness sanity
138/// check, not a success signal: failed requests stay in-band as result lines.
139#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
140pub struct RequestCounts {
141    #[serde(default)]
142    pub total: i64,
143    #[serde(default)]
144    pub completed: i64,
145    #[serde(default)]
146    pub failed: i64,
147}
148
149/// Batch usage. `cost` is dollars as a plain f64: the API emits
150/// scientific-notation floats (`1.33e-05`), so fixed-point types are wrong.
151#[derive(Debug, Clone, Copy, Default, Deserialize)]
152pub struct BatchUsage {
153    #[serde(default)]
154    pub prompt_tokens: Option<i64>,
155    #[serde(default)]
156    pub completion_tokens: Option<i64>,
157    #[serde(default)]
158    pub total_tokens: Option<i64>,
159    #[serde(default)]
160    pub cost: Option<f64>,
161    #[serde(default)]
162    pub is_byok: Option<bool>,
163}
164
165/// One completed request's result line, decoded per line: a line that does
166/// not match the documented envelope arrives as `Undecodable` with the raw
167/// JSON preserved, so one malformed line never fails a completed batch and
168/// callers can persist it and mark the matched request retryable
169/// (context/05). Unknown fields on a parsed line survive a decode/encode
170/// round trip (pinned by the decode tests).
171#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(untagged)]
173pub enum BatchResultLine {
174    Parsed(ResultLine),
175    Undecodable(serde_json::Value),
176}
177
178/// The documented result envelope: exactly one of `response` or `error` is
179/// populated.
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub struct ResultLine {
182    #[serde(default)]
183    pub id: Option<String>,
184    pub custom_id: String,
185    #[serde(default)]
186    pub response: Option<ResultResponse>,
187    #[serde(default)]
188    pub error: Option<serde_json::Value>,
189    #[serde(flatten)]
190    pub extra: serde_json::Map<String, serde_json::Value>,
191}
192
193/// A per-line failure of one completed request. The `Display` strings are
194/// the durable `last_error` texts operators see in Postgres.
195#[derive(Debug, thiserror::Error)]
196pub enum BatchLineError {
197    /// The line did not match the envelope; `BatchResultLine::raw` still
198    /// carries the provider payload for storage.
199    #[error("result envelope decode: line does not match the documented result envelope")]
200    Undecodable,
201    #[error("line carries an error instead of a response")]
202    NoResponse,
203    #[error("provider returned status {0}")]
204    ProviderStatus(u16),
205    #[error("response body has no assistant message content")]
206    NoContent,
207}
208
209impl BatchResultLine {
210    /// The deterministic join key — also recovered from an undecodable
211    /// line, so its failure can still be routed to the matching request.
212    pub fn custom_id(&self) -> Option<&str> {
213        match self {
214            Self::Parsed(line) => Some(&line.custom_id),
215            Self::Undecodable(raw) => raw.get("custom_id").and_then(serde_json::Value::as_str),
216        }
217    }
218
219    /// The provider payload exactly as received: the Lineage value for
220    /// `raw_provider_results.raw_line`.
221    pub fn raw(&self) -> serde_json::Value {
222        match self {
223            Self::Parsed(line) => {
224                serde_json::to_value(line).expect("result line re-serialization cannot fail")
225            }
226            Self::Undecodable(raw) => raw.clone(),
227        }
228    }
229
230    /// The chat-completions outcome of one line: the assistant message
231    /// content for local schema validation, or a typed per-line error.
232    pub fn chat_content(&self) -> Result<&str, BatchLineError> {
233        let line = match self {
234            Self::Parsed(line) => line,
235            Self::Undecodable(_) => return Err(BatchLineError::Undecodable),
236        };
237        let response = line.response.as_ref().ok_or(BatchLineError::NoResponse)?;
238        if response.status_code != 200 {
239            return Err(BatchLineError::ProviderStatus(response.status_code));
240        }
241        chat_completion_content(&response.body).ok_or(BatchLineError::NoContent)
242    }
243}
244
245/// The assistant message content of one chat-completions response body.
246/// Shared by the batch result decode and the synchronous client, so the two
247/// paths cannot drift on where the content lives.
248pub fn chat_completion_content(body: &serde_json::Value) -> Option<&str> {
249    body.get("choices")
250        .and_then(|choices| choices.get(0))
251        .and_then(|choice| choice.get("message"))
252        .and_then(|message| message.get("content"))
253        .and_then(serde_json::Value::as_str)
254}
255
256/// The per-request success envelope: an HTTP status plus the endpoint's
257/// response body (`chat.completion` or embedding list).
258#[derive(Debug, Clone, Serialize, Deserialize)]
259pub struct ResultResponse {
260    pub status_code: u16,
261    #[serde(default)]
262    pub request_id: Option<String>,
263    pub body: serde_json::Value,
264    #[serde(flatten)]
265    pub extra: serde_json::Map<String, serde_json::Value>,
266}
267
268/// One synchronous chat-completions response: the provider payload exactly
269/// as received, plus the request ID header when present. The eval harness
270/// decodes it with `chat_completion_content` — the same traversal the batch
271/// result lines use.
272#[derive(Debug, Clone)]
273pub struct RawChatResponse {
274    pub raw: serde_json::Value,
275    pub provider_request_id: Option<String>,
276}
277
278impl RawChatResponse {
279    /// The assistant message content, or a typed decode failure.
280    pub fn content(&self) -> Result<&str, BatchLineError> {
281        chat_completion_content(&self.raw).ok_or(BatchLineError::NoContent)
282    }
283}
284
285/// The batch object returned by both submit (202, `validating`) and poll.
286#[derive(Debug, Clone, Deserialize)]
287pub struct BatchObject {
288    pub id: String,
289    #[serde(default)]
290    pub endpoint: Option<String>,
291    #[serde(default)]
292    pub model: Option<String>,
293    pub status: BatchStatus,
294    #[serde(default)]
295    pub request_counts: Option<RequestCounts>,
296    #[serde(default)]
297    pub usage: Option<BatchUsage>,
298    /// Inline result lines, decoded per line (see `BatchResultLine`);
299    /// `None` for every non-`completed` status. There is no separate
300    /// results-download endpoint.
301    #[serde(default)]
302    pub results: Option<Vec<BatchResultLine>>,
303    #[serde(default)]
304    pub error: Option<serde_json::Value>,
305    #[serde(default)]
306    pub created_at: Option<i64>,
307    #[serde(default)]
308    pub finalized_at: Option<i64>,
309    /// Provider request identifier from the HTTP response header.
310    #[serde(skip)]
311    pub provider_request_id: Option<String>,
312}
313
314impl BatchObject {
315    /// Seconds since `created_at` — the operator-visible batch age
316    /// (context/11 `provider_batch_age_seconds`).
317    pub fn batch_age_seconds(&self) -> Option<i64> {
318        let created = self.created_at?;
319        let end = self
320            .finalized_at
321            .unwrap_or_else(|| finalized_or_now(None).unwrap_or(0));
322        Some(end - created)
323    }
324}
325
326fn finalized_or_now(finalized_at: Option<i64>) -> Option<i64> {
327    finalized_at.or_else(|| {
328        std::time::SystemTime::now()
329            .duration_since(std::time::UNIX_EPOCH)
330            .ok()
331            .map(|d| d.as_secs() as i64)
332    })
333}
334
335/// Typed transport failure. Retry classification lives here so every caller
336/// shares one policy (context/11): timeouts, connect failures, 408/429, and
337/// 5xx are retryable; other 4xx and decode errors are not.
338#[derive(Debug, thiserror::Error)]
339pub enum TransportError {
340    /// The API answered with a non-2xx status. `message` is the provider's
341    /// `error.message` when present, else the truncated body.
342    #[error("provider rejected the request ({status}): {message}")]
343    Api {
344        status: u16,
345        message: String,
346        provider_request_id: Option<String>,
347    },
348    /// A request key exceeded the provider's 64-character cap; the batch
349    /// never left the process. Non-retryable: the key scheme is a code fix.
350    #[error("custom_id exceeds the provider's 64-character cap: {custom_id}")]
351    CustomIdTooLong { custom_id: String },
352    #[error("http transport failure: {source}")]
353    Http {
354        #[source]
355        source: reqwest::Error,
356        provider_request_id: Option<String>,
357    },
358    #[error("response decode failure: {source}")]
359    Decode {
360        #[source]
361        source: serde_json::Error,
362        provider_request_id: Option<String>,
363    },
364}
365
366impl From<reqwest::Error> for TransportError {
367    fn from(source: reqwest::Error) -> Self {
368        Self::Http {
369            source,
370            provider_request_id: None,
371        }
372    }
373}
374
375impl From<serde_json::Error> for TransportError {
376    fn from(source: serde_json::Error) -> Self {
377        Self::Decode {
378            source,
379            provider_request_id: None,
380        }
381    }
382}
383
384impl TransportError {
385    pub fn retryable(&self) -> bool {
386        fn status_retryable(status: u16) -> bool {
387            status == 408 || status == 429 || status >= 500
388        }
389        match self {
390            Self::Api { status, .. } => status_retryable(*status),
391            Self::CustomIdTooLong { .. } => false,
392            Self::Http { source, .. } => {
393                source.is_timeout()
394                    || source.is_connect()
395                    || source.is_body()
396                    || source.is_decode()
397                    || source
398                        .status()
399                        .is_some_and(|status| status_retryable(status.as_u16()))
400            }
401            Self::Decode { .. } => false,
402        }
403    }
404
405    /// Return the provider request ID when an HTTP response supplied one.
406    pub fn provider_request_id(&self) -> Option<&str> {
407        match self {
408            Self::CustomIdTooLong { .. } => None,
409            Self::Api {
410                provider_request_id,
411                ..
412            } => provider_request_id.as_deref(),
413            Self::Http {
414                provider_request_id,
415                ..
416            }
417            | Self::Decode {
418                provider_request_id,
419                ..
420            } => provider_request_id.as_deref(),
421        }
422    }
423
424    /// Whether this is a 404: on poll it means the provider accepted the
425    /// batch but does not expose it yet (eventual consistency after the
426    /// 202, observed live 2026-08-30), not that the batch is gone.
427    pub fn not_found(&self) -> bool {
428        matches!(self, Self::Api { status: 404, .. })
429    }
430}
431
432/// The transport seam: submit and poll, nothing else. Implementations must
433/// treat any 2xx with a parseable batch body as success (submit answers 202).
434#[async_trait]
435pub trait BatchTransport: Send + Sync {
436    async fn submit(
437        &self,
438        batch: &SubmitBatch,
439        idempotency_key: &str,
440    ) -> Result<BatchObject, TransportError>;
441    async fn get(&self, batch_id: &str) -> Result<BatchObject, TransportError>;
442}
443
444/// The synchronous embedding I/O seam.
445#[async_trait]
446pub trait EmbeddingTransport: Send + Sync {
447    async fn embed(
448        &self,
449        model: &str,
450        inputs: &[String],
451        idempotency_key: &str,
452    ) -> Result<RawEmbeddingResponse, TransportError>;
453}