hyperhive/swarm-controller/src/matrix_account/agent_token.rs

//! Each agent's own account on the swarm's homeserver: created here with the
//! **swarm's** appservice token, stored at `swarm/agents/<agent>/matrix/main`,
//! and pulled from there by the agent's matrix daemon itself.
//!
//! The appservice token is minted inside the matrix container and published
//! to `swarm_secret_client::matrix::swarm_appservice_token_path`, which only
//! matrix-ctl and this daemon may read. No hive holds it, and no hive mints an
//! agent's account any more.
//!
//! ⚠️ **Every mint replaces the agent's live token.** Both calls pin the
//! device id `hyperhive-<agent>`, so a login for an existing account replaces
//! that device's token: the one in the store before, or the file token a hive
//! minted. That is why the decision reads the stored token back with `whoami`
//! ([`classify`]) instead of only checking that the store holds something, and
//! why a failed read never mints ([`Observed::Unknown`]).
//!
//! Two callers insert the same `MintAgentMatrixAccount` job node: agent
//! creation, and [`spawn`]'s pass at start and every [`RECONCILE_INTERVAL`]
//! over every agent that holds a store identity — the same roster and shape as
//! `crate::forge::agent_token`.

use std::sync::Arc;

use anyhow::{Context, Result, bail};
use swarm_matrix_client::{self as homeserver, Whoami};
use swarm_secret_client::{SecretStore, client::DEFAULT_CERT_MOUNT, matrix, policy};

/// The account name every agent's own matrix credential is stored under.
///
/// `main` is what `nix/agent-modules/matrix.nix` declares per agent and what
/// the agent's daemon reads first. ⚠️ `PUT .../matrix-accounts/{account}`
/// refuses this exact name (`super::is_reserved_account`): that route stores
/// an **external** account an operator supplies, and one called `main` would
/// overwrite this. Two writers, one reserved name, and the reservation is what
/// keeps them apart.
pub const ACCOUNT: &str = "main";

/// How often [`spawn`] re-checks every agent's account.
const RECONCILE_INTERVAL: std::time::Duration = std::time::Duration::from_mins(5);

/// Why an account's token has to be (re)minted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MintReason {
    /// The store holds nothing for this agent: a new agent, or one whose
    /// account a hive created before the swarm did this.
    NotStored,
    /// The homeserver does not know the stored token: another login on the
    /// same device replaced it.
    Revoked,
    /// The stored token is live but belongs to another account.
    OtherUser,
}

/// What one look at an agent's stored token found.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Probe {
    /// Nothing at the agent's path.
    NotStored,
    /// The homeserver's verdict on the stored token.
    Whoami(Whoami),
}

/// What to do about one agent's account.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decision {
    /// The stored token is a live token for this agent's own account.
    Keep,
    /// Mint and store a new one.
    Mint(MintReason),
}

/// What one pass found for one agent.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Observed {
    /// Both reads worked, and this is what [`classify`] made of them.
    Decided(Decision),
    /// A read failed. Nothing is known, so nothing is done.
    Unknown,
}

/// The localpart of a full user id, `@<localpart>:<server>`.
fn localpart(user_id: &str) -> Option<&str> {
    user_id.strip_prefix('@')?.split_once(':').map(|(l, _)| l)
}

/// Decide what to do about `agent`'s account from what its stored token is.
pub fn classify(agent: &str, probe: &Probe) -> Decision {
    match probe {
        Probe::NotStored => Decision::Mint(MintReason::NotStored),
        Probe::Whoami(Whoami::UnknownToken) => Decision::Mint(MintReason::Revoked),
        Probe::Whoami(Whoami::User(user)) if localpart(user) == Some(agent) => Decision::Keep,
        Probe::Whoami(Whoami::User(_)) => Decision::Mint(MintReason::OtherUser),
    }
}

/// The agents a pass mints for.
///
/// [`Observed::Unknown`] never mints: a homeserver or store outage must not
/// turn into a new token for every agent, each of which kills the one the
/// agent is running on.
pub fn plan(observed: &[(String, Observed)]) -> Vec<String> {
    observed
        .iter()
        .filter(|(_, o)| matches!(o, Observed::Decided(Decision::Mint(_))))
        .map(|(agent, _)| agent.clone())
        .collect()
}

/// The swarm appservice token, or an error naming why there is none.
pub(super) async fn appservice_token(store: &SecretStore) -> Result<String> {
    let path = matrix::swarm_appservice_token_path()?;
    let stored: Option<matrix::Credential> = store
        .read_optional(&path)
        .await
        .with_context(|| format!("reading {path}"))?;
    match stored {
        Some(c) if !c.value.trim().is_empty() => Ok(c.value.trim().to_owned()),
        _ => bail!(
            "nothing at {path}: the matrix container has not published the swarm \
             appservice token (swarm-matrix-appservice-publish)"
        ),
    }
}

/// What the store and the homeserver say about `agent`'s stored token.
async fn probe(
    store: &SecretStore,
    http: &reqwest::Client,
    base: &str,
    agent: &str,
) -> Result<Probe> {
    probe_at(store, http, base, &matrix::account_path(agent, ACCOUNT)?).await
}

/// What the store and the homeserver say about the token stored at `path`.
pub(super) async fn probe_at(
    store: &SecretStore,
    http: &reqwest::Client,
    base: &str,
    path: &str,
) -> Result<Probe> {
    let stored: Option<matrix::Credential> = store
        .read_optional(path)
        .await
        .with_context(|| format!("reading {path}"))?;
    let Some(stored) = stored else {
        return Ok(Probe::NotStored);
    };
    let verdict = homeserver::whoami(http, base, stored.value.trim())
        .await
        .with_context(|| format!("asking the homeserver about the token at {path}"))?;
    Ok(Probe::Whoami(verdict))
}

/// Make sure `agent` holds a live token for its own account on the swarm's
/// homeserver at `base`, creating the account or logging in to it when it
/// does not. The whole job of the `MintAgentMatrixAccount` node.
///
/// Reads the new token back with `whoami` before reporting success, so the
/// node does not report success on a write nobody has read. A crash between
/// the login and the write leaves a dead token in the store, which the next
/// pass classifies as [`MintReason::Revoked`].
///
/// # Errors
/// When the store or the homeserver refuses a step, the swarm appservice
/// token has not been published, or the new token authenticates as someone
/// else.
pub async fn ensure_agent_matrix_account(base: &str, agent: &str) -> Result<()> {
    let store = crate::store::connect()
        .await
        .context("logging in to the swarm secret store")?;
    let as_token = appservice_token(&store).await?;
    let http = homeserver::client()?;
    let reason = match classify(agent, &probe(&store, &http, base, agent).await?) {
        Decision::Keep => {
            tracing::debug!(agent, "agent matrix account is current; left as it is");
            return Ok(());
        }
        Decision::Mint(reason) => reason,
    };

    let path = matrix::account_path(agent, ACCOUNT)?;
    mint_at(&store, &http, base, agent, &as_token, &path).await?;
    tracing::info!(agent, ?reason, %path, "agent matrix account token minted and stored");
    Ok(())
}

/// Create the account `account` (a localpart), or log in to it as the appservice when it
/// exists, store the token at `path`, and read it back with `whoami`.
///
/// # Errors
/// When the store or the homeserver refuses a step, or the new token
/// authenticates as someone else.
pub(super) async fn mint_at(
    store: &SecretStore,
    http: &reqwest::Client,
    base: &str,
    account: &str,
    as_token: &str,
    path: &str,
) -> Result<()> {
    let token = match homeserver::register(http, base, account, as_token).await? {
        homeserver::Registered::Token(token) => token,
        // An account minted before, or one a hive created itself: log in as
        // the appservice on the same device instead.
        homeserver::Registered::AlreadyExists => {
            homeserver::appservice_login(http, base, account, as_token).await?
        }
    };
    store
        .write(
            path,
            &matrix::Credential {
                value: token.clone(),
                homeserver: Some(base.to_owned()),
            },
        )
        .await
        .with_context(|| format!("storing @{account}'s matrix token at {path}"))?;

    match homeserver::whoami(http, base, &token)
        .await
        .with_context(|| format!("using @{account}'s new matrix token"))?
    {
        Whoami::User(user) if localpart(&user) == Some(account) => Ok(()),
        _ => bail!("@{account}'s new matrix token does not authenticate as @{account}"),
    }
}

/// One pass: every agent holding a store identity, observed.
///
/// Fails as a whole, rather than per agent, when the swarm appservice token is
/// not published: every mint would fail on it, and one message says so.
async fn observe_all(base: &str) -> Result<Vec<(String, Observed)>> {
    let store = crate::store::connect()
        .await
        .context("logging in to the swarm secret store")?;
    appservice_token(&store).await?;
    let http = homeserver::client()?;
    let roles = store
        .list_cert_roles(DEFAULT_CERT_MOUNT)
        .await
        .context("listing the store's cert-auth roles")?;
    let mut observed = Vec::new();
    for agent in policy::agents_from_role_names(&roles) {
        let o = match probe(&store, &http, base, &agent).await {
            Ok(p) => Observed::Decided(classify(&agent, &p)),
            Err(e) => {
                tracing::warn!(agent, error = %format!("{e:#}"), "agent matrix account: read failed");
                Observed::Unknown
            }
        };
        observed.push((agent, o));
    }
    Ok(observed)
}

/// Check every agent's account now and every [`RECONCILE_INTERVAL`] after, and
/// hand the agents that need a token to `enqueue`, which inserts a
/// `MintAgentMatrixAccount` node for each.
///
/// The roster is the store's `hive-agent-*` cert-auth roles: exactly the
/// agents that can read what the node writes. A pass that fails is logged and
/// retried on the next tick; it never stops the daemon.
pub fn spawn(base: Arc<str>, enqueue: impl Fn(Vec<String>) + Send + 'static) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(RECONCILE_INTERVAL);
        loop {
            ticker.tick().await;
            match observe_all(&base).await {
                Ok(observed) => {
                    let agents = plan(&observed);
                    if agents.is_empty() {
                        tracing::debug!(
                            checked = observed.len(),
                            "agent matrix accounts: all current"
                        );
                    } else {
                        tracing::info!(
                            checked = observed.len(),
                            minting = agents.len(),
                            "agent matrix accounts: queueing mints"
                        );
                        enqueue(agents);
                    }
                }
                Err(e) => tracing::warn!(
                    error = %format!("{e:#}"),
                    retry_in_s = RECONCILE_INTERVAL.as_secs(),
                    "agent matrix accounts: pass failed; retrying next tick"
                ),
            }
        }
    });
}

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

    fn user(id: &str) -> Probe {
        Probe::Whoami(Whoami::User(id.to_owned()))
    }

    #[test]
    fn a_live_token_for_the_agent_is_kept() {
        assert_eq!(classify("atlas", &user("@atlas:t.local")), Decision::Keep);
    }

    #[test]
    fn nothing_stored_mints() {
        assert_eq!(
            classify("atlas", &Probe::NotStored),
            Decision::Mint(MintReason::NotStored)
        );
    }

    #[test]
    fn a_replaced_token_mints() {
        // The shape a hive re-login on the same device leaves behind.
        assert_eq!(
            classify("atlas", &Probe::Whoami(Whoami::UnknownToken)),
            Decision::Mint(MintReason::Revoked)
        );
    }

    #[test]
    fn a_token_for_another_account_mints() {
        // Localpart compared whole, so a prefix of the agent's name is not it.
        for other in ["@argus:t.local", "@atlas2:t.local", "@atla:t.local"] {
            assert_eq!(
                classify("atlas", &user(other)),
                Decision::Mint(MintReason::OtherUser),
                "{other}"
            );
        }
    }

    #[test]
    fn a_pass_mints_only_what_it_knows_needs_one() {
        let observed = [
            ("a".to_owned(), Observed::Decided(Decision::Keep)),
            (
                "b".to_owned(),
                Observed::Decided(Decision::Mint(MintReason::NotStored)),
            ),
            ("c".to_owned(), Observed::Unknown),
            (
                "d".to_owned(),
                Observed::Decided(Decision::Mint(MintReason::Revoked)),
            ),
        ];
        assert_eq!(plan(&observed), ["b", "d"]);
    }

    #[test]
    fn an_outage_plans_nothing() {
        // Each mint replaces the token an agent is running on, so a pass that
        // could not read must not mint for everyone.
        let observed = [
            ("a".to_owned(), Observed::Unknown),
            ("b".to_owned(), Observed::Unknown),
        ];
        assert!(plan(&observed).is_empty());
    }

    #[test]
    fn the_published_path_is_the_hiveless_one_the_agent_reads() {
        // The literal is the point, and the absence of a hive segment is the
        // half of it that was decided: an agent's account follows the agent.
        assert_eq!(
            matrix::account_path("atlas", ACCOUNT).expect("a plain name is legal"),
            "swarm/agents/atlas/matrix/main"
        );
    }

    #[test]
    fn the_account_name_is_the_one_the_agent_module_declares() {
        // `nix/agent-modules/matrix.nix` renders this as a `matrixAccounts`
        // key and nothing wires an override across.
        assert_eq!(ACCOUNT, "main");
    }
}