zc2 0.0.29

//! Wiring for `zc agent run` (spec §6.1): one tokio runtime and three loops
//! (the local API, the hub poll, the reconcile tick), then a clean shutdown.

use crate::agent::config::AgentConfig;
use crate::agent::core::Core;
use crate::agent::{files, hub, server};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::watch;

/// How long `shutdown` gives the axum server's graceful drain before it gives
/// up waiting and aborts the task outright. `axum::serve`'s graceful shutdown
/// has no timeout of its own, so this is what keeps a stuck connection from
/// stalling the whole exit.
const SERVER_SHUTDOWN_BOUND: Duration = Duration::from_secs(2);
/// How long `shutdown` waits for the reconcile task to notice the stop
/// request and return. In the ordinary case this is near-instant — the
/// `Core`'s stop flag makes an in-flight `terminate` wait give up right
/// away — this is only a safety net against something unexpected blocking it.
const RECONCILE_SHUTDOWN_BOUND: Duration = Duration::from_secs(5);

pub struct RunningAgent {
    pub port: u16,
    pub token: String,
    pub core: Arc<Core>,
    stop: watch::Sender<bool>,
    server: tokio::task::JoinHandle<()>,
    hub: tokio::task::JoinHandle<()>,
    reconcile: tokio::task::JoinHandle<()>,
}

impl RunningAgent {
    /// Stop the loops, then SIGTERM every child and wait (spec §6.5,
    /// launchd's `ExitTimeOut = 60`). Returns `false` if some part of the
    /// shutdown could not be completed cleanly within its bound (a task had
    /// to be force-aborted, or terminating the children panicked), so the
    /// caller can report a failing exit code instead of a silent success.
    pub async fn shutdown(self) -> bool {
        let _ = self.stop.send(true);
        let mut clean = true;

        // The hub task may be blocked inside an HTTP call to an unreachable
        // hub for up to hub_timeout; there is nothing worth waiting for
        // there, so it is aborted rather than joined.
        self.hub.abort();

        // Tell the supervisor to give up early on any wait an in-flight
        // reconcile tick is blocked in (stopping the broker, killing a
        // worker), so the tick releases the lock `shutdown_children` needs
        // instead of running out its full grace period. `reconcile_once`
        // also checks this flag itself, before touching the supervisor, so
        // no tick started after this point can respawn anything.
        self.core.request_stop();
        let reconcile_abort = self.reconcile.abort_handle();
        if tokio::time::timeout(RECONCILE_SHUTDOWN_BOUND, self.reconcile)
            .await
            .is_err()
        {
            reconcile_abort.abort();
            clean = false;
        }

        // axum::serve's graceful shutdown has no timeout of its own: bound
        // it, and abort the task if it is still draining after that.
        let server_abort = self.server.abort_handle();
        if tokio::time::timeout(SERVER_SHUTDOWN_BOUND, self.server)
            .await
            .is_err()
        {
            server_abort.abort();
            clean = false;
        }

        // Killing child process groups can take up to the configured grace
        // periods (tens of seconds), so it never runs directly on an async task.
        let core = self.core.clone();
        if tokio::task::spawn_blocking(move || core.shutdown_children())
            .await
            .is_err()
        {
            clean = false;
        }
        clean
    }
}

/// Pure decision, injected so it's testable without real processes or
/// sockets: does `agent.json` (`pid`, `port`) describe another zc agent for
/// this same zc dir that is genuinely still running?
///
/// This never looks at the preferred port at all -- that is the point (review
/// A2): once an agent can live on a fallback port, "is 4720 busy" no longer
/// implies "is another agent already running", and vice versa. `agent.json`'s
/// OWN recorded pid and port are the only thing that can answer that,
/// whatever port a NEW `start()` call happens to prefer.
///
/// Both signals must agree: `pid_alive(pid)` (and `pid != us`) says the
/// process is still there, `answers_like_agent(port)` says something at that
/// port still speaks the agent API (200, or 401 with the JSON error
/// envelope). Either alone is not enough -- a stale `agent.json` with a dead
/// pid must not refuse, and an unrelated 401-emitting service on a recycled
/// pid's port must not either.
pub(crate) fn existing_agent_conflict(
    agent_json: Option<(u32, u16)>,
    us: u32,
    pid_alive: impl Fn(u32) -> bool,
    answers_like_agent: impl Fn(u16) -> bool,
) -> Option<(u32, u16)> {
    let (pid, port) = agent_json?;
    if pid == us || !pid_alive(pid) {
        return None;
    }
    if !answers_like_agent(port) {
        return None;
    }
    Some((pid, port))
}

