Watch
0
0
Fork
You've already forked hyperhive
0
hyperhive/swarm-controller/src/agent_input.rs

621 lines
22 KiB
Rust

//! `POST /api/agents/{name}/input` — text the operator types in the swarm UI,
//! handed to one agent's harness as an inbox turn. No hive is in the path: the
//! agent subscribes to [`swarm_queue_client::agent_input_subject`] itself.
//!
//! `POST /api/agents/broadcast` sends one text to several agents through the
//! same send and answers per agent, in request order: one agent not getting it
//! does not fail the rest. A name must be in the swarm roster (what
//! `GET /api/agents` lists), or nothing is sent.
//!
//! **The sender is the `Remote-User` header**, which the gateway sets on
//! `/api/` from authelia's auth-request answer and so overwrites whatever the
//! client sent (`nix/host-modules/swarm-ui.nix`). A request without one is
//! refused rather than sent anonymously.
//!
//! **Sent as a NATS request** so that an agent with no subscriber is a 503
//! instead of text that silently goes nowhere. The agent does not reply to
//! input (it does to a command, see [`crate::agent_control`]). A request that
//! times out therefore means a subscriber held the message, and that is all
//! "sent" means. Core NATS, no stream: an agent that restarts before its turn
//! loses the text.
use std::collections::HashSet;
use std::time::Duration;
use axum::Json;
use axum::extract::{Path, State};
use axum::http::{HeaderMap, StatusCode};
use futures_util::StreamExt as _;
use problem_details::ProblemDetails;
use serde::{Deserialize, Serialize};
use swarm_queue_client::AgentInput;
use utoipa::ToSchema;
use crate::routes::{AppState, error_problem};
/// The header the gateway names the session's user in.
const SENDER_HEADER: &str = "remote-user";
/// How long to wait for the server's no-responders answer. It comes back
/// within a round trip when nothing subscribes; a subscriber sends nothing,
/// so a send to a live agent always waits this long.
const NO_RESPONDERS_WINDOW: Duration = Duration::from_secs(2);
/// The most agents one `POST /api/agents/broadcast` may name.
const MAX_BROADCAST_AGENTS: usize = 64;
/// How many sends of one broadcast are in flight at once. A send to a live
/// agent takes [`NO_RESPONDERS_WINDOW`], so a broadcast to
/// [`MAX_BROADCAST_AGENTS`] live agents answers after four windows.
const BROADCAST_CONCURRENCY: usize = 16;
/// Text for one agent.
#[derive(Deserialize, ToSchema)]
pub struct PostAgentInputRequest {
/// What the agent's turn shows as the message body, verbatim. Not logged.
body: String,
}
/// One text for several agents.
#[derive(Deserialize, ToSchema)]
pub struct BroadcastRequest {
/// The agents to send to, each named once.
agents: Vec<String>,
/// What each agent's turn shows as the message body, verbatim. Not logged.
body: String,
}
/// What became of a broadcast: one entry per requested agent, in request
/// order.
#[derive(Debug, Serialize, ToSchema)]
pub struct BroadcastResponse {
results: Vec<BroadcastResult>,
}
/// What became of the text for one agent of a broadcast.
#[derive(Debug, PartialEq, Eq, Serialize, ToSchema)]
pub struct BroadcastResult {
agent: String,
/// Whether the agent's harness was subscribed and handed the message.
sent: bool,
/// Why not, in the words `POST /api/agents/{name}/input` answers with.
/// Present only when `sent` is false.
#[serde(skip_serializing_if = "Option::is_none")]
reason: Option<String>,
}
/// Why [`send_one`] did not hand a message over, with the text to show for it.
#[derive(Debug, PartialEq, Eq)]
enum NotSent {
NotRunning(String),
Failed(String),
}
/// What became of one message.
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Delivery {
/// A subscriber held it.
Sent,
/// Nothing subscribes to the agent's input subject.
NotRunning,
/// The queue refused or failed the request.
Failed(String),
}
/// Where a message goes: the queue in production, a fake in the tests.
pub(crate) trait InputSink {
fn deliver(&self, subject: String, payload: Vec<u8>) -> impl Future<Output = Delivery> + Send;
}
impl InputSink for async_nats::Client {
async fn deliver(&self, subject: String, payload: Vec<u8>) -> Delivery {
use async_nats::RequestErrorKind;
let request = async_nats::Request::new()
.payload(payload.into())
.timeout(Some(NO_RESPONDERS_WINDOW));
match self.send_request(subject, request).await {
Ok(_) => Delivery::Sent,
Err(e) => match e.kind() {
RequestErrorKind::TimedOut => Delivery::Sent,
RequestErrorKind::NoResponders => Delivery::NotRunning,
_ => Delivery::Failed(e.to_string()),
},
}
}
}
/// Send text to an agent as the logged-in operator.
#[utoipa::path(
post,
path = "/api/agents/{name}/input",
params(("name" = String, Path, description = "agent to send the text to")),
request_body = PostAgentInputRequest,
responses(
(status = 202, description = "the agent's harness is subscribed and was handed the \
message; it becomes a turn from the sender when the harness next picks one up"),
(status = 400, description = "the agent name is not shaped like an identifier, or \
the body is empty (problem+json)", body = String),
(status = 401, description = "the request carries no `Remote-User`, so there is no \
sender to name (problem+json)", body = String),
(status = 503, description = "the agent's harness is not subscribed — not running, \
or not connected to the swarm queue — or no swarm queue is configured on this \
host or it could not be reached (problem+json)", body = String),
(status = 500, description = "the queue refused or failed the request (problem+json)",
body = String),
),
tag = "agents"
)]
pub(crate) async fn post_agent_input(
State(state): State<AppState>,
Path(agent): Path<String>,
headers: HeaderMap,
Json(req): Json<PostAgentInputRequest>,
) -> Result<StatusCode, ProblemDetails> {
let (agent, input) =
message(&headers, &agent, req.body).map_err(|(s, d)| error_problem(s, d))?;
let client = queue(&state).map_err(|d| error_problem(StatusCode::SERVICE_UNAVAILABLE, &d))?;
send(&client, &agent, &input).await
}
/// Send one text to several agents as the logged-in operator.
#[utoipa::path(
post,
path = "/api/agents/broadcast",
request_body = BroadcastRequest,
responses(
(status = 200, description = "one entry per requested agent, in request order: \
`sent`, and when false a `reason` in the words the single-agent input route \
answers with. One agent not getting it does not stop the others",
body = BroadcastResponse),
(status = 400, description = "`body` is empty, `agents` is empty, names an agent \
twice, names more than 64 agents, or names one not in the swarm roster or not \
shaped like an identifier; nothing was sent (problem+json)", body = String),
(status = 401, description = "the request carries no `Remote-User`, so there is no \
sender to name (problem+json)", body = String),
(status = 503, description = "the swarm roster could not be read, or no swarm queue \
is configured on this host or it could not be reached; nothing was sent \
(problem+json)", body = String),
),
tag = "agents"
)]
pub(crate) async fn post_agents_broadcast(
State(state): State<AppState>,
headers: HeaderMap,
Json(req): Json<BroadcastRequest>,
) -> Result<Json<BroadcastResponse>, ProblemDetails> {
let (agents, input) =
broadcast(&headers, req.agents, req.body).map_err(|(s, d)| error_problem(s, &d))?;
let roster = roster(&state).await?;
known(&agents, &roster).map_err(|(s, d)| error_problem(s, &d))?;
let client = queue(&state).map_err(|d| error_problem(StatusCode::SERVICE_UNAVAILABLE, &d))?;
Ok(Json(BroadcastResponse {
results: send_all(&client, agents, &input).await,
}))
}
/// The connected queue client, or why there is none: the caller's 503.
fn queue(state: &AppState) -> Result<async_nats::Client, String> {
let Some(status) = state.status.as_ref() else {
return Err("no swarm queue is configured on this host".to_owned());
};
let client = status.queue_client();
swarm_queue_client::ensure_connected(&client).map_err(|e| swarm_queue_client::chain(&e))?;
Ok(client)
}
/// The swarm roster, or the 503 `GET /api/agents` answers when it has none.
pub(crate) async fn roster(state: &AppState) -> Result<Vec<String>, ProblemDetails> {
let Some(bridge) = state.auth.as_ref() else {
return Err(error_problem(
StatusCode::SERVICE_UNAVAILABLE,
"no identity bridge is configured on this host",
));
};
bridge.list_agent_identities().await.map_err(|e| {
let detail = format!("reading the agent roster failed: {e:#}");
tracing::warn!(error = %detail, "broadcast refused");
error_problem(StatusCode::SERVICE_UNAVAILABLE, &detail)
})
}
/// Who the gateway says is asking, or the 401 for a request that did not come
/// through it. Shared with [`crate::agent_control`].
pub(crate) fn sender(headers: &HeaderMap) -> Result<String, (StatusCode, &'static str)> {
headers
.get(SENDER_HEADER)
.and_then(|v| v.to_str().ok())
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_owned)
.ok_or((
StatusCode::UNAUTHORIZED,
"no Remote-User header: the gateway names the sender, and this request \
did not come through it",
))
}
/// The agent and the message, or the refusal. Everything decided before the
/// queue is touched.
fn message(
headers: &HeaderMap,
agent: &str,
body: String,
) -> Result<(String, AgentInput), (StatusCode, &'static str)> {
let from = sender(headers)?;
let agent = agent_name(agent).map_err(|reason| (StatusCode::BAD_REQUEST, reason))?;
Ok((agent, input(from, body)?))
}
/// The agents and the one message for all of them, or the refusal of the
/// whole request. Everything decided before the roster and the queue are
/// touched.
fn broadcast(
headers: &HeaderMap,
agents: Vec<String>,
body: String,
) -> Result<(Vec<String>, AgentInput), (StatusCode, String)> {
let refused = |detail: String| (StatusCode::BAD_REQUEST, detail);
let from = sender(headers).map_err(|(s, d)| (s, d.to_owned()))?;
let input = input(from, body).map_err(|(s, d)| (s, d.to_owned()))?;
if agents.is_empty() {
return Err(refused("agents is empty".to_owned()));
}
if agents.len() > MAX_BROADCAST_AGENTS {
return Err(refused(format!(
"{} agents named; one broadcast may name at most {MAX_BROADCAST_AGENTS}",
agents.len()
)));
}
let mut seen = HashSet::new();
let mut names = Vec::with_capacity(agents.len());
for agent in agents {
let name = agent_name(&agent).map_err(|reason| refused(format!("{agent:?}: {reason}")))?;
if !seen.insert(name.clone()) {
return Err(refused(format!("{name} is named twice")));
}
names.push(name);
}
Ok((names, input))
}
/// Refuses the first of `agents` not in `roster`.
fn known(agents: &[String], roster: &[String]) -> Result<(), (StatusCode, String)> {
match agents.iter().find(|a| !roster.contains(a)) {
Some(agent) => Err((
StatusCode::BAD_REQUEST,
format!("{agent} is not an agent of this swarm"),
)),
None => Ok(()),
}
}
fn agent_name(agent: &str) -> Result<String, &'static str> {
Ok(hive_types::Ident::parse(agent)?.into_string())
}
fn input(from: String, body: String) -> Result<AgentInput, (StatusCode, &'static str)> {
if body.trim().is_empty() {
return Err((StatusCode::BAD_REQUEST, "body is empty"));
}
Ok(AgentInput { from, body })
}
/// `POST /api/agents/{name}/input`'s answer for one send.
async fn send(
sink: &impl InputSink,
agent: &str,
input: &AgentInput,
) -> Result<StatusCode, ProblemDetails> {
match send_one(sink, agent, input).await {
Ok(()) => Ok(StatusCode::ACCEPTED),
Err(NotSent::NotRunning(detail)) => {
Err(error_problem(StatusCode::SERVICE_UNAVAILABLE, &detail))
}
Err(NotSent::Failed(detail)) => {
Err(error_problem(StatusCode::INTERNAL_SERVER_ERROR, &detail))
}
}
}
/// Every agent's result, in the order of `agents`; at most
/// [`BROADCAST_CONCURRENCY`] sends are in flight.
async fn send_all<S: InputSink + Sync>(
sink: &S,
agents: Vec<String>,
input: &AgentInput,
) -> Vec<BroadcastResult> {
futures_util::stream::iter(agents)
.map(|agent| async move {
let reason = match send_one(sink, &agent, input).await {
Ok(()) => None,
Err(NotSent::NotRunning(detail) | NotSent::Failed(detail)) => Some(detail),
};
BroadcastResult {
agent,
sent: reason.is_none(),
reason,
}
})
.buffered(BROADCAST_CONCURRENCY)
.collect()
.await
}
/// Hand `input` to `agent`'s harness. The one send both routes make.
async fn send_one(sink: &impl InputSink, agent: &str, input: &AgentInput) -> Result<(), NotSent> {
let subject = swarm_queue_client::agent_input_subject(agent);
let payload = serde_json::to_vec(input)
.map_err(|e| NotSent::Failed(format!("encoding the input failed: {e}")))?;
match sink.deliver(subject.clone(), payload).await {
Delivery::Sent => {
tracing::info!(%agent, from = %input.from, "agent input sent");
Ok(())
}
Delivery::NotRunning => Err(NotSent::NotRunning(format!(
"{agent} is not listening for input: its harness is not running or not connected to the swarm queue"
))),
Delivery::Failed(e) => {
tracing::warn!(%agent, %subject, error = %e, "agent input failed");
Err(NotSent::Failed(format!("sending to {subject} failed: {e}")))
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use swarm_queue_client::AgentInput;
use super::{
BroadcastResult, Delivery, InputSink, MAX_BROADCAST_AGENTS, broadcast, known, message,
send, send_all,
};
/// Answers every message with one fixed outcome and records what it got.
struct FakeSink {
answer: fn() -> Delivery,
got: Mutex<Vec<(String, Vec<u8>)>>,
}
impl FakeSink {
fn new(answer: fn() -> Delivery) -> Self {
Self {
answer,
got: Mutex::default(),
}
}
}
impl InputSink for FakeSink {
async fn deliver(&self, subject: String, payload: Vec<u8>) -> Delivery {
self.got
.lock()
.expect("unpoisoned")
.push((subject, payload));
(self.answer)()
}
}
fn from(user: &str) -> HeaderMap {
let mut h = HeaderMap::new();
h.insert("Remote-User", HeaderValue::from_str(user).expect("ascii"));
h
}
fn status_of(r: Result<(String, AgentInput), (StatusCode, &'static str)>) -> StatusCode {
r.expect_err("refused").0
}
fn names(n: &[&str]) -> Vec<String> {
n.iter().map(|&a| a.to_owned()).collect()
}
fn broadcast_refusal(agents: &[&str], body: &str) -> (StatusCode, String) {
broadcast(&from("mara"), names(agents), body.to_owned()).expect_err("refused")
}
/// Sends to `atlas` and `iris` succeed, everyone else has no subscriber.
struct SomeListen {
got: Mutex<Vec<String>>,
}
impl InputSink for SomeListen {
async fn deliver(&self, subject: String, _payload: Vec<u8>) -> Delivery {
let agent = subject.trim_start_matches("$SWARM.input.").to_owned();
self.got.lock().expect("unpoisoned").push(agent.clone());
if agent == "atlas" || agent == "iris" {
Delivery::Sent
} else {
Delivery::NotRunning
}
}
}
#[test]
fn the_sender_is_the_remote_user_header() {
let (agent, input) = message(&from("mara"), "atlas", "hi".to_owned()).expect("valid");
assert_eq!(agent, "atlas");
assert_eq!(
input,
AgentInput {
from: "mara".to_owned(),
body: "hi".to_owned()
}
);
}
#[test]
fn no_sender_is_refused_before_anything_else() {
// A bad agent and an empty body as well: the missing sender is the
// answer, not whichever check happens to run first.
assert_eq!(
status_of(message(&HeaderMap::new(), "../x", String::new())),
StatusCode::UNAUTHORIZED
);
assert_eq!(
status_of(message(&from(" "), "atlas", "hi".to_owned())),
StatusCode::UNAUTHORIZED
);
}
#[test]
fn a_bad_agent_or_an_empty_body_is_refused() {
assert_eq!(
status_of(message(&from("mara"), "../x", "hi".to_owned())),
StatusCode::BAD_REQUEST
);
assert_eq!(
status_of(message(&from("mara"), "atlas", " \n".to_owned())),
StatusCode::BAD_REQUEST
);
}
#[test]
fn a_broadcast_names_its_agents_in_request_order_with_one_message() {
let (agents, input) =
broadcast(&from("mara"), names(&["iris", "atlas"]), "hi".to_owned()).expect("valid");
assert_eq!(agents, names(&["iris", "atlas"]));
assert_eq!(
input,
AgentInput {
from: "mara".to_owned(),
body: "hi".to_owned()
}
);
}
#[test]
fn a_broadcast_is_refused_whole_with_the_single_route_s_words() {
assert_eq!(
broadcast(&HeaderMap::new(), names(&["atlas"]), "hi".to_owned())
.expect_err("no sender")
.0,
StatusCode::UNAUTHORIZED
);
assert_eq!(
broadcast_refusal(&["atlas"], " \n"),
(StatusCode::BAD_REQUEST, "body is empty".to_owned())
);
assert_eq!(
broadcast_refusal(&[], "hi"),
(StatusCode::BAD_REQUEST, "agents is empty".to_owned())
);
let (status, detail) = broadcast_refusal(&["atlas", "../x"], "hi");
assert_eq!(status, StatusCode::BAD_REQUEST);
assert!(detail.starts_with("\"../x\": "), "{detail}");
assert_eq!(
broadcast_refusal(&["atlas", "iris", "atlas"], "hi"),
(StatusCode::BAD_REQUEST, "atlas is named twice".to_owned())
);
}
#[test]
fn a_broadcast_may_name_up_to_the_cap_and_no_more() {
let at_cap: Vec<String> = (0..MAX_BROADCAST_AGENTS).map(|i| format!("a{i}")).collect();
let mut over = at_cap.clone();
over.push("one-more".to_owned());
broadcast(&from("mara"), at_cap, "hi".to_owned()).expect("the cap itself is allowed");
let (status, detail) =
broadcast(&from("mara"), over, "hi".to_owned()).expect_err("over the cap");
assert_eq!(status, StatusCode::BAD_REQUEST);
assert!(detail.contains("at most 64"), "{detail}");
}
#[test]
fn a_name_outside_the_roster_refuses_the_broadcast() {
let roster = names(&["atlas", "iris"]);
known(&names(&["iris", "atlas"]), &roster).expect("both in the roster");
assert_eq!(
known(&names(&["atlas", "ghost"]), &roster).expect_err("ghost is unknown"),
(
StatusCode::BAD_REQUEST,
"ghost is not an agent of this swarm".to_owned()
)
);
}
#[tokio::test]
async fn one_agent_not_listening_does_not_stop_the_others() {
let sink = SomeListen {
got: Mutex::default(),
};
let input = AgentInput {
from: "mara".to_owned(),
body: "hi".to_owned(),
};
let results = send_all(&sink, names(&["iris", "argus", "atlas"]), &input).await;
assert_eq!(
results,
vec![
BroadcastResult {
agent: "iris".to_owned(),
sent: true,
reason: None
},
BroadcastResult {
agent: "argus".to_owned(),
sent: false,
reason: Some(
"argus is not listening for input: its harness is not running or not \
connected to the swarm queue"
.to_owned()
)
},
BroadcastResult {
agent: "atlas".to_owned(),
sent: true,
reason: None
},
]
);
let mut got = sink.got.lock().expect("unpoisoned").clone();
got.sort();
assert_eq!(got, names(&["argus", "atlas", "iris"]));
let json = serde_json::to_value(&results).expect("serializes");
assert_eq!(
json[0],
serde_json::json!({ "agent": "iris", "sent": true })
);
assert_eq!(json[1]["sent"], false);
assert!(json[1]["reason"].is_string());
}
#[tokio::test]
async fn a_held_message_is_202_on_the_agent_s_own_subject() {
let sink = FakeSink::new(|| Delivery::Sent);
let input = AgentInput {
from: "mara".to_owned(),
body: "hi".to_owned(),
};
let status = send(&sink, "atlas", &input).await.expect("sent");
assert_eq!(status, StatusCode::ACCEPTED);
let got = sink.got.lock().expect("unpoisoned");
assert_eq!(got.len(), 1);
assert_eq!(got[0].0, "$SWARM.input.atlas");
let decoded: AgentInput = serde_json::from_slice(&got[0].1).expect("an AgentInput");
assert_eq!(decoded, input);
}
#[tokio::test]
async fn no_subscriber_is_503_and_a_queue_failure_is_500() {
let input = AgentInput {
from: "mara".to_owned(),
body: "hi".to_owned(),
};
let not_running = send(&FakeSink::new(|| Delivery::NotRunning), "atlas", &input)
.await
.expect_err("nobody listens");
assert_eq!(not_running.status, Some(StatusCode::SERVICE_UNAVAILABLE));
let failed = send(
&FakeSink::new(|| Delivery::Failed("refused".to_owned())),
"atlas",
&input,
)
.await
.expect_err("the queue failed");
assert_eq!(failed.status, Some(StatusCode::INTERNAL_SERVER_ERROR));
}
}