openrouter-batch 0.1.2

//! Wire types and the transport seam. The only async I/O surface of the
//! batch path; e2e wire fakes sit below the trait, unit fakes above it.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::fmt;

/// The API shape every request in one normalization batch uses.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BatchEndpoint {
    ChatCompletions,
}

impl BatchEndpoint {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::ChatCompletions => "/v1/chat/completions",
        }
    }
}

impl fmt::Display for BatchEndpoint {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

/// One request inside a batch: a deterministic join key plus the endpoint
/// body. `custom_id` must be unique within the batch.
#[derive(Debug, Clone, Serialize)]
pub struct BatchItem {
    pub custom_id: String,
    pub body: serde_json::Value,
}

/// A batch submission.
///
/// Field declaration order is wire order: the API stream-parses the body and
/// rejects it with 400 when `requests` is serialized before `endpoint` and
/// `model` (quickstart warning; pinned by a serialization test).
#[derive(Debug, Clone, Serialize)]
pub struct SubmitBatch {
    endpoint: &'static str,
    model: String,
    requests: Vec<BatchItem>,
}

/// One synchronous embedding request with array input.
#[derive(Debug, Clone, Serialize)]
pub struct EmbeddingRequest {
    model: String,
    input: Vec<String>,
    encoding_format: &'static str,
}

impl EmbeddingRequest {
    pub fn new(model: impl Into<String>, input: Vec<String>) -> Self {
        Self {
            model: model.into(),
            input,
            encoding_format: "float",
        }
    }
}

/// One indexed vector in a synchronous embedding response.
#[derive(Debug, Clone, Deserialize)]
pub struct EmbeddingData {
    pub index: i64,
    pub embedding: Vec<f32>,
}

/// The typed fields needed from a synchronous embedding response. Internal
/// to the client: the public decode is `facade::embedding_vectors` over the
/// stored raw payload.
#[derive(Debug, Clone, Deserialize)]
pub struct EmbeddingResponseBody {
    pub data: Vec<EmbeddingData>,
}

/// A successful HTTP response before typed result validation.
#[derive(Debug, Clone)]
pub struct RawEmbeddingResponse {
    pub raw: serde_json::Value,
    pub provider_request_id: Option<String>,
}

impl SubmitBatch {
    pub fn new(
        endpoint: BatchEndpoint,
        model: impl Into<String>,
        requests: Vec<BatchItem>,
    ) -> Self {
        Self {
            endpoint: endpoint.as_str(),
            model: model.into(),
            requests,
        }
    }

    /// The batch's request items, for tests and callers that must inspect
    /// the exact wire payload after submission.
    pub fn requests(&self) -> &[BatchItem] {
        &self.requests
    }
}

/// Batch status lifecycle: `validating -> in_progress -> finalizing ->
/// completed`, with `failed`, `expired`, and `cancelled` as the other
/// terminal outcomes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BatchStatus {
    Validating,
    InProgress,
    Finalizing,
    Completed,
    Failed,
    Expired,
    Cancelling,
    Cancelled,
    /// A status added after this client was written. Never terminal: the
    /// caller keeps polling instead of guessing.
    #[serde(other)]
    Unknown,
}

impl BatchStatus {
    pub fn is_terminal(self) -> bool {
        matches!(
            self,
            Self::Completed | Self::Failed | Self::Expired | Self::Cancelled
        )
    }
}

/// `request_counts.completed + failed == total` is a completeness sanity
/// check, not a success signal: failed requests stay in-band as result lines.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
pub struct RequestCounts {
    #[serde(default)]
    pub total: i64,
    #[serde(default)]
    pub completed: i64,
    #[serde(default)]
    pub failed: i64,
}

/// Batch usage. `cost` is dollars as a plain f64: the API emits
/// scientific-notation floats (`1.33e-05`), so fixed-point types are wrong.
#[derive(Debug, Clone, Copy, Default, Deserialize)]
pub struct BatchUsage {
    #[serde(default)]
    pub prompt_tokens: Option<i64>,
    #[serde(default)]
    pub completion_tokens: Option<i64>,
    #[serde(default)]
    pub total_tokens: Option<i64>,
    #[serde(default)]
    pub cost: Option<f64>,
    #[serde(default)]
    pub is_byok: Option<bool>,
}

/// One completed request's result line, decoded per line: a line that does
/// not match the documented envelope arrives as `Undecodable` with the raw
/// JSON preserved, so one malformed line never fails a completed batch and
/// callers can persist it and mark the matched request retryable
/// (context/05). Unknown fields on a parsed line survive a decode/encode
/// round trip (pinned by the decode tests).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum BatchResultLine {
    Parsed(ResultLine),
    Undecodable(serde_json::Value),
}

/// The documented result envelope: exactly one of `response` or `error` is
/// populated.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResultLine {
    #[serde(default)]
    pub id: Option<String>,
    pub custom_id: String,
    #[serde(default)]
    pub response: Option<ResultResponse>,
    #[serde(default)]
    pub error: Option<serde_json::Value>,
    #[serde(flatten)]
    pub extra: serde_json::Map<String, serde_json::Value>,
}

/// A per-line failure of one completed request. The `Display` strings are
/// the durable `last_error` texts operators see in Postgres.
#[derive(Debug, thiserror::Error)]
pub enum BatchLineError {
    /// The line did not match the envelope; `BatchResultLine::raw` still
    /// carries the provider payload for storage.
    #[error("result envelope decode: line does not match the documented result envelope")]
    Undecodable,
    #[error("line carries an error instead of a response")]
    NoResponse,
    #[error("provider returned status {0}")]
    ProviderStatus(u16),
    #[error("response body has no assistant message content")]
    NoContent,
}

impl BatchResultLine {
    /// The deterministic join key — also recovered from an undecodable
    /// line, so its failure can still be routed to the matching request.
    pub fn custom_id(&self) -> Option<&str> {
        match self {
            Self::Parsed(line) => Some(&line.custom_id),
            Self::Undecodable(raw) => raw.get("custom_id").and_then(serde_json::Value::as_str),
        }
    }

    /// The provider payload exactly as received: the Lineage value for
    /// `raw_provider_results.raw_line`.
    pub fn raw(&self) -> serde_json::Value {
        match self {
            Self::Parsed(line) => {
                serde_json::to_value(line).expect("result line re-serialization cannot fail")
            }
            Self::Undecodable(raw) => raw.clone(),
        }
    }

    /// The chat-completions outcome of one line: the assistant message
    /// content for local schema validation, or a typed per-line error.
    pub fn chat_content(&self) -> Result<&str, BatchLineError> {
        let line = match self {
            Self::Parsed(line) => line,
            Self::Undecodable(_) => return Err(BatchLineError::Undecodable),
        };
        let response = line.response.as_ref().ok_or(BatchLineError::NoResponse)?;
        if response.status_code != 200 {
            return Err(BatchLineError::ProviderStatus(response.status_code));
        }
        chat_completion_content(&response.body).ok_or(BatchLineError::NoContent)
    }
}

/// The assistant message content of one chat-completions response body.
/// Shared by the batch result decode and the synchronous client, so the two
/// paths cannot drift on where the content lives.
pub fn chat_completion_content(body: &serde_json::Value) -> Option<&str> {
    body.get("choices")
        .and_then(|choices| choices.get(0))
        .and_then(|choice| choice.get("message"))
        .and_then(|message| message.get("content"))
        .and_then(serde_json::Value::as_str)
}

/// The per-request success envelope: an HTTP status plus the endpoint's
/// response body (`chat.completion` or embedding list).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResultResponse {
    pub status_code: u16,
    #[serde(default)]
    pub request_id: Option<String>,
    pub body: serde_json::Value,
    #[serde(flatten)]
    pub extra: serde_json::Map<String, serde_json::Value>,
}

/// One synchronous chat-completions response: the provider payload exactly
/// as received, plus the request ID header when present. The eval harness
/// decodes it with `chat_completion_content` — the same traversal the batch
/// result lines use.
#[derive(Debug, Clone)]
pub struct RawChatResponse {
    pub raw: serde_json::Value,
    pub provider_request_id: Option<String>,
}

impl RawChatResponse {
    /// The assistant message content, or a typed decode failure.
    pub fn content(&self) -> Result<&str, BatchLineError> {
        chat_completion_content(&self.raw).ok_or(BatchLineError::NoContent)
    }
}

/// The batch object returned by both submit (202, `validating`) and poll.
#[derive(Debug, Clone, Deserialize)]
pub struct BatchObject {
    pub id: String,
    #[serde(default)]
    pub endpoint: Option<String>,
    #[serde(default)]
    pub model: Option<String>,
    pub status: BatchStatus,
    #[serde(default)]
    pub request_counts: Option<RequestCounts>,
    #[serde(default)]
    pub usage: Option<BatchUsage>,
    /// Inline result lines, decoded per line (see `BatchResultLine`);
    /// `None` for every non-`completed` status. There is no separate
    /// results-download endpoint.
    #[serde(default)]
    pub results: Option<Vec<BatchResultLine>>,
    #[serde(default)]
    pub error: Option<serde_json::Value>,
    #[serde(default)]
    pub created_at: Option<i64>,
    #[serde(default)]
    pub finalized_at: Option<i64>,
    /// Provider request identifier from the HTTP response header.
    #[serde(skip)]
    pub provider_request_id: Option<String>,
}

