ai-dispatch 10.52.1

// Live model price feed (https://llm-prices.agent-tools.org/v1/prices.json),
// cached under the aid home. Only exact ids and aliases price a model; the
// cache is refreshed out of band so the dispatch path never blocks on the network.
// Exports: Feed, feed_lookup, maybe_refresh
// Deps: crate::paths, crate::feed_refresh, chrono, serde, serde_json

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::time::Duration;

const FEED_URL: &str = "https://llm-prices.agent-tools.org/v1/prices.json";
const CACHE_FILE: &str = "prices.json";
/// How long a locally cached feed is trusted before a refresh is attempted.
/// 24h matches the update-check cache: model prices move on a daily cadence and
/// a day-old price is never materially wrong for an estimate. `stale` and
/// `age_seconds` from the server guard the *content* of the cache — a stale
/// feed is never written over a fresh one.
const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
/// An excessively old feed (`age_seconds` from the server) is not trusted.
const MAX_SERVER_AGE: i64 = 24 * 60 * 60;

/// A canonical priced model from the feed.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct FeedModel {
    pub id: String,
    #[serde(default)]
    pub aliases: Vec<String>,
    pub input_per_mtok: f64,
    pub output_per_mtok: f64,
    pub cached_input_per_mtok: Option<f64>,
    pub context_length: Option<i64>,
    pub source: Option<String>,
}

/// The feed envelope. `stale` and `age_seconds` are produced by the server and
/// we trust them over re-deriving freshness locally.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct Feed {
    pub built_at: String,
    pub age_seconds: Option<i64>,
    pub stale: Option<bool>,
    pub count: Option<i64>,
    pub models: Vec<FeedModel>,
}

impl Feed {
    /// Index the feed once: canonical ids and per-CLI aliases -> canonical id.
    pub fn index(&self) -> HashMap<String, usize> {
        let mut index = HashMap::with_capacity(self.models.len() * 2);
        for (i, model) in self.models.iter().enumerate() {
            index.insert(model.id.clone(), i);
            for alias in &model.aliases {
                index.insert(alias.clone(), i);
            }
        }
        index
    }

    /// The feed is usable when the server says so and the server-reported age
    /// is within bounds. A stale or excessively old feed must never be written
    /// over a fresh cache.
    pub fn usable(&self) -> bool {
        if self.stale.unwrap_or(true) {
            return false;
        }
        self.age_seconds.unwrap_or(i64::MAX) <= MAX_SERVER_AGE
    }
}

/// Resolve `model` to a priced feed entry by exact canonical id or alias. A
/// model the feed does not carry is a miss: no basename, substring, or
/// near-match fallback.
pub fn feed_lookup<'a>(feed: &'a Feed, index: &HashMap<String, usize>, model: &str) -> Option<&'a FeedModel> {
    index.get(model).and_then(|&i| feed.models.get(i))
}

/// Where the refresh result is written. Cache lives under the aid home so a
/// cold or stale cache never delays a run — serving is a local file read.
fn cache_path() -> std::path::PathBuf {
    crate::paths::aid_dir().join(CACHE_FILE)
}

/// Read the cached feed, if any. Errors are swallowed: a missing or corrupt
/// cache is an offline first run: only catalog and override prices apply.
pub fn load_cache() -> Option<Feed> {
    let bytes = fs::read(cache_path()).ok()?;
    serde_json::from_slice(&bytes).ok()
}

/// The feed is fresh enough to serve without a refresh attempt.
pub fn cache_fresh(feed: &Feed) -> bool {
    let Ok(built) = DateTime::parse_from_rfc3339(&feed.built_at) else {
        return false;
    };
    Utc::now().signed_duration_since(built.with_timezone(&Utc)) < chrono::Duration::from_std(CACHE_TTL).unwrap_or_default()
}

