hyperhive/swarm-controller/src/inference.rs

//! `/api/agents/{name}/inference` — the model, effort and subagent models the
//! swarm sets for one agent.
//!
//! 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.

use std::collections::BTreeSet;

use axum::Json;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use problem_details::ProblemDetails;
use serde::{Deserialize, Serialize};
use swarm_queue_client::agent_inference::{self, AgentInference, Effort};
use utoipa::ToSchema;

use crate::schedules::{live, queue, read, 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.
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"));
    }
    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,
        })
    }
}

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 agent's settings; `{}` when none is set",
            body = InferenceSettings),
        (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<InferenceSettings>, ProblemDetails> {
    let agent = agent_name(&name).map_err(refused)?;
    let client = queue(&state).map_err(refused)?;
    let store = open(&client).await?;
    let (current, _) = read::<AgentInference>(&store, &agent)
        .await
        .map_err(|e| queue_problem(&e))?;
    Ok(Json(current.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, 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) = read::<AgentInference>(&store, &agent)
            .await
            .map_err(|e| queue_problem(&e))?;
        if current == value
            || 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());
    }

    #[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"]}"#,
        ] {
            assert!(settings(json).unwrap().into_stored().is_err(), "{json}");
        }
    }
}