hyperhive/hive-c0re/src/workers/swarm_permissions.rs

//! Reading the permissions the swarm grants this hive's agents.
//!
//! swarm-controller keeps one key per agent in the `agent-permissions` bucket
//! ([`swarm_queue_client::agent_permissions`]). Every [`READ_INTERVAL`] this
//! reads the key of each agent this hive has and stores what it finds with
//! [`crate::permissions`], where the container's bind mounts are worked out.
//! Those are fixed when the container starts, so an agent whose
//! `manage_root_agent` changed is restarted, after its current turn. The
//! agent applies the rest of its permissions itself; after every pass each
//! agent's copy of its stored entry is refreshed
//! ([`crate::permissions::hand_at`]), which it holds until it reads its key.
//!
//! A failed read, or a bucket that does not answer, leaves the stored
//! permissions as they are and restarts nothing.

use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use hive_types::Permission;
use swarm_queue_client::agent_permissions::{self, AgentPermissions, Read};

use crate::coordinator::Coordinator;
use crate::permissions::{self, Stored};
use crate::stats::sweep_health::{self, SweepHealth};

/// How often the permissions are read.
pub const READ_INTERVAL: Duration = Duration::from_secs(30);

/// At [`READ_INTERVAL`] this is ~90 seconds of failing reads.
const FAILURES_BEFORE_BANNER: u32 = 3;

/// Send the meta repo's grants once, then read the permissions now and every
/// [`READ_INTERVAL`] until shutdown.
pub async fn follow(
    client: async_nats::Client,
    coord: Arc<Coordinator>,
    hive: String,
    mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
    let mut health = SweepHealth::new("swarm_permissions_read", "warn", FAILURES_BEFORE_BANNER);
    loop {
        let import = super::permissions_import::send_pending(&client, &hive).await;
        let pass = if import.read_allowed() {
            sync(&client).await
        } else {
            Err(import
                .into_error()
                .unwrap_or_else(|| anyhow::anyhow!("import unsettled")))
        };
        match pass {
            Ok(changed) => {
                health.record_ok();
                restart(&coord, changed).await;
            }
            Err(e) => {
                tracing::warn!(error = ?e, "swarm permissions: the pass failed");
                let err = format!("{e:#}");
                health.record_err(|ctx| {
                    let age = ctx.since_last_ok.map_or_else(
                        || "no success this session".to_owned(),
                        |d| format!("last ok {} ago", sweep_health::fmt_age(d)),
                    );
                    format!(
                        "syncing the swarm's agent permissions is failing ({} consecutive, \
                         {age}) — agents keep the permissions last read: {err}",
                        ctx.consecutive
                    )
                });
            }
        }
        hand_all().await;
        tokio::select! {
            () = tokio::time::sleep(READ_INTERVAL) => {}
            _ = shutdown.changed() => {
                tracing::info!("swarm permissions: shutdown signal received");
                return;
            }
        }
    }
}

/// Read every agent's key and store the result. Returns the agents whose
/// `manage_root_agent` changed.
async fn sync(client: &async_nats::Client) -> Result<Vec<String>> {
    // An unconnected client does not fail a JetStream request, it hangs on it.
    swarm_queue_client::ensure_connected(client)?;
    let agents: Vec<String> = crate::lifecycle::agents_for_meta_listing()
        .await
        .context("enumerating this hive's agents")?
        .into_iter()
        .map(|spec| spec.name)
        .collect();
    sync_with(&agents, &permissions::path(), |agent| async move {
        agent_permissions::read(client, &agent)
            .await
            .map_err(|e| anyhow::anyhow!(e))
            .with_context(|| format!("reading the swarm's permissions for {agent}"))
    })
    .await
}

/// [`sync`] with the read and the file given. Nothing is written unless every
/// read succeeds.
async fn sync_with<F, Fut>(agents: &[String], path: &Path, read: F) -> Result<Vec<String>>
where
    F: Fn(String) -> Fut,
    Fut: std::future::Future<Output = Result<Read>>,
{
    let mut read_now = Stored::new();
    for agent in agents {
        match read(agent.clone()).await? {
            Read::Current(record) => {
                if record != AgentPermissions::default() {
                    read_now.insert(agent.clone(), record);
                }
            }
            Read::NoBucket => {
                tracing::debug!("swarm permissions: the agent-permissions bucket does not answer");
                return Ok(Vec::new());
            }
        }
    }
    let stored = permissions::read_at(path).context("reading the stored swarm permissions")?;
    if read_now == stored {
        return Ok(Vec::new());
    }
    permissions::write_at(path, &read_now).context("storing the swarm permissions")?;
    tracing::info!(permissions = ?read_now, "swarm permissions: stored the swarm's permissions");
    Ok(agents
        .iter()
        .filter(|a| manages(a, &stored) != manages(a, &read_now))
        .cloned()
        .collect())
}