/// Out-of-band refresh: fetch the feed and replace the cache only when the new
/// feed is fresh. A network failure keeps the old cache ("keep the old cache",
/// never "no prices"). Never blocks or fails a run.
pub fn maybe_refresh() {
    crate::feed_refresh::maybe_refresh(
        &cache_path(),
        FEED_URL,
        |body| {
            serde_json::from_slice::<Feed>(body)
                .is_ok_and(|feed| !feed.models.is_empty() && feed.usable())
        },
        || load_cache().is_some_and(|feed| cache_fresh(&feed)),
    );
}

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

    fn sample_feed() -> Feed {
        Feed {
            built_at: Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
            age_seconds: Some(60),
            stale: Some(false),
            count: Some(2),
            models: vec![
                FeedModel {
                    id: "openai/gpt-5.3-codex".to_string(),
                    aliases: vec!["gpt-5.3-codex-high".to_string()],
                    input_per_mtok: 1.75,
                    output_per_mtok: 14.0,
                    cached_input_per_mtok: None,
                    context_length: None,
                    source: Some("openrouter".to_string()),
                },
                FeedModel {
                    id: "gpt-5.6-sol".to_string(),
                    aliases: vec![],
                    input_per_mtok: 5.0,
                    output_per_mtok: 30.0,
                    cached_input_per_mtok: None,
                    context_length: None,
                    source: Some("openrouter".to_string()),
                },
            ],
        }
    }

    #[test]
    fn lookup_exact_id_and_alias() {
        let feed = sample_feed();
        let index = feed.index();
        assert_eq!(feed_lookup(&feed, &index, "openai/gpt-5.3-codex").unwrap().id, "openai/gpt-5.3-codex");
        assert_eq!(feed_lookup(&feed, &index, "gpt-5.3-codex-high").unwrap().id, "openai/gpt-5.3-codex");
        assert_eq!(feed_lookup(&feed, &index, "gpt-5.6-sol").unwrap().id, "gpt-5.6-sol");
    }

    #[test]
    fn vendor_prefixed_name_is_not_a_basename_match() {
        let feed = sample_feed();
        let index = feed.index();
        assert!(feed_lookup(&feed, &index, "mimo/gpt-5.6-sol").is_none());
    }

    #[test]
    fn lookup_miss_is_none() {
        let feed = sample_feed();
        let index = feed.index();
        assert!(feed_lookup(&feed, &index, "nobody/has-this-model").is_none());
        assert!(feed_lookup(&feed, &index, "gpt-5.3-codex-high-extra").is_none());
    }

    #[test]
    fn stale_feed_is_not_usable() {
        let mut feed = sample_feed();
        feed.stale = Some(true);
        assert!(!feed.usable());
        feed.stale = Some(false);
        feed.age_seconds = Some(MAX_SERVER_AGE + 1);
        assert!(!feed.usable());
    }

    #[test]
    fn cache_round_trip_under_aid_home() {
        let temp = tempfile::tempdir().unwrap();
        let _guard = AidHomeGuard::set(temp.path());
        std::fs::create_dir_all(crate::paths::aid_dir()).unwrap();
        let feed = sample_feed();
        fs::write(cache_path(), serde_json::to_vec(&feed).unwrap()).unwrap();
        assert_eq!(load_cache().unwrap(), feed);
        assert!(cache_fresh(&feed));
    }

    #[test]
    fn missing_cache_loads_none() {
        let temp = tempfile::tempdir().unwrap();
        let _guard = AidHomeGuard::set(temp.path());
        assert!(load_cache().is_none());
    }

    #[test]
    #[ignore = "live network probe"]
    fn live_refresh_writes_cache() {
        let temp = tempfile::tempdir().unwrap();
        let _guard = AidHomeGuard::set(temp.path());
        maybe_refresh();
        std::thread::sleep(std::time::Duration::from_secs(20));
        maybe_refresh();
        let feed = load_cache().expect("cache should be adopted after live refresh");
        assert!(feed.usable());
    }

    #[test]
    fn ready_refresh_preserves_cache_on_invalid_empty_or_server_stale_feed() {
        let temp = tempfile::tempdir().expect("home");
        let _guard = AidHomeGuard::set(temp.path());
        let old = sample_feed();
        let bytes = serde_json::to_vec(&old).expect("old");
        fs::write(cache_path(), &bytes).expect("cache");
        let ready = cache_path().with_extension("json.ready");
        let mut empty = old.clone();
        empty.models.clear();
        let mut stale = old.clone();
        stale.stale = Some(true);
        let mut aged = old;
        aged.age_seconds = Some(MAX_SERVER_AGE + 1);
        for response in [
            b"invalid JSON".to_vec(),
            serde_json::to_vec(&empty).expect("empty"),
            serde_json::to_vec(&stale).expect("stale"),
            serde_json::to_vec(&aged).expect("aged"),
        ] {
            fs::write(&ready, response).expect("ready");
            maybe_refresh();
            assert!(!ready.exists());
            assert_eq!(fs::read(cache_path()).expect("cache"), bytes);
        }
    }

    #[test]
    fn valid_ready_replaces_even_fresh_price_cache() {
        let temp = tempfile::tempdir().expect("home");
        let _guard = AidHomeGuard::set(temp.path());
        let mut feed = sample_feed();
        fs::write(cache_path(), serde_json::to_vec(&feed).expect("feed")).expect("cache");
        feed.models[0].input_per_mtok = 9.0;
        let ready = cache_path().with_extension("json.ready");
        fs::write(&ready, serde_json::to_vec(&feed).expect("feed")).expect("ready");
        maybe_refresh();
        assert_eq!(load_cache(), Some(feed));
        assert!(!ready.exists());
    }
}