/// A short-timeout `GET /v1/summary` that answers 200 or 401 (with a JSON
/// error envelope) looks like the agent API; anything else (including no
/// answer at all) does not.
fn answers_like_agent(port: u16) -> bool {
    let http = ureq::Agent::new_with_config(
        ureq::Agent::config_builder()
            .timeout_global(Some(Duration::from_millis(300)))
            .http_status_as_error(false)
            .build(),
    );
    let resp = http
        .get(format!("http://127.0.0.1:{port}/v1/summary"))
        .call();
    match resp {
        Ok(r) if r.status().as_u16() == 200 => true,
        Ok(r) if r.status().as_u16() == 401 => r
            .into_body()
            .read_to_string()
            .map(|b| b.contains("\"error\""))
            .unwrap_or(false),
        _ => false,
    }
}

/// Bind `127.0.0.1:<port>` (0 = any free port), write `agent.json`, start the loops.
///
/// Before touching any socket, `<zc dir>/agent/agent.json` is checked: if it
/// names a pid that is alive and isn't us, AND that pid's own recorded port
/// still answers like the agent API, this refuses -- two agents for one zc
/// dir must never run, and (once an agent can live on a fallback port)
/// checking only the *preferred* port is not enough to guarantee that (review
/// A2). Otherwise, the preferred port (`cfg.port`, 4720 or
/// `ZAKURO_AGENT_PORT`) is tried; if it's held by anything else, an
/// OS-assigned port is used instead (logged), so a stray process on 4720
/// never blocks the agent.
pub async fn start(cfg: AgentConfig, rotate_token: bool) -> Result<RunningAgent, String> {
    let agent_json_path = cfg.dir.join(files::AGENT_FILE);
    let recorded: Option<(u32, u16)> =
        files::load_json::<files::AgentFile>(&agent_json_path).map(|f| (f.pid, f.port));
    let us = std::process::id();
    let conflict = tokio::task::spawn_blocking(move || {
        existing_agent_conflict(
            recorded,
            us,
            crate::agent::supervisor::is_alive,
            answers_like_agent,
        )
    })
    .await
    .map_err(|e| e.to_string())?;
    if let Some((pid, port)) = conflict {
        return Err(format!(
            "another zc agent is already running on 127.0.0.1:{port} (pid {pid})"
        ));
    }

    let listener = match tokio::net::TcpListener::bind(("127.0.0.1", cfg.port)).await {
        Ok(l) => l,
        Err(e) if e.kind() == std::io::ErrorKind::AddrInUse => {
            let fallback = tokio::net::TcpListener::bind(("127.0.0.1", 0))
                .await
                .map_err(|e| format!("bind_error: cannot bind 127.0.0.1:0: {e}"))?;
            let actual = fallback.local_addr().map_err(|e| e.to_string())?.port();
            eprintln!(
                "  [AGENT] 127.0.0.1:{} is held by another program; the agent API is on 127.0.0.1:{actual}",
                cfg.port
            );
            fallback
        }
        Err(e) => {
            return Err(format!(
                "bind_error: cannot bind 127.0.0.1:{}: {e}",
                cfg.port
            ))
        }
    };
    let port = listener.local_addr().map_err(|e| e.to_string())?.port();
    let agent_file =
        files::load_or_init_agent_file(&cfg.dir, port, rotate_token).map_err(|e| e.to_string())?;
    let core = Core::new(cfg).map_err(|e| e.to_string())?;
    let (stop, stop_rx) = watch::channel(false);

    let app = server::router(core.clone(), agent_file.token.clone(), port);
    let mut rx = stop_rx.clone();
    let server = tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.wait_for(|stopped| *stopped).await;
            })
            .await;
    });

    let c = core.clone();
    let mut rx = stop_rx.clone();
    let hub_task = tokio::spawn(async move {
        loop {
            c.refresh_hub().await;
            let delay = hub::poll_delay(c.cfg.timings.hub_poll, c.hub_failures());
            tokio::select! {
                _ = tokio::time::sleep(delay) => {}
                _ = c.wake_hub.notified() => {}
                _ = rx.wait_for(|stopped| *stopped) => break,
            }
        }
    });

    let c = core.clone();
    let mut rx = stop_rx;
    let reconcile = tokio::spawn(async move {
        loop {
            let tick = c.clone();
            // reconcile_once() spawns and kills child processes with grace
            // periods of up to several seconds; it must never block the
            // async executor directly.
            let _ = tokio::task::spawn_blocking(move || tick.reconcile_once()).await;
            let pause = if c.is_draining() {
                c.cfg.timings.drain_poll
            } else {
                c.cfg.timings.reconcile
            };
            tokio::select! {
                // The stop branch goes first and `biased` turns off random
                // selection: a `wake_reconcile` permit left over from
                // whatever the last tick did must never win a race against
                // an already-requested stop and pull in one more tick.
                biased;
                _ = rx.wait_for(|stopped| *stopped) => break,
                _ = tokio::time::sleep(pause) => {}
                _ = c.wake_reconcile.notified() => {}
            }
        }
    });

    Ok(RunningAgent {
        port,
        token: agent_file.token,
        core,
        stop,
        server,
        hub: hub_task,
        reconcile,
    })
}

