hyperhive/hive-agent/src/swarm_inference.rs

//! Applies the model and effort the swarm sets for this agent.
//!
//! swarm-controller keeps the agent's settings under its key in the
//! `agent-inference` KV bucket. This task reads that key every
//! [`POLL_INTERVAL`] with the one subject the agent's own queue credential is
//! granted there (`swarm_queue_client::agent_inference::read_own`).
//!
//! **A set field wins**: it replaces the Bus value, so the next turn runs on
//! it and the web UI's badge follows. It is not persisted: `hyperhive-model`
//! and `hyperhive-effort` keep the agent's local value.
//!
//! **An unset field changes nothing**, so an agent the swarm sets nothing for
//! runs on its local value ([`crate::harness_state::local_model`] and
//! [`crate::harness_state::local_effort`]). A field that was set and then
//! cleared goes back to that local value, read again at that moment.
//!
//! Each value is applied with its [`Source`], which the agent's turn-state
//! header (`crate::swarm_agent_state`) carries to the swarm.
//!
//! A failed read changes nothing. Until the first successful read after a
//! start, the agent runs on its local values.
//!
//! `subagent_models` and `max_subagents` are not read here.

use std::time::Duration;

use swarm_queue_client::agent_inference::{AgentInference, Effort, Read, Source, read_own};

use crate::events::Bus;
use crate::harness_state::{configured_effort, configured_model, local_effort, local_model};

/// How often the key is read. Also the longest a change can take to reach
/// the agent's next turn.
const POLL_INTERVAL: Duration = Duration::from_secs(10);

/// Where the settings are applied, and where the local values come from.
trait Target {
    fn model(&self) -> (String, Source);
    fn effort(&self) -> (String, Source);
    fn apply_model(&self, value: String, source: Source);
    fn apply_effort(&self, value: String, source: Source);
    fn local_model(&self) -> (String, Source);
    fn local_effort(&self) -> (String, Source);
}

/// The harness's Bus, with the nix values read once at start.
struct Harness {
    bus: Bus,
    nix_model: Option<&'static str>,
    nix_effort: Option<&'static str>,
}

impl Target for Harness {
    fn model(&self) -> (String, Source) {
        (self.bus.model(), self.bus.model_source())
    }
    fn effort(&self) -> (String, Source) {
        (self.bus.effort(), self.bus.effort_source())
    }
    fn apply_model(&self, value: String, source: Source) {
        self.bus.apply_model(value, source);
    }
    fn apply_effort(&self, value: String, source: Source) {
        self.bus.apply_effort(value, source);
    }
    fn local_model(&self) -> (String, Source) {
        local_model(self.nix_model)
    }
    fn local_effort(&self) -> (String, Source) {
        local_effort(self.nix_effort)
    }
}

/// Which fields currently carry the swarm's value.
#[derive(Debug, Default)]
struct Overridden {
    model: bool,
    effort: bool,
}

/// What one field and its source should become: the swarm's value while it
/// sets one, the local value on the read after it stops, and otherwise
/// `None` (left alone). Also `None` when both are already what they should
/// be; a swarm value equal to the current one still changes the source.
fn next(
    swarm: Option<&str>,
    overridden: &mut bool,
    current: &(String, Source),
    local: impl FnOnce() -> (String, Source),
) -> Option<(String, Source)> {
    let wanted = match swarm {
        Some(value) => {
            *overridden = true;
            (value.to_owned(), Source::Swarm)
        }
        None if std::mem::take(overridden) => local(),
        None => return None,
    };
    (wanted != *current).then_some(wanted)
}

impl Overridden {
    /// Apply one read of the key to `target`.
    fn apply(&mut self, read: &AgentInference, target: &impl Target) {
        if let Some((model, source)) = next(
            read.model.as_deref(),
            &mut self.model,
            &target.model(),
            || target.local_model(),
        ) {
            tracing::info!(%model, ?source, "inference: switching model");
            target.apply_model(model, source);
        }
        if let Some((effort, source)) = next(
            read.effort.map(Effort::as_str),
            &mut self.effort,
            &target.effort(),
            || target.local_effort(),
        ) {
            tracing::info!(%effort, ?source, "inference: switching effort");
            target.apply_effort(effort, source);
        }
    }
}

