nolgia 0.1.0

//! Hand-written. Every request the SDK sends, generated or hand-written, goes
//! through [`execute`]: retries with backoff, the idempotency key, and the
//! `Content-Length: 0` the production load balancer insists on.
//!
//! The generated client (`api.rs`) calls `ClientHooks::exec` for every
//! operation. progenitor implements that trait for `&Client`; implementing it
//! for `Client` itself here wins method resolution (auto-ref specialization,
//! documented on `progenitor_client::ClientHooks`), so no generated file is
//! edited to add this behaviour.

use std::time::Duration;

use progenitor_client::{ClientHooks, ClientInfo, OperationInfo};
use reqwest::{
    Method, Request, Response, StatusCode,
    header::{CONTENT_LENGTH, HeaderName, HeaderValue, RETRY_AFTER},
};

use crate::api;

/// The `Idempotency-Key` request header.
pub(crate) const IDEMPOTENCY_KEY: HeaderName = HeaderName::from_static("idempotency-key");

/// How the SDK retries a request the server refused for now (`429`) or failed
/// to answer (`500`, `502`, `503`, `504`, or a dropped connection).
///
/// Delays grow exponentially from [`initial_backoff`](Self::initial_backoff)
/// with jitter, capped at [`max_backoff`](Self::max_backoff); a `Retry-After`
/// header from the server takes precedence (capped at
/// [`max_retry_after`](Self::max_retry_after)).
///
/// **What is retried.** `429` is retried for every method: the server refused
/// the request before doing anything. A `5xx` or a dropped connection is
/// retried only when repeating the request cannot do the work twice: `GET`,
/// `HEAD`, `PUT`, `DELETE` and `OPTIONS`; any request carrying an
/// `Idempotency-Key`; and the generation submits, which the API itself
/// deduplicates (a repeated identical submit within five minutes is refused
/// with `409` naming the job that already exists, never billed twice).
/// A connection that could not be established is always retried.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetryPolicy {
    /// Retries after the first attempt. `0` disables retries. Default `2`.
    pub max_retries: u32,
    /// The first backoff delay. Default 500 ms.
    pub initial_backoff: Duration,
    /// The longest computed backoff delay. Default 8 s.
    pub max_backoff: Duration,
    /// The longest `Retry-After` the SDK will honour by sleeping; a longer one
    /// surfaces the `429` to the caller instead. Default 60 s.
    pub max_retry_after: Duration,
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_retries: 2,
            initial_backoff: Duration::from_millis(500),
            max_backoff: Duration::from_secs(8),
            max_retry_after: Duration::from_secs(60),
        }
    }
}

impl RetryPolicy {
    /// A policy that never retries.
    pub fn none() -> Self {
        Self {
            max_retries: 0,
            ..Self::default()
        }
    }

    fn backoff(&self, retry: u32) -> Duration {
        let exp = self
            .initial_backoff
            .saturating_mul(2u32.saturating_pow(retry.min(16)));
        let capped = exp.min(self.max_backoff);
        // Equal jitter: half fixed, half random, so concurrent clients spread
        // out without any delay collapsing to zero.
        let half = capped / 2;
        half + half.mul_f64(jitter())
    }
}

/// A number in `[0, 1)`. Not cryptographic; it only spreads retries apart.
fn jitter() -> f64 {
    use std::{
        collections::hash_map::RandomState,
        hash::{BuildHasher, Hasher},
    };
    let mut hasher = RandomState::new().build_hasher();
    hasher.write_u128(
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or_default(),
    );
    (hasher.finish() >> 11) as f64 / (1u64 << 53) as f64
}

/// Configuration the generated client carries for the transport (progenitor's
/// "inner" value).
#[derive(Debug, Clone, Default)]
pub struct Inner {
    pub(crate) retry: RetryPolicy,
    /// Sent as `Idempotency-Key` on every `POST` this client makes, when set.
    pub(crate) idempotency_key: Option<HeaderValue>,
}

/// Operations the API deduplicates on its own (body fingerprint, five
/// minutes), or that only read, so a repeated `POST` is harmless. progenitor's
/// snake_case operation ids.
const POST_SAFE_OPERATIONS: &[&str] = &[
    "generate_image",
    "generate_audio",
    "generate_video",
    "generate3_d",
    "generate_3d",
    "restore_video",
    "remove_background_video",
    "quote_job_cost",
    "validate_mask",
    "validate_motion",
    "estimate_preset_intake_run",
    "assemble_preset_prompt",
    "quote_multicam_run",
];

/// Whether this attempt may be repeated after a `5xx` or a lost response.
fn replay_safe(request: &Request, operation: &str) -> bool {
    matches!(
        *request.method(),
        Method::GET | Method::HEAD | Method::PUT | Method::DELETE | Method::OPTIONS
    ) || request.headers().contains_key(IDEMPOTENCY_KEY)
        || POST_SAFE_OPERATIONS.contains(&operation)
}

/// How many retries a response took, stamped into its extensions so the
/// hand-written layer can tell "my own retry collapsed onto my first submit"
/// from "a duplicate of an earlier call".
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) struct Attempts(pub u32);

