hyperhive/swarm-queue-client/src/agent_inference.rs

//! The per-agent inference-settings KV bucket: one key per agent, holding the
//! model, effort and subagent model list 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 [`crate::agent_schedules`] gives: the direct get names the key
//! in its subject ([`crate::agent_inference::direct_get_subject`]), which is
//! what lets the queue grant scope the read to one agent.
//!
//! Keyed by the agent alone, like [`crate::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>>,
}

/// 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",
        }
    }
}

/// Open the inference bucket for writing, creating it if nothing has yet.
///
/// **Controller-side only.** `history: 1`: an agent wants its current value.
#[cfg(feature = "kv")]
pub async fn open_or_create(
    client: &async_nats::Client,
) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    crate::kv::open_or_create(
        &js,
        BUCKET,
        "agent-inference",
        "Inference settings the swarm controller sets for each agent".to_owned(),
    )
    .await
}

/// What reading one agent's key found.
#[cfg(feature = "agent-inference")]
#[derive(Debug)]
pub enum Read {
    /// The bucket answered; an absent or deleted key has every field unset.
    Current(AgentInference),
    /// Nothing answers for the bucket's stream: no setting has been written
    /// yet, or `JetStream` is not serving it right now. Not an unset value,
    /// since the second case says nothing about the agent's settings.
    NoBucket,
}

/// Read `agent`'s current settings with the one subject its grant covers.
///
/// # Errors
///
/// A failed request or an undecodable value. A request the queue refused for
/// permissions reaches the client as a timeout, not as a refusal.
#[cfg(feature = "agent-inference")]
pub async fn read_own(
    client: &async_nats::Client,
    agent: &str,
) -> Result<Read, Box<dyn std::error::Error + Send + Sync>> {
    use async_nats::jetstream::stream::DirectGetErrorKind;

    let js = async_nats::jetstream::new(client.clone());
    let stream = js.get_stream_no_info(STREAM).await?;
    let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
        Ok(message) => message,
        Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => {
            return Ok(Read::Current(AgentInference::default()));
        }
        Err(e) if is_no_responders(&e) => return Ok(Read::NoBucket),
        Err(e) => return Err(e.into()),
    };
    // A KV delete or purge is stored as an empty marker carrying this header.
    let removed = message
        .headers
        .get("KV-Operation")
        .is_some_and(|op| matches!(op.as_str(), "DEL" | "PURGE"));
    if removed {
        return Ok(Read::Current(AgentInference::default()));
    }
    Ok(Read::Current(serde_json::from_slice(&message.payload)?))
}

#[cfg(feature = "agent-inference")]
fn is_no_responders(e: &async_nats::jetstream::stream::DirectGetError) -> bool {
    matches!(
        e.kind(),
        async_nats::jetstream::stream::DirectGetErrorKind::ErrorResponse(status, _)
            if status == async_nats::StatusCode::NO_RESPONDERS
    )
}

#[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()]),
        };
        let encoded = serde_json::to_string(&value).unwrap();
        assert_eq!(
            encoded,
            r#"{"model":"opus","effort":"xhigh","subagent_models":["haiku","sonnet"]}"#
        );
        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()
        );
    }

    #[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"));
    }
}