After a rollback, a record a newer build wrote (a new enum variant, a changed required field) stopped decoding, and the strict `schedules::read` turned that into an error on every read and write of the key until someone deleted it by hand. - agent_limits / inference (one document per agent): an undecodable document reads as unset, with a WARN naming the bucket, key and error. The PUT's "equal to stored, skip the write" shortcut no longer applies to an undecodable document, so a PUT of the defaults replaces it too. - schedule definitions (one document, every schedule): salvaged per definition, as the review-classifier run list is. An undecodable definition is kept as its raw JSON, hidden from the list, edits and projections, written back by every commit after the readable ones, and its id is never minted again. A document with no list of schedules stays an error: reading it as empty would empty every agent's key. - per-agent agent-schedules: an undecodable key reads as no schedules and is rewritten on the next projection, even when the agent has none. `schedules::read` is replaced by `read_or_default`; `wanted` keeps its own strict read. Closes #5193
541 lines
19 KiB
Rust
541 lines
19 KiB
Rust
//! `/api/agents/{name}/inference` — the model, effort and subagent models the
|
|
//! swarm sets for one agent, and what the agent reports running on.
|
|
//!
|
|
//! Each agent's settings are one value under its own key of `agent-inference`
|
|
//! (`swarm_queue_client::agent_inference`), which this daemon alone writes.
|
|
//! The agent reads its key and applies each set field over its local value;
|
|
//! an unset field leaves the local value in force.
|
|
//!
|
|
//! What the agent runs on comes from its turn-state header
|
|
//! (`$SWARM.agent-state.<agent>`), which [`Reported`] keeps the latest of.
|
|
//! Agents publish a header when it changes and re-send it once a minute, and
|
|
//! the queue keeps none, so an agent has no entry for up to a minute after
|
|
//! this daemon starts.
|
|
|
|
use std::collections::{BTreeSet, HashMap};
|
|
use std::sync::{Arc, Mutex, PoisonError};
|
|
use std::time::Duration;
|
|
|
|
use axum::Json;
|
|
use axum::extract::{Path, State};
|
|
use axum::http::StatusCode;
|
|
use futures_util::StreamExt as _;
|
|
use problem_details::ProblemDetails;
|
|
use serde::{Deserialize, Serialize};
|
|
use swarm_queue_client::agent_inference::{self, AgentInference, Effective, Effort, Source};
|
|
use utoipa::ToSchema;
|
|
|
|
use crate::schedules::{live, queue, read_or_default, refused, write};
|
|
use crate::{AppState, error_problem, wanted_error_status};
|
|
|
|
/// How many times a losing writer re-reads and re-applies, as in
|
|
/// `wanted::MAX_ATTEMPTS`.
|
|
const MAX_ATTEMPTS: usize = 5;
|
|
|
|
/// One agent's inference settings. An absent field is not set at swarm level,
|
|
/// and the agent keeps its own value for it.
|
|
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
|
|
#[serde(deny_unknown_fields)]
|
|
pub struct InferenceSettings {
|
|
/// The model the agent's turns run on, as claude's `--model` takes it.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
model: Option<String>,
|
|
/// One of `low`, `medium`, `high`, `xhigh`, `max`.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
#[schema(value_type = Option<String>)]
|
|
effort: Option<Effort>,
|
|
/// The models the agent may start its subagents on. At least one.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
subagent_models: Option<Vec<String>>,
|
|
}
|
|
|
|
impl From<AgentInference> for InferenceSettings {
|
|
fn from(v: AgentInference) -> Self {
|
|
Self {
|
|
model: v.model,
|
|
effort: v.effort,
|
|
subagent_models: v.subagent_models,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A model name as the agent passes it to claude: one non-empty word that
|
|
/// cannot be read as a command-line flag.
|
|
pub(crate) fn check_model(name: &str, what: &str) -> Result<(), String> {
|
|
if name.is_empty() {
|
|
return Err(format!("the {what} is empty"));
|
|
}
|
|
if name.chars().any(|c| c.is_whitespace() || c.is_control()) {
|
|
return Err(format!("the {what} {name:?} contains whitespace"));
|
|
}
|
|
if name.starts_with('-') {
|
|
return Err(format!("the {what} {name:?} starts with \"-\""));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
impl InferenceSettings {
|
|
/// The value to store, or why it is refused.
|
|
fn into_stored(self) -> Result<AgentInference, String> {
|
|
if let Some(model) = &self.model {
|
|
check_model(model, "model")?;
|
|
}
|
|
if let Some(models) = &self.subagent_models {
|
|
if models.is_empty() {
|
|
return Err(
|
|
"subagent_models is empty; leave it out to keep the agent's own list"
|
|
.to_owned(),
|
|
);
|
|
}
|
|
let mut seen = BTreeSet::new();
|
|
for model in models {
|
|
check_model(model, "subagent model")?;
|
|
if !seen.insert(model) {
|
|
return Err(format!("subagent model {model:?} is listed twice"));
|
|
}
|
|
}
|
|
}
|
|
Ok(AgentInference {
|
|
model: self.model,
|
|
effort: self.effort,
|
|
subagent_models: self.subagent_models,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// What `GET` answers for one agent.
|
|
#[derive(Debug, Serialize, ToSchema)]
|
|
pub struct AgentInferenceView {
|
|
/// The settings the swarm sets; `null` when it has none for the agent (no
|
|
/// key, or a cleared one), `{}` for a stored setting with no fields.
|
|
configured: Option<InferenceSettings>,
|
|
/// What the agent reported running on in its latest turn-state header;
|
|
/// `null` when this daemon has received none with these fields since it
|
|
/// started.
|
|
effective: Option<EffectiveInference>,
|
|
}
|
|
|
|
/// The model and effort an agent runs on, and where each comes from.
|
|
#[derive(Debug, PartialEq, Eq, Serialize, ToSchema)]
|
|
pub struct EffectiveInference {
|
|
model: String,
|
|
effort: String,
|
|
/// `swarm`, `agent_nix`, `local_file` or `default`.
|
|
#[schema(value_type = String)]
|
|
model_source: Source,
|
|
/// As `model_source`.
|
|
#[schema(value_type = String)]
|
|
effort_source: Source,
|
|
}
|
|
|
|
impl From<Effective> for EffectiveInference {
|
|
fn from(v: Effective) -> Self {
|
|
Self {
|
|
model: v.model,
|
|
effort: v.effort,
|
|
model_source: v.model_source,
|
|
effort_source: v.effort_source,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The latest [`Effective`] each agent reported, by agent.
|
|
#[derive(Debug, Default)]
|
|
pub(crate) struct Reported {
|
|
latest: Mutex<HashMap<String, Effective>>,
|
|
}
|
|
|
|
impl Reported {
|
|
/// Start recording every agent's turn-state headers from `client`.
|
|
pub(crate) fn spawn(client: async_nats::Client) -> Arc<Self> {
|
|
let reported = Arc::new(Self::default());
|
|
let recording = Arc::clone(&reported);
|
|
tokio::spawn(async move { recording.follow(&client).await });
|
|
reported
|
|
}
|
|
|
|
/// Subscribe to every agent's header, and subscribe again if the
|
|
/// subscription ends.
|
|
async fn follow(&self, client: &async_nats::Client) {
|
|
const RETRY: Duration = Duration::from_secs(5);
|
|
let subject = swarm_queue_client::agent_state_subject_wildcard();
|
|
loop {
|
|
match client.subscribe(subject.clone()).await {
|
|
Ok(mut headers) => {
|
|
tracing::info!(%subject, "inference: following agent turn-state headers");
|
|
while let Some(msg) = headers.next().await {
|
|
self.record(&msg.subject, &msg.payload);
|
|
}
|
|
tracing::warn!(%subject, "inference: header subscription ended");
|
|
}
|
|
Err(e) => tracing::warn!(%subject, error = %e, "inference: subscribe failed"),
|
|
}
|
|
tokio::time::sleep(RETRY).await;
|
|
}
|
|
}
|
|
|
|
/// Keep one header's inference fields. A header without them, from an
|
|
/// agent that does not report them, removes what was kept.
|
|
fn record(&self, subject: &str, payload: &[u8]) {
|
|
let Some(agent) = swarm_queue_client::agent_from_state_subject(subject) else {
|
|
return;
|
|
};
|
|
let mut latest = self.latest.lock().unwrap_or_else(PoisonError::into_inner);
|
|
match serde_json::from_slice::<Effective>(payload) {
|
|
Ok(effective) => {
|
|
latest.insert(agent.to_owned(), effective);
|
|
}
|
|
Err(e) => {
|
|
tracing::debug!(%agent, error = %e, "inference: header has no inference fields");
|
|
latest.remove(agent);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn get(&self, agent: &str) -> Option<Effective> {
|
|
self.latest
|
|
.lock()
|
|
.unwrap_or_else(PoisonError::into_inner)
|
|
.get(agent)
|
|
.cloned()
|
|
}
|
|
}
|
|
|
|
fn agent_name(name: &str) -> Result<String, (StatusCode, String)> {
|
|
hive_types::Ident::parse(name)
|
|
.map(hive_types::Ident::into_string)
|
|
.map_err(|reason| (StatusCode::BAD_REQUEST, reason.to_owned()))
|
|
}
|
|
|
|
fn queue_problem(e: &anyhow::Error) -> ProblemDetails {
|
|
tracing::warn!(error = %format!("{e:#}"), "inference: queue failure");
|
|
error_problem(wanted_error_status(e), &format!("{e:#}"))
|
|
}
|
|
|
|
async fn open(
|
|
client: &async_nats::Client,
|
|
) -> Result<async_nats::jetstream::kv::Store, ProblemDetails> {
|
|
agent_inference::open_or_create(client)
|
|
.await
|
|
.map_err(|e| queue_problem(&anyhow::Error::new(e)))
|
|
}
|
|
|
|
#[utoipa::path(
|
|
get,
|
|
path = "/api/agents/{name}/inference",
|
|
params(("name" = String, Path, description = "agent name")),
|
|
responses(
|
|
(status = 200, description = "the settings the swarm sets for the agent, and what \
|
|
it reported running on", body = AgentInferenceView),
|
|
(status = 400, description = "not an identifier (problem+json)", body = String),
|
|
(status = 503, description = "no swarm queue is configured on this host, or it is \
|
|
not connected (problem+json)", body = String),
|
|
(status = 500, description = "the settings could not be read (problem+json)",
|
|
body = String),
|
|
),
|
|
tag = "agents"
|
|
)]
|
|
pub(crate) async fn get_agent_inference(
|
|
State(state): State<AppState>,
|
|
Path(name): Path<String>,
|
|
) -> Result<Json<AgentInferenceView>, ProblemDetails> {
|
|
let agent = agent_name(&name).map_err(refused)?;
|
|
let client = queue(&state).map_err(refused)?;
|
|
let store = open(&client).await?;
|
|
// `Option`'s default is the `None` an absent, deleted or undecodable key
|
|
// reads as; a stored `{}` decodes to `Some`.
|
|
let (current, _, _) = read_or_default::<Option<AgentInference>>(&store, &agent)
|
|
.await
|
|
.map_err(|e| queue_problem(&e))?;
|
|
Ok(Json(AgentInferenceView {
|
|
configured: current.map(Into::into),
|
|
effective: state
|
|
.reported_inference
|
|
.as_ref()
|
|
.and_then(|r| r.get(&agent))
|
|
.map(Into::into),
|
|
}))
|
|
}
|
|
|
|
#[utoipa::path(
|
|
put,
|
|
path = "/api/agents/{name}/inference",
|
|
params(("name" = String, Path, description = "agent name")),
|
|
request_body = InferenceSettings,
|
|
responses(
|
|
(status = 200, description = "the settings as stored; they replace the previous \
|
|
ones whole", body = InferenceSettings),
|
|
(status = 400, description = "not an identifier, not declared on any hive, or a \
|
|
model name is empty, has whitespace, starts with `-`, or is listed twice \
|
|
(problem+json)",
|
|
body = String),
|
|
(status = 422, description = "an unknown field, or an effort that is not one of \
|
|
the five levels"),
|
|
(status = 503, description = "no swarm queue is configured on this host, or it is \
|
|
not connected (problem+json)", body = String),
|
|
(status = 500, description = "the settings could not be written (problem+json)",
|
|
body = String),
|
|
),
|
|
tag = "agents"
|
|
)]
|
|
pub(crate) async fn put_agent_inference(
|
|
State(state): State<AppState>,
|
|
Path(name): Path<String>,
|
|
Json(req): Json<InferenceSettings>,
|
|
) -> Result<Json<InferenceSettings>, ProblemDetails> {
|
|
let agent = agent_name(&name).map_err(refused)?;
|
|
let value = req
|
|
.into_stored()
|
|
.map_err(|why| error_problem(StatusCode::BAD_REQUEST, &why))?;
|
|
let client = queue(&state).map_err(refused)?;
|
|
if !live(&state).await?.contains(&agent) {
|
|
return Err(error_problem(
|
|
StatusCode::BAD_REQUEST,
|
|
&format!("no hive's wanted state declares an agent {agent:?} that is not destroyed"),
|
|
));
|
|
}
|
|
let store = open(&client).await?;
|
|
for _ in 0..MAX_ATTEMPTS {
|
|
let (current, revision, undecodable) = read_or_default::<AgentInference>(&store, &agent)
|
|
.await
|
|
.map_err(|e| queue_problem(&e))?;
|
|
if (current == value && !undecodable)
|
|
|| write(&store, &agent, &value, revision)
|
|
.await
|
|
.map_err(|e| queue_problem(&e))?
|
|
{
|
|
tracing::info!(%agent, ?value, "set inference settings");
|
|
return Ok(Json(value.into()));
|
|
}
|
|
}
|
|
Err(queue_problem(&anyhow::anyhow!(
|
|
"gave up writing the inference settings of {agent} after {MAX_ATTEMPTS} conflicting \
|
|
writes"
|
|
)))
|
|
}
|
|
|
|
#[utoipa::path(
|
|
delete,
|
|
path = "/api/agents/{name}/inference",
|
|
params(("name" = String, Path, description = "agent name")),
|
|
responses(
|
|
(status = 204, description = "cleared; the agent goes back to its own values"),
|
|
(status = 400, description = "not an identifier (problem+json)", body = String),
|
|
(status = 503, description = "no swarm queue is configured on this host, or it is \
|
|
not connected (problem+json)", body = String),
|
|
(status = 500, description = "the settings could not be cleared (problem+json)",
|
|
body = String),
|
|
),
|
|
tag = "agents"
|
|
)]
|
|
pub(crate) async fn delete_agent_inference(
|
|
State(state): State<AppState>,
|
|
Path(name): Path<String>,
|
|
) -> Result<StatusCode, ProblemDetails> {
|
|
let agent = agent_name(&name).map_err(refused)?;
|
|
let client = queue(&state).map_err(refused)?;
|
|
let store = open(&client).await?;
|
|
store
|
|
.delete(&agent)
|
|
.await
|
|
.map_err(|e| queue_problem(&anyhow::Error::new(e)))?;
|
|
tracing::info!(%agent, "cleared inference settings");
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn settings(json: &str) -> serde_json::Result<InferenceSettings> {
|
|
serde_json::from_str(json)
|
|
}
|
|
|
|
#[test]
|
|
fn a_full_put_is_stored_as_sent() {
|
|
let stored = settings(r#"{"model":"opus","effort":"high","subagent_models":["haiku"]}"#)
|
|
.unwrap()
|
|
.into_stored()
|
|
.unwrap();
|
|
assert_eq!(
|
|
stored,
|
|
AgentInference {
|
|
model: Some("opus".to_owned()),
|
|
effort: Some(Effort::High),
|
|
subagent_models: Some(vec!["haiku".to_owned()]),
|
|
}
|
|
);
|
|
assert_eq!(
|
|
serde_json::to_value(InferenceSettings::from(stored)).unwrap(),
|
|
serde_json::json!({"model": "opus", "effort": "high", "subagent_models": ["haiku"]})
|
|
);
|
|
}
|
|
|
|
/// Each field is optional, and one left out stays out of the stored value.
|
|
#[test]
|
|
fn a_partial_put_leaves_the_other_fields_unset() {
|
|
let stored = settings(r#"{"effort":"max"}"#)
|
|
.unwrap()
|
|
.into_stored()
|
|
.unwrap();
|
|
assert_eq!(
|
|
stored,
|
|
AgentInference {
|
|
effort: Some(Effort::Max),
|
|
..AgentInference::default()
|
|
}
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn an_unknown_field_or_effort_does_not_decode() {
|
|
for json in [
|
|
r#"{"modle":"opus"}"#,
|
|
r#"{"available_models":["opus"]}"#,
|
|
r#"{"effort":"extreme"}"#,
|
|
r#"{"effort":"HIGH"}"#,
|
|
] {
|
|
assert!(settings(json).is_err(), "{json}");
|
|
}
|
|
// Control: the same shape with a known field and level decodes.
|
|
assert!(settings(r#"{"model":"opus","effort":"xhigh"}"#).is_ok());
|
|
}
|
|
|
|
const NEW_HEADER: &[u8] = br#"{"turn_state":"idle","model":"opus","effort":"high",
|
|
"model_source":"swarm","effort_source":"local_file"}"#;
|
|
|
|
/// A header from an agent that does not report effort or sources.
|
|
const OLD_HEADER: &[u8] = br#"{"turn_state":"idle","model":"sonnet"}"#;
|
|
|
|
fn iris() -> String {
|
|
swarm_queue_client::agent_state_subject("iris")
|
|
}
|
|
|
|
#[test]
|
|
fn a_header_is_kept_under_its_agent() {
|
|
let reported = Reported::default();
|
|
reported.record(&iris(), NEW_HEADER);
|
|
assert_eq!(
|
|
reported.get("iris"),
|
|
Some(Effective {
|
|
model: "opus".to_owned(),
|
|
effort: "high".to_owned(),
|
|
model_source: Source::Swarm,
|
|
effort_source: Source::LocalFile,
|
|
})
|
|
);
|
|
assert_eq!(reported.get("atlas"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn a_header_without_the_inference_fields_leaves_no_entry() {
|
|
let reported = Reported::default();
|
|
reported.record(&iris(), NEW_HEADER);
|
|
reported.record(&iris(), OLD_HEADER);
|
|
assert_eq!(reported.get("iris"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn a_subject_naming_no_agent_is_skipped() {
|
|
let reported = Reported::default();
|
|
reported.record("$SWARM.agent-state.iris.sub.scout", NEW_HEADER);
|
|
reported.record("$SWARM.term.iris", NEW_HEADER);
|
|
assert!(reported.latest.lock().unwrap().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn the_view_answers_both_halves() {
|
|
let view = |effective: Option<Effective>| {
|
|
serde_json::to_value(AgentInferenceView {
|
|
configured: Some(InferenceSettings::default()),
|
|
effective: effective.map(Into::into),
|
|
})
|
|
.unwrap()
|
|
};
|
|
assert_eq!(
|
|
view(None),
|
|
serde_json::json!({"configured": {}, "effective": null})
|
|
);
|
|
let unset = AgentInferenceView {
|
|
configured: None,
|
|
effective: None,
|
|
};
|
|
assert_eq!(
|
|
serde_json::to_value(unset).unwrap(),
|
|
serde_json::json!({"configured": null, "effective": null})
|
|
);
|
|
let reported = Reported::default();
|
|
reported.record(&iris(), NEW_HEADER);
|
|
assert_eq!(
|
|
view(reported.get("iris"))["effective"],
|
|
serde_json::json!({"model": "opus", "effort": "high",
|
|
"model_source": "swarm", "effort_source": "local_file"})
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn bad_model_names_are_refused() {
|
|
for json in [
|
|
r#"{"model":""}"#,
|
|
r#"{"model":"opus 4"}"#,
|
|
r#"{"subagent_models":[]}"#,
|
|
r#"{"subagent_models":["haiku",""]}"#,
|
|
r#"{"subagent_models":["haiku","haiku"]}"#,
|
|
r#"{"model":"--dangerously-skip-permissions"}"#,
|
|
r#"{"model":"-x"}"#,
|
|
r#"{"subagent_models":["haiku","--print"]}"#,
|
|
] {
|
|
assert!(settings(json).unwrap().into_stored().is_err(), "{json}");
|
|
}
|
|
}
|
|
|
|
/// Control for the leading-dash refusal: full model ids have dashes inside.
|
|
#[test]
|
|
fn a_dash_inside_a_model_name_is_accepted() {
|
|
let json = r#"{"model":"claude-opus-4-8","subagent_models":["claude-haiku-4-5"]}"#;
|
|
assert!(settings(json).unwrap().into_stored().is_ok());
|
|
}
|
|
|
|
/// What `get_agent_inference` reads with: an absent or deleted key is
|
|
/// `Option`'s default, `None`; a stored `{}` is a setting with no fields.
|
|
#[test]
|
|
fn no_key_and_an_empty_setting_read_apart() {
|
|
assert_eq!(Option::<AgentInference>::default(), None);
|
|
assert_eq!(
|
|
serde_json::from_slice::<Option<AgentInference>>(b"{}").unwrap(),
|
|
Some(AgentInference::default())
|
|
);
|
|
}
|
|
|
|
/// A stored value this build cannot decode, such as an effort a newer
|
|
/// build added, reads as no settings, and a `PUT`, even of no fields,
|
|
/// replaces it.
|
|
#[tokio::test]
|
|
async fn an_undecodable_value_reads_as_unset_until_a_put() {
|
|
let nats = crate::test_nats::TestNats::start();
|
|
let client = nats.client().await;
|
|
let state = crate::agent_limits::tests::declared(&client).await;
|
|
let store = agent_inference::open_or_create(&client)
|
|
.await
|
|
.expect("bucket");
|
|
store
|
|
.put("atlas", br#"{"effort":"later"}"#.to_vec().into())
|
|
.await
|
|
.expect("written");
|
|
|
|
let view = get_agent_inference(State(state.clone()), Path("atlas".to_owned()))
|
|
.await
|
|
.expect("read");
|
|
assert_eq!(view.0.configured, None);
|
|
let stored = put_agent_inference(
|
|
State(state),
|
|
Path("atlas".to_owned()),
|
|
Json(InferenceSettings::default()),
|
|
)
|
|
.await
|
|
.expect("set");
|
|
assert_eq!(stored.0, InferenceSettings::default());
|
|
let entry = store.get("atlas").await.expect("read").expect("a value");
|
|
assert_eq!(&entry[..], b"{}");
|
|
}
|
|
}
|