hyperhive/swarm-controller/src/inference.rs

//! `/api/agents/{name}/inference` — the model, effort, subagent models and
//! subagent cap 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_wire::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.
//!
//! [`Reported`] also keeps each agent's newest header as it arrived, so the
//! state stream (`crate::agent_state_stream`) can hand a client that attaches
//! mid-idle the last header at once, rather than leaving it waiting for the
//! next re-send.

use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, Instant};

use axum::Json;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use futures_util::StreamExt as _;
use problem_details::ProblemDetails;
use swarm_controller_api::inference::{
    AgentInferenceView, EffectiveInference, InferenceSettings, MAX_SUBAGENTS,
};
use swarm_queue_client::agent_inference;
use swarm_wire::agent_inference::{AgentInference, Effective};

use crate::routes::{AppState, error_problem, wanted_error_status};
use crate::schedules::{live, queue, read_or_default, refused, write_revision};

/// How many times a losing writer re-reads and re-applies, as in
/// `wanted::MAX_ATTEMPTS`.
const MAX_ATTEMPTS: usize = 5;

/// The stored settings as the API shows them.
fn settings_view(v: AgentInference) -> InferenceSettings {
    InferenceSettings {
        model: v.model,
        effort: v.effort,
        subagent_models: v.subagent_models,
        max_subagents: v.max_subagents,
    }
}

/// 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(())
}

/// The settings to store, or why they are refused.
pub(crate) fn into_stored(settings: InferenceSettings) -> Result<AgentInference, String> {
    if let Some(model) = &settings.model {
        check_model(model, "model")?;
    }
    if let Some(models) = &settings.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"));
            }
        }
    }
    if let Some(max) = settings.max_subagents
        && max > MAX_SUBAGENTS
    {
        return Err(format!(
            "max_subagents is {max}; it is from 0 (no subagents) to {MAX_SUBAGENTS}, or left out for no cap"
        ));
    }
    Ok(AgentInference {
        model: settings.model,
        effort: settings.effort,
        subagent_models: settings.subagent_models,
        max_subagents: settings.max_subagents,
    })
}

/// The agent's reported values as the API shows them.
fn effective_view(v: Effective) -> EffectiveInference {
    EffectiveInference {
        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>>,
    /// Each agent's newest header exactly as published, with when it came.
    raw: Mutex<HashMap<String, (Instant, Vec<u8>)>>,
}

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;
        }
    }

    /// `agent`'s newest header, as published, when it arrived within `max_age`.
    /// An older one describes an agent that may have stopped since.
    pub(crate) fn header_within(&self, agent: &str, max_age: Duration) -> Option<Vec<u8>> {
        let raw = self.raw.lock().unwrap_or_else(PoisonError::into_inner);
        let (at, payload) = raw.get(agent)?;
        (at.elapsed() <= max_age).then(|| payload.clone())
    }

    /// 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;
        };
        self.raw
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .insert(agent.to_owned(), (Instant::now(), payload.to_vec()));
        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(settings_view),
        effective: state
            .reported_inference
            .as_ref()
            .and_then(|r| r.get(&agent))
            .map(effective_view),
    }))
}

#[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, or `max_subagents` is above 32 (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 = into_stored(req).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"),
        ));
    }
    store(&client, &agent, &value).await?;
    Ok(Json(settings_view(value)))
}

/// What [`store`] replaced, or `None` when it wrote nothing because the
/// settings were already stored.
pub(crate) type Replaced = Option<crate::schedules::Replaced<AgentInference>>;

/// Replace `agent`'s stored inference settings with `value`.
pub(crate) async fn store(
    client: &async_nats::Client,
    agent: &str,
    value: &AgentInference,
) -> Result<Replaced, ProblemDetails> {
    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 {
            tracing::info!(%agent, ?value, "set inference settings");
            return Ok(None);
        }
        // A deleted key reads as present at its marker's revision, so ask
        // whether it holds a value before keeping one to put back.
        let held = store
            .get(agent)
            .await
            .map_err(|e| queue_problem(&anyhow::Error::new(e)))?
            .is_some();
        if let Some(written) = write_revision(&store, agent, value, revision)
            .await
            .map_err(|e| queue_problem(&e))?
        {
            tracing::info!(%agent, ?value, "set inference settings");
            return Ok(Some(crate::schedules::Replaced::new(
                held.then_some(current),
                written,
            )));
        }
    }
    Err(queue_problem(&anyhow::anyhow!(
        "gave up writing the inference settings of {agent} after {MAX_ATTEMPTS} conflicting \
         writes"
    )))
}

