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
|
|
@ -16,8 +16,11 @@ swarmctl agent inference clear iris
|
||||||
- `set` changes only the flags you pass; the other fields keep what's
|
- `set` changes only the flags you pass; the other fields keep what's
|
||||||
stored. A hive must declare the agent.
|
stored. A hive must declare the agent.
|
||||||
- `--effort` is one of `low`, `medium`, `high`, `xhigh`, `max`.
|
- `--effort` is one of `low`, `medium`, `high`, `xhigh`, `max`.
|
||||||
- `--subagent-models` takes a comma-separated list. It's stored and shown,
|
- `--subagent-models` takes a comma-separated list. It replaces the
|
||||||
but agents don't apply it yet.
|
agent's `availableModels` as the models its subagent `start` and
|
||||||
|
`continue` calls may name, read on each call
|
||||||
|
([`subagent.md`](../tools/subagent.md#which-models-a-subagent-may-run-on)).
|
||||||
|
The web UI's model picker doesn't use it.
|
||||||
- `get` shows each field, naming the unset ones.
|
- `get` shows each field, naming the unset ones.
|
||||||
- `clear` removes every field; the agent goes back to its own values.
|
- `clear` removes every field; the agent goes back to its own values.
|
||||||
|
|
||||||
|
|
@ -65,4 +68,8 @@ The harness (`hive-agent`'s `swarm_inference`) applies a set field to the
|
||||||
in-memory model or effort without writing `hyperhive-model` or
|
in-memory model or effort without writing `hyperhive-model` or
|
||||||
`hyperhive-effort`, which keep the agent's own value.
|
`hyperhive-effort`, which keep the agent's own value.
|
||||||
|
|
||||||
|
The subagent daemon (`hive-subagent-mcp`'s `swarm_models`) reads
|
||||||
|
`subagent_models` with the same direct get, under the same credential, on
|
||||||
|
each `start` or `continue` that names a model.
|
||||||
|
|
||||||
</details>
|
</details>
|
||||||
|
|
|
||||||
|
|
@ -307,6 +307,28 @@ agent that needs a stable or non-default port.
|
||||||
Own systemd unit, defined alongside the other per-agent MCP daemons in
|
Own systemd unit, defined alongside the other per-agent MCP daemons in
|
||||||
`nix/agent-modules/mcp.nix`.
|
`nix/agent-modules/mcp.nix`.
|
||||||
|
|
||||||
|
## Which models a subagent may run on
|
||||||
|
|
||||||
|
The daemon refuses a `start` or `continue` that names a `model` outside the
|
||||||
|
agent's list before anything spawns, and the error names the list and where
|
||||||
|
it came from. One without a `model` runs on claude's own default, which the
|
||||||
|
daemon doesn't check. The list is:
|
||||||
|
|
||||||
|
- the swarm's `subagent_models` for this agent, when it sets one
|
||||||
|
([`inference.md`](../swarm/inference.md)). It replaces the agent's own
|
||||||
|
list, so it can name models that list leaves out.
|
||||||
|
- otherwise `services.hyperhive.agent.availableModels`, which reaches the
|
||||||
|
daemon as `HIVE_AVAILABLE_MODELS`. An empty list restricts nothing.
|
||||||
|
|
||||||
|
The daemon reads the swarm's value on each such call, so a change applies
|
||||||
|
to the next one. When the read fails, or doesn't answer within 2 seconds,
|
||||||
|
the daemon logs a warning and checks against `availableModels`. The same
|
||||||
|
happens until the daemon's first connection to the swarm queue succeeds.
|
||||||
|
The match ignores case. On an ACP agent the names are the agent's own, such
|
||||||
|
as `<provider>/<model>` on `opencode`.
|
||||||
|
|
||||||
|
The web UI's model picker still offers `availableModels`.
|
||||||
|
|
||||||
## Runtime
|
## Runtime
|
||||||
|
|
||||||
A subagent runs on its parent agent's runtime
|
A subagent runs on its parent agent's runtime
|
||||||
|
|
|
||||||
|
|
@ -29,7 +29,8 @@ schemars.workspace = true
|
||||||
serde.workspace = true
|
serde.workspace = true
|
||||||
serde_json.workspace = true
|
serde_json.workspace = true
|
||||||
# `subagent_term`: the subject each subagent's rows are published on.
|
# `subagent_term`: the subject each subagent's rows are published on.
|
||||||
swarm-queue-client.workspace = true
|
# `agent-inference`: this agent's `subagent_models`, read with one direct get.
|
||||||
|
swarm-queue-client = { workspace = true, features = ["agent-inference"] }
|
||||||
# The agent's own queue credential, read under its store identity.
|
# The agent's own queue credential, read under its store identity.
|
||||||
swarm-secret-client.workspace = true
|
swarm-secret-client.workspace = true
|
||||||
tokio.workspace = true
|
tokio.workspace = true
|
||||||
|
|
|
||||||
|
|
@ -21,4 +21,5 @@ pub mod mcp_config;
|
||||||
pub mod paths;
|
pub mod paths;
|
||||||
pub mod role;
|
pub mod role;
|
||||||
pub mod session;
|
pub mod session;
|
||||||
|
pub mod swarm_models;
|
||||||
pub mod swarm_term;
|
pub mod swarm_term;
|
||||||
|
|
|
||||||
|
|
@ -56,7 +56,8 @@ async fn main() -> Result<()> {
|
||||||
let state = Arc::new(
|
let state = Arc::new(
|
||||||
hive_subagent_mcp::session::State::new(todo_socket, signal_base)
|
hive_subagent_mcp::session::State::new(todo_socket, signal_base)
|
||||||
.on_runtime(runtime)
|
.on_runtime(runtime)
|
||||||
.publishing_to(hive_subagent_mcp::swarm_term::Publisher::spawn()),
|
.publishing_to(hive_subagent_mcp::swarm_term::Publisher::spawn())
|
||||||
|
.reading_models_from(hive_subagent_mcp::swarm_models::Reader::spawn()),
|
||||||
);
|
);
|
||||||
|
|
||||||
// Serve the MCP tools over streamable-http forever. No background poll
|
// Serve the MCP tools over streamable-http forever. No background poll
|
||||||
|
|
|
||||||
|
|
@ -216,7 +216,7 @@ impl SubagentMcp {
|
||||||
you do not have refuses the whole call rather than spawning without it. See the \
|
you do not have refuses the whole call rather than spawning without it. See the \
|
||||||
`base:claude-subagents` skill for when to reach for this."
|
`base:claude-subagents` skill for when to reach for this."
|
||||||
)]
|
)]
|
||||||
fn start(&self, Parameters(args): Parameters<StartArgs>) -> String {
|
async fn start(&self, Parameters(args): Parameters<StartArgs>) -> String {
|
||||||
match session::start(
|
match session::start(
|
||||||
&self.state,
|
&self.state,
|
||||||
session::StartRequest {
|
session::StartRequest {
|
||||||
|
|
@ -231,7 +231,9 @@ impl SubagentMcp {
|
||||||
max_turns: args.max_turns,
|
max_turns: args.max_turns,
|
||||||
report_file: args.report_file,
|
report_file: args.report_file,
|
||||||
},
|
},
|
||||||
) {
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
Ok(msg) => msg,
|
Ok(msg) => msg,
|
||||||
Err(e) => format!("start error: {e:#}"),
|
Err(e) => format!("start error: {e:#}"),
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -386,6 +386,9 @@ pub struct State {
|
||||||
/// Where every subagent's output lines go up to the swarm, when this
|
/// Where every subagent's output lines go up to the swarm, when this
|
||||||
/// agent can publish there ([`State::publishing_to`]).
|
/// agent can publish there ([`State::publishing_to`]).
|
||||||
term: Option<crate::swarm_term::Publisher>,
|
term: Option<crate::swarm_term::Publisher>,
|
||||||
|
/// Where the swarm's `subagent_models` for this agent is read from, when
|
||||||
|
/// this agent can read it ([`State::reading_models_from`]).
|
||||||
|
models: Option<crate::swarm_models::Reader>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Which opaque URL segment belongs to which session — the whole of a
|
/// Which opaque URL segment belongs to which session — the whole of a
|
||||||
|
|
@ -439,6 +442,7 @@ impl State {
|
||||||
signal_tokens: Mutex::new(SignalTokens::default()),
|
signal_tokens: Mutex::new(SignalTokens::default()),
|
||||||
runtime: SubagentRuntime::Claude,
|
runtime: SubagentRuntime::Claude,
|
||||||
term: None,
|
term: None,
|
||||||
|
models: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -449,6 +453,27 @@ impl State {
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Check every named `model` against the swarm's `subagent_models` as
|
||||||
|
/// `models` reads it, where the swarm sets one.
|
||||||
|
#[must_use]
|
||||||
|
pub fn reading_models_from(mut self, models: Option<crate::swarm_models::Reader>) -> Self {
|
||||||
|
self.models = models;
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Refuse a `model` outside the list in effect for this agent (see
|
||||||
|
/// `check_model`). The swarm key is read only for a named model.
|
||||||
|
async fn allow_model(&self, model: Option<&str>) -> anyhow::Result<()> {
|
||||||
|
if model.is_none() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let read = match &self.models {
|
||||||
|
Some(reader) => Some(reader.read().await),
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
check_model(model, &model_list(read, available_models()))
|
||||||
|
}
|
||||||
|
|
||||||
/// Run every subagent on `runtime` — the daemon passes the parent
|
/// Run every subagent on `runtime` — the daemon passes the parent
|
||||||
/// agent's, read from its environment at startup. On ACP this resolves
|
/// agent's, read from its environment at startup. On ACP this resolves
|
||||||
/// the harness dir, which only exists inside a container.
|
/// the harness dir, which only exists inside a container.
|
||||||
|
|
@ -918,28 +943,69 @@ fn available_models() -> Option<String> {
|
||||||
std::env::var(AVAILABLE_MODELS_ENV).ok()
|
std::env::var(AVAILABLE_MODELS_ENV).ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The list a named `model` is checked against.
|
||||||
|
#[derive(Debug, PartialEq, Eq)]
|
||||||
|
enum ModelList {
|
||||||
|
/// The swarm's `subagent_models` for this agent.
|
||||||
|
Swarm(Vec<String>),
|
||||||
|
/// [`AVAILABLE_MODELS_ENV`]; `None` when the var is absent.
|
||||||
|
Local(Option<String>),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The swarm's list when it sets one, else `local`. `read` is `None` when
|
||||||
|
/// this daemon has no swarm queue to read from; a failed read is logged and
|
||||||
|
/// falls back to `local`.
|
||||||
|
fn model_list(read: Option<crate::swarm_models::Read>, local: Option<String>) -> ModelList {
|
||||||
|
use crate::swarm_models::Read;
|
||||||
|
match read {
|
||||||
|
Some(Read::Set(models)) => ModelList::Swarm(models),
|
||||||
|
Some(Read::Failed(error)) => {
|
||||||
|
tracing::warn!(
|
||||||
|
%error,
|
||||||
|
"reading this agent's subagent_models from the swarm failed; checking against \
|
||||||
|
{AVAILABLE_MODELS_ENV}"
|
||||||
|
);
|
||||||
|
ModelList::Local(local)
|
||||||
|
}
|
||||||
|
Some(Read::Unset) | None => ModelList::Local(local),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Refuse a `model` the operator hasn't made available to this agent.
|
/// Refuse a `model` the operator hasn't made available to this agent.
|
||||||
///
|
///
|
||||||
/// Default-open on purpose: this is a safety rail, not a security boundary,
|
/// Default-open on purpose: this is a safety rail, not a security boundary,
|
||||||
/// so an absent `allowed` (or one naming nothing) restricts nothing and an
|
/// so an absent list (or one naming nothing) restricts nothing and an
|
||||||
/// agent deployed before this gate existed keeps working unchanged. A `None`
|
/// agent deployed before this gate existed keeps working unchanged. A `None`
|
||||||
/// model isn't a choice of model either — it lets claude fall back to its own
|
/// model isn't a choice of model either — it lets claude fall back to its own
|
||||||
/// default — so there is nothing to check.
|
/// default — so there is nothing to check.
|
||||||
///
|
///
|
||||||
/// # Errors
|
/// # Errors
|
||||||
///
|
///
|
||||||
/// A named model absent from `allowed`; the message names it and every
|
/// A named model absent from `list`; the message names it, every permitted
|
||||||
/// permitted one, since it is all the caller sees.
|
/// one and where the list came from, since it is all the caller sees.
|
||||||
fn check_model(model: Option<&str>, allowed: Option<&str>) -> anyhow::Result<()> {
|
fn check_model(model: Option<&str>, list: &ModelList) -> anyhow::Result<()> {
|
||||||
let (Some(model), Some(allowed)) = (model, allowed) else {
|
let Some(model) = model else {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
};
|
};
|
||||||
// Same split/trim/skip-empty rules as `ToolGroup::parse_list`.
|
let (allowed, source): (Vec<&str>, &str) = match list {
|
||||||
let allowed: Vec<&str> = allowed
|
ModelList::Swarm(models) => (
|
||||||
.split(',')
|
models
|
||||||
.map(str::trim)
|
.iter()
|
||||||
.filter(|token| !token.is_empty())
|
.map(|m| m.trim())
|
||||||
.collect();
|
.filter(|m| !m.is_empty())
|
||||||
|
.collect(),
|
||||||
|
"the swarm's subagent_models for this agent",
|
||||||
|
),
|
||||||
|
ModelList::Local(None) => return Ok(()),
|
||||||
|
// Same split/trim/skip-empty rules as `ToolGroup::parse_list`.
|
||||||
|
ModelList::Local(Some(env)) => (
|
||||||
|
env.split(',')
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|token| !token.is_empty())
|
||||||
|
.collect(),
|
||||||
|
AVAILABLE_MODELS_ENV,
|
||||||
|
),
|
||||||
|
};
|
||||||
if allowed.is_empty()
|
if allowed.is_empty()
|
||||||
|| allowed
|
|| allowed
|
||||||
.iter()
|
.iter()
|
||||||
|
|
@ -948,7 +1014,7 @@ fn check_model(model: Option<&str>, allowed: Option<&str>) -> anyhow::Result<()>
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
anyhow::bail!(
|
anyhow::bail!(
|
||||||
"model `{model}` is not available to this agent — {AVAILABLE_MODELS_ENV} permits: {}",
|
"model `{model}` is not available to this agent — {source} permits: {}",
|
||||||
allowed.join(", ")
|
allowed.join(", ")
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
@ -1207,10 +1273,10 @@ pub struct StartRequest {
|
||||||
/// case is the only one that can happen *after* commit-to-run, and it's
|
/// case is the only one that can happen *after* commit-to-run, and it's
|
||||||
/// exactly why nothing is registered in `running` until spawn actually
|
/// exactly why nothing is registered in `running` until spawn actually
|
||||||
/// succeeds.
|
/// succeeds.
|
||||||
pub fn start(state: &Arc<State>, req: StartRequest) -> anyhow::Result<String> {
|
pub async fn start(state: &Arc<State>, req: StartRequest) -> anyhow::Result<String> {
|
||||||
let name = req.name.as_str();
|
let name = req.name.as_str();
|
||||||
validate_name(name)?;
|
validate_name(name)?;
|
||||||
check_model(req.model.as_deref(), available_models().as_deref())?;
|
state.allow_model(req.model.as_deref()).await?;
|
||||||
// Before `reserve`, so a role this agent does not have costs the caller
|
// Before `reserve`, so a role this agent does not have costs the caller
|
||||||
// an error and nothing else: no name claimed, no session archived, no
|
// an error and nothing else: no name claimed, no session archived, no
|
||||||
// process spawned.
|
// process spawned.
|
||||||
|
|
@ -1388,7 +1454,7 @@ pub async fn continue_(
|
||||||
dir: Option<&str>,
|
dir: Option<&str>,
|
||||||
) -> anyhow::Result<String> {
|
) -> anyhow::Result<String> {
|
||||||
validate_name(name)?;
|
validate_name(name)?;
|
||||||
check_model(model.as_deref(), available_models().as_deref())?;
|
state.allow_model(model.as_deref()).await?;
|
||||||
if !state.reserve(name) {
|
if !state.reserve(name) {
|
||||||
anyhow::bail!(
|
anyhow::bail!(
|
||||||
"subagent `{name}` is already running — use `interrupt` first if you meant to redirect it"
|
"subagent `{name}` is already running — use `interrupt` first if you meant to redirect it"
|
||||||
|
|
@ -2626,15 +2692,17 @@ mod tests {
|
||||||
assert!(err.to_string().contains("task instructions"), "{err}");
|
assert!(err.to_string().contains("task instructions"), "{err}");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn a_start_naming_an_unknown_role_errors_and_reserves_nothing() {
|
async fn a_start_naming_an_unknown_role_errors_and_reserves_nothing() {
|
||||||
// The error path end to end: the refusal happens before anything is
|
// The error path end to end: the refusal happens before anything is
|
||||||
// claimed, so an unknown role costs a message and nothing else —
|
// claimed, so an unknown role costs a message and nothing else —
|
||||||
// no reservation to release, no archived prior session, no spawn.
|
// no reservation to release, no archived prior session, no spawn.
|
||||||
let state = Arc::new(State::new(PathBuf::from("/dev/null"), signal_url()));
|
let state = Arc::new(State::new(PathBuf::from("/dev/null"), signal_url()));
|
||||||
let mut req = start_request("unknown-role-run");
|
let mut req = start_request("unknown-role-run");
|
||||||
req.role = Some("no-such-role-ships-anywhere".to_owned());
|
req.role = Some("no-such-role-ships-anywhere".to_owned());
|
||||||
let err = start(&state, req).expect_err("a role this agent lacks must refuse the start");
|
let err = start(&state, req)
|
||||||
|
.await
|
||||||
|
.expect_err("a role this agent lacks must refuse the start");
|
||||||
assert!(
|
assert!(
|
||||||
err.to_string().contains("no-such-role-ships-anywhere"),
|
err.to_string().contains("no-such-role-ships-anywhere"),
|
||||||
"the refusal names the role that was asked for: {err}"
|
"the refusal names the role that was asked for: {err}"
|
||||||
|
|
@ -2646,12 +2714,14 @@ mod tests {
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn a_start_naming_an_unusable_role_refuses_before_touching_the_filesystem() {
|
async fn a_start_naming_an_unusable_role_refuses_before_touching_the_filesystem() {
|
||||||
let state = Arc::new(State::new(PathBuf::from("/dev/null"), signal_url()));
|
let state = Arc::new(State::new(PathBuf::from("/dev/null"), signal_url()));
|
||||||
let mut req = start_request("bad-role-name");
|
let mut req = start_request("bad-role-name");
|
||||||
req.role = Some("../../etc/passwd".to_owned());
|
req.role = Some("../../etc/passwd".to_owned());
|
||||||
let err = start(&state, req).expect_err("a role name that is not an identifier refuses");
|
let err = start(&state, req)
|
||||||
|
.await
|
||||||
|
.expect_err("a role name that is not an identifier refuses");
|
||||||
assert!(err.to_string().contains("invalid role name"), "{err}");
|
assert!(err.to_string().contains("invalid role name"), "{err}");
|
||||||
assert_eq!(state.occupancy("bad-role-name"), None);
|
assert_eq!(state.occupancy("bad-role-name"), None);
|
||||||
}
|
}
|
||||||
|
|
@ -2809,8 +2879,8 @@ mod tests {
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn start_does_not_corrupt_the_remembered_dir_when_the_name_is_already_running() {
|
async fn start_does_not_corrupt_the_remembered_dir_when_the_name_is_already_running() {
|
||||||
// The exact scenario argus caught in review: a rejected concurrent
|
// The exact scenario argus caught in review: a rejected concurrent
|
||||||
// `start` used to commit its `dir` before the reservation check
|
// `start` used to commit its `dir` before the reservation check
|
||||||
// refused it, corrupting what a later bare `continue` would resolve
|
// refused it, corrupting what a later bare `continue` would resolve
|
||||||
|
|
@ -2829,7 +2899,8 @@ mod tests {
|
||||||
dir: Some("/tmp/rejected".to_owned()),
|
dir: Some("/tmp/rejected".to_owned()),
|
||||||
..start_request("dup")
|
..start_request("dup")
|
||||||
},
|
},
|
||||||
);
|
)
|
||||||
|
.await;
|
||||||
assert!(
|
assert!(
|
||||||
result.is_err(),
|
result.is_err(),
|
||||||
"a name already reserved must be refused, not spawned"
|
"a name already reserved must be refused, not spawned"
|
||||||
|
|
@ -4000,8 +4071,8 @@ mod tests {
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn a_fresh_start_inherits_nothing_from_the_run_before_it() {
|
async fn a_fresh_start_inherits_nothing_from_the_run_before_it() {
|
||||||
// `start` archives the prior session precisely so this is a fresh
|
// `start` archives the prior session precisely so this is a fresh
|
||||||
// start; a leftover stop reason or report path would make the new run
|
// start; a leftover stop reason or report path would make the new run
|
||||||
// report the old one's ending, into the old one's file.
|
// report the old one's ending, into the old one's file.
|
||||||
|
|
@ -4009,7 +4080,7 @@ mod tests {
|
||||||
state.record_stop("n", StopReason::NeedHelp("old block".to_owned()));
|
state.record_stop("n", StopReason::NeedHelp("old block".to_owned()));
|
||||||
state.set_report_file("n", Some("/tmp/old-report.md"));
|
state.set_report_file("n", Some("/tmp/old-report.md"));
|
||||||
state.reserve("n");
|
state.reserve("n");
|
||||||
let refused = start(&state, start_request("n"));
|
let refused = start(&state, start_request("n")).await;
|
||||||
assert!(refused.is_err(), "the reserved name is refused as before");
|
assert!(refused.is_err(), "the reserved name is refused as before");
|
||||||
|
|
||||||
// The same clearing `start` does once it owns the name.
|
// The same clearing `start` does once it owns the name.
|
||||||
|
|
@ -4139,22 +4210,22 @@ mod tests {
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn an_unset_model_whitelist_allows_any_model() {
|
fn an_unset_model_whitelist_allows_any_model() {
|
||||||
assert!(check_model(Some("opus"), None).is_ok());
|
assert!(check_model(Some("opus"), &local(None)).is_ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn an_unset_model_whitelist_allows_an_omitted_model() {
|
fn an_unset_model_whitelist_allows_an_omitted_model() {
|
||||||
assert!(check_model(None, None).is_ok());
|
assert!(check_model(None, &local(None)).is_ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn a_listed_model_is_allowed() {
|
fn a_listed_model_is_allowed() {
|
||||||
assert!(check_model(Some("sonnet"), Some("haiku,sonnet,opus")).is_ok());
|
assert!(check_model(Some("sonnet"), &local(Some("haiku,sonnet,opus"))).is_ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn an_unlisted_model_is_rejected_naming_what_is_permitted() {
|
fn an_unlisted_model_is_rejected_naming_what_is_permitted() {
|
||||||
let err = check_model(Some("opus"), Some("haiku,sonnet"))
|
let err = check_model(Some("opus"), &local(Some("haiku,sonnet")))
|
||||||
.expect_err("a model outside the list must be refused")
|
.expect_err("a model outside the list must be refused")
|
||||||
.to_string();
|
.to_string();
|
||||||
assert!(
|
assert!(
|
||||||
|
|
@ -4167,7 +4238,7 @@ mod tests {
|
||||||
#[test]
|
#[test]
|
||||||
fn a_set_whitelist_still_allows_an_omitted_model() {
|
fn a_set_whitelist_still_allows_an_omitted_model() {
|
||||||
assert!(
|
assert!(
|
||||||
check_model(None, Some("haiku")).is_ok(),
|
check_model(None, &local(Some("haiku"))).is_ok(),
|
||||||
"omitting `model` lets claude pick its own default — it is not a \
|
"omitting `model` lets claude pick its own default — it is not a \
|
||||||
choice of model, so there is nothing to gate"
|
choice of model, so there is nothing to gate"
|
||||||
);
|
);
|
||||||
|
|
@ -4175,14 +4246,111 @@ mod tests {
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn whitespace_and_empty_entries_are_ignored() {
|
fn whitespace_and_empty_entries_are_ignored() {
|
||||||
assert!(check_model(Some("sonnet"), Some(" haiku , sonnet ,, ")).is_ok());
|
assert!(check_model(Some("sonnet"), &local(Some(" haiku , sonnet ,, "))).is_ok());
|
||||||
assert!(check_model(Some("opus"), Some(" haiku , sonnet ,, ")).is_err());
|
assert!(check_model(Some("opus"), &local(Some(" haiku , sonnet ,, "))).is_err());
|
||||||
assert!(
|
assert!(
|
||||||
check_model(Some("opus"), Some(" ,, ")).is_ok(),
|
check_model(Some("opus"), &local(Some(" ,, "))).is_ok(),
|
||||||
"a var naming nothing restricts nothing — default open"
|
"a var naming nothing restricts nothing — default open"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn local(env: Option<&str>) -> ModelList {
|
||||||
|
ModelList::Local(env.map(str::to_owned))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn swarm(models: &[&str]) -> crate::swarm_models::Read {
|
||||||
|
crate::swarm_models::Read::Set(models.iter().map(|m| (*m).to_owned()).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The swarm's list replaces the local one, so it can widen it as well
|
||||||
|
/// as narrow it.
|
||||||
|
#[test]
|
||||||
|
fn a_swarm_list_replaces_the_local_one() {
|
||||||
|
let env = || Some("haiku,sonnet".to_owned());
|
||||||
|
let list = model_list(Some(swarm(&["opus"])), env());
|
||||||
|
assert_eq!(list, ModelList::Swarm(vec!["opus".to_owned()]));
|
||||||
|
assert!(check_model(Some("opus"), &list).is_ok());
|
||||||
|
assert!(check_model(Some("sonnet"), &list).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn without_a_swarm_list_the_local_one_applies() {
|
||||||
|
let env = || Some("haiku,sonnet".to_owned());
|
||||||
|
for read in [None, Some(crate::swarm_models::Read::Unset)] {
|
||||||
|
assert_eq!(model_list(read, env()), local(Some("haiku,sonnet")));
|
||||||
|
}
|
||||||
|
assert_eq!(
|
||||||
|
model_list(Some(crate::swarm_models::Read::Unset), None),
|
||||||
|
local(None)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_failed_read_falls_back_to_the_local_list_and_warns() {
|
||||||
|
let logs = CapturedLogs::default();
|
||||||
|
let subscriber = tracing_subscriber::fmt()
|
||||||
|
.with_writer(logs.clone())
|
||||||
|
.with_ansi(false)
|
||||||
|
.finish();
|
||||||
|
let list = tracing::subscriber::with_default(subscriber, || {
|
||||||
|
model_list(
|
||||||
|
Some(crate::swarm_models::Read::Failed("no answer".to_owned())),
|
||||||
|
Some("haiku".to_owned()),
|
||||||
|
)
|
||||||
|
});
|
||||||
|
assert_eq!(list, local(Some("haiku")));
|
||||||
|
let logged = logs.text();
|
||||||
|
assert!(
|
||||||
|
logged.contains("WARN") && logged.contains("no answer"),
|
||||||
|
"the fallback is logged as a warning with the cause: {logged}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn a_model_outside_the_swarm_list_is_refused_naming_the_list_and_its_source() {
|
||||||
|
let list = model_list(Some(swarm(&["haiku", "sonnet"])), Some("opus".to_owned()));
|
||||||
|
let err = check_model(Some("opus"), &list)
|
||||||
|
.expect_err("a model outside the swarm's list must be refused")
|
||||||
|
.to_string();
|
||||||
|
assert!(
|
||||||
|
err.contains("opus")
|
||||||
|
&& err.contains("haiku, sonnet")
|
||||||
|
&& err.contains("swarm's subagent_models"),
|
||||||
|
"{err}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
check_model(None, &list).is_ok(),
|
||||||
|
"an omitted model is not checked against the swarm's list either"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A `tracing` writer whose output a test can read back.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct CapturedLogs(Arc<Mutex<Vec<u8>>>);
|
||||||
|
|
||||||
|
impl CapturedLogs {
|
||||||
|
fn text(&self) -> String {
|
||||||
|
String::from_utf8_lossy(&self.0.lock().unwrap()).into_owned()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::io::Write for CapturedLogs {
|
||||||
|
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||||
|
self.0.lock().unwrap().extend_from_slice(buf);
|
||||||
|
Ok(buf.len())
|
||||||
|
}
|
||||||
|
fn flush(&mut self) -> std::io::Result<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturedLogs {
|
||||||
|
type Writer = Self;
|
||||||
|
fn make_writer(&'a self) -> Self::Writer {
|
||||||
|
self.clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// An ACP agent in plain `sh`, after `hive-runtime`'s own test agent. It
|
/// An ACP agent in plain `sh`, after `hive-runtime`'s own test agent. It
|
||||||
/// has one session, `s1`, appends every method it is sent to
|
/// has one session, `s1`, appends every method it is sent to
|
||||||
/// `<log>.methods` and every `session/prompt` line to `<log>`, and
|
/// `<log>.methods` and every `session/prompt` line to `<log>`, and
|
||||||
|
|
|
||||||
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(_)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -149,9 +149,9 @@ impl Publisher {
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Where to publish, and what to connect with.
|
/// Where to publish, and what to connect with.
|
||||||
struct Target {
|
pub(crate) struct Target {
|
||||||
url: String,
|
pub(crate) url: String,
|
||||||
agent: String,
|
pub(crate) agent: String,
|
||||||
store: StoreSource,
|
store: StoreSource,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -168,7 +168,9 @@ impl Target {
|
||||||
///
|
///
|
||||||
/// # Errors
|
/// # Errors
|
||||||
/// A store address without the agent name or the identity beside it.
|
/// A store address without the agent name or the identity beside it.
|
||||||
fn from_lookup(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<Option<Self>> {
|
pub(crate) fn from_lookup(
|
||||||
|
get: impl Fn(&str) -> Option<String>,
|
||||||
|
) -> anyhow::Result<Option<Self>> {
|
||||||
let Some(url) = get(ENV_NATS_URL).filter(|v| !v.is_empty()) else {
|
let Some(url) = get(ENV_NATS_URL).filter(|v| !v.is_empty()) else {
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
};
|
};
|
||||||
|
|
@ -225,7 +227,7 @@ impl StoreSource {
|
||||||
|
|
||||||
/// Connect presenting the agent's own token, read from the store on every
|
/// Connect presenting the agent's own token, read from the store on every
|
||||||
/// connection attempt so a re-minted secret is picked up on reconnect.
|
/// connection attempt so a re-minted secret is picked up on reconnect.
|
||||||
async fn connect(target: &Arc<Target>) -> anyhow::Result<async_nats::Client> {
|
pub(crate) async fn connect(target: &Arc<Target>) -> anyhow::Result<async_nats::Client> {
|
||||||
let url = target.url.clone();
|
let url = target.url.clone();
|
||||||
let target = Arc::clone(target);
|
let target = Arc::clone(target);
|
||||||
let agent = target.agent.clone();
|
let agent = target.agent.clone();
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue