hyperhive/swarm-wire/src/agent_inference.rs

//! The per-agent inference-settings KV bucket: one key per agent, holding the
//! model, effort, subagent model list and subagent cap the swarm sets for it.
//!
//! **swarm-controller is the only writer**, and creates the bucket on its
//! first write. Every field is optional: an absent field leaves the agent on
//! its own local value, so an agent with no key runs exactly as it would with
//! no swarm.
//!
//! **An agent reads with one `DIRECT.GET` on its own key, not a watch**, for
//! the reason `swarm_queue_client::agent_schedules` gives: the direct get names the key
//! in its subject ([`direct_get_subject`]), which is
//! what lets the queue grant scope the read to one agent.
//!
//! Keyed by the agent alone, like `swarm_queue_client::agent_schedules`: an agent can
//! move between hives.

/// The KV bucket inference settings are published into, keyed by agent name.
pub const BUCKET: &str = "agent-inference";

/// The stream backing [`BUCKET`]. `create_key_value` prefixes `KV_`.
pub const STREAM: &str = "KV_agent-inference";

/// The subject `agent`'s entry is stored under.
#[must_use]
pub fn key_subject(agent: &str) -> String {
    format!("$KV.{BUCKET}.{agent}")
}

/// The request subject that reads `agent`'s current entry, and so the one
/// subject the queue grants that agent in this bucket.
#[must_use]
pub fn direct_get_subject(agent: &str) -> String {
    format!("$JS.API.DIRECT.GET.{STREAM}.{}", key_subject(agent))
}

/// One agent's inference settings: the value under `BUCKET`/`<agent>`.
/// `None` is "not set at swarm level", never "use a default".
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AgentInference {
    /// The model the agent's turns run on.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    /// The reasoning effort the agent's turns run at.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub effort: Option<Effort>,
    /// The models the agent may start its subagents on.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub subagent_models: Option<Vec<String>>,
    /// How many subagents the agent may have running at once. `None` is no
    /// cap.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_subagents: Option<u32>,
}

/// A claude `--effort` level, ascending.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Effort {
    Low,
    Medium,
    High,
    Xhigh,
    Max,
}

impl Effort {
    /// The level as claude's `--effort` takes it, and as it is stored.
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Low => "low",
            Self::Medium => "medium",
            Self::High => "high",
            Self::Xhigh => "xhigh",
            Self::Max => "max",
        }
    }
}

/// Where the model or effort an agent runs on comes from.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Source {
    /// The agent's key in [`BUCKET`].
    Swarm,
    /// `services.hyperhive.agent.model` or `effortLevel` in its agent.nix.
    AgentNix,
    /// A `hyperhive-model` or `hyperhive-effort` file in its harness dir, which
    /// the agent only reads.
    LocalFile,
    /// The harness's built-in default.
    Default,
}

/// The model and effort an agent runs on now, and where each comes from.
///
/// Part of the turn-state header the agent publishes on
/// `swarm_queue_client::agent_state_subject`; swarm-controller decodes these fields
/// out of that header.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Effective {
    /// The alias `claude --model` is invoked with (e.g. `"sonnet"`).
    pub model: String,
    /// The level `claude --effort` is invoked with.
    pub effort: String,
    pub model_source: Source,
    pub effort_source: Source,
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The names the grant in `swarm-nats.nix` spells with `{agent}`.
    #[test]
    fn names_match_the_queue_grant() {
        assert_eq!(STREAM, format!("KV_{BUCKET}"));
        assert_eq!(key_subject("iris"), "$KV.agent-inference.iris");
        assert_eq!(
            direct_get_subject("iris"),
            "$JS.API.DIRECT.GET.KV_agent-inference.$KV.agent-inference.iris"
        );
    }

    #[test]
    fn a_full_value_round_trips() {
        let value = AgentInference {
            model: Some("opus".to_owned()),
            effort: Some(Effort::Xhigh),
            subagent_models: Some(vec!["haiku".to_owned(), "sonnet".to_owned()]),
            max_subagents: Some(3),
        };
        let encoded = serde_json::to_string(&value).unwrap();
        assert_eq!(
            encoded,
            r#"{"model":"opus","effort":"xhigh","subagent_models":["haiku","sonnet"],"max_subagents":3}"#
        );
        assert_eq!(
            serde_json::from_str::<AgentInference>(&encoded).unwrap(),
            value
        );
    }

    /// An unset field is left out on the wire, and `{}` is every field unset.
    #[test]
    fn unset_fields_are_absent() {
        let value = AgentInference {
            effort: Some(Effort::Low),
            ..AgentInference::default()
        };
        assert_eq!(
            serde_json::to_string(&value).unwrap(),
            r#"{"effort":"low"}"#
        );
        assert_eq!(
            serde_json::from_str::<AgentInference>("{}").unwrap(),
            AgentInference::default()
        );
    }

    /// A value written before `max_subagents` existed decodes with no cap.
    #[test]
    fn a_value_without_max_subagents_has_no_cap() {
        let decoded: AgentInference =
            serde_json::from_str(r#"{"model":"opus","effort":"high","subagent_models":["haiku"]}"#)
                .unwrap();
        assert_eq!(decoded.max_subagents, None);
        assert_eq!(decoded.model.as_deref(), Some("opus"));
    }

    #[test]
    fn every_effort_is_stored_as_its_cli_spelling() {
        for effort in [
            Effort::Low,
            Effort::Medium,
            Effort::High,
            Effort::Xhigh,
            Effort::Max,
        ] {
            assert_eq!(
                serde_json::to_string(&effort).unwrap(),
                format!("\"{}\"", effort.as_str())
            );
        }
    }

    #[test]
    fn an_unknown_effort_is_refused() {
        for doc in [
            r#"{"effort":"extreme"}"#,
            r#"{"effort":"High"}"#,
            r#"{"effort":""}"#,
        ] {
            assert!(
                serde_json::from_str::<AgentInference>(doc).is_err(),
                "{doc}"
            );
        }
    }

    /// Fields a newer controller adds must not cost an older agent its value.
    #[test]
    fn decoding_tolerates_unknown_fields() {
        let decoded: AgentInference =
            serde_json::from_str(r#"{"model":"opus","temperature":1}"#).unwrap();
        assert_eq!(decoded.model.as_deref(), Some("opus"));
    }

    #[test]
    fn every_source_is_sent_snake_case() {
        for (source, wire) in [
            (Source::Swarm, "swarm"),
            (Source::AgentNix, "agent_nix"),
            (Source::LocalFile, "local_file"),
            (Source::Default, "default"),
        ] {
            assert_eq!(serde_json::to_value(source).unwrap(), wire);
        }
    }

    /// The fields are read out of a whole turn-state header; the rest of it
    /// is skipped.
    #[test]
    fn effective_decodes_out_of_a_header() {
        let decoded: Effective = serde_json::from_str(
            r#"{"turn_state":"idle","model":"opus","effort":"high",
                "model_source":"swarm","effort_source":"agent_nix"}"#,
        )
        .unwrap();
        assert_eq!(
            decoded,
            Effective {
                model: "opus".to_owned(),
                effort: "high".to_owned(),
                model_source: Source::Swarm,
                effort_source: Source::AgentNix,
            }
        );
    }
}