/// Run the poller for the harness's lifetime, if this agent has a label to
/// name its key with.
pub async fn run(bus: Bus) {
    let agent = crate::identity::label();
    if agent.is_empty() {
        tracing::warn!("this agent has no label; not reading its swarm inference settings");
        return;
    }
    let target = Harness {
        bus,
        nix_model: configured_model(),
        nix_effort: configured_effort(),
    };
    let client = crate::swarm_queue::client().await;
    tracing::info!(%agent, "reading swarm inference settings from the swarm queue");
    let mut overridden = Overridden::default();
    loop {
        if let Some(read) = read(&client, &agent).await {
            overridden.apply(&read, &target);
        }
        tokio::time::sleep(POLL_INTERVAL).await;
    }
}

/// The key's current value, or `None` when it could not be read.
async fn read(client: &async_nats::Client, agent: &str) -> Option<AgentInference> {
    // An unconnected client does not fail a JetStream request, it hangs on it.
    swarm_queue_client::ensure_connected(client).ok()?;
    match read_own(client, agent).await {
        Ok(Read::Current(current)) => Some(current),
        Ok(Read::NoBucket) => {
            tracing::debug!("inference: the agent-inference bucket does not answer yet");
            None
        }
        Err(e) => {
            tracing::warn!(error = %e, "inference: reading this agent's settings failed");
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use std::cell::RefCell;

    use super::*;

    /// A Bus stand-in recording every apply. Starts on its local values:
    /// the model from agent.nix, the effort from the file.
    struct Fake {
        model: RefCell<String>,
        effort: RefCell<String>,
        sources: RefCell<(Source, Source)>,
        local_model: RefCell<String>,
        local_effort: String,
        applies: RefCell<usize>,
    }

    impl Fake {
        fn new() -> Self {
            Self {
                model: RefCell::new("sonnet".to_owned()),
                effort: RefCell::new("medium".to_owned()),
                sources: RefCell::new((Source::AgentNix, Source::LocalFile)),
                local_model: RefCell::new("sonnet".to_owned()),
                local_effort: "medium".to_owned(),
                applies: RefCell::new(0),
            }
        }
        fn now(&self) -> (String, String) {
            (self.model.borrow().clone(), self.effort.borrow().clone())
        }
    }

    impl Target for Fake {
        fn model(&self) -> (String, Source) {
            (self.model.borrow().clone(), self.sources.borrow().0)
        }
        fn effort(&self) -> (String, Source) {
            (self.effort.borrow().clone(), self.sources.borrow().1)
        }
        fn apply_model(&self, value: String, source: Source) {
            *self.applies.borrow_mut() += 1;
            *self.model.borrow_mut() = value;
            self.sources.borrow_mut().0 = source;
        }
        fn apply_effort(&self, value: String, source: Source) {
            *self.applies.borrow_mut() += 1;
            *self.effort.borrow_mut() = value;
            self.sources.borrow_mut().1 = source;
        }
        fn local_model(&self) -> (String, Source) {
            (self.local_model.borrow().clone(), Source::AgentNix)
        }
        fn local_effort(&self) -> (String, Source) {
            (self.local_effort.clone(), Source::LocalFile)
        }
    }

    fn set(model: Option<&str>, effort: Option<Effort>) -> AgentInference {
        AgentInference {
            model: model.map(str::to_owned),
            effort,
            subagent_models: None,
            max_subagents: None,
        }
    }

    fn pair(model: &str, effort: &str) -> (String, String) {
        (model.to_owned(), effort.to_owned())
    }

    #[test]
    fn a_set_field_replaces_the_bus_value() {
        let fake = Fake::new();
        let mut o = Overridden::default();
        o.apply(&set(Some("opus"), Some(Effort::High)), &fake);
        assert_eq!(fake.now(), pair("opus", "high"));
        // Re-reading the same value applies nothing again.
        o.apply(&set(Some("opus"), Some(Effort::High)), &fake);
        assert_eq!(*fake.applies.borrow(), 2);
    }

    /// An absent key, and a key with every field unset, leave the agent on
    /// whatever it runs now — even when that differs from its local value.
    #[test]
    fn an_absent_key_leaves_the_bus_alone() {
        let fake = Fake::new();
        *fake.model.borrow_mut() = "haiku".to_owned();
        let mut o = Overridden::default();
        for _ in 0..3 {
            o.apply(&AgentInference::default(), &fake);
        }
        assert_eq!(fake.now(), pair("haiku", "medium"));
        assert_eq!(*fake.applies.borrow(), 0);
    }

    /// Only the field the swarm sets moves; the other keeps its value.
    #[test]
    fn an_unset_field_beside_a_set_one_is_left_alone() {
        let fake = Fake::new();
        *fake.effort.borrow_mut() = "max".to_owned();
        let mut o = Overridden::default();
        o.apply(&set(Some("opus"), None), &fake);
        assert_eq!(fake.now(), pair("opus", "max"));
    }

    #[test]
    fn a_cleared_field_goes_back_to_the_local_value() {
        let fake = Fake::new();
        let mut o = Overridden::default();
        o.apply(&set(Some("opus"), Some(Effort::Xhigh)), &fake);
        assert_eq!(fake.now(), pair("opus", "xhigh"));
        // The effort is cleared, the model still set.
        o.apply(&set(Some("opus"), None), &fake);
        assert_eq!(fake.now(), pair("opus", "medium"));
        // The local value is read at the clear, not at start.
        "haiku".clone_into(&mut fake.local_model.borrow_mut());
        o.apply(&AgentInference::default(), &fake);
        assert_eq!(fake.now(), pair("haiku", "medium"));
        // Once reverted, a later unset read changes nothing.
        *fake.model.borrow_mut() = "sonnet".to_owned();
        o.apply(&AgentInference::default(), &fake);
        assert_eq!(fake.now(), pair("sonnet", "medium"));
    }

    /// A swarm value equal to what the agent runs is still applied once, so
    /// the source moves to the swarm and the header is republished.
    #[test]
    fn a_swarm_value_equal_to_the_current_one_is_applied_once() {
        let fake = Fake::new();
        let mut o = Overridden::default();
        o.apply(&set(Some("sonnet"), Some(Effort::Medium)), &fake);
        assert_eq!(fake.now(), pair("sonnet", "medium"));
        assert_eq!(*fake.applies.borrow(), 2);
        o.apply(&set(Some("sonnet"), Some(Effort::Medium)), &fake);
        assert_eq!(*fake.applies.borrow(), 2);
    }

    #[test]
    fn the_source_follows_the_swarm_value() {
        let fake = Fake::new();
        let mut o = Overridden::default();
        o.apply(&set(Some("opus"), Some(Effort::High)), &fake);
        assert_eq!(*fake.sources.borrow(), (Source::Swarm, Source::Swarm));
        o.apply(&set(None, Some(Effort::High)), &fake);
        assert_eq!(*fake.sources.borrow(), (Source::AgentNix, Source::Swarm));
        o.apply(&AgentInference::default(), &fake);
        assert_eq!(
            *fake.sources.borrow(),
            (Source::AgentNix, Source::LocalFile)
        );
    }

    /// On the real Bus: the swarm's value reaches the Bus but not the
    /// persisted files, so clearing it, or a restart, finds the local value.
    #[test]
    fn the_swarm_value_is_not_persisted_over_the_local_one() {
        let (swarm, swarm_sources, files, cleared, cleared_sources) =
            crate::paths::with_harness_dir(|| {
                let bus = Bus::new();
                let dir = crate::paths::harness_dir();
                std::fs::write(dir.join("hyperhive-model"), "sonnet\n").unwrap();
                std::fs::write(dir.join("hyperhive-effort"), "low\n").unwrap();
                let target = Harness {
                    bus: bus.clone(),
                    nix_model: None,
                    nix_effort: None,
                };
                let mut o = Overridden::default();
                o.apply(&set(Some("opus"), Some(Effort::Max)), &target);
                let swarm = (bus.model(), bus.effort());
                let swarm_sources = (bus.model_source(), bus.effort_source());
                let files = (
                    crate::harness_state::load_model(),
                    crate::harness_state::load_effort(),
                );
                o.apply(&AgentInference::default(), &target);
                let cleared = (bus.model(), bus.effort());
                let cleared_sources = (bus.model_source(), bus.effort_source());
                (swarm, swarm_sources, files, cleared, cleared_sources)
            });
        assert_eq!(swarm, pair("opus", "max"));
        assert_eq!(files, (Some("sonnet".to_owned()), Some("low".to_owned())));
        assert_eq!(cleared, pair("sonnet", "low"));
        assert_eq!(swarm_sources, (Source::Swarm, Source::Swarm));
        assert_eq!(cleared_sources, (Source::LocalFile, Source::LocalFile));
    }
}