mara/hyperhive
mara/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Activity

hive-bash-mcp: publish each task's status and output to the swarm queue

Browse source
This commit is contained in:
damocles 2026-11-15 17:48:13 +01:00
commit c0b9c1778c
14 changed files with 841 additions and 211 deletions
Download patch file Download diff file Expand all files Collapse all files

6
Cargo.lock generated
View file

@ -1761,6 +1761,7 @@ name = "hive-bash-mcp"
version = "0.1.0"
dependencies = [
"anyhow",
"async-nats",
"axum",
"chrono",
"clap",
@ -1774,6 +1775,9 @@ dependencies = [
"schemars",
"serde",
"serde_json",
"swarm-queue-client",
"swarm-secret-client",
"swarm-wire",
"tempfile",
"tokio",
"tracing",
@ -5298,6 +5302,7 @@ dependencies = [
name = "swarm-queue-client"
version = "0.1.0"
dependencies = [
"anyhow",
"async-nats",
"axum",
"data-encoding",
@ -5306,6 +5311,7 @@ dependencies = [
"reqwest",
"serde",
"serde_json",
"swarm-agent-verify",
"swarm-secret-client",
"swarm-wire",
"tempfile",

23
docs/4-implementation/agent-tools/bash.md
View file

@ -78,6 +78,29 @@ unread rooms. An inline `wait_seconds` / `status` observation that already
delivered the result instead clears the keyed todo, so no redundant
loose-end follows.
### Status and output the swarm sees
The daemon also tells the swarm about each task, so the swarm UI can list the
running ones and read their output without calling the agent
(`swarm_wire::bash_task`; the bucket, the stream and the grants are in
[`controller.md`](../swarm/controller.md)). It publishes as the agent, over the
connection `swarm_queue_client::agent_publisher` makes with the agent's own
queue credential, and publishes nothing in a container without a queue address
or a secret store. Everything is best-effort: a failed connect or publish is a
log line, and a task never waits on the queue.
- **Status.** The runner writes a task file at every state change, and each
write also publishes the task's state to `$KV.bash-task-status.<agent>.<task>`.
While a task is `pending` or `running`, the publish task repeats it every 60
seconds, so a daemon that died reads as stale. A finished task leaves that
repeat.
- **Output.** A task's output goes to files, so while it runs the daemon reads
what the files gain twice a second and publishes it as rows on
`$SWARM.term.<agent>.bash.<task>`, each with the `Swarm-Stream` header
(`stdout` or `stderr`). A chunk never splits a character, and one larger than
the queue's payload limit goes out in pieces. When the process exits, the daemon
publishes what the files still hold before it writes the final status.
## Relationship to the `execution` permission
`Permission::Execution` (in `Permission::DEFAULT`) is the real gate on bash:

3
docs/4-implementation/swarm/controller.md
View file

@ -202,8 +202,7 @@ goes to `$SWARM.term.<agent>.bash.<task>`, with a `Swarm-Stream` header of `stdo
`stderr`, into the stream `term-bash-<agent>` that swarm-controller creates for every declared
agent, for 24 hours or 64 MiB like the subagent rows. The agent's credential may publish
`$KV.bash-task-status.<agent>.*` and `$SWARM.term.<agent>.bash.>` and nothing else of
JetStream. Only the names, the grants and the two stores exist so far; the bash MCP doesn't
publish yet.
JetStream. The bash daemon publishes both (see [`bash.md`](../agent-tools/bash.md)).
The queue would refuse a row too large for its `max_payload` outright and
take the connection down with it, so the harness drops such a row's body before

6
hive-bash-mcp/Cargo.toml
View file

@ -19,6 +19,12 @@ hive-types.workspace = true
libc.workspace = true
rmcp.workspace = true
rusqlite.workspace = true
# The agent's own queue connection, and what a task publishes over it
# (`swarm_term`).
async-nats.workspace = true
swarm-queue-client = { workspace = true, features = ["agent-publisher"] }
swarm-secret-client.workspace = true
swarm-wire.workspace = true
schemars.workspace = true
serde.workspace = true
serde_json.workspace = true

1
hive-bash-mcp/src/lib.rs
View file

@ -9,5 +9,6 @@ pub mod paths;
pub mod protocol;
pub mod runner;
pub mod stats;
pub mod swarm_term;
#[cfg(test)]
mod test_util;

3
hive-bash-mcp/src/main.rs
View file

@ -39,6 +39,9 @@ async fn main() -> Result<()> {
"hive-bash-daemon starting"
);
// Tell the swarm about every task, when this container has a queue.
hive_bash_mcp::swarm_term::install();
let cgroups = hive_agent_sock::run_cgroup::RunCgroups::init("tasks");
// Start the background runner loop — scans for pending tasks and

10
hive-bash-mcp/src/runner.rs
View file

@ -197,7 +197,10 @@ fn write_task(task: &TaskFile) -> std::io::Result<()> {
let dest = paths::task_json(&task.id);
let tmp = dest.with_extension("json.tmp");
std::fs::write(&tmp, &json)?;
std::fs::rename(&tmp, &dest)
std::fs::rename(&tmp, &dest)?;
// Every state change is written here, so the swarm hears of each one.
crate::swarm_term::status(task);
Ok(())
}
/// Read a task file. Returns `None` if the file doesn't exist or is
@ -626,6 +629,7 @@ async fn run_task(mut task: TaskFile, socket: &Path, cgroups: Option<&RunCgroups
);
let cgroup = cgroups.and_then(|cgroups| task_cgroup(cgroups, &id));
let follower = crate::swarm_term::follow(&id, out_path.clone(), err_path.clone());
let outcome = exec_cmd(
&task.cmd,
&out_path,
@ -637,6 +641,10 @@ async fn run_task(mut task: TaskFile, socket: &Path, cgroups: Option<&RunCgroups
)
.await;
drop(cgroup);
// The last output goes out before the terminal status does.
if let Some(follower) = follower {
follower.finish().await;
}
// Deregister before writing terminal state — the task is no longer killable.
running().lock().unwrap().remove(&id);

534
hive-bash-mcp/src/swarm_term.rs Normal file
View file

@ -0,0 +1,534 @@
//! Publishing each bash task's status and output onto the swarm queue, as the
//! agent, so the swarm UI can list the tasks and read what they print without
//! calling the agent (`swarm_wire::bash_task`).
//!
//! **Status** goes to `$KV.bash-task-status.<agent>.<task>`: every state change
//! (the runner writes a task file on each, and [`status`] hears about it there)
//! and, while the task is `pending` or `running`, every
//! [`bash_task::HEARTBEAT_S`], so a daemon that died reads as stale.
//!
//! **Output** goes to `$SWARM.term.<agent>.bash.<task>`, one row per chunk with
//! the `Swarm-Stream` header. The task's output lands in files, so [`follow`]
//! reads what they gain while the task runs and hands it over.
//!
//! **Best-effort.** A failed connect or publish is a log line; the task never
//! waits on the queue. Without a queue address or a secret store nothing is
//! published. The connection is the agent's own
//! (`swarm_queue_client::agent_publisher`).
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use hive_sh4re::bash_task::{TaskFile, TaskStatus};
use swarm_queue_client::agent_publisher::{Queue, Target};
use swarm_wire::bash_task::{self, State, Status};
use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _};
use tokio::sync::{Notify, mpsc};
/// Output chunks waiting for the publish task. Past this, a chunk is dropped
/// and counted rather than held back from the task's reader.
const QUEUE_DEPTH: usize = 1024;
/// Room left under the server's payload limit for subject, headers and framing.
const HEADROOM: usize = 1024;
/// How often [`follow`] looks for new output, and the most it reads at once.
const FOLLOW_EVERY: Duration = Duration::from_millis(500);
const READ_AT_MOST: usize = 64 * 1024;
/// Which of a task's two output streams a chunk came from; the header value.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Stream {
Stdout,
Stderr,
}
impl Stream {
fn as_str(self) -> &'static str {
match self {
Self::Stdout => "stdout",
Self::Stderr => "stderr",
}
}
}
enum Event {
Status {
task: String,
status: Status,
},
Output {
task: String,
stream: Stream,
text: String,
},
}
/// The handle on the publish task.
pub struct Publisher {
status: mpsc::UnboundedSender<Event>,
output: mpsc::Sender<Event>,
dropped: Arc<AtomicU64>,
}
static PUBLISHER: OnceLock<Publisher> = OnceLock::new();
/// Start publishing, when this container has a queue address and a store to
/// read the agent's credential from. Logs which, once. Call once at startup.
pub fn install() {
let target = match Target::from_lookup(|k| std::env::var(k).ok()) {
Ok(Some(target)) => target,
Ok(None) => {
tracing::info!(
"no swarm queue address or no secret store for this agent; bash tasks are not \
published"
);
return;
}
Err(e) => {
tracing::warn!(
error = format!("{e:#}"),
"this agent's secret store coordinates are incomplete; bash tasks are not \
published"
);
return;
}
};
swarm_secret_client::install_crypto_provider();
tracing::info!(url = %target.url, "publishing bash tasks to the swarm queue as this agent");
let (status, status_rx) = mpsc::unbounded_channel();
let (output, output_rx) = mpsc::channel(QUEUE_DEPTH);
let dropped = Arc::new(AtomicU64::new(0));
tokio::spawn(run(status_rx, output_rx, target, Arc::clone(&dropped)));
let _ = PUBLISHER.set(Publisher {
status,
output,
dropped,
});
}
/// Tell the swarm `task`'s state, as the runner has just written it. A no-op
/// when nothing is being published.
pub fn status(task: &TaskFile) {
if let Some(p) = PUBLISHER.get() {
// A closed receiver means the task is gone; there is nobody to tell.
let _ = p.status.send(Event::Status {
task: task.id.clone(),
status: status_of(task, unix_now()),
});
}
}
fn offer_output(task: &str, stream: Stream, text: String) {
if let Some(p) = PUBLISHER.get() {
let event = Event::Output {
task: task.to_owned(),
stream,
text,
};
if p.output.try_send(event).is_err() {
p.dropped.fetch_add(1, Ordering::Relaxed);
}
}
}
/// Unix seconds now.
fn unix_now() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX))
}
fn state_of(status: &TaskStatus) -> State {
match status {
TaskStatus::Pending => State::Pending,
TaskStatus::Running => State::Running,
TaskStatus::Done => State::Done,
TaskStatus::TimedOut => State::TimedOut,
TaskStatus::Interrupted => State::Interrupted,
TaskStatus::Killed => State::Killed,
}
}
/// `task` as the swarm sees it, stamped `now`.
fn status_of(task: &TaskFile, now: i64) -> Status {
let created = task.created_at.timestamp();
let since = match task.status {
TaskStatus::Pending => created,
TaskStatus::Running => task.started_at.map_or(created, |t| t.timestamp()),
_ => task.completed_at.map_or(now, |t| t.timestamp()),
};
Status {
state: state_of(&task.status),
cmd: bash_task::clip_cmd(&task.cmd),
created,
since,
exit: task.exit_code,
ts: now,
}
}
/// Fold a status into the live tasks the heartbeat keeps going: a live one is
/// kept, a finished one leaves.
fn apply(live: &mut HashMap<String, Status>, task: &str, status: &Status) {
if status.state.is_live() {
live.insert(task.to_owned(), status.clone());
} else {
live.remove(task);
}
}
/// Stamp every live task's status with `ts`, the heartbeat, and return them.
fn refresh(live: &mut HashMap<String, Status>, ts: i64) -> Vec<(String, Status)> {
live.iter_mut()
.map(|(task, status)| {
status.ts = ts;
(task.clone(), status.clone())
})
.collect()
}
async fn run(
mut status_rx: mpsc::UnboundedReceiver<Event>,
mut output_rx: mpsc::Receiver<Event>,
target: Target,
dropped: Arc<AtomicU64>,
) {
let target = Arc::new(target);
let mut queue = Queue::new("bash tasks");
let mut live: HashMap<String, Status> = HashMap::new();
let period = Duration::from_secs(bash_task::HEARTBEAT_S.unsigned_abs());
let mut beat = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
loop {
tokio::select! {
Some(event) = status_rx.recv() => {
if let Event::Status { task, status } = event {
apply(&mut live, &task, &status);
publish_status(&mut queue, &target, &task, &status).await;
}
}
event = output_rx.recv() => {
let Some(Event::Output { task, stream, text }) = event else { return };
let missed = dropped.swap(0, Ordering::Relaxed);
if missed > 0 {
tracing::warn!(missed, "bash task output: queue full, chunks dropped");
}
publish_output(&mut queue, &target, &task, stream, &text).await;
}
_ = beat.tick() => {
for (task, status) in refresh(&mut live, unix_now()) {
publish_status(&mut queue, &target, &task, &status).await;
}
}
}
}
}
async fn publish_status(queue: &mut Queue, target: &Arc<Target>, task: &str, status: &Status) {
let Some(client) = queue.ready(target).await else {
return;
};
if let Err(e) = swarm_queue_client::ensure_connected(&client) {
tracing::warn!(error = %swarm_queue_client::chain(&e), "bash task status: publish skipped");
return;
}
let payload = match serde_json::to_vec(status) {
Ok(payload) => payload,
Err(e) => {
tracing::warn!(error = %e, "bash task status: serialising failed, not published");
return;
}
};
let subject = bash_task::status_subject(&target.agent, task);
if let Err(e) = client.publish(subject, payload.into()).await {
tracing::warn!(error = %e, task, "bash task status: publish failed");
}
}
async fn publish_output(
queue: &mut Queue,
target: &Arc<Target>,
task: &str,
stream: Stream,
text: &str,
) {
let Some(client) = queue.ready(target).await else {
return;
};
if let Err(e) = swarm_queue_client::ensure_connected(&client) {
tracing::warn!(error = %swarm_queue_client::chain(&e), "bash task output: publish skipped");
return;
}
let limit = swarm_queue_client::max_payload(&client).saturating_sub(HEADROOM);
let subject = bash_task::output_subject(&target.agent, task);
let mut headers = async_nats::HeaderMap::new();
headers.insert(bash_task::HEADER_STREAM, stream.as_str());
for chunk in chunks(text, limit) {
let sent = client
.publish_with_headers(
subject.clone(),
headers.clone(),
chunk.to_owned().into_bytes().into(),
)
.await;
if let Err(e) = sent {
tracing::warn!(error = %e, task, "bash task output: publish failed, chunk dropped");
}
}
}
/// `text` cut into pieces of at most `limit` bytes, each on a character
/// boundary. A `limit` smaller than one character yields nothing rather than
/// loop.
fn chunks(text: &str, limit: usize) -> Vec<&str> {
let mut out = Vec::new();
let mut rest = text;
while !rest.is_empty() {
let mut end = limit.min(rest.len());
while end > 0 && !rest.is_char_boundary(end) {
end -= 1;
}
if end == 0 {
break;
}
let (piece, tail) = rest.split_at(end);
out.push(piece);
rest = tail;
}
out
}
/// What [`follow`] returns: finish it once the task's process has exited, so
/// the last output is published before the task's final status.
pub struct Follower {
stop: Arc<Notify>,
done: tokio::task::JoinHandle<()>,
}
impl Follower {
/// Publish whatever the files still hold, then stop.
pub async fn finish(self) {
self.stop.notify_one();
let _ = self.done.await;
}
}
/// Publish what `task`'s output files gain while it runs, or `None` when
/// nothing is being published.
#[must_use]
pub fn follow(task: &str, out: PathBuf, err: PathBuf) -> Option<Follower> {
PUBLISHER.get()?;
let stop = Arc::new(Notify::new());
let task = task.to_owned();
let done = tokio::spawn({
let stop = Arc::clone(&stop);
async move {
let mut files = [
Tail::new(out, Stream::Stdout),
Tail::new(err, Stream::Stderr),
];
loop {
let last = tokio::select! {
() = stop.notified() => true,
() = tokio::time::sleep(FOLLOW_EVERY) => false,
};
// On the stop, read until both files are empty: the process
// is gone, so what is there is all there will be.
loop {
let mut any = false;
for file in &mut files {
if let Some(text) = file.read_new(last).await {
any = true;
offer_output(&task, file.stream, text);
}
}
if !last || !any {
break;
}
}
if last {
return;
}
}
}
});
Some(Follower { stop, done })
}
/// One output file being followed.
struct Tail {
path: PathBuf,
stream: Stream,
offset: u64,
/// The end of the last read when it cut a character in two, held until the
/// rest arrives.
partial: Vec<u8>,
}
impl Tail {
fn new(path: PathBuf, stream: Stream) -> Self {
Self {
path,
stream,
offset: 0,
partial: Vec::new(),
}
}
/// The text the file has gained since the last read, or `None` when there is
/// none. `last` flushes a held partial character (as the replacement
/// character) instead of waiting for the rest.
async fn read_new(&mut self, last: bool) -> Option<String> {
let mut file = tokio::fs::File::open(&self.path).await.ok()?;
file.seek(std::io::SeekFrom::Start(self.offset))
.await
.ok()?;
let mut buf = Vec::new();
file.take(READ_AT_MOST as u64)
.read_to_end(&mut buf)
.await
.ok()?;
self.offset += buf.len() as u64;
self.partial.extend_from_slice(&buf);
let text = take_text(&mut self.partial, last);
(!text.is_empty()).then_some(text)
}
}
/// The text in `bytes`, leaving a trailing incomplete character in `bytes`
/// unless `flush`. Invalid bytes become the replacement character.
fn take_text(bytes: &mut Vec<u8>, flush: bool) -> String {
let keep = if flush {
0
} else {
match std::str::from_utf8(bytes) {
Err(e) if e.error_len().is_none() => bytes.len() - e.valid_up_to(),
_ => 0,
}
};
let cut = bytes.len() - keep;
let text = String::from_utf8_lossy(&bytes[..cut]).into_owned();
bytes.drain(..cut);
text
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{TimeZone as _, Utc};
fn task(status: TaskStatus) -> TaskFile {
TaskFile {
id: "build".to_owned(),
cmd: "make -j8".to_owned(),
timeout_secs: None,
status,
created_at: Utc.timestamp_opt(1_000, 0).unwrap(),
started_at: Some(Utc.timestamp_opt(1_010, 0).unwrap()),
completed_at: Some(Utc.timestamp_opt(1_090, 0).unwrap()),
exit_code: Some(2),
stdout_tail: None,
stderr_tail: None,
}
}
/// `since` is when the task entered its state, and the state keeps the names
/// the swarm reads.
#[test]
fn a_status_dates_each_state_from_the_task_file() {
let pending = status_of(&task(TaskStatus::Pending), 5_000);
assert_eq!((pending.state, pending.since), (State::Pending, 1_000));
let running = status_of(&task(TaskStatus::Running), 5_000);
assert_eq!((running.state, running.since), (State::Running, 1_010));
let done = status_of(&task(TaskStatus::Done), 5_000);
assert_eq!(
(done.state, done.since, done.exit),
(State::Done, 1_090, Some(2))
);
assert_eq!((done.created, done.ts), (1_000, 5_000));
for (status, state) in [
(TaskStatus::TimedOut, State::TimedOut),
(TaskStatus::Interrupted, State::Interrupted),
(TaskStatus::Killed, State::Killed),
] {
assert_eq!(status_of(&task(status), 5_000).state, state);
}
}
#[test]
fn the_command_is_cut_to_the_length_the_swarm_shows() {
let mut long = task(TaskStatus::Running);
long.cmd = "x".repeat(bash_task::CMD_MAX_CHARS + 50);
assert_eq!(
status_of(&long, 0).cmd.chars().count(),
bash_task::CMD_MAX_CHARS
);
}
/// A live task is heartbeated until it finishes, and the heartbeat moves `ts`
/// only.
#[test]
fn the_heartbeat_keeps_live_tasks_and_drops_finished_ones() {
let mut live = HashMap::new();
apply(&mut live, "a", &status_of(&task(TaskStatus::Running), 100));
apply(&mut live, "b", &status_of(&task(TaskStatus::Pending), 100));
let mut beat = refresh(&mut live, 160);
beat.sort_by(|x, y| x.0.cmp(&y.0));
assert_eq!(beat.len(), 2);
assert!(beat.iter().all(|(_, s)| s.ts == 160 && s.since != 160));
apply(&mut live, "a", &status_of(&task(TaskStatus::Done), 170));
let beat = refresh(&mut live, 220);
assert_eq!(beat.len(), 1);
assert_eq!(beat[0].0, "b");
}
#[test]
fn a_chunk_is_cut_on_a_character_boundary_within_the_limit() {
assert_eq!(chunks("abcdef", 4), vec!["abcd", "ef"]);
// "é" is two bytes: a limit of 3 cannot split the second one.
assert_eq!(chunks("ééé", 3), vec!["é", "é", "é"]);
assert!(
chunks("é", 1).is_empty(),
"a limit under one character yields nothing"
);
assert!(chunks("", 10).is_empty());
}
/// Output read in pieces joins up: a character split across two reads comes
/// out whole, and the held bytes are flushed at the end.
#[test]
fn a_character_split_between_reads_is_held_until_it_is_whole() {
let bytes = "aé".as_bytes();
let mut held = bytes[..2].to_vec();
assert_eq!(take_text(&mut held, false), "a");
assert_eq!(held, vec![bytes[1]]);
held.push(bytes[2]);
assert_eq!(take_text(&mut held, false), "é");
assert!(held.is_empty());
let mut cut = vec![0xc3];
assert_eq!(take_text(&mut cut, true), "\u{fffd}");
assert!(cut.is_empty());
}
#[tokio::test]
async fn a_tail_reads_only_what_a_file_gained() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("t.out");
let mut tail = Tail::new(path.clone(), Stream::Stdout);
assert_eq!(tail.read_new(false).await, None, "no file yet");
std::fs::write(&path, "one\n").unwrap();
assert_eq!(tail.read_new(false).await.as_deref(), Some("one\n"));
assert_eq!(tail.read_new(false).await, None, "nothing new");
let mut f = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.unwrap();
std::io::Write::write_all(&mut f, b"two\n").unwrap();
assert_eq!(tail.read_new(false).await.as_deref(), Some("two\n"));
}
}

1
hive-subagent-mcp/Cargo.toml
View file

@ -32,6 +32,7 @@ serde_json.workspace = true
swarm-queue-client = { workspace = true, features = [
"agent-inference",
"agent-permissions",
"agent-publisher",
] }
# The spelling of the token the queue connection presents.
swarm-agent-verify.workspace = true

2
hive-subagent-mcp/src/swarm_models.rs
View file

@ -30,7 +30,7 @@ use swarm_queue_client::agent_inference::{self, read_own};
use swarm_queue_client::agent_permissions::{self, BEFORE_FIRST_READ, held};
use swarm_wire::agent_inference::AgentInference;
use crate::swarm_term::{Target, connect};
use swarm_queue_client::agent_publisher::{Target, connect};
/// The last permission set [`Reader::refresh_permissions`] read; `None` until
/// one succeeds.

218
hive-subagent-mcp/src/swarm_term.rs
View file

@ -32,38 +32,18 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context as _, anyhow};
use hive_sh4re::term_msg::{ClassifyCtx, Level, TermMsg, fit, iso8601_utc};
use swarm_secret_client::{
SecretStore,
client::{DEFAULT_CERT_MOUNT, ENV_ADDR, ENV_CACERT, Settings},
policy, queue,
};
use swarm_queue_client::agent_publisher::{Queue, Target};
use swarm_wire::subagent_status::{self, State as RunState, Status};
use swarm_wire::subagent_term;
use tokio::sync::mpsc;
/// Names the agent. The store role, the store path, the queue token and the
/// subjects are all built from it.
pub(crate) const ENV_AGENT_NAME: &str = "HIVE_AGENT_NAME";
/// Where the swarm queue listens, as this container reaches it.
const ENV_NATS_URL: &str = "HIVE_AGENT_NATS_URL";
/// Lines waiting for the publish task. Past this, `offer` drops the line
/// rather than block a subagent's output reader.
const QUEUE_DEPTH: usize = 4096;
/// How long one read of the store may take; it runs inside a queue
/// connection attempt.
const STORE_READ_TIMEOUT: Duration = Duration::from_secs(10);
/// How long after a failed connect the task tries again. Rows that arrive
/// before a first connect succeeds are dropped.
const RETRY_AFTER: Duration = Duration::from_mins(1);
/// Room left under the server's payload limit for subject, headers and
/// framing, as in `hive-agent`'s `swarm_term`.
const HEADROOM: usize = 1024;
@ -245,136 +225,6 @@ fn unix_now() -> i64 {
.map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX))
}
/// Where to publish, and what to connect with.
pub(crate) struct Target {
pub(crate) url: String,
pub(crate) agent: String,
store: StoreSource,
}
/// This agent's queue secret as the store holds it.
struct StoreSource {
settings: Settings,
role: String,
path: String,
}
impl Target {
/// `None` when the queue address or the store address is absent, which
/// is a container without a swarm queue or without a store.
///
/// # Errors
/// A store address without the agent name or the identity beside it.
pub(crate) fn from_lookup(
get: impl Fn(&str) -> Option<String>,
) -> anyhow::Result<Option<Self>> {
let Some(url) = get(ENV_NATS_URL).filter(|v| !v.is_empty()) else {
return Ok(None);
};
if get(ENV_ADDR).is_none_or(|v| v.is_empty()) {
return Ok(None);
}
let agent = get(ENV_AGENT_NAME)
.filter(|v| !v.is_empty())
.with_context(|| format!("{ENV_AGENT_NAME} is unset or empty"))?;
let settings =
Settings::from_lookup(|k| get(k).filter(|v| k != ENV_CACERT || ca_is_usable(v)))
.context("reading the swarm secret store's coordinates from the environment")?;
let store = StoreSource {
settings,
role: policy::agent_object_name(&agent)?,
path: queue::agent_queue_path(&agent)?,
};
Ok(Some(Self { url, agent, store }))
}
fn subject(&self, subagent: &str) -> String {
subagent_term::subject(&self.agent, subagent)
}
}
/// Whether the CA bundle at `path` is a file with bytes in it. Absent means
/// the container's own trust store.
fn ca_is_usable(path: &str) -> bool {
std::fs::metadata(path).is_ok_and(|m| m.len() > 0)
}
impl StoreSource {
/// The stored secret, read with a fresh login. Errors name the path and
/// never the value.
async fn read(&self) -> anyhow::Result<String> {
let fetch = async {
let store = SecretStore::connect(&self.settings, &self.role, DEFAULT_CERT_MOUNT)
.await
.context("logging in to the swarm secret store as this agent")?;
let credential: Option<queue::AgentCredential> = store
.read_optional(&self.path)
.await
.with_context(|| format!("reading {} from the store", self.path))?;
credential
.map(|c| c.value.trim().to_owned())
.filter(|v| !v.is_empty())
.with_context(|| format!("nothing is stored at {}", self.path))
};
tokio::time::timeout(STORE_READ_TIMEOUT, fetch)
.await
.map_err(|_| anyhow!("the store did not answer within {STORE_READ_TIMEOUT:?}"))?
}
}
/// Connect presenting the agent's own token, read from the store on every
/// connection attempt so a re-minted secret is picked up on reconnect.
pub(crate) async fn connect(target: &Arc<Target>) -> anyhow::Result<async_nats::Client> {
let url = target.url.clone();
let target = Arc::clone(target);
let agent = target.agent.clone();
swarm_queue_client::connect_with_token(&url, None, &agent, move || {
let target = Arc::clone(&target);
// Spawned so the future handed back is only a `JoinHandle`, which is
// `Sync` as async-nats's auth callback requires; the store read is not.
let attempt = tokio::spawn(async move {
let value = target.store.read().await.map_err(|e| format!("{e:#}"))?;
swarm_agent_verify::agent_token::format_agent_token(&target.agent, &value)
.map_err(|e| e.to_string())
});
async move { attempt.await.map_err(|e| e.to_string())? }
})
.await
.map_err(|e| anyhow!(swarm_queue_client::chain(&e)))
}
/// The connection, retried at most once per [`RETRY_AFTER`].
#[derive(Default)]
struct Queue {
client: Option<async_nats::Client>,
connect_after: Option<Instant>,
}
impl Queue {
/// A client to publish on, connecting first when a connect is due.
async fn ready(&mut self, target: &Arc<Target>) -> Option<async_nats::Client> {
let now = Instant::now();
if self.client.is_none() && self.connect_after.is_none_or(|t| now >= t) {
match connect(target).await {
Ok(client) => {
tracing::info!(url = %target.url, "subagent terminal: connected as this agent");
self.client = Some(client);
}
Err(e) => {
tracing::warn!(
error = format!("{e:#}"),
retry_in = ?RETRY_AFTER,
"subagent terminal: connecting with this agent's own queue credential \
failed; rows are dropped until a connect succeeds"
);
self.connect_after = Some(now + RETRY_AFTER);
}
}
}
self.client.clone()
}
}
async fn run(
mut rx: mpsc::Receiver<Line>,
mut status_rx: mpsc::UnboundedReceiver<StatusEvent>,
@ -383,7 +233,7 @@ async fn run(
) {
let target = Arc::new(target);
let mut ctxs: HashMap<String, ClassifyCtx> = HashMap::new();
let mut queue = Queue::default();
let mut queue = Queue::new("subagent terminal");
let mut runs: HashMap<String, Status> = HashMap::new();
let period = Duration::from_secs(subagent_status::HEARTBEAT_S.unsigned_abs());
let mut beat = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
@ -425,7 +275,7 @@ async fn publish_line(
let Some(client) = queue.ready(target).await else {
return;
};
let subject = target.subject(&line.subagent);
let subject = subagent_term::subject(&target.agent, &line.subagent);
let headers = line.info.headers();
for row in rows {
publish(&client, &subject, &headers, row).await;
@ -552,64 +402,18 @@ async fn publish(
mod tests {
use super::*;
fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
let pairs: Vec<(String, String)> = pairs
.iter()
.map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
.collect();
move |k| pairs.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone())
}
const STORE: [(&str, &str); 3] = [
("BAO_ADDR", "https://bao.t.local:8200"),
("BAO_CLIENT_CERT", "/run/credentials/c"),
("BAO_CLIENT_KEY", "/run/credentials/k"),
];
fn full() -> Vec<(&'static str, &'static str)> {
let mut pairs = STORE.to_vec();
pairs.push((ENV_NATS_URL, "tls://nats.t.local:4222"));
pairs.push((ENV_AGENT_NAME, "iris"));
pairs
}
/// Each subagent publishes on its own subject under this agent's family,
/// the one the queue grants the agent.
#[test]
fn a_subagent_publishes_under_the_agents_own_family() {
let target = Target::from_lookup(env(&full()))
.expect("complete")
.expect("configured");
assert_eq!(target.subject("scout"), "$SWARM.term.iris.sub.scout");
assert_eq!(target.subject("w-2"), "$SWARM.term.iris.sub.w-2");
assert_eq!(target.store.path, "swarm/agents/iris/queue");
}
#[test]
fn without_a_queue_address_or_a_store_nothing_is_published() {
let no_url: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != ENV_NATS_URL)
.collect();
assert!(Target::from_lookup(env(&no_url)).expect("legal").is_none());
let no_store: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != "BAO_ADDR")
.collect();
assert!(
Target::from_lookup(env(&no_store))
.expect("legal")
.is_none()
assert_eq!(
subagent_term::subject("iris", "scout"),
"$SWARM.term.iris.sub.scout"
);
assert_eq!(
subagent_term::subject("iris", "w-2"),
"$SWARM.term.iris.sub.w-2"
);
}
#[test]
fn a_store_without_the_agent_name_is_an_error() {
let no_name: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != ENV_AGENT_NAME)
.collect();
assert!(Target::from_lookup(env(&no_name)).is_err());
}
/// The sink must never wait on the publisher: a full queue drops the line

6
swarm-queue-client/Cargo.toml
View file

@ -33,6 +33,9 @@ notices = ["async-nats/jetstream"]
# `subagent_term::open_or_create`, for swarm-controller, which creates every
# agent's subagent stream. Explicit for the same reason as `notices`.
subagent-term = ["async-nats/jetstream"]
# `agent_publisher`: the connection a daemon in an agent's container publishes
# over as the agent, with the agent's own queue credential from the secret store.
agent-publisher = ["dep:anyhow", "dep:swarm-agent-verify"]
# `bash_task::open_status` and `open_output`, for swarm-controller, which creates
# the status bucket and every agent's task output stream. Needs `kv` for the bucket.
bash-task = ["kv"]
@ -84,6 +87,7 @@ router-todo = ["async-nats/jetstream"]
[dependencies]
# Bare (no `kv`/`jetstream`) unless a consumer opts into the `kv` feature
# above - the connect itself needs none of them.
anyhow = { workspace = true, optional = true }
async-nats.workspace = true
hive-types = { workspace = true, optional = true }
# `agent_status::roster` collects the stream's subject listing.
@ -106,6 +110,8 @@ swarm-secret-client.workspace = true
# The inference bucket's stored types are in `swarm-wire`; this crate holds
# the connection side.
swarm-wire.workspace = true
# The spelling of the token `agent_publisher` presents.
swarm-agent-verify = { workspace = true, optional = true }
# A library, so its errors are a matchable enum rather than an opaque
# `anyhow::Error`. The binaries that consume this keep anyhow; `?` converts.
thiserror.workspace = true

236
swarm-queue-client/src/agent_publisher.rs Normal file
View file

@ -0,0 +1,236 @@
//! The connection a daemon publishes to the swarm queue over, **as the agent**.
//!
//! The agent's own queue credential is read from the swarm secret store under
//! the agent's store identity on every connection attempt, so a re-minted
//! secret is picked up on reconnect. The hive's shared client is never tried:
//! the queue grants it nothing under the agent's own subjects.
//!
//! [`Target::from_lookup`] reads where to connect from the environment and is
//! `None` in a container without a queue address or without a store;
//! [`Queue`] connects lazily and at most once per [`RETRY_AFTER`]. Callers
//! publish; they never subscribe.
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::{Context as _, anyhow};
use swarm_secret_client::{
SecretStore,
client::{DEFAULT_CERT_MOUNT, ENV_ADDR, ENV_CACERT, Settings},
policy, queue,
};
/// Names the agent. The store role, the store path, the queue token and the
/// subjects are all built from it.
pub const ENV_AGENT_NAME: &str = "HIVE_AGENT_NAME";
/// Where the swarm queue listens, as this container reaches it.
pub const ENV_NATS_URL: &str = "HIVE_AGENT_NATS_URL";
/// How long one read of the store may take; it runs inside a queue
/// connection attempt.
const STORE_READ_TIMEOUT: Duration = Duration::from_secs(10);
/// How long after a failed connect the task tries again. What arrives before a first connect
/// succeeds is dropped.
pub const RETRY_AFTER: Duration = Duration::from_mins(1);
/// Where to publish, and what to connect with.
pub struct Target {
pub url: String,
pub agent: String,
store: StoreSource,
}
/// This agent's queue secret as the store holds it.
struct StoreSource {
settings: Settings,
role: String,
path: String,
}
impl Target {
/// `None` when the queue address or the store address is absent, which
/// is a container without a swarm queue or without a store.
///
/// # Errors
/// A store address without the agent name or the identity beside it.
pub fn from_lookup(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<Option<Self>> {
let Some(url) = get(ENV_NATS_URL).filter(|v| !v.is_empty()) else {
return Ok(None);
};
if get(ENV_ADDR).is_none_or(|v| v.is_empty()) {
return Ok(None);
}
let agent = get(ENV_AGENT_NAME)
.filter(|v| !v.is_empty())
.with_context(|| format!("{ENV_AGENT_NAME} is unset or empty"))?;
let settings =
Settings::from_lookup(|k| get(k).filter(|v| k != ENV_CACERT || ca_is_usable(v)))
.context("reading the swarm secret store's coordinates from the environment")?;
let store = StoreSource {
settings,
role: policy::agent_object_name(&agent)?,
path: queue::agent_queue_path(&agent)?,
};
Ok(Some(Self { url, agent, store }))
}
}
/// Whether the CA bundle at `path` is a file with bytes in it. Absent means
/// the container's own trust store.
fn ca_is_usable(path: &str) -> bool {
std::fs::metadata(path).is_ok_and(|m| m.len() > 0)
}
impl StoreSource {
/// The stored secret, read with a fresh login. Errors name the path and
/// never the value.
async fn read(&self) -> anyhow::Result<String> {
let fetch = async {
let store = SecretStore::connect(&self.settings, &self.role, DEFAULT_CERT_MOUNT)
.await
.context("logging in to the swarm secret store as this agent")?;
let credential: Option<queue::AgentCredential> = store
.read_optional(&self.path)
.await
.with_context(|| format!("reading {} from the store", self.path))?;
credential
.map(|c| c.value.trim().to_owned())
.filter(|v| !v.is_empty())
.with_context(|| format!("nothing is stored at {}", self.path))
};
tokio::time::timeout(STORE_READ_TIMEOUT, fetch)
.await
.map_err(|_| anyhow!("the store did not answer within {STORE_READ_TIMEOUT:?}"))?
}
}
/// Connect presenting the agent's own token, read from the store on every
/// connection attempt so a re-minted secret is picked up on reconnect.
pub async fn connect(target: &Arc<Target>) -> anyhow::Result<async_nats::Client> {
let url = target.url.clone();
let target = Arc::clone(target);
let agent = target.agent.clone();
crate::connect_with_token(&url, None, &agent, move || {
let target = Arc::clone(&target);
// Spawned so the future handed back is only a `JoinHandle`, which is
// `Sync` as async-nats's auth callback requires; the store read is not.
let attempt = tokio::spawn(async move {
let value = target.store.read().await.map_err(|e| format!("{e:#}"))?;
swarm_agent_verify::agent_token::format_agent_token(&target.agent, &value)
.map_err(|e| e.to_string())
});
async move { attempt.await.map_err(|e| e.to_string())? }
})
.await
.map_err(|e| anyhow!(crate::chain(&e)))
}
/// The connection, retried at most once per [`RETRY_AFTER`].
pub struct Queue {
client: Option<async_nats::Client>,
connect_after: Option<Instant>,
/// What the log lines call the publisher, e.g. `subagent terminal`.
label: &'static str,
}
impl Queue {
/// A queue that has not connected yet; `label` names the publisher in
/// the log.
#[must_use]
pub fn new(label: &'static str) -> Self {
Self {
client: None,
connect_after: None,
label,
}
}
/// A client to publish on, connecting first when a connect is due.
pub async fn ready(&mut self, target: &Arc<Target>) -> Option<async_nats::Client> {
let now = Instant::now();
if self.client.is_none() && self.connect_after.is_none_or(|t| now >= t) {
match connect(target).await {
Ok(client) => {
tracing::info!(url = %target.url, label = self.label, "connected to the swarm queue as this agent");
self.client = Some(client);
}
Err(e) => {
tracing::warn!(
error = format!("{e:#}"),
retry_in = ?RETRY_AFTER,
label = self.label,
"connecting with this agent's own queue credential failed; nothing is \
published until a connect succeeds"
);
self.connect_after = Some(now + RETRY_AFTER);
}
}
}
self.client.clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
let pairs: Vec<(String, String)> = pairs
.iter()
.map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
.collect();
move |k| pairs.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone())
}
const STORE: [(&str, &str); 3] = [
("BAO_ADDR", "https://bao.t.local:8200"),
("BAO_CLIENT_CERT", "/run/credentials/c"),
("BAO_CLIENT_KEY", "/run/credentials/k"),
];
fn full() -> Vec<(&'static str, &'static str)> {
let mut pairs = STORE.to_vec();
pairs.push((ENV_NATS_URL, "tls://nats.t.local:4222"));
pairs.push((ENV_AGENT_NAME, "iris"));
pairs
}
#[test]
fn the_target_names_the_agent_and_its_queue_secret() {
let target = Target::from_lookup(env(&full()))
.expect("complete")
.expect("configured");
assert_eq!(target.agent, "iris");
assert_eq!(target.url, "tls://nats.t.local:4222");
assert_eq!(target.store.path, "swarm/agents/iris/queue");
}
#[test]
fn without_a_queue_address_or_a_store_nothing_is_published() {
let no_url: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != ENV_NATS_URL)
.collect();
assert!(Target::from_lookup(env(&no_url)).expect("legal").is_none());
let no_store: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != "BAO_ADDR")
.collect();
assert!(
Target::from_lookup(env(&no_store))
.expect("legal")
.is_none()
);
}
#[test]
fn a_store_without_the_agent_name_is_an_error() {
let no_name: Vec<_> = full()
.into_iter()
.filter(|(k, _)| *k != ENV_AGENT_NAME)
.collect();
assert!(Target::from_lookup(env(&no_name)).is_err());
}
}

3
swarm-queue-client/src/lib.rs
View file

@ -484,6 +484,9 @@ pub mod notices;
#[cfg(feature = "agent-term")]
pub mod agent_term;
#[cfg(feature = "agent-publisher")]
pub mod agent_publisher;
/// Creating the per-agent stream that lists an agent's subagents. The names,
/// subjects and grants are `swarm_wire::subagent_term`. Behind the
/// `subagent-term` feature.