A separate $SWARM.control.<agent> subject carries commands, apart from
the text subject. swarm-controller exposes POST /api/agents/{name}/cancel
as a NATS request; the harness answers whether a turn was actually
stopped, through the same cancel path as the dashboard button. The agent
credential gets subscribe on its control subject and allow_responses to
answer; only the controller may publish there.
Refs #3803
223 lines
8.2 KiB
Rust
223 lines
8.2 KiB
Rust
//! 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 arrives on [`swarm_queue_client::agent_control_subject`], never on
|
|
//! the input subject, so nothing a person types can be read as a command. The
|
|
//! reply says whether a turn was actually stopped; 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, |command| match command {
|
|
swarm_queue_client::AgentControl::Cancel => {
|
|
crate::web_ui::cancel_turn(&bus, &interrupted, "swarm UI: cancel")
|
|
}
|
|
}) else {
|
|
continue;
|
|
};
|
|
let payload = serde_json::to_vec(&answer).expect("a reply encodes");
|
|
if let Err(e) = client.publish(reply, payload.into()).await {
|
|
tracing::warn!(error = %e, "swarm control: answering failed");
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The answer to a command payload, running `cancel` for a cancel; `None`
|
|
/// (logged) when the payload is not a command, which gets no answer.
|
|
fn act(
|
|
payload: &[u8],
|
|
cancel: impl FnOnce(swarm_queue_client::AgentControl) -> bool,
|
|
) -> Option<swarm_queue_client::AgentControlReply> {
|
|
match serde_json::from_slice::<swarm_queue_client::AgentControl>(payload) {
|
|
Ok(command) => Some(swarm_queue_client::AgentControlReply {
|
|
acted: cancel(command),
|
|
}),
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "swarm control: undecodable command dropped");
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
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::{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;
|
|
}
|
|
|
|
#[test]
|
|
fn a_cancel_is_answered_with_whether_it_acted() {
|
|
for acted in [true, false] {
|
|
let got = act(br#""cancel""#, |c| {
|
|
assert_eq!(c, AgentControl::Cancel);
|
|
acted
|
|
});
|
|
assert_eq!(got, Some(AgentControlReply { acted }));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
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"));
|
|
assert_eq!(got, None);
|
|
}
|
|
}
|