hyperhive/hive-subagent-mcp/src/swarm_models.rs

//! Reads the subagent settings and permissions the swarm sets for this agent.
//!
//! swarm-controller keeps the agent's settings under its key in the
//! `agent-inference` KV bucket. Its `subagent_models`, when set, is the list
//! a `start` or `continue` naming a `model` is checked against, in place of
//! `HIVE_AVAILABLE_MODELS` (`crate::session::check_model`). Its
//! `max_subagents`, when set, caps how many subagents run at once
//! (`crate::session::State::reserve_under`). The agent's permissions, under
//! its key in `agent-permissions`, decide a new subagent's tools
//! ([`permissions_now`]).
//!
//! **Read on every `start` and `continue`**, never cached, with the one subject the
//! agent's own queue credential is granted in that bucket
//! (`swarm_queue_client::agent_inference::read_own`). The connection is made
//! the way [`crate::swarm_term`] makes its own: as the agent, never as the
//! hive.
//!
//! **A subagent start never waits on the queue.** Before the first connect
//! succeeds, and when a read fails or outlasts `READ_TIMEOUT`, the read is
//! [`Read::Failed`]: the caller checks against `HIVE_AVAILABLE_MODELS` and
//! applies no cap.

use std::path::Path;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;

use hive_sh4re::permissions::Permission;
use swarm_queue_client::agent_inference::{self, read_own};
use swarm_queue_client::agent_permissions::{self, HIVE_COPY_PATH, HiveCopy};
use swarm_wire::agent_inference::AgentInference;

use crate::swarm_term::{Target, connect};

/// The last permission set [`Reader::refresh_permissions`] read; `None` until
/// one succeeds.
static LAST_PERMISSIONS: std::sync::RwLock<Option<Vec<Permission>>> = std::sync::RwLock::new(None);

/// The permissions this agent holds now, for the subagent being spawned: the
/// last set read, or what hive-c0re stored for the agent before any read
/// succeeds ([`HiveCopy::before_read`]).
#[must_use]
pub fn permissions_now() -> Vec<Permission> {
    let last = LAST_PERMISSIONS
        .read()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .clone();
    last.unwrap_or_else(|| HiveCopy::read_at(Path::new(HIVE_COPY_PATH)).before_read())
}

/// How long after a failed connect the task tries again.
const RETRY_AFTER: Duration = Duration::from_mins(1);

/// How long one read of the key may take. A read the queue refuses for
/// permissions reaches the client as a timeout, so this is also how long a
/// misgranted agent's `start` waits.
const READ_TIMEOUT: Duration = Duration::from_secs(2);

/// What one read of this agent's key found.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Read {
    /// The settings the swarm holds; an absent key or bucket has every field
    /// unset.
    Current(AgentInference),
    /// Nothing could be read, and why.
    Failed(String),
}

/// The handle `start`/`continue` read through.
pub struct Reader {
    source: Source,
}

enum Source {
    Queue {
        agent: String,
        /// Set once the first connect succeeds; async-nats reconnects it after.
        client: Arc<OnceLock<async_nats::Client>>,
    },
    /// Answers every read with `read`, counting them.
    #[cfg(test)]
    Fixed { read: Read, reads: Arc<AtomicUsize> },
}

impl Reader {
    /// Start connecting, or return `None` when this daemon has no queue
    /// address or no store to read the agent's credential from. Logs which,
    /// once.
    #[must_use]
    pub fn spawn() -> Option<Self> {
        let target = match Target::from_lookup(|k| std::env::var(k).ok()) {
            Ok(Some(target)) => target,
            Ok(None) => {
                tracing::info!(
                    "no swarm queue address or no secret store for this agent; subagent models \
                     are checked against HIVE_AVAILABLE_MODELS only, and no subagent cap applies"
                );
                return None;
            }
            Err(e) => {
                tracing::warn!(
                    error = format!("{e:#}"),
                    "this agent's secret store coordinates are incomplete; subagent models are \
                     checked against HIVE_AVAILABLE_MODELS only, and no subagent cap applies"
                );
                return None;
            }
        };
        swarm_secret_client::install_crypto_provider();
        let agent = target.agent.clone();
        let client = Arc::new(OnceLock::new());
        tokio::spawn(connect_until_up(Arc::new(target), Arc::clone(&client)));
        Some(Self {
            source: Source::Queue { agent, client },
        })
    }

