hive-subagent-mcp: check subagent models against the swarm's subagent_models
A `start` or `continue` naming a `model` now reads this agent's key in the `agent-inference` bucket and, when `subagent_models` is set, checks the model against that list instead of HIVE_AVAILABLE_MODELS (agent.nix `availableModels`). An unset field, an absent key or an absent bucket keeps today's check. A read that fails or takes over 2 s falls back to HIVE_AVAILABLE_MODELS with a warning, so a subagent start never waits on the queue. A call without a `model` is not checked and does not read. The read uses the agent's own queue credential through swarm_term's connect, and the one direct-get subject the agent is granted in that bucket. The refusal names the list and that it came from the swarm. Refs #5034 (S5)
This commit is contained in:
parent
971426a1c0
commit
90e6b984eb
9 changed files with 394 additions and 46 deletions
144
hive-subagent-mcp/src/swarm_models.rs
Normal file
144
hive-subagent-mcp/src/swarm_models.rs
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
//! 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(_)));
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue