diff --git a/hive-subagent-mcp/src/session.rs b/hive-subagent-mcp/src/session.rs index 452ff4a3..75031cd1 100644 --- a/hive-subagent-mcp/src/session.rs +++ b/hive-subagent-mcp/src/session.rs @@ -389,6 +389,9 @@ pub struct State { /// Where the swarm's `subagent_models` for this agent is read from, when /// this agent can read it ([`State::reading_models_from`]). models: Option, + /// The agent's own model list, checked when the swarm sets none: + /// [`available_models`] outside tests. + local_models: fn() -> Option, } /// Which opaque URL segment belongs to which session — the whole of a @@ -443,6 +446,7 @@ impl State { runtime: SubagentRuntime::Claude, term: None, models: None, + local_models: available_models, } } @@ -471,7 +475,7 @@ impl State { Some(reader) => Some(reader.read().await), None => None, }; - check_model(model, &model_list(read, available_models())) + check_model(model, &model_list(read, (self.local_models)())) } /// Run every subagent on `runtime` — the daemon passes the parent @@ -4324,6 +4328,82 @@ mod tests { ); } + /// A state whose swarm read answers `read` and whose own list is `haiku`, + /// with `name` already reserved: a call that gets past the model check + /// stops at `reserve` and spawns nothing. Also the reader's read count. + fn state_reading( + read: crate::swarm_models::Read, + name: &str, + ) -> (Arc, Arc) { + let (reader, reads) = crate::swarm_models::Reader::fixed(read); + let mut state = + State::new(PathBuf::from("/dev/null"), signal_url()).reading_models_from(Some(reader)); + state.local_models = || Some("haiku".to_owned()); + assert!(state.reserve(name)); + (Arc::new(state), reads) + } + + fn start_on(name: &str, model: Option<&str>) -> StartRequest { + StartRequest { + model: model.map(str::to_owned), + ..start_request(name) + } + } + + async fn continue_on( + state: &Arc, + name: &str, + model: Option<&str>, + ) -> anyhow::Result { + continue_( + state, + name, + "go".to_owned(), + model.map(str::to_owned), + None, + None, + ) + .await + } + + /// The refusal comes from the model check, ahead of `reserve`'s. + #[tokio::test] + async fn start_and_continue_refuse_a_model_outside_the_swarm_list_before_reserving() { + let (state, reads) = state_reading(swarm(&["sonnet"]), "n"); + let started = start(&state, start_on("n", Some("opus"))).await; + let continued = continue_on(&state, "n", Some("opus")).await; + for result in [started, continued] { + let err = result + .expect_err("opus is outside the swarm's list") + .to_string(); + assert!(err.contains("swarm's subagent_models"), "{err}"); + } + assert_eq!(reads.load(std::sync::atomic::Ordering::Relaxed), 2); + } + + #[tokio::test] + async fn start_and_continue_take_a_swarm_listed_model_the_local_list_lacks() { + let (state, _) = state_reading(swarm(&["opus"]), "n"); + let started = start(&state, start_on("n", Some("opus"))).await; + let continued = continue_on(&state, "n", Some("opus")).await; + for result in [started, continued] { + let err = result.expect_err("the name is reserved").to_string(); + assert!(err.contains("already running"), "{err}"); + } + } + + #[tokio::test] + async fn start_and_continue_without_a_model_do_not_read_the_swarm() { + let (state, reads) = state_reading(swarm(&["sonnet"]), "n"); + let started = start(&state, start_on("n", None)).await; + let continued = continue_on(&state, "n", None).await; + for result in [started, continued] { + let err = result.expect_err("the name is reserved").to_string(); + assert!(err.contains("already running"), "{err}"); + } + assert_eq!(reads.load(std::sync::atomic::Ordering::Relaxed), 0); + } + /// A `tracing` writer whose output a test can read back. #[derive(Clone, Default)] struct CapturedLogs(Arc>>); diff --git a/hive-subagent-mcp/src/swarm_models.rs b/hive-subagent-mcp/src/swarm_models.rs index 875b750c..2c4595c0 100644 --- a/hive-subagent-mcp/src/swarm_models.rs +++ b/hive-subagent-mcp/src/swarm_models.rs @@ -15,6 +15,8 @@ //! 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; @@ -31,7 +33,7 @@ const RETRY_AFTER: Duration = Duration::from_mins(1); const READ_TIMEOUT: Duration = Duration::from_secs(2); /// What one read of this agent's key said about `subagent_models`. -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum Read { /// The swarm sets the list. Set(Vec), @@ -43,9 +45,18 @@ pub enum Read { /// The handle `start`/`continue` read through. pub struct Reader { - agent: String, - /// Set once the first connect succeeds; async-nats reconnects it after. - client: Arc>, + 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 { @@ -73,36 +84,55 @@ impl Reader { } }; swarm_queue_client::install_crypto_provider(); - let reader = Self { - agent: target.agent.clone(), - client: Arc::new(OnceLock::new()), + 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), }; - tokio::spawn(connect_until_up( - Arc::new(target), - Arc::clone(&reader.client), - )); - Some(reader) + (Self { source }, reads) } /// 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 &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() + } } - 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) - } + } +} + +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) } } } @@ -135,10 +165,7 @@ mod tests { #[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(_))); + let client = OnceLock::new(); + assert!(matches!(read_queue("iris", &client).await, Read::Failed(_))); } }