car-inference 0.60.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,
}

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;
        }
        let base = api_base.trim_end_matches('/');
        let base = base.strip_suffix("/v1/chat/completions").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,
        })
    }

    /// 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`].
#[derive(Debug, Default)]
pub struct LateCostProbe(Mutex<Option<LateCost>>);

impl LateCostProbe {
    /// The lookup for the scope's call, if the provider offers one.
    pub fn take(&self) -> Option<LateCost> {
        self.0
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .take()
    }

    /// The lookup for the call served by `model`, taken: what a settling
    /// spend ticket runs for its own call. 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 slot = self
            .0
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if slot.as_ref().is_some_and(|late| late.serves(model)) {
            slot.take()
        } else {
            None
        }
    }

    /// Leave `late` for the scope's caller, replacing any earlier one.
    pub fn set(&self, late: LateCost) {
        *self
            .0
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(late);
    }
}

tokio::task_local! {
    static LATE_COST: Arc<LateCostProbe>;
}

/// 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);
    }

    /// 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}");
    }
}