car-inference 0.63.0

//! A provider's final cost for a call, looked up after the call ended.
//!
//! OpenRouter reports a call's dollar cost on the response itself
//! (`usage.cost`, [`crate::TokenUsage::reported_cost_usd`]). A call that ends
//! without that figure (a stream cut off before its final usage chunk, a
//! response that omitted it) was still billed, and OpenRouter answers for it
//! by generation id at `GET /v1/generation?id=<id>` once the generation is
//! recorded. A meter that charged such a call at its list price or estimate
//! uses the answer to correct the charge upward (`SpendMeter::correct`).
//!
//! The engine learns the generation id where the response arrives, inside a
//! caller's [`scope_late_cost`]: it leaves a ready-to-run [`LateCost`] in the
//! caller's [`LateCostProbe`]. The lookup itself needs the provider credential
//! the call used, which only the engine holds, so the [`LateCost`] carries it
//! and never prints it.
//!
//! The call's spend ticket runs the lookup when it settles
//! (`spend::SpendTicket`), on every way the call ends — a success, an error,
//! a cancelled or broken stream, a candidate the router fell back from — so
//! a caller only provides the scope.

use std::future::Future;
use std::sync::{Arc, Mutex};
use std::time::Duration;

/// How long a late cost is looked for, from the end of the call. After this
/// the call stays as first charged.
pub const LATE_COST_WINDOW: Duration = Duration::from_secs(10 * 60);

/// The wait after a lookup that found nothing yet, doubling up to
/// [`MAX_RETRY_DELAY`].
const FIRST_RETRY_DELAY: Duration = Duration::from_secs(2);
const MAX_RETRY_DELAY: Duration = Duration::from_secs(60);
/// The bound on one lookup request.
const LOOKUP_TIMEOUT: Duration = Duration::from_secs(30);

/// One lookup's answer.
#[derive(Debug, Clone, PartialEq)]
pub enum LateFetch {
    /// The provider's final dollar figure.
    Ready(f64),
    /// Not recorded yet, or a failure worth retrying.
    NotYet,
    /// An answer retrying cannot fix: the credential was refused, the
    /// generation does not belong to it, or the figure is not a dollar amount
    /// (negative, not finite). Treated as no figure.
    Refused,
}

/// A ready-to-run lookup of one call's final cost.
#[derive(Clone)]
pub struct LateCost {
    /// The served model's catalog id and display name, either of which a
    /// caller may hold as the model that served the call.
    model_id: String,
    model_name: String,
    url: String,
    headers: Vec<(String, String)>,
    client: reqwest::Client,
    /// The dispatch that made the call ([`attempt_scope`]), when one was
    /// bound: what ties the lookup to its own spend ticket (car#3630).
    attempt: Option<u64>,
}

impl std::fmt::Debug for LateCost {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // The headers carry the provider credential.
        f.debug_struct("LateCost")
            .field("model_id", &self.model_id)
            .field("url", &self.url)
            .finish_non_exhaustive()
    }
}

/// Whether `id` is shaped like a generation id. Anything else is not put in a
/// URL.
fn valid_generation_id(id: &str) -> bool {
    !id.is_empty()
        && id.len() <= 128
        && id
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
}

impl LateCost {
    /// The OpenRouter lookup of generation `generation_id`, at the same API
    /// base and with the same auth headers as the call itself. `None` when
    /// the id is not a generation id.
    pub fn openrouter(
        client: reqwest::Client,
        api_base: &str,
        generation_id: &str,
        headers: Vec<(String, String)>,
        model_id: &str,
        model_name: &str,
    ) -> Option<Self> {
        if !valid_generation_id(generation_id) {
            return None;
        }
        // The call's base comes as configured: `…/api`, `…/api/v1` or the
        // full chat URL. Each names the same API, so each must reach
        // `…/api/v1/generation`, never `…/v1/v1/generation` (car#3483).
        let base = api_base.trim_end_matches('/');
        let base = base.strip_suffix("/chat/completions").unwrap_or(base);
        let base = base.trim_end_matches('/');
        let base = base.strip_suffix("/v1").unwrap_or(base);
        Some(Self {
            model_id: model_id.to_string(),
            model_name: model_name.to_string(),
            url: format!("{base}/v1/generation?id={generation_id}"),
            headers: headers
                .into_iter()
                .filter(|(name, _)| !name.eq_ignore_ascii_case("content-type"))
                .collect(),
            client,
            attempt: current_attempt(),
        })
    }

    /// This lookup as made by dispatch `attempt`: for a call whose lookup is
    /// built in a task the engine spawned (a stream reader), where the
    /// dispatch's scope does not reach.
    pub(crate) fn with_attempt(mut self, attempt: Option<u64>) -> Self {
        if attempt.is_some() {
            self.attempt = attempt;
        }
        self
    }

    /// Whether this lookup is for the call served by `model` (a catalog id or
    /// a display name).
    pub fn serves(&self, model: &str) -> bool {
        model == self.model_id || model == self.model_name
    }

    /// One lookup. OpenRouter answers 404 until the generation is recorded,
    /// then `{"data": {"total_cost": <usd>, …}}`.
    pub async fn fetch(&self) -> LateFetch {
        let mut request = self.client.get(&self.url).timeout(LOOKUP_TIMEOUT);
        for (name, value) in &self.headers {
            request = request.header(name.as_str(), value.as_str());
        }
        let response = match request.send().await {
            Ok(response) => response,
            Err(_) => return LateFetch::NotYet,
        };
        let status = response.status().as_u16();
        if matches!(status, 401..=403) {
            return LateFetch::Refused;
        }
        if !response.status().is_success() {
            return LateFetch::NotYet;
        }
        let Ok(body) = response.json::<serde_json::Value>().await else {
            return LateFetch::NotYet;
        };
        match body
            .pointer("/data/total_cost")
            .and_then(serde_json::Value::as_f64)
        {
            Some(cost) if cost.is_finite() && cost >= 0.0 => LateFetch::Ready(cost),
            Some(_) => LateFetch::Refused,
            None => LateFetch::NotYet,
        }
    }

    /// Look the cost up until it is recorded, for at most
    /// [`LATE_COST_WINDOW`]. `None` when it never was, or the lookup was
    /// refused.
    pub async fn resolve(&self) -> Option<f64> {
        resolve_late_cost(LATE_COST_WINDOW, || self.fetch()).await
    }
}

/// Run `fetch` until it answers [`LateFetch::Ready`] or [`LateFetch::Refused`],
/// waiting 2 s after the first miss and doubling up to 60 s, and give up when
/// `window` has passed: no lookup starts or runs past it.
pub async fn resolve_late_cost<F, Fut>(window: Duration, mut fetch: F) -> Option<f64>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = LateFetch>,
{
    let deadline = tokio::time::Instant::now() + window;
    let mut delay = FIRST_RETRY_DELAY;
    loop {
        match tokio::time::timeout_at(deadline, fetch()).await {
            Ok(LateFetch::Ready(cost)) => return Some(cost),
            Ok(LateFetch::Refused) | Err(_) => return None,
            Ok(LateFetch::NotYet) => {}
        }
        if tokio::time::Instant::now() + delay >= deadline {
            return None;
        }
        tokio::time::sleep(delay).await;
        delay = (delay * 2).min(MAX_RETRY_DELAY);
    }
}

/// Where the engine leaves a [`LateCost`] for the caller of
/// [`scope_late_cost`]: one per call, so two calls running at once in one
/// scope never take each other's (car#3630).
#[derive(Debug, Default)]
pub struct LateCostProbe(Mutex<Vec<LateCost>>);

impl LateCostProbe {
    fn lookups(&self) -> std::sync::MutexGuard<'_, Vec<LateCost>> {
        self.0
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    /// The latest call's lookup, if the provider offers one.
    pub fn take(&self) -> Option<LateCost> {
        self.lookups().pop()
    }

    /// The lookup for the call served by `model`, taken: the latest one that
    /// serves it. A lookup for another model is left in place, since it
    /// belongs to another call.
    pub fn take_serving(&self, model: &str) -> Option<LateCost> {
        let mut lookups = self.lookups();
        let at = lookups.iter().rposition(|late| late.serves(model))?;
        Some(lookups.remove(at))
    }

    /// The lookup the dispatch `attempt` left for its call served by `model`:
    /// what a settling spend ticket runs for its own call. A lookup another
    /// dispatch left is never taken, whatever model served it. One made
    /// outside any dispatch (no key) is matched by model, as before.
    pub(crate) fn take_for_attempt(&self, attempt: u64, model: &str) -> Option<LateCost> {
        let mut lookups = self.lookups();
        let at = lookups
            .iter()
            .rposition(|late| late.attempt == Some(attempt) && late.serves(model))
            .or_else(|| {
                lookups
                    .iter()
                    .rposition(|late| late.attempt.is_none() && late.serves(model))
            })?;
        Some(lookups.remove(at))
    }

    /// Leave `late` for the scope's caller, replacing an earlier lookup from
    /// the same dispatch.
    pub fn set(&self, late: LateCost) {
        let mut lookups = self.lookups();
        if late.attempt.is_some() {
            lookups.retain(|earlier| earlier.attempt != late.attempt);
        }
        lookups.push(late);
    }
}

tokio::task_local! {
    static LATE_COST: Arc<LateCostProbe>;
    /// The dispatch the current call belongs to ([`attempt_scope`]).
    static LATE_ATTEMPT: u64;
}

/// A fresh dispatch key for [`attempt_scope`].
pub(crate) fn next_attempt() -> u64 {
    static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
    NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
}

/// Run `future` as dispatch `attempt`: a [`LateCost`] made inside it carries
/// that key, so the spend ticket of that dispatch takes it and no other.
pub(crate) async fn attempt_scope<F: Future>(attempt: u64, future: F) -> F::Output {
    LATE_ATTEMPT.scope(attempt, future).await
}

/// The current dispatch's key, to carry into a task the engine spawns (a
/// stream reader), as [`current_probe`] carries the probe.
pub(crate) fn current_attempt() -> Option<u64> {
    LATE_ATTEMPT.try_with(|attempt| *attempt).ok()
}

/// Run `future`, leaving in `probe` the [`LateCost`] of a call it makes to a
/// provider that settles cost late (OpenRouter). The last such call wins.
pub async fn scope_late_cost<F: Future>(probe: Arc<LateCostProbe>, future: F) -> F::Output {
    LATE_COST.scope(probe, future).await
}

/// The probe of the current scope, to carry into a task the engine spawns
/// (a stream reader).
pub(crate) fn current_probe() -> Option<Arc<LateCostProbe>> {
    LATE_COST.try_with(Arc::clone).ok()
}

/// Leave `late` in `probe`.
pub(crate) fn note(probe: Option<&Arc<LateCostProbe>>, late: Option<LateCost>) {
    if let (Some(probe), Some(late)) = (probe, late) {
        probe.set(late);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};

    /// A cost that is recorded after a few misses is found; the misses are
    /// retried on a backoff, never in a tight loop.
    #[tokio::test(start_paused = true)]
    async fn a_cost_recorded_later_is_found() {
        let calls = AtomicU32::new(0);
        let started = tokio::time::Instant::now();
        let found = resolve_late_cost(LATE_COST_WINDOW, || {
            let n = calls.fetch_add(1, Ordering::SeqCst);
            async move {
                if n < 3 {
                    LateFetch::NotYet
                } else {
                    LateFetch::Ready(0.0042)
                }
            }
        })
        .await;
        assert_eq!(found, Some(0.0042));
        assert_eq!(calls.load(Ordering::SeqCst), 4);
        // 2 s + 4 s + 8 s of waiting between the four lookups.
        assert_eq!(started.elapsed(), Duration::from_secs(14));
    }

    /// A cost never recorded is given up on within ten minutes, and no lookup
    /// runs past them, however slow it is.
    #[tokio::test(start_paused = true)]
    async fn the_lookup_gives_up_within_ten_minutes() {
        let calls = AtomicU32::new(0);
        let started = tokio::time::Instant::now();
        let found = resolve_late_cost(LATE_COST_WINDOW, || {
            calls.fetch_add(1, Ordering::SeqCst);
            async { LateFetch::NotYet }
        })
        .await;
        assert_eq!(found, None);
        assert!(
            started.elapsed() <= LATE_COST_WINDOW,
            "{:?}",
            started.elapsed()
        );
        assert!(calls.load(Ordering::SeqCst) > 5);

        let started = tokio::time::Instant::now();
        let hung = resolve_late_cost(LATE_COST_WINDOW, || async {
            tokio::time::sleep(Duration::from_secs(3600)).await;
            LateFetch::Ready(1.0)
        })
        .await;
        assert_eq!(hung, None);
        assert_eq!(started.elapsed(), LATE_COST_WINDOW);
    }

    /// A refused lookup stops at once.
    #[tokio::test(start_paused = true)]
    async fn a_refused_lookup_stops() {
        let calls = AtomicU32::new(0);
        let found = resolve_late_cost(LATE_COST_WINDOW, || {
            calls.fetch_add(1, Ordering::SeqCst);
            async { LateFetch::Refused }
        })
        .await;
        assert_eq!(found, None);
        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    /// Two calls to one model at once in one scope each get their own lookup,
    /// whichever settles first (car#3630).
    #[tokio::test]
    async fn concurrent_calls_in_one_scope_keep_their_own_lookups() {
        let probe = Arc::new(LateCostProbe::default());
        let lookup = |generation: &'static str| async move {
            LateCost::openrouter(
                reqwest::Client::new(),
                "https://openrouter.ai/api",
                generation,
                Vec::new(),
                "openrouter/same",
                "same",
            )
        };
        let (a, b) = (next_attempt(), next_attempt());
        note(Some(&probe), attempt_scope(a, lookup("gen-a")).await);
        note(Some(&probe), attempt_scope(b, lookup("gen-b")).await);
        let taken_a = probe
            .take_for_attempt(a, "openrouter/same")
            .expect("a's lookup");
        assert!(taken_a.url.ends_with("id=gen-a"), "{}", taken_a.url);
        assert!(probe.take_for_attempt(a, "openrouter/same").is_none());
        let taken_b = probe.take_for_attempt(b, "same").expect("b's lookup");
        assert!(taken_b.url.ends_with("id=gen-b"), "{}", taken_b.url);
        // A dispatch's later lookup replaces its earlier one, not another's.
        note(Some(&probe), attempt_scope(a, lookup("gen-a1")).await);
        note(Some(&probe), attempt_scope(a, lookup("gen-a2")).await);
        note(Some(&probe), attempt_scope(b, lookup("gen-b1")).await);
        assert!(probe
            .take_for_attempt(a, "same")
            .unwrap()
            .url
            .ends_with("id=gen-a2"));
        assert!(probe.take().unwrap().url.ends_with("id=gen-b1"));
        assert!(probe.take().is_none());
    }

    /// Every configured shape of the OpenRouter base reaches the one
    /// generation URL (car#3483).
    #[test]
    fn every_openrouter_base_shape_reaches_one_generation_url() {
        for base in [
            "https://openrouter.ai/api",
            "https://openrouter.ai/api/",
            "https://openrouter.ai/api/v1",
            "https://openrouter.ai/api/v1/",
            "https://openrouter.ai/api/v1/chat/completions",
            "https://openrouter.ai/api/chat/completions",
        ] {
            let lookup = LateCost::openrouter(
                reqwest::Client::new(),
                base,
                "gen-1",
                Vec::new(),
                "openrouter/test",
                "test",
            )
            .expect("a generation id");
            assert_eq!(
                lookup.url, "https://openrouter.ai/api/v1/generation?id=gen-1",
                "{base}"
            );
        }
    }

    /// The lookup asks OpenRouter's generation endpoint for the call's id with
    /// the call's own credential, reads `data.total_cost`, and treats a 404 as
    /// not recorded yet.
    #[tokio::test]
    async fn the_openrouter_lookup_reads_total_cost_by_generation_id() {
        use wiremock::matchers::{header, method, path, query_param};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .and(path("/v1/generation"))
            .and(query_param("id", "gen-recorded"))
            .and(header("authorization", "Bearer k"))
            .respond_with(ResponseTemplate::new(200).set_body_json(
                serde_json::json!({"data": {"id": "gen-recorded", "total_cost": 0.0125}}),
            ))
            .mount(&server)
            .await;
        Mock::given(method("GET"))
            .and(path("/v1/generation"))
            .and(query_param("id", "gen-negative"))
            .respond_with(ResponseTemplate::new(200).set_body_json(
                serde_json::json!({"data": {"id": "gen-negative", "total_cost": -0.5}}),
            ))
            .mount(&server)
            .await;
        Mock::given(method("GET"))
            .and(path("/v1/generation"))
            .and(query_param("id", "gen-pending"))
            .respond_with(ResponseTemplate::new(404))
            .mount(&server)
            .await;
        let headers = vec![
            ("Authorization".to_string(), "Bearer k".to_string()),
            ("Content-Type".to_string(), "application/json".to_string()),
        ];
        let lookup = |id: &str| {
            LateCost::openrouter(
                crate::tls_client::test_seam::trust_store_independent_client(),
                &server.uri(),
                id,
                headers.clone(),
                "openrouter/m",
                "m",
            )
            .unwrap()
        };
        assert_eq!(
            lookup("gen-recorded").fetch().await,
            LateFetch::Ready(0.0125)
        );
        assert_eq!(lookup("gen-pending").fetch().await, LateFetch::NotYet);
        // A negative figure is no figure, and asking again will not change it.
        assert_eq!(lookup("gen-negative").fetch().await, LateFetch::Refused);
        assert!(lookup("gen-recorded").serves("m"));
        assert!(lookup("gen-recorded").serves("openrouter/m"));
        assert!(!lookup("gen-recorded").serves("other"));
        assert!(LateCost::openrouter(
            crate::tls_client::test_seam::trust_store_independent_client(),
            &server.uri(),
            "gen&x=1",
            headers.clone(),
            "openrouter/m",
            "m"
        )
        .is_none());
        let printed = format!("{:?}", lookup("gen-recorded"));
        assert!(!printed.contains("Bearer"), "{printed}");
    }
}