/// Refresh the copy of its stored permissions each agent with a runtime dir
/// holds, so it follows this pass's reads and the import's progress.
async fn hand_all() {
    let agents = match crate::lifecycle::agents_for_meta_listing().await {
        Ok(agents) => agents,
        Err(e) => {
            tracing::warn!(error = ?e, "swarm permissions: enumerating agents to hand their copies failed");
            return;
        }
    };
    for spec in agents {
        let dir = crate::paths::agent_runtime_dir(&spec.name);
        if dir.is_dir() {
            crate::lifecycle::hand_permissions(&spec.name, &dir);
        }
    }
}

/// Whether `agent` holds the one permission hive-c0re enforces under `stored`.
fn manages(agent: &str, stored: &Stored) -> bool {
    permissions::effective(stored.get(agent)).contains(&Permission::ManageRootAgent)
}

/// Restart each running agent in `changed` once its current turn is done, so
/// its container starts with the bind mounts its new permissions give. A
/// stopped agent gets them at its next start.
async fn restart(coord: &Arc<Coordinator>, changed: Vec<String>) {
    let mut running = Vec::new();
    for agent in changed {
        if crate::lifecycle::is_running(&agent).await {
            running.push(agent);
        }
    }
    if running.is_empty() {
        return;
    }
    tracing::info!(agents = ?running, "swarm permissions: manage_root_agent changed, restarting");
    if let Err(e) = crate::job_queue::power::restart_many(coord, &running, true).await {
        tracing::error!(agents = ?running, error = ?e, "swarm permissions: restart insert failed");
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use super::{Read, sync_with};
    use swarm_queue_client::agent_permissions::AgentPermissions;

    fn record(json: &str) -> AgentPermissions {
        serde_json::from_str(json).expect("decodes")
    }

    fn agents() -> Vec<String> {
        vec!["atlas".to_owned(), "ruth".to_owned()]
    }

    /// Sync `agents()` against `path` with each agent read as `answer` gives.
    async fn sync(
        path: &Path,
        answer: impl Fn(&str) -> anyhow::Result<Read>,
    ) -> anyhow::Result<Vec<String>> {
        sync_with(&agents(), path, |agent| {
            let out = answer(&agent);
            async move { out }
        })
        .await
    }

    /// No key anywhere is everyone on the default: nothing stored, and no
    /// restart, since that is what the agents already run with.
    #[tokio::test]
    async fn missing_keys_are_the_default_and_restart_nothing() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.json");
        let changed = sync(&path, |_| Ok(Read::Current(AgentPermissions::default())))
            .await
            .expect("read");
        assert!(changed.is_empty());
        assert!(crate::permissions::read_at(&path).expect("read").is_empty());
        assert!(
            !crate::permissions::holds(&path, "ruth", hive_types::Permission::ManageRootAgent)
                .expect("read")
        );
    }

    /// Granting `manage_root_agent` to agents on the default restarts them;
    /// clearing it again restarts them back.
    #[tokio::test]
    async fn a_changed_manage_root_agent_asks_for_a_restart() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.json");
        let granted = || {
            Ok(Read::Current(record(
                r#"{"permissions":["manage_root_agent"]}"#,
            )))
        };
        let changed = sync(&path, |_| granted()).await.expect("read");
        assert_eq!(changed, agents());

        // The same answer again changes nothing.
        let unchanged = sync(&path, |_| granted()).await.expect("read");
        assert!(unchanged.is_empty());

        let cleared = sync(&path, |_| Ok(Read::Current(AgentPermissions::default())))
            .await
            .expect("read");
        assert_eq!(cleared, agents());
    }

    /// A change that leaves `manage_root_agent` where it was is stored and
    /// restarts nothing: the agent applies tool permissions itself.
    #[tokio::test]
    async fn a_tool_only_change_is_stored_without_a_restart() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.json");
        let changed = sync(&path, |agent| {
            Ok(Read::Current(if agent == "ruth" {
                record(r#"{"permissions":["meta"]}"#)
            } else {
                record(r#"{"permissions":["messaging","web_tools"]}"#)
            }))
        })
        .await
        .expect("read");
        assert!(changed.is_empty());
        assert!(
            crate::permissions::read_at(&path)
                .expect("read")
                .contains_key("atlas")
        );
    }

    /// One failed read, or a bucket that does not answer, keeps what was
    /// stored and restarts nothing.
    #[tokio::test]
    async fn a_failed_read_keeps_the_last_value() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.json");
        // atlas stays on the default, so only ruth changes.
        let ruth_granted = |agent: &str| {
            Ok(Read::Current(record(if agent == "ruth" {
                r#"{"permissions":["manage_root_agent"]}"#
            } else {
                "{}"
            })))
        };
        assert_eq!(
            sync(&path, ruth_granted).await.expect("read"),
            vec!["ruth".to_owned()]
        );
        let before = std::fs::read(&path).expect("stored");

        sync(&path, |agent| {
            if agent == "ruth" {
                Err(anyhow::anyhow!("timed out"))
            } else {
                Ok(Read::Current(AgentPermissions::default()))
            }
        })
        .await
        .expect_err("a failed read fails the pass");
        assert_eq!(std::fs::read(&path).expect("stored"), before);

        let changed = sync(&path, |_| Ok(Read::NoBucket))
            .await
            .expect("no bucket");
        assert!(changed.is_empty());
        assert_eq!(std::fs::read(&path).expect("stored"), before);
    }
}