zc2 0.0.28

//! The agent's hub client (spec §6.7): credentials re-read on change, the
//! provider summary, price writes, poll backoff and new-worker price seeding.

use crate::agent::hubwire::{HubSummary, HubWorker};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};

#[derive(Debug, Clone, PartialEq)]
pub struct Credentials {
    pub api_key: String,
    pub api_url: Option<String>,
}

/// `~/.zakuro/credentials` (`api_key=…`, `api_url=…`), as `zc login` writes it.
pub fn read_credentials(path: &Path) -> Option<Credentials> {
    let m = crate::credentials::parse(&std::fs::read_to_string(path).ok()?);
    let api_key = m.get("api_key")?.trim().to_string();
    if api_key.is_empty() {
        return None;
    }
    Some(Credentials {
        api_key,
        api_url: m.get("api_url").cloned().filter(|u| !u.trim().is_empty()),
    })
}

/// Credentials re-read only when the file's modification time changes, so a
/// fresh `zc login` takes effect without restarting the agent.
pub struct CredWatch {
    path: PathBuf,
    mtime: Option<SystemTime>,
    cached: Option<Credentials>,
    loaded: bool,
}

impl CredWatch {
    pub fn new(path: PathBuf) -> Self {
        Self {
            path,
            mtime: None,
            cached: None,
            loaded: false,
        }
    }