    /// A reader answering every read with `read`, and the count of reads.
    #[cfg(test)]
    pub(crate) fn fixed(read: Read) -> (Self, Arc<AtomicUsize>) {
        let reads = Arc::new(AtomicUsize::new(0));
        let source = Source::Fixed {
            read,
            reads: Arc::clone(&reads),
        };
        (Self { source }, reads)
    }

    /// The agent whose key is read; `"test"` for a fixed reader.
    #[must_use]
    pub fn agent(&self) -> &str {
        match &self.source {
            Source::Queue { agent, .. } => agent,
            #[cfg(test)]
            Source::Fixed { .. } => "test",
        }
    }

    /// This agent's settings as the swarm holds them now.
    pub async fn read(&self) -> Read {
        match &self.source {
            Source::Queue { agent, client } => read_queue(agent, client).await,
            #[cfg(test)]
            Source::Fixed { read, reads } => {
                reads.fetch_add(1, Ordering::Relaxed);
                read.clone()
            }
        }
    }

    /// Read the permissions the swarm grants this agent and keep them for
    /// [`permissions_now`], judged by [`HiveCopy::after_read`]. A read that
    /// fails keeps the last value.
    pub async fn refresh_permissions(&self) {
        let read = match &self.source {
            Source::Queue { agent, client } => read_permissions(agent, client).await,
            #[cfg(test)]
            Source::Fixed { .. } => None,
        };
        if let Some(perms) = read {
            *LAST_PERMISSIONS
                .write()
                .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(perms);
        }
    }
}

async fn read_permissions(
    agent: &str,
    client: &OnceLock<async_nats::Client>,
) -> Option<Vec<Permission>> {
    let client = client.get()?;
    swarm_queue_client::ensure_connected(client).ok()?;
    match tokio::time::timeout(READ_TIMEOUT, agent_permissions::read(client, agent)).await {
        Ok(Ok(agent_permissions::Read::Current(record))) => {
            Some(HiveCopy::read_at(Path::new(HIVE_COPY_PATH)).after_read(&record))
        }
        Ok(Ok(agent_permissions::Read::NoBucket)) => None,
        Ok(Err(e)) => {
            tracing::warn!(error = %e, "subagent permissions: reading this agent's key failed");
            None
        }
        Err(_) => {
            tracing::warn!("subagent permissions: no answer within {READ_TIMEOUT:?}");
            None
        }
    }
}

async fn read_queue(agent: &str, client: &OnceLock<async_nats::Client>) -> Read {
    let Some(client) = client.get() else {
        return Read::Failed("not connected to the swarm queue yet".to_owned());
    };
    // An unconnected client does not fail a JetStream request, it hangs on it.
    if let Err(e) = swarm_queue_client::ensure_connected(client) {
        return Read::Failed(swarm_queue_client::chain(&e));
    }
    match tokio::time::timeout(READ_TIMEOUT, read_own(client, agent)).await {
        Err(_) => Read::Failed(format!("no answer within {READ_TIMEOUT:?}")),
        Ok(Err(e)) => Read::Failed(e.to_string()),
        Ok(Ok(agent_inference::Read::NoBucket)) => {
            tracing::debug!("subagent settings: the agent-inference bucket does not answer");
            Read::Current(AgentInference::default())
        }
        Ok(Ok(agent_inference::Read::Current(current))) => Read::Current(current),
    }
}

/// Connect, retrying every [`RETRY_AFTER`] until one succeeds.
async fn connect_until_up(target: Arc<Target>, client: Arc<OnceLock<async_nats::Client>>) {
    loop {
        match connect(&target).await {
            Ok(connected) => {
                tracing::info!(url = %target.url, "subagent settings: connected as this agent");
                let _already_set = client.set(connected);
                return;
            }
            Err(e) => {
                tracing::warn!(
                    error = format!("{e:#}"),
                    retry_in = ?RETRY_AFTER,
                    "subagent settings: connecting with this agent's own queue credential failed; \
                     HIVE_AVAILABLE_MODELS and no subagent cap apply until a connect succeeds"
                );
                tokio::time::sleep(RETRY_AFTER).await;
            }
        }
    }
}

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

    #[tokio::test]
    async fn a_reader_not_yet_connected_reads_as_failed() {
        let client = OnceLock::new();
        assert!(matches!(read_queue("iris", &client).await, Read::Failed(_)));
    }
}