/// `zc agent run`: blocks until SIGTERM or Ctrl-C (SIGINT), then shuts down.
/// Returns non-zero if the agent failed to start, or if shutdown could not
/// complete cleanly.
pub fn run(cfg: AgentConfig) -> i32 {
    let rt = match tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .enable_all()
        .build()
    {
        Ok(rt) => rt,
        Err(e) => {
            eprintln!("zc agent: {e}");
            return 1;
        }
    };
    rt.block_on(async move {
        // Installed before `start()` spawns anything: a SIGTERM/SIGINT that
        // arrives in that window must still be caught, not take the default
        // action while children could already be running.
        let mut term = match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
        {
            Ok(s) => s,
            Err(e) => {
                eprintln!("zc agent: install SIGTERM handler: {e}");
                return 1;
            }
        };
        let mut int = match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
        {
            Ok(s) => s,
            Err(e) => {
                eprintln!("zc agent: install SIGINT handler: {e}");
                return 1;
            }
        };
        let agent = match start(cfg, false).await {
            Ok(a) => a,
            Err(e) => {
                eprintln!("zc agent: {e}");
                return 1;
            }
        };
        eprintln!("  [AGENT] listening on 127.0.0.1:{}", agent.port);
        tokio::select! {
            _ = term.recv() => {}
            _ = int.recv() => {}
        }
        eprintln!("  [AGENT] stopping: SIGTERM to every child");
        if agent.shutdown().await {
            0
        } else {
            eprintln!("  [AGENT] shutdown did not complete cleanly");
            1
        }
    })
}

/// An agent on its own thread and runtime, for the integration tests.
/// Sending on the returned channel (or dropping it) shuts it down; join the
/// returned handle to wait for that shutdown to actually finish, so this
/// agent's children never overlap whatever the caller starts next.
#[cfg(test)]
pub fn start_background(
    cfg: AgentConfig,
) -> Result<
    (
        u16,
        String,
        std::sync::mpsc::Sender<()>,
        std::thread::JoinHandle<()>,
    ),
    String,