    pub fn current(&mut self) -> Option<Credentials> {
        let mtime = std::fs::metadata(&self.path)
            .and_then(|m| m.modified())
            .ok();
        if !self.loaded || mtime != self.mtime {
            self.loaded = true;
            self.mtime = mtime;
            self.cached = read_credentials(&self.path);
        }
        self.cached.clone()
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum HubError {
    /// 401: the key is missing, wrong or revoked → `not_logged_in`.
    Unauthorized,
    /// 404: the hub predates the route → `hub_too_old`.
    NotFound,
    Status(u16, String),
    Transport(String),
}

impl HubError {
    pub fn message(&self) -> String {
        match self {
            HubError::Unauthorized => "the hub rejected this Mac's key (401)".into(),
            HubError::NotFound => "the hub does not have this route yet (404)".into(),
            HubError::Status(s, body) => format!("hub answered HTTP {s}: {body}"),
            HubError::Transport(e) => format!("could not reach the hub: {e}"),
        }
    }
}

/// How much of a hub error body `hub.error` and a 502 message carry.
const HUB_ERROR_MAX_CHARS: usize = 200;

/// The readable part of a non-2xx hub body: FastAPI's `detail` when it is a
/// string, otherwise the body itself. Either way whitespace runs collapse to
/// one space and the text is cut to `HUB_ERROR_MAX_CHARS` characters (on a
/// char boundary, marked `…`), so a proxy's HTML error page never lands whole
/// in the app.
fn readable_error(body: &str) -> String {
    let detail = serde_json::from_str::<serde_json::Value>(body)
        .ok()
        .and_then(|v| v.get("detail")?.as_str().map(str::to_string));
    let text = detail
        .as_deref()
        .unwrap_or(body)
        .split_whitespace()
        .collect::<Vec<_>>()
        .join(" ");
    match text.char_indices().nth(HUB_ERROR_MAX_CHARS) {
        Some((cut, _)) => format!("{}…", &text[..cut]),
        None => text,
    }
}

fn classify(status: u16, body: String) -> HubError {
    match status {
        401 => HubError::Unauthorized,
        404 => HubError::NotFound,
        s => HubError::Status(s, readable_error(&body)),
    }
}

#[derive(Clone)]
pub struct HubClient {
    http: reqwest::Client,
    base: String,
    key: String,
}

impl HubClient {
    pub fn new(base: &str, key: &str, timeout: Duration) -> Self {
        Self {
            http: reqwest::Client::builder()
                .timeout(timeout)
                .build()
                .expect("reqwest client builds with rustls"),
            base: base.trim_end_matches('/').to_string(),
            key: key.to_string(),
        }
    }

    async fn send(&self, req: reqwest::RequestBuilder) -> Result<String, HubError> {
        let resp = req
            .bearer_auth(&self.key)
            .send()
            .await
            .map_err(|e| HubError::Transport(e.to_string()))?;
        let status = resp.status().as_u16();
        let body = resp.text().await.unwrap_or_default();
        if (200..300).contains(&status) {
            Ok(body)
        } else {
            Err(classify(status, body))
        }
    }

    /// `GET /api/provide/summary?tz=` (spec §5.3).
    pub async fn summary(&self, tz: &str) -> Result<HubSummary, HubError> {
        let body = self
            .send(
                self.http
                    .get(format!("{}/api/provide/summary", self.base))
                    .query(&[("tz", tz)]),
            )
            .await?;
        serde_json::from_str(&body)
            .map_err(|e| HubError::Transport(format!("bad summary JSON: {e}")))
    }

    /// `PUT /api/workers/{id}/price`: `null` means "use the account default".
    pub async fn put_worker_price(&self, id: i64, price: Option<f64>) -> Result<(), HubError> {
        self.send(
            self.http
                .put(format!("{}/api/workers/{id}/price", self.base))
                .json(&serde_json::json!({ "price_per_hour": price })),
        )
        .await
        .map(|_| ())
    }

    /// `PUT /api/workers/global-price`: `null` means "use the mesh price".
    pub async fn put_default_price(&self, price: Option<f64>) -> Result<(), HubError> {
        self.send(
            self.http
                .put(format!("{}/api/workers/global-price", self.base))
                .json(&serde_json::json!({ "default_price_per_hour": price })),
        )
        .await
        .map(|_| ())
    }
}

/// Every `base` normally; after failures 60 s → 2 min → 5 min (base 60 s).
pub fn poll_delay(base: Duration, consecutive_failures: u32) -> Duration {
    match consecutive_failures {
        0 | 1 => base,
        2 => base * 2,
        _ => base * 5,
    }
}

/// Spec §6.7: PUT the price on every one of this Mac's worker rows, all at
/// once, then retry each failure once (again all at once). The error lists
/// the rows that still failed. However many workers this Mac has, a device
/// price takes about two hub timeouts at worst.
pub async fn set_device_price(
    client: &HubClient,
    ids: &[i64],
    price: Option<f64>,
) -> Result<(), Vec<(i64, String)>> {
    let failed: Vec<i64> = put_prices(client, ids, price)
        .await
        .into_iter()
        .map(|(id, _)| id)
        .collect();
    let still = put_prices(client, &failed, price).await;
    if still.is_empty() {
        Ok(())
    } else {
        Err(still)
    }
}

/// PUT `price` on every row in `ids` concurrently. Returns the rows that
/// failed, with their errors, in `ids` order.
async fn put_prices(client: &HubClient, ids: &[i64], price: Option<f64>) -> Vec<(i64, String)> {
    let mut writes = tokio::task::JoinSet::new();
    for &id in ids {
        let c = client.clone();
        writes.spawn(async move { (id, c.put_worker_price(id, price).await) });
    }
    // Every row counts as failed until its write reports back, so a write
    // task that panicked is still reported instead of silently dropped.
    let mut errors: std::collections::HashMap<i64, String> = ids
        .iter()
        .map(|&id| (id, "the price write did not finish".to_string()))
        .collect();
    while let Some(joined) = writes.join_next().await {
        match joined {
            Ok((id, Ok(()))) => {
                errors.remove(&id);
            }
            Ok((id, Err(e))) => {
                errors.insert(id, e.message());
            }
            Err(_) => {}
        }
    }
    ids.iter()
        .filter_map(|id| errors.remove(id).map(|e| (*id, e)))
        .collect()
}

/// Spec §6.7 seeding: a worker row the agent has not seen before gets the
/// price every known sibling agrees on (an explicit value), as long as that
/// price is at least the account's minimum. A price below the minimum means
/// the owner disabled those siblings, and that is never copied onto a new
/// worker. Otherwise nothing: the hub stays the only price owner.
pub fn seed_targets(this_mac: &[&HubWorker], known: &[i64], min: f64) -> Vec<(i64, f64)> {
    let known_prices: Vec<Option<f64>> = this_mac
        .iter()
        .filter(|w| known.contains(&w.id))
        .map(|w| w.price_per_hour)
        .collect();
    let agreed = match known_prices.first() {
        Some(Some(p)) if known_prices.iter().all(|q| *q == Some(*p)) => *p,
        _ => return vec![],
    };
    if agreed < min {
        return vec![];
    }
    this_mac
        .iter()
        .filter(|w| !known.contains(&w.id) && w.price_per_hour != Some(agreed))
        .map(|w| (w.id, agreed))
        .collect()
}

/// IANA zone for the summary's `?tz=`: `$TZ`, else the `/etc/localtime` link.
pub fn system_tz() -> String {
    tz_from(
        std::env::var("TZ").ok(),
        std::fs::read_link("/etc/localtime").ok(),
    )
}

pub fn tz_from(tz_env: Option<String>, localtime_link: Option<PathBuf>) -> String {
    if let Some(tz) = tz_env.filter(|t| !t.is_empty() && !t.starts_with(':')) {
        return tz;
    }
    localtime_link
        .and_then(|p| {
            let s = p.to_string_lossy().to_string();
            s.split_once("zoneinfo/").map(|(_, z)| z.to_string())
        })
        .filter(|z| !z.is_empty())
        .unwrap_or_else(|| "UTC".to_string())
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::agent::files::tests::tmp;
    use crate::agent::hubwire::tests::HUB_JSON;
    use std::sync::{Arc, Mutex};

    /// (method, url, body, authorization header) of every request.
    pub(crate) type Seen = Arc<Mutex<Vec<(String, String, String, Option<String>)>>>;

    /// A hub on a free loopback port. `handler(method, url, nth_time_this_url)`
    /// returns (status, JSON body).
    pub(crate) fn mock_hub<F>(handler: F) -> (String, Seen)
    where
        F: Fn(&str, &str, usize) -> (u16, String) + Send + 'static,
    {
        let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
        let port = server.server_addr().to_ip().unwrap().port();
        let seen: Seen = Default::default();
        let log = seen.clone();
        std::thread::spawn(move || {
            for mut req in server.incoming_requests() {
                let mut body = String::new();
                let _ = std::io::Read::read_to_string(req.as_reader(), &mut body);
                let method = req.method().as_str().to_string();
                let url = req.url().to_string();
                let auth = req
                    .headers()
                    .iter()
                    .find(|h| h.field.equiv("Authorization"))
                    .map(|h| h.value.as_str().to_string());
                let nth = {
                    let mut l = log.lock().unwrap();
                    l.push((method.clone(), url.clone(), body, auth));
                    l.iter().filter(|e| e.1 == url).count()
                };
                let (status, reply) = handler(&method, &url, nth);
                let _ = req.respond(
                    tiny_http::Response::from_string(reply)
                        .with_status_code(status)
                        .with_header(
                            tiny_http::Header::from_bytes("Content-Type", "application/json")
                                .unwrap(),
                        ),
                );
            }
        });
        (format!("http://127.0.0.1:{port}"), seen)
    }

    fn block_on<F: std::future::Future>(f: F) -> F::Output {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap()
            .block_on(f)
    }

    #[test]
    fn credentials_are_reread_only_when_the_file_changes() {
        let dir = tmp("creds");
        std::fs::create_dir_all(&dir).unwrap();
        let p = dir.join("credentials");
        std::fs::write(
            &p,
            "api_key=zk_1_aaa\napi_url=https://stg.api.zakuro-ai.com\n",
        )
        .unwrap();
        let mut w = CredWatch::new(p.clone());
        let first = w.current().unwrap();
        assert_eq!(first.api_key, "zk_1_aaa");
        assert_eq!(
            first.api_url.as_deref(),
            Some("https://stg.api.zakuro-ai.com")
        );

        std::fs::write(&p, "api_key=zk_1_bbb\n").unwrap();
        std::fs::File::options()
            .write(true)
            .open(&p)
            .unwrap()
            .set_modified(std::time::SystemTime::now() + Duration::from_secs(5))
            .unwrap();
        assert_eq!(
            w.current().unwrap().api_key,
            "zk_1_bbb",
            "zc login applies without a restart"
        );

        std::fs::write(&p, "api_url=x\n").unwrap();
        assert_eq!(read_credentials(&p), None, "no key, no credentials");
    }

    #[test]
    fn summary_sends_the_key_and_tz_and_classifies_errors() {
        let (hub, seen) = mock_hub(|_, url, _| match url {
            u if u.starts_with("/api/provide/summary") => (200, HUB_JSON.to_string()),
            "/unauthorized" => (401, "{}".into()),
            _ => (404, "{}".into()),
        });
        let c = HubClient::new(&hub, "zk_1_x", Duration::from_secs(5));
        let s = block_on(c.summary("Asia/Tokyo")).unwrap();
        assert_eq!(s.account.username, "jean");
        let (method, url, _, auth) = seen.lock().unwrap()[0].clone();
        assert_eq!(method, "GET");
        assert_eq!(url, "/api/provide/summary?tz=Asia%2FTokyo");
        assert_eq!(auth.as_deref(), Some("Bearer zk_1_x"));

        assert_eq!(classify(401, "{}".into()), HubError::Unauthorized);
        assert_eq!(classify(404, "{}".into()), HubError::NotFound);
        assert_eq!(
            classify(500, "boom".into()),
            HubError::Status(500, "boom".into())
        );

        let dead = HubClient::new("http://127.0.0.1:1", "k", Duration::from_secs(2));
        assert!(matches!(
            block_on(dead.summary("UTC")),
            Err(HubError::Transport(_))
        ));
    }

    #[test]
    fn device_price_fans_out_and_retries_each_failure_once() {
        let (hub, seen) = mock_hub(|_, url, nth| match (url, nth) {
            ("/api/workers/12/price", 1) => (500, r#"{"detail":"flaky"}"#.into()),
            ("/api/workers/13/price", _) => (500, r#"{"detail":"down"}"#.into()),
            _ => (200, "{}".into()),
        });
        let c = HubClient::new(&hub, "zk_1_x", Duration::from_secs(5));
        assert_eq!(
            block_on(set_device_price(&c, &[11, 12], Some(18.0))),
            Ok(())
        );
        let bodies: Vec<(String, String, String)> = seen
            .lock()
            .unwrap()
            .iter()
            .map(|(m, u, b, _)| (m.clone(), u.clone(), b.clone()))
            .collect();
        assert_eq!(bodies.len(), 3, "11, 12 (failed), 12 (retry)");
        assert!(bodies
            .iter()
            .all(|(m, _, b)| m == "PUT" && b == r#"{"price_per_hour":18.0}"#));

        let err = block_on(set_device_price(&c, &[13], None)).unwrap_err();
        assert_eq!(err.len(), 1);
        assert_eq!(err[0].0, 13);
        assert!(err[0].1.contains("down"));
    }

    /// A non-2xx hub answer ends up in `hub.error` and in the 502s the app
    /// shows: FastAPI's `detail` string when there is one, otherwise the body
    /// trimmed and cut to about 200 characters, never a proxy's whole HTML page.
    #[test]
    fn hub_error_text_is_the_detail_or_a_short_trimmed_body() {
        let msg = |status: u16, body: &str| classify(status, body.to_string()).message();
        assert_eq!(
            msg(403, r#"{"detail":"Not your worker"}"#),
            "hub answered HTTP 403: Not your worker"
        );
        assert_eq!(msg(500, "  boom \n"), "hub answered HTTP 500: boom");

        let page = format!(
            "<!DOCTYPE html>\n<html><head><title>502 Bad Gateway</title></head>\n<body>{}</body></html>\n",
            "<p>nginx</p>".repeat(100)
        );
        let m = msg(502, &page);
        assert!(
            m.starts_with("hub answered HTTP 502: <!DOCTYPE html> <html>"),
            "{m}"
        );
        assert!(!m.contains("</html>"), "never the whole page: {m}");
        assert!(m.chars().count() < 240, "{} chars: {m}", m.chars().count());

        let m = msg(
            422,
            r#"{"detail":[{"loc":["body","price_per_hour"],"msg":"bad"}]}"#,
        );
        assert!(
            m.contains("price_per_hour"),
            "a non-string detail falls back to the body: {m}"
        );

        let m = msg(500, &"é".repeat(400));
        assert!(m.ends_with('…'), "{m}");
        assert_eq!(
            m.chars().filter(|c| *c == 'é').count(),
            200,
            "cut on a char boundary"
        );
    }

    /// Every row's PUT is in flight at the same time, so a device price takes
    /// about two hub timeouts at worst, however many workers this Mac has.
    #[test]
    fn device_price_puts_are_sent_concurrently() {
        use std::sync::atomic::{AtomicUsize, Ordering};
        let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
        let port = server.server_addr().to_ip().unwrap().port();
        let in_flight = Arc::new(AtomicUsize::new(0));
        let peak = Arc::new(AtomicUsize::new(0));
        let (now, most) = (in_flight.clone(), peak.clone());
        // Each request is held on its own thread, so the hub itself never
        // serializes them.
        std::thread::spawn(move || {
            for req in server.incoming_requests() {
                let (now, most) = (now.clone(), most.clone());
                std::thread::spawn(move || {
                    most.fetch_max(now.fetch_add(1, Ordering::SeqCst) + 1, Ordering::SeqCst);
                    std::thread::sleep(Duration::from_millis(600));
                    now.fetch_sub(1, Ordering::SeqCst);
                    let _ = req.respond(tiny_http::Response::from_string("{}"));
                });
            }
        });
        let c = HubClient::new(
            &format!("http://127.0.0.1:{port}"),
            "zk_1_x",
            Duration::from_secs(5),
        );
        let ids = [11, 12, 13, 14];
        assert_eq!(block_on(set_device_price(&c, &ids, Some(18.0))), Ok(()));
        assert_eq!(
            peak.load(Ordering::SeqCst),
            ids.len(),
            "every worker's PUT was in flight at once"
        );
    }

    #[test]
    fn default_price_writes_the_global_price() {
        let (hub, seen) = mock_hub(|_, _, _| (200, "{}".into()));
        let c = HubClient::new(&hub, "zk_1_x", Duration::from_secs(5));
        block_on(c.put_default_price(None)).unwrap();
        let (m, u, b, _) = seen.lock().unwrap()[0].clone();
        assert_eq!(
            (m.as_str(), u.as_str(), b.as_str()),
            (
                "PUT",
                "/api/workers/global-price",
                r#"{"default_price_per_hour":null}"#
            )
        );
    }

    #[test]
    fn poll_backs_off_60s_2min_5min() {
        let base = Duration::from_secs(60);
        let secs: Vec<u64> = (0..5).map(|n| poll_delay(base, n).as_secs()).collect();
        assert_eq!(secs, vec![60, 60, 120, 300, 300]);
    }

    #[test]
    fn new_workers_get_the_price_their_known_siblings_agree_on() {
        let w = |id: i64, price: Option<f64>| crate::agent::hubwire::HubWorker {
            id,
            worker_id: format!("fp-w{id}"),
            status: "online".into(),
            last_seen: None,
            price_per_hour: price,
            reported_price_per_hour: None,
            effective_price_per_hour: price,
            disabled: false,
        };
        let (a, b, new) = (w(11, Some(18.0)), w(12, Some(18.0)), w(13, None));
        assert_eq!(
            seed_targets(&[&a, &b, &new], &[11, 12], 1.0),
            vec![(13, 18.0)]
        );
        let c = w(12, Some(20.0));
        assert_eq!(
            seed_targets(&[&a, &c, &new], &[11, 12], 1.0),
            vec![],
            "siblings disagree"
        );
        let (d, e) = (w(11, None), w(12, None));
        assert_eq!(
            seed_targets(&[&d, &e, &new], &[11, 12], 1.0),
            vec![],
            "inherited: nothing to copy"
        );
        assert_eq!(
            seed_targets(&[&a, &new], &[], 1.0),
            vec![],
            "no known sibling yet"
        );
    }

    #[test]
    fn seeding_never_copies_a_price_below_the_minimum() {
        let w = |id: i64, price: Option<f64>| crate::agent::hubwire::HubWorker {
            id,
            worker_id: format!("fp-w{id}"),
            status: "online".into(),
            last_seen: None,
            price_per_hour: price,
            reported_price_per_hour: None,
            effective_price_per_hour: price,
            disabled: false,
        };
        let (a, b, new) = (w(11, Some(-1.0)), w(12, Some(-1.0)), w(13, None));
        assert_eq!(
            seed_targets(&[&a, &b, &new], &[11, 12], 1.0),
            vec![],
            "an owner-disabled price is not seeded onto a new worker"
        );
    }

    #[test]
    fn tz_prefers_tz_then_the_localtime_link() {
        assert_eq!(tz_from(Some("Europe/Paris".into()), None), "Europe/Paris");
        assert_eq!(
            tz_from(None, Some("/var/db/timezone/zoneinfo/Asia/Tokyo".into())),
            "Asia/Tokyo"
        );
        assert_eq!(
            tz_from(Some(String::new()), Some("/usr/share/zoneinfo/UTC".into())),
            "UTC"
        );
        assert_eq!(tz_from(None, None), "UTC");
    }
}