//! 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`. #[cfg(test)] use std::sync::atomic::{AtomicUsize, Ordering}; 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, Clone, PartialEq, Eq)] pub enum Read { /// The swarm sets the list. Set(Vec), /// 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 { source: Source, } enum Source { Queue { agent: String, /// Set once the first connect succeeds; async-nats reconnects it after. client: Arc>, }, /// Answers every read with `read`, counting them. #[cfg(test)] Fixed { read: Read, reads: Arc }, } 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 { 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 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) { let reads = Arc::new(AtomicUsize::new(0)); let source = Source::Fixed { read, reads: Arc::clone(&reads), }; (Self { source }, reads) } /// This agent's `subagent_models` as the swarm holds it 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() } } } } async fn read_queue(agent: &str, client: &OnceLock) -> 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 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, client: Arc>) { 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 client = OnceLock::new(); assert!(matches!(read_queue("iris", &client).await, Read::Failed(_))); } }