hyperhive/hive-agent/src/swarm_input.rs

//! Operator input from the swarm UI, taken straight off the swarm queue.
//!
//! swarm-controller sends each message on
//! [`swarm_queue_client::agent_input_subject`] for this agent. This task
//! subscribes to it and pushes each message into the serve loop's local
//! message channel, where it becomes a turn from the operator it names.
//! hive-c0re is not in the path, so the message has no broker row: the hive
//! dashboard's history and `recv` never see it, and `send` has no recipient
//! by the sender's name.
//!
//! **`from` is taken as given.** The controller is the only client granted
//! publish on the subject (`swarm-nats-auth`'s reader grant), and it names
//! the sender from the swarm gateway's authelia answer.
//!
//! **Own credential only.** The queue grants subscribe on this subject to the
//! agent's own credential alone.
//!
//! **Core NATS, no reply to input.** A message sent while this task is not
//! subscribed is gone; the controller sends it as a request so the operator
//! gets a 503 for that instead. Input is not answered: the controller reads a
//! request that times out as "a subscriber held it".
//!
//! **Commands are a separate subject and are answered.** The swarm UI's
//! cancel and `AGENTS.md` edits arrive on
//! [`swarm_queue_client::agent_control_subject`], never on the input subject,
//! so nothing a person types can be read as a command. A cancel's reply says
//! whether a turn was actually stopped, an edit's whether it was applied
//! ([`crate::swarm_agents_md`]); the agent's grant allows answering only the
//! request it was just sent (`allow_responses`).

use futures_util::{Stream, StreamExt as _};
use hive_sh4re::inbox::DeliveredMessage;
use tokio::sync::mpsc::UnboundedSender;

use crate::events::Bus;

/// Start the subscribers, if this agent has a label to name its subjects with.
pub fn spawn(
    tx: UnboundedSender<DeliveredMessage>,
    bus: Bus,
    interrupted: std::sync::Arc<std::sync::atomic::AtomicBool>,
) {
    let agent = crate::identity::label();
    if agent.is_empty() {
        tracing::warn!("this agent has no label; not taking input from the swarm queue");
        return;
    }
    tokio::spawn(run(tx, agent.clone()));
    tokio::spawn(run_control(agent, bus, interrupted));
}

/// Answer each control command for `agent` until the subscription ends.
async fn run_control(
    agent: String,
    bus: Bus,
    interrupted: std::sync::Arc<std::sync::atomic::AtomicBool>,
) {
    let client = crate::swarm_queue::client().await;
    let subject = swarm_queue_client::agent_control_subject(&agent);
    let mut subscriber = match client.subscribe(subject.clone()).await {
        Ok(subscriber) => subscriber,
        Err(e) => {
            tracing::warn!(subject, error = %e, "swarm control: subscribe failed");
            return;
        }
    };
    tracing::info!(subject, "taking control commands from the swarm queue");
    while let Some(msg) = subscriber.next().await {
        let Some(reply) = msg.reply else {
            // Sent without a reply subject: nobody to tell, and the
            // controller always asks for one.
            tracing::warn!("swarm control: command without a reply subject dropped");
            continue;
        };
        let Some(answer) = act(
            &msg.payload,
            || crate::web_ui::cancel_turn(&bus, &interrupted, "swarm UI: cancel"),
            crate::swarm_agents_md::handle,
        )
        .await
        else {
            continue;
        };
        if let Err(e) = client.publish(reply, answer.into()).await {
            tracing::warn!(error = %e, "swarm control: answering failed");
        }
    }
}

/// The encoded answer to a command payload: `cancel` runs for a cancel and
/// `apply` for an `AGENTS.md` edit. `None` (logged) when the payload is not a
/// command, which gets no answer.
async fn act<F>(
    payload: &[u8],
    cancel: impl FnOnce() -> bool,
    apply: impl FnOnce(swarm_queue_client::agents_md::Apply) -> F,
) -> Option<Vec<u8>>
where
    F: Future<Output = swarm_queue_client::agents_md::ApplyReply>,
{
    use swarm_queue_client::AgentControl;

    let encoded = match serde_json::from_slice::<AgentControl>(payload) {
        Ok(AgentControl::Cancel) => {
            serde_json::to_vec(&swarm_queue_client::AgentControlReply { acted: cancel() })
        }
        Ok(AgentControl::ApplyAgentsMd(edit)) => serde_json::to_vec(&apply(edit).await),
        Err(e) => {
            tracing::warn!(error = %e, "swarm control: undecodable command dropped");
            return None;
        }
    };
    Some(encoded.expect("a reply encodes"))
}