/// Send `request` with `http`, retrying under `inner`'s policy.
pub(crate) async fn execute(
    http: &reqwest::Client,
    inner: &Inner,
    mut request: Request,
    operation: &str,
) -> reqwest::Result<Response> {
    prepare(inner, &mut request);
    let policy = &inner.retry;
    let mut retry = 0u32;
    loop {
        // A streamed body cannot be replayed; such a request gets one attempt.
        let replay = if retry < policy.max_retries {
            request.try_clone()
        } else {
            None
        };
        let outcome = http.execute(request).await;
        let Some(next) = replay else {
            return stamp(outcome, retry);
        };
        let delay = match &outcome {
            Ok(response) => match retry_delay(policy, response, &next, operation, retry) {
                Some(delay) => delay,
                None => return stamp(outcome, retry),
            },
            Err(err) if err.is_connect() || (err.is_timeout() && replay_safe(&next, operation)) => {
                policy.backoff(retry)
            }
            Err(err) if err.is_request() && replay_safe(&next, operation) => policy.backoff(retry),
            Err(_) => return stamp(outcome, retry),
        };
        tokio::time::sleep(delay).await;
        request = next;
        retry += 1;
    }
}

fn stamp(outcome: reqwest::Result<Response>, retries: u32) -> reqwest::Result<Response> {
    outcome.map(|mut response| {
        response.extensions_mut().insert(Attempts(retries));
        response
    })
}

fn retry_delay(
    policy: &RetryPolicy,
    response: &Response,
    request: &Request,
    operation: &str,
    retry: u32,
) -> Option<Duration> {
    let status = response.status();
    let retryable = status == StatusCode::TOO_MANY_REQUESTS
        || (matches!(status.as_u16(), 500 | 502 | 503 | 504) && replay_safe(request, operation));
    if !retryable {
        return None;
    }
    match retry_after(response) {
        Some(wait) if wait > policy.max_retry_after => None,
        Some(wait) => Some(wait),
        None => Some(policy.backoff(retry)),
    }
}

/// `Retry-After` in delta-seconds. (The API never sends an HTTP date.)
pub(crate) fn retry_after(response: &Response) -> Option<Duration> {
    response
        .headers()
        .get(RETRY_AFTER)?
        .to_str()
        .ok()?
        .trim()
        .parse::<u64>()
        .ok()
        .map(Duration::from_secs)
}

/// Request fix-ups applied once, before the first attempt.
fn prepare(inner: &Inner, request: &mut Request) {
    let method = request.method().clone();
    // `GET /jobs/{id}/wait` holds the connection for `timeout_seconds` (up to
    // 900). Give that one request room beyond the client-wide timeout, so a
    // generated `wait_for_job().timeout_seconds(600)` call is not cut off.
    if method == Method::GET
        && request.timeout().is_none()
        && request.url().path().ends_with("/wait")
        && let Some(window) = request
            .url()
            .query_pairs()
            .find(|(k, _)| k == "timeout_seconds")
            .and_then(|(_, v)| v.parse::<u64>().ok())
    {
        *request.timeout_mut() = Some(Duration::from_secs(window.min(900) + 30));
    }
    if method == Method::POST
        && let Some(key) = &inner.idempotency_key
        && !request.headers().contains_key(IDEMPOTENCY_KEY)
    {
        request.headers_mut().insert(IDEMPOTENCY_KEY, key.clone());
    }
    // reqwest omits Content-Length on a bodyless POST/PUT/PATCH, and the
    // production load balancer answers such a request `411 Length Required`
    // before the API sees it (several operations, e.g. `POST /jobs/{id}/cancel`
    // and `POST /assets/uploads/{id}/complete`, declare no body).
    if matches!(method, Method::POST | Method::PUT | Method::PATCH) && request.body().is_none() {
        *request.body_mut() = Some(reqwest::Body::from(Vec::<u8>::new()));
        request
            .headers_mut()
            .insert(CONTENT_LENGTH, HeaderValue::from_static("0"));
    }
}

impl ClientHooks<Inner> for api::Client {
    async fn exec(&self, request: Request, info: &OperationInfo) -> reqwest::Result<Response> {
        execute(self.client(), self.inner(), request, info.operation_id).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backoff_grows_and_is_capped() {
        let policy = RetryPolicy::default();
        for retry in 0..10 {
            let delay = policy.backoff(retry);
            assert!(delay <= policy.max_backoff, "{delay:?}");
            assert!(delay >= policy.initial_backoff.min(policy.max_backoff) / 2);
        }
        assert!(policy.backoff(5) >= Duration::from_secs(4));
    }

    #[test]
    fn replay_safety_follows_method_key_and_operation() {
        let client = reqwest::Client::new();
        let get = client.get("http://x/v1/jobs").build().unwrap();
        assert!(replay_safe(&get, "list_jobs"));
        let post = client.post("http://x/v1/agent/messages").build().unwrap();
        assert!(!replay_safe(&post, "submit_agent_message"));
        assert!(replay_safe(&post, "generate_image"));
        let keyed = client
            .post("http://x/v1/agent/messages")
            .header(IDEMPOTENCY_KEY, "k")
            .build()
            .unwrap();
        assert!(replay_safe(&keyed, "submit_agent_message"));
    }

    #[test]
    fn bodyless_post_gets_content_length_and_key() {
        let inner = Inner {
            idempotency_key: Some(HeaderValue::from_static("abc")),
            ..Inner::default()
        };
        let mut request = reqwest::Client::new()
            .post("http://x/v1/jobs/1/cancel")
            .build()
            .unwrap();
        prepare(&inner, &mut request);
        assert_eq!(request.headers()[CONTENT_LENGTH], "0");
        assert_eq!(request.headers()[IDEMPOTENCY_KEY], "abc");
        let mut get = reqwest::Client::new()
            .get("http://x/v1/jobs")
            .build()
            .unwrap();
        prepare(&inner, &mut get);
        assert!(!get.headers().contains_key(IDEMPOTENCY_KEY));
        assert!(!get.headers().contains_key(CONTENT_LENGTH));
    }
}