90 lines
3.4 KiB
Rust
90 lines
3.4 KiB
Rust
//! The connection side of an agent's bash tasks: the status bucket and the
|
|
//! output stream. The names, keys, subjects and the status value are
|
|
//! `swarm_wire::bash_task`.
|
|
//!
|
|
//! `swarm-controller` creates the bucket once and one output stream per agent
|
|
//! with the fixed config `output_config` spells.
|
|
|
|
use swarm_wire::bash_task::{OUTPUT_MAX_AGE, output_stream_name, output_stream_subjects};
|
|
|
|
/// How large one agent's output stream may grow before the discard policy
|
|
/// reclaims space, as for the subagent rows.
|
|
const OUTPUT_MAX_BYTES: i64 = 64 * 1024 * 1024;
|
|
|
|
/// The config [`open_output`] creates `agent`'s output stream with.
|
|
///
|
|
/// `max_message_size` is the queue's live `max_payload` ([`crate::max_payload`]),
|
|
/// so the stream never rejects a row the queue itself accepted.
|
|
pub(crate) fn output_config(
|
|
agent: &str,
|
|
max_message_size: i32,
|
|
) -> async_nats::jetstream::stream::Config {
|
|
async_nats::jetstream::stream::Config {
|
|
name: output_stream_name(agent),
|
|
description: Some(format!("Output of {agent}'s bash tasks")),
|
|
subjects: vec![output_stream_subjects(agent)],
|
|
max_age: OUTPUT_MAX_AGE,
|
|
max_bytes: OUTPUT_MAX_BYTES,
|
|
discard: async_nats::jetstream::stream::DiscardPolicy::Old,
|
|
max_message_size,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
/// Open the status bucket, creating it if nothing has yet. `history: 1`: a
|
|
/// reader wants each task's current state, not every state it has been in. A
|
|
/// key expires [`swarm_wire::bash_task::STATUS_MAX_AGE`] after its last write,
|
|
/// since nothing deletes the key of a finished task.
|
|
pub async fn open_status(
|
|
client: &async_nats::Client,
|
|
) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
|
|
let js = async_nats::jetstream::new(client.clone());
|
|
crate::kv::open_or_create_expiring(
|
|
&js,
|
|
swarm_wire::bash_task::BUCKET,
|
|
"bash-task-status",
|
|
"The state of each bash task, as its agent's bash MCP last published it".to_owned(),
|
|
Some(swarm_wire::bash_task::STATUS_MAX_AGE),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Open `agent`'s output stream, creating it if it does not exist yet.
|
|
pub async fn open_output(
|
|
client: &async_nats::Client,
|
|
agent: &str,
|
|
) -> Result<async_nats::jetstream::stream::Stream, crate::Error> {
|
|
let js = async_nats::jetstream::new(client.clone());
|
|
let name = output_stream_name(agent);
|
|
match js.get_stream(&name).await {
|
|
Ok(stream) => Ok(stream),
|
|
Err(e) => {
|
|
tracing::info!(stream = %name, reason = %e, "bash task output stream not available, creating it");
|
|
let max_message_size = i32::try_from(crate::max_payload(client)).unwrap_or(i32::MAX);
|
|
js.create_stream(output_config(agent, max_message_size))
|
|
.await
|
|
.map_err(|source| crate::Error::CreateAgentStream {
|
|
stream: name,
|
|
source,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn the_output_stream_caps_bytes_and_discards_the_oldest() {
|
|
let cfg = output_config("iris", 8_388_608);
|
|
assert_eq!(cfg.name, "term-bash-iris");
|
|
assert_eq!(cfg.subjects, vec!["$SWARM.term.iris.bash.>".to_owned()]);
|
|
assert_eq!(cfg.max_bytes, OUTPUT_MAX_BYTES);
|
|
assert_eq!(
|
|
cfg.discard,
|
|
async_nats::jetstream::stream::DiscardPolicy::Old
|
|
);
|
|
assert_eq!(cfg.max_message_size, 8_388_608);
|
|
}
|
|
}
|