async fn run(tx: UnboundedSender<DeliveredMessage>, agent: String) {
    let client = crate::swarm_queue::client().await;
    let subject = swarm_queue_client::agent_input_subject(&agent);
    let subscriber = match client.subscribe(subject.clone()).await {
        Ok(subscriber) => subscriber,
        Err(e) => {
            tracing::warn!(subject, error = %e, "swarm input: subscribe failed");
            return;
        }
    };
    tracing::info!(subject, "taking operator input from the swarm queue");
    forward(subscriber.map(|msg| msg.payload), &tx).await;
}

/// Push each decodable payload into `tx` until either end closes.
async fn forward<P: AsRef<[u8]>>(
    payloads: impl Stream<Item = P>,
    tx: &UnboundedSender<DeliveredMessage>,
) {
    let mut payloads = std::pin::pin!(payloads);
    while let Some(payload) = payloads.next().await {
        let Some(dm) = delivered(payload.as_ref()) else {
            continue;
        };
        if tx.send(dm).is_err() {
            // The serve loop is gone: the harness is shutting down.
            return;
        }
    }
}

/// The message a payload becomes, or `None` (logged) for one that is not an
/// [`swarm_queue_client::AgentInput`].
fn delivered(payload: &[u8]) -> Option<DeliveredMessage> {
    match serde_json::from_slice::<swarm_queue_client::AgentInput>(payload) {
        Ok(input) => {
            tracing::info!(from = %input.from, "swarm input: message received");
            Some(DeliveredMessage {
                from: input.from,
                body: input.body,
                id: 0,
                redelivered: false,
                in_reply_to: None,
            })
        }
        Err(e) => {
            tracing::warn!(error = %e, "swarm input: undecodable message dropped");
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{act, forward};
    use swarm_queue_client::agents_md::{Apply, ApplyReply};
    use swarm_queue_client::{AgentControl, AgentControlReply, AgentInput};

    fn payload(from: &str, body: &str) -> Vec<u8> {
        serde_json::to_vec(&AgentInput {
            from: from.to_owned(),
            body: body.to_owned(),
        })
        .expect("encodes")
    }

    #[tokio::test]
    async fn each_message_reaches_the_serve_loop_as_from_its_sender() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        let payloads = vec![
            payload("mara", "first"),
            b"not json".to_vec(),
            payload("argus", "second"),
        ];
        forward(futures_util::stream::iter(payloads), &tx).await;
        drop(tx);

        let mut got = Vec::new();
        while let Some(dm) = rx.recv().await {
            got.push(dm);
        }
        // The undecodable one is dropped, the ones either side of it are not.
        assert_eq!(got.len(), 2, "{got:?}");
        assert_eq!(
            (got[0].from.as_str(), got[0].body.as_str()),
            ("mara", "first")
        );
        assert_eq!(
            (got[1].from.as_str(), got[1].body.as_str()),
            ("argus", "second")
        );
        // No broker row behind it, so no `[msg #id]` tag in the prompt.
        assert!(got.iter().all(|dm| dm.id == 0 && dm.in_reply_to.is_none()));
    }

    #[tokio::test]
    async fn a_closed_serve_loop_stops_the_forwarder() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        drop(rx);
        // An endless stream: returning at all proves the closed channel ended it.
        let payloads = futures_util::stream::repeat(payload("mara", "hi"));
        forward(payloads, &tx).await;
    }

    #[tokio::test]
    async fn a_cancel_is_answered_with_whether_it_acted() {
        for acted in [true, false] {
            let got = act(
                br#""cancel""#,
                || acted,
                |_| async { panic!("applied an edit for a cancel") },
            )
            .await
            .expect("answered");
            let got: AgentControlReply = serde_json::from_slice(&got).expect("a cancel reply");
            assert_eq!(got, AgentControlReply { acted });
        }
    }

    #[tokio::test]
    async fn an_edit_is_answered_with_the_apply_s_reply() {
        let edit = Apply {
            base_revision: "abc".to_owned(),
            diff: "a diff".to_owned(),
        };
        let payload =
            serde_json::to_vec(&AgentControl::ApplyAgentsMd(edit.clone())).expect("encodes");
        let got = act(
            &payload,
            || panic!("cancelled for an edit"),
            |got| async move {
                assert_eq!(got, edit);
                ApplyReply::Applied {
                    revision: "def".to_owned(),
                }
            },
        )
        .await
        .expect("answered");
        let got: ApplyReply = serde_json::from_slice(&got).expect("an apply reply");
        assert_eq!(
            got,
            ApplyReply::Applied {
                revision: "def".to_owned()
            }
        );
    }

    #[tokio::test]
    async fn a_payload_that_is_not_a_command_is_not_answered_or_acted_on() {
        let got = act(
            b"pause",
            || panic!("ran a command that was not one"),
            |_| async { panic!("ran a command that was not one") },
        )
        .await;
        assert_eq!(got, None);
    }
}