zc2 0.0.27

P2P compute broker with credit-based billing, WAL, and broker mesh support
//! A blocking client of the running agent, for the CLI (spec §6.1 hand-off).

use crate::agent::files::{load_json, AgentFile, AGENT_FILE};
use std::path::Path;
use std::time::Duration;

pub struct AgentClient {
    base: String,
    port: u16,
    token: String,
}

impl AgentClient {
    pub fn from_dir(agent_dir: &Path) -> Option<Self> {
        let f: AgentFile = load_json(&agent_dir.join(AGENT_FILE))?;
        Some(Self {
            base: format!("http://127.0.0.1:{}", f.port),
            port: f.port,
            token: f.token,
        })
    }

    /// The port `agent.json` named when this client was built.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// Spec §6.1: running = `agent.json` exists AND `/v1/summary` answers within 1 s.
    pub fn running() -> Option<Self> {
        let c = Self::from_dir(&crate::credentials::dir()?.join("agent"))?;
        c.summary(Duration::from_secs(1)).ok()?;
        Some(c)
    }

    fn http(timeout: Duration) -> ureq::Agent {
        ureq::Agent::new_with_config(
            ureq::Agent::config_builder()
                .http_status_as_error(false)
                .timeout_global(Some(timeout))
                .build(),
        )
    }

    fn read(resp: ureq::http::Response<ureq::Body>) -> Result<serde_json::Value, String> {
        let status = resp.status().as_u16();
        let body: serde_json::Value = resp.into_body().read_json().unwrap_or_default();
        if status == 200 {
            Ok(body)
        } else {
            Err(body["error"]["message"]
                .as_str()
                .map(str::to_string)
                .unwrap_or_else(|| format!("agent answered HTTP {status}")))
        }
    }

    pub fn summary(&self, timeout: Duration) -> Result<serde_json::Value, String> {
        let resp = Self::http(timeout)
            .get(&format!("{}/v1/summary", self.base))
            .header("Authorization", format!("Bearer {}", self.token).as_str())
            .call()
            .map_err(|e| e.to_string())?;
        Self::read(resp)
    }

    pub fn put(
        &self,
        path: &str,
        body: serde_json::Value,
        timeout: Duration,
    ) -> Result<serde_json::Value, String> {
        let resp = Self::http(timeout)
            .put(&format!("{}{path}", self.base))
            .header("Authorization", format!("Bearer {}", self.token).as_str())
            .send_json(body)
            .map_err(|e| e.to_string())?;
        Self::read(resp)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::agent::files::tests::tmp;

    /// Owns the stop sender and join handle from `run::start_background` so
    /// dropping it always stops the agent thread and waits for it to finish.
    /// No agent thread may outlive the test that started it.
    pub(crate) struct AgentGuard {
        stop: Option<std::sync::mpsc::Sender<()>>,
        thread: Option<std::thread::JoinHandle<()>>,
    }

    impl AgentGuard {
        pub(crate) fn new(
            stop: std::sync::mpsc::Sender<()>,
            thread: std::thread::JoinHandle<()>,
        ) -> Self {
            Self {
                stop: Some(stop),
                thread: Some(thread),
            }
        }
    }

    impl Drop for AgentGuard {
        fn drop(&mut self) {
            if let Some(stop) = self.stop.take() {
                let _ = stop.send(());
            }
            if let Some(thread) = self.thread.take() {
                let _ = thread.join();
            }
        }
    }

    /// A test agent's config, before it starts: free ports, harmless `sleep`
    /// children, a fresh state dir under `tag`. Never leave the default
    /// `<test binary> broker 9000` in a test agent: libtest would read
    /// "broker" as a test filter and run the suite.
    pub(crate) fn sleeper_cfg(tag: &str) -> crate::agent::config::AgentConfig {
        let free = || {
            std::net::TcpListener::bind("127.0.0.1:0")
                .unwrap()
                .local_addr()
                .unwrap()
                .port()
        };
        let sleeper = vec![
            "sh".to_string(),
            "-c".to_string(),
            "exec sleep 30".to_string(),
        ];
        let mut cfg = crate::agent::config::AgentConfig::for_dirs(tmp(tag));
        cfg.port = 0;
        cfg.broker_port = free();
        cfg.max_workers = 3;
        cfg.broker_argv = sleeper.clone();
        cfg.worker_template = sleeper;
        cfg.worker_template_overridden = true;
        cfg.timings.shutdown_wait = Duration::from_secs(2);
        cfg
    }

    /// Start `cfg` on its own thread; dropping the guard stops it.
    pub(crate) fn start_agent(cfg: crate::agent::config::AgentConfig) -> (AgentClient, AgentGuard) {
        let dir = cfg.dir.clone();
        let (_, _, stop, thread) = crate::agent::run::start_background(cfg).expect("agent starts");
        (
            AgentClient::from_dir(&dir).expect("agent.json written"),
            AgentGuard::new(stop, thread),
        )
    }

    /// An agent on free ports whose broker and worker commands are harmless
    /// `sleep`s.
    pub(crate) fn idle_agent() -> (AgentClient, AgentGuard) {
        start_agent(sleeper_cfg("client"))
    }

    #[test]
    fn talks_to_a_running_agent_and_surfaces_api_errors() {
        let (c, _guard) = idle_agent();
        let s = c.summary(Duration::from_secs(2)).unwrap();
        assert_eq!(s["v"], 1);
        let s = c
            .put(
                "/v1/workers",
                serde_json::json!({"count": 2}),
                Duration::from_secs(2),
            )
            .unwrap();
        assert_eq!(s["this_mac"]["workers"]["desired"], 2);
        let err = c
            .put(
                "/v1/workers",
                serde_json::json!({"count": 9}),
                Duration::from_secs(2),
            )
            .unwrap_err();
        assert!(err.contains("count must be between 0 and 3"), "{err}");
    }

    #[test]
    fn a_stale_agent_file_is_not_a_running_agent() {
        let dir = tmp("stale");
        crate::agent::files::ensure_dir(&dir).unwrap();
        crate::agent::files::save_json(
            &dir.join(crate::agent::files::AGENT_FILE),
            &crate::agent::files::AgentFile {
                v: 1,
                port: 1,
                token: "t".into(),
                pid: 1,
                version: "x".into(),
            },
        )
        .unwrap();
        let c = AgentClient::from_dir(&dir).unwrap();
        assert!(c.summary(Duration::from_secs(1)).is_err());
        assert!(AgentClient::from_dir(&tmp("absent")).is_none());
    }
}