/// Undo what [`store`] did, unless another write has landed since. Best
/// effort: a failure is logged, since the caller is already on its way to an
/// error of its own.
pub(crate) async fn restore(client: &async_nats::Client, agent: &str, replaced: Replaced) {
    let Some(replaced) = replaced else { return };
    let undone = match open(client).await {
        Ok(store) => replaced.undo(&store, agent).await,
        Err(e) => Err(anyhow::anyhow!("{e:?}")),
    };
    match undone {
        Ok(true) => {}
        Ok(false) => {
            tracing::info!(%agent, "left the agent's inference settings: another write landed since");
        }
        Err(e) => {
            tracing::warn!(%agent, error = %format!("{e:#}"), "could not put back the agent's inference settings");
        }
    }
}

#[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::*;
    use swarm_controller_api::inference::{Effort, Source};

    fn settings(json: &str) -> serde_json::Result<InferenceSettings> {
        serde_json::from_str(json)
    }

    #[test]
    fn a_full_put_is_stored_as_sent() {
        let stored = into_stored(
            settings(r#"{"model":"opus","effort":"high","subagent_models":["haiku"]}"#).unwrap(),
        )
        .unwrap();
        assert_eq!(
            stored,
            AgentInference {
                model: Some("opus".to_owned()),
                effort: Some(Effort::High),
                subagent_models: Some(vec!["haiku".to_owned()]),
                max_subagents: None,
            }
        );
        assert_eq!(
            serde_json::to_value(settings_view(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 = into_stored(settings(r#"{"effort":"max"}"#).unwrap()).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 the_newest_header_is_kept_as_published_while_it_is_fresh() {
        let reported = Reported::default();
        reported.record(&iris(), OLD_HEADER);
        reported.record(&iris(), NEW_HEADER);
        assert_eq!(
            reported.header_within("iris", Duration::from_mins(1)),
            Some(NEW_HEADER.to_vec())
        );
        assert_eq!(
            reported.header_within("atlas", Duration::from_mins(1)),
            None
        );
        // A header with no inference fields still counts as the newest one.
        reported.record(&iris(), OLD_HEADER);
        assert_eq!(
            reported.header_within("iris", Duration::from_mins(1)),
            Some(OLD_HEADER.to_vec())
        );
    }

    #[test]
    fn a_header_older_than_the_limit_is_not_handed_out() {
        let reported = Reported::default();
        reported.record(&iris(), NEW_HEADER);
        std::thread::sleep(Duration::from_millis(20));
        assert_eq!(
            reported.header_within("iris", Duration::from_millis(1)),
            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(effective_view),
            })
            .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!(into_stored(settings(json).unwrap()).is_err(), "{json}");
        }
    }

    #[test]
    fn max_subagents_outside_its_range_is_refused() {
        assert!(into_stored(settings(r#"{"max_subagents":33}"#).unwrap()).is_err());
        assert!(settings(r#"{"max_subagents":-1}"#).is_err());
    }

    #[test]
    fn max_subagents_in_range_is_stored() {
        for max in [0, 1, 3, MAX_SUBAGENTS] {
            let stored =
                into_stored(settings(&format!(r#"{{"max_subagents":{max}}}"#)).unwrap()).unwrap();
            assert_eq!(stored.max_subagents, Some(max));
        }
    }

    /// 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!(into_stored(settings(json).unwrap()).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"{}");
    }

    /// `restore` puts back the settings a `store` replaced, or takes the key
    /// out again when there was none.
    #[tokio::test]
    async fn restore_undoes_a_store() {
        let nats = crate::test_nats::TestNats::start();
        let client = nats.client().await;
        let bucket = agent_inference::open_or_create(&client)
            .await
            .expect("bucket");
        let value = |json: &str| into_stored(settings(json).unwrap()).unwrap();

        let first = store(&client, "atlas", &value(r#"{"model":"opus"}"#))
            .await
            .expect("stored");
        restore(&client, "atlas", first).await;
        assert!(bucket.get("atlas").await.expect("read").is_none());

        store(&client, "atlas", &value(r#"{"model":"opus"}"#))
            .await
            .expect("stored");
        let second = store(&client, "atlas", &value(r#"{"model":"haiku"}"#))
            .await
            .expect("replaced");
        restore(&client, "atlas", second).await;
        let entry = bucket.get("atlas").await.expect("read").expect("a value");
        assert_eq!(&entry[..], br#"{"model":"opus"}"#);
    }

    /// A write that landed after the one being undone is newer, so `restore`
    /// leaves it.
    #[tokio::test]
    async fn restore_leaves_a_write_that_landed_since() {
        let nats = crate::test_nats::TestNats::start();
        let client = nats.client().await;
        let bucket = agent_inference::open_or_create(&client)
            .await
            .expect("bucket");
        let value = |json: &str| into_stored(settings(json).unwrap()).unwrap();

        let made = store(&client, "atlas", &value(r#"{"model":"opus"}"#))
            .await
            .expect("stored");
        store(&client, "atlas", &value(r#"{"model":"haiku"}"#))
            .await
            .expect("a PUT");
        restore(&client, "atlas", made).await;
        let entry = bucket.get("atlas").await.expect("read").expect("a value");
        assert_eq!(&entry[..], br#"{"model":"haiku"}"#);
    }
}