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

//! Reads the subagent model list 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`).
//!
//! **Read on every such call**, 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`] and the caller checks against `HIVE_AVAILABLE_MODELS`.

use std::sync::{Arc, OnceLock};
use std::time::Duration;

use swarm_queue_client::agent_inference::{self, read_own};

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

/// 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 said about `subagent_models`.
#[derive(Debug, PartialEq, Eq)]
pub enum Read {
    /// The swarm sets the list.
    Set(Vec<String>),
    /// The key, the field, or the bucket is absent.
    Unset,
    /// Nothing could be read, and why.
    Failed(String),
}

/// The handle `start`/`continue` read through.
pub struct Reader {
    agent: String,
    /// Set once the first connect succeeds; async-nats reconnects it after.
    client: Arc<OnceLock<async_nats::Client>>,
}

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"
                );
                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"
                );
                return None;
            }
        };
        swarm_queue_client::install_crypto_provider();
        let reader = Self {
            agent: target.agent.clone(),
            client: Arc::new(OnceLock::new()),
        };
        tokio::spawn(connect_until_up(
            Arc::new(target),
            Arc::clone(&reader.client),
        ));
        Some(reader)
    }

    /// This agent's `subagent_models` as the swarm holds it now.
    pub async fn read(&self) -> Read {
        let Some(client) = self.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, &self.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 models: the agent-inference bucket does not answer");
                Read::Unset
            }
            Ok(Ok(agent_inference::Read::Current(current))) => {
                current.subagent_models.map_or(Read::Unset, Read::Set)
            }
        }
    }
}

/// 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 models: connected as this agent");
                let _already_set = client.set(connected);
                return;
            }
            Err(e) => {
                tracing::warn!(
                    error = format!("{e:#}"),
                    retry_in = ?RETRY_AFTER,
                    "subagent models: connecting with this agent's own queue credential failed; \
                     HIVE_AVAILABLE_MODELS applies 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 reader = Reader {
            agent: "iris".to_owned(),
            client: Arc::new(OnceLock::new()),
        };
        assert!(matches!(reader.read().await, Read::Failed(_)));
    }
}