impl BatchObject {
    /// Seconds since `created_at` — the operator-visible batch age
    /// (context/11 `provider_batch_age_seconds`).
    pub fn batch_age_seconds(&self) -> Option<i64> {
        let created = self.created_at?;
        let end = self
            .finalized_at
            .unwrap_or_else(|| finalized_or_now(None).unwrap_or(0));
        Some(end - created)
    }
}

fn finalized_or_now(finalized_at: Option<i64>) -> Option<i64> {
    finalized_at.or_else(|| {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .ok()
            .map(|d| d.as_secs() as i64)
    })
}

/// Typed transport failure. Retry classification lives here so every caller
/// shares one policy (context/11): timeouts, connect failures, 408/429, and
/// 5xx are retryable; other 4xx and decode errors are not.
#[derive(Debug, thiserror::Error)]
pub enum TransportError {
    /// The API answered with a non-2xx status. `message` is the provider's
    /// `error.message` when present, else the truncated body.
    #[error("provider rejected the request ({status}): {message}")]
    Api {
        status: u16,
        message: String,
        provider_request_id: Option<String>,
    },
    /// A request key exceeded the provider's 64-character cap; the batch
    /// never left the process. Non-retryable: the key scheme is a code fix.
    #[error("custom_id exceeds the provider's 64-character cap: {custom_id}")]
    CustomIdTooLong { custom_id: String },
    #[error("http transport failure: {source}")]
    Http {
        #[source]
        source: reqwest::Error,
        provider_request_id: Option<String>,
    },
    #[error("response decode failure: {source}")]
    Decode {
        #[source]
        source: serde_json::Error,
        provider_request_id: Option<String>,
    },
}

impl From<reqwest::Error> for TransportError {
    fn from(source: reqwest::Error) -> Self {
        Self::Http {
            source,
            provider_request_id: None,
        }
    }
}

impl From<serde_json::Error> for TransportError {
    fn from(source: serde_json::Error) -> Self {
        Self::Decode {
            source,
            provider_request_id: None,
        }
    }
}

impl TransportError {
    pub fn retryable(&self) -> bool {
        fn status_retryable(status: u16) -> bool {
            status == 408 || status == 429 || status >= 500
        }
        match self {
            Self::Api { status, .. } => status_retryable(*status),
            Self::CustomIdTooLong { .. } => false,
            Self::Http { source, .. } => {
                source.is_timeout()
                    || source.is_connect()
                    || source.is_body()
                    || source.is_decode()
                    || source
                        .status()
                        .is_some_and(|status| status_retryable(status.as_u16()))
            }
            Self::Decode { .. } => false,
        }
    }

    /// Return the provider request ID when an HTTP response supplied one.
    pub fn provider_request_id(&self) -> Option<&str> {
        match self {
            Self::CustomIdTooLong { .. } => None,
            Self::Api {
                provider_request_id,
                ..
            } => provider_request_id.as_deref(),
            Self::Http {
                provider_request_id,
                ..
            }
            | Self::Decode {
                provider_request_id,
                ..
            } => provider_request_id.as_deref(),
        }
    }

    /// Whether this is a 404: on poll it means the provider accepted the
    /// batch but does not expose it yet (eventual consistency after the
    /// 202, observed live 2026-08-30), not that the batch is gone.
    pub fn not_found(&self) -> bool {
        matches!(self, Self::Api { status: 404, .. })
    }
}

/// The transport seam: submit and poll, nothing else. Implementations must
/// treat any 2xx with a parseable batch body as success (submit answers 202).
#[async_trait]
pub trait BatchTransport: Send + Sync {
    async fn submit(
        &self,
        batch: &SubmitBatch,
        idempotency_key: &str,
    ) -> Result<BatchObject, TransportError>;
    async fn get(&self, batch_id: &str) -> Result<BatchObject, TransportError>;
}

/// The synchronous embedding I/O seam.
#[async_trait]
pub trait EmbeddingTransport: Send + Sync {
    async fn embed(
        &self,
        model: &str,
        inputs: &[String],
        idempotency_key: &str,
    ) -> Result<RawEmbeddingResponse, TransportError>;
}