> {
    let (ready_tx, ready_rx) = std::sync::mpsc::channel();
    let (stop_tx, stop_rx) = std::sync::mpsc::channel::<()>();
    let thread = std::thread::spawn(move || {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap();
        rt.block_on(async move {
            match start(cfg, false).await {
                Ok(agent) => {
                    let _ = ready_tx.send(Ok((agent.port, agent.token.clone())));
                    let _ = tokio::task::spawn_blocking(move || stop_rx.recv()).await;
                    agent.shutdown().await;
                }
                Err(e) => {
                    let _ = ready_tx.send(Err(e));
                }
            }
        });
    });
    let ready = ready_rx
        .recv_timeout(std::time::Duration::from_secs(10))
        .map_err(|_| "the agent thread ended before it became ready".to_string())?;
    let (port, token) = ready?;
    Ok((port, token, stop_tx, thread))
}

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

    /// The pure decision, with injected closures: a live pid on a FALLBACK
    /// port refuses even though the caller's preferred port (never passed to
    /// this function at all) is free -- that is the whole point of A2's fix.
    #[test]
    fn existing_agent_conflict_decision() {
        // No agent.json at all.
        assert_eq!(existing_agent_conflict(None, 111, |_| true, |_| true), None);
        // The recorded pid is us.
        assert_eq!(
            existing_agent_conflict(Some((111, 4720)), 111, |_| true, |_| true),
            None
        );
        // The recorded pid is dead.
        assert_eq!(
            existing_agent_conflict(Some((222, 4720)), 111, |_| false, |_| true),
            None
        );
        // The pid is alive, but its recorded port doesn't answer like an agent.
        assert_eq!(
            existing_agent_conflict(Some((222, 4720)), 111, |_| true, |_| false),
            None
        );
        // A live pid recorded on a FALLBACK port (50324, not 4720) that does
        // answer -- refuses. The function never even sees a "preferred" port.
        assert_eq!(
            existing_agent_conflict(Some((222, 50324)), 111, |_| true, |_| true),
            Some((222, 50324))
        );
    }

    fn write_agent_json(state: &std::path::Path, port: u16, pid: u32) {
        std::fs::create_dir_all(state.join("agent")).unwrap();
        std::fs::write(
            state.join("agent").join(crate::agent::files::AGENT_FILE),
            format!(r#"{{"v":1,"port":{port},"token":"x","pid":{pid},"version":"0.0.0"}}"#),
        )
        .unwrap();
    }

    /// Bind a real TCP listener that speaks just enough HTTP to identify as
    /// the agent API once (a 401 JSON error envelope), on the given port.
    /// Returns immediately; the listener answers on a background thread.
    fn serve_mock_agent_once(listener: std::net::TcpListener) {
        std::thread::spawn(move || {
            if let Ok((mut stream, _)) = listener.accept() {
                use std::io::{Read, Write};
                let mut buf = [0u8; 1024];
                let _ = stream.read(&mut buf);
                let body = r#"{"error":{"code":"unauthorized","message":"nope"}}"#;
                let resp = format!(
                    "HTTP/1.1 401 Unauthorized\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    body.len(),
                    body
                );
                let _ = stream.write_all(resp.as_bytes());
            }
        });
    }

    /// The brief's first-named case: a free, NON-ZERO preferred port is used
    /// directly (`agent.port == cfg.port` exactly), with no fallback and no
    /// `agent.json` conflict check getting in the way (there is no
    /// `agent.json` here at all).
    ///
    /// `cfg.port = 0` would not prove this: it is always "free" (it means
    /// "any port"), so `start()` would pass this test even if it ignored
    /// `cfg.port` entirely and always bound an OS-assigned port -- e.g. if
    /// the direct-bind line were changed to bind `("127.0.0.1", 0)` outright.
    /// Instead: bind a real ephemeral port, read it, drop it immediately,
    /// then assert the agent bound EXACTLY that port.
    #[tokio::test]
    async fn start_uses_the_preferred_port_when_it_is_free() {
        let mut last = (0, 0);
        for attempt in 0..5 {
            // Retried: once the probe is dropped, a test running in parallel
            // can bind its port before `start()` does (seen on CI).
            let probe = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
            let preferred = probe.local_addr().unwrap().port();
            drop(probe);

            let mut cfg = AgentConfig::for_dirs(tmp(&format!("run-preferred-free-{attempt}")));
            cfg.port = preferred;
            let agent = start(cfg, false)
                .await
                .expect("a free preferred port must bind directly");
            let got = agent.port;
            agent.shutdown().await;
            if got == preferred {
                return;
            }
            last = (preferred, got);
        }
        panic!(
            "a free, non-zero preferred port must be used exactly, not silently replaced: \
             after 5 attempts the last wanted port {} and the agent got {}",
            last.0, last.1
        );
    }

    /// A held port with no `agent.json` at all falls back to a random port
    /// instead of failing, and still writes `agent.json` normally.
    #[tokio::test]
    async fn start_on_a_held_port_with_no_agent_falls_back_to_a_random_port() {
        let held = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let held_port = held.local_addr().unwrap().port();
        let state = tmp("run-port-in-use");
        let mut cfg = AgentConfig::for_dirs(state.clone());
        cfg.port = held_port;

        let agent = start(cfg, false)
            .await
            .expect("a non-agent on the port must not block startup");
        assert_ne!(agent.port, held_port);
        let agent_json = state.join("agent").join(crate::agent::files::AGENT_FILE);
        assert!(agent_json.exists());
        agent.shutdown().await;
        drop(held);
    }

    /// A dead pid does not refuse, EVEN THOUGH its recorded port genuinely
    /// answers like the agent API: only the pid check can be why this falls
    /// back (review A5 -- the old version of this test could pass however
    /// `other_agent_pid` treated liveness, because its held port never
    /// answered HTTP at all).
    #[tokio::test]
    async fn start_falls_back_when_agent_json_names_a_dead_pid_even_though_its_port_answers_like_the_agent(
    ) {
        let mut child = std::process::Command::new("true").spawn().unwrap();
        let dead_pid = child.id();
        let _ = child.wait(); // reaped: definitely dead now

        let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let held_port = listener.local_addr().unwrap().port();
        serve_mock_agent_once(listener);
        std::thread::sleep(std::time::Duration::from_millis(50));

        let state = tmp("run-dead-pid-answers");
        write_agent_json(&state, held_port, dead_pid);
        let mut cfg = AgentConfig::for_dirs(state.clone());
        cfg.port = held_port;

        let agent = start(cfg, false).await.expect(
            "a dead pid must not refuse, even though its recorded port answers like an agent",
        );
        assert_ne!(agent.port, held_port);
        agent.shutdown().await;
    }

    /// Both signals present -- a live, non-self pid in `agent.json` AND its
    /// recorded port answering like the agent API -- refuses to start,
    /// naming the pid. Here the recorded port IS the preferred port.
    #[tokio::test]
    async fn start_refuses_when_a_live_agent_already_holds_the_preferred_port() {
        // Bind the port once, for the mock agent itself -- no separate
        // throwaway listener that would have to be dropped and immediately
        // re-bound, which races with any other test grabbing the same
        // freshly-freed ephemeral port.
        let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let held_port = listener.local_addr().unwrap().port();
        serve_mock_agent_once(listener);
        // Give the mock listener a moment to bind before `start` connects.
        std::thread::sleep(std::time::Duration::from_millis(50));

        let mut other = std::process::Command::new("sleep")
            .arg("5")
            .spawn()
            .unwrap();
        let other_pid = other.id();
        let state = tmp("run-live-agent");
        write_agent_json(&state, held_port, other_pid);
        let mut cfg = AgentConfig::for_dirs(state.clone());
        cfg.port = held_port;

        let err = start(cfg, false).await.err().expect("must refuse");
        assert!(err.contains("already running"), "{err}");
        assert!(err.contains(&other_pid.to_string()), "{err}");

        let _ = other.kill();
        let _ = other.wait();
    }

    /// The invariant A2 restores: a live agent recorded on a FALLBACK port
    /// refuses a second `start()` even though the second call's preferred
    /// port is completely free -- the old check only ever looked at the
    /// preferred port, so this exact scenario used to start a second agent.
    #[tokio::test]
    async fn start_refuses_a_live_agent_on_a_fallback_port_even_when_the_preferred_port_is_free() {
        let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap();
        let recorded_port = listener.local_addr().unwrap().port();
        serve_mock_agent_once(listener);
        std::thread::sleep(std::time::Duration::from_millis(50));

        let mut other = std::process::Command::new("sleep")
            .arg("5")
            .spawn()
            .unwrap();
        let other_pid = other.id();
        let state = tmp("run-live-agent-fallback-port");
        write_agent_json(&state, recorded_port, other_pid);
        let mut cfg = AgentConfig::for_dirs(state.clone());
        // A distinct "preferred" port, never actually bound: the conflict is
        // caught before any bind is attempted, so its value is irrelevant
        // except that it differs from `recorded_port`, proving the refusal
        // does not depend on the preferred port being busy.
        cfg.port = recorded_port.wrapping_add(1);

        let err = start(cfg, false)
            .await
            .err()
            .expect("must refuse even though the preferred port is free");
        assert!(err.contains("already running"), "{err}");
        assert!(err.contains(&recorded_port.to_string()), "{err}");
        assert!(err.contains(&other_pid.to_string()), "{err}");

        let _ = other.kill();
        let _ = other.wait();
    }
}