433 lines
16 KiB
Rust
433 lines
16 KiB
Rust
//! Agent archetypes: the `agent-archetypes` record this daemon keeps from the
|
|
//! swarm builder's evaluation, and `GET /api/archetypes`, which the create
|
|
//! screen prefills from.
|
|
//!
|
|
//! Every [`REFRESH_INTERVAL`], starting once the queue is connected, the pass
|
|
//! asks the builder for the `meta` of its hyperhive source's archetypes, checks
|
|
//! each one the way `POST /api/agents` checks its settings, and writes the
|
|
//! result under [`BUILTIN_SOURCE`] when it differs from what is stored. A
|
|
//! failed or unanswered evaluation keeps the stored record.
|
|
|
|
use std::time::Duration;
|
|
|
|
use axum::Json;
|
|
use axum::extract::State;
|
|
use hive_types::Permission;
|
|
use problem_details::ProblemDetails;
|
|
use serde::{Deserialize, Serialize};
|
|
use swarm_queue_client::archetypes::{ArchetypeMeta, BUILTIN_SOURCE, SourceRecord};
|
|
use swarm_wire::builder::{
|
|
ARCHETYPES_DEADLINE, ARCHETYPES_SUBJECT, ArchetypesReply, EvalArchetypes,
|
|
};
|
|
use utoipa::ToSchema;
|
|
|
|
use crate::routes::{AppState, error_problem};
|
|
use crate::schedules::{queue, read_or_default, refused, write};
|
|
use crate::{agent_limits, agent_permissions, inference};
|
|
|
|
/// How often the builder is asked. Its source only changes when its host is
|
|
/// redeployed, so this bounds how long a redeploy takes to reach the create
|
|
/// screen.
|
|
const REFRESH_INTERVAL: Duration = Duration::from_mins(10);
|
|
|
|
/// One archetype's `meta` as a source flake declares it. Keys this daemon
|
|
/// does not know, `ui` among them, are ignored, so an archetype written for a
|
|
/// newer swarm still lists.
|
|
#[derive(Deserialize)]
|
|
struct DeclaredMeta {
|
|
description: String,
|
|
#[serde(default)]
|
|
permissions: Option<Vec<Permission>>,
|
|
#[serde(default)]
|
|
limits: Option<agent_limits::Caps>,
|
|
#[serde(default)]
|
|
inference: Option<inference::InferenceSettings>,
|
|
}
|
|
|
|
/// `meta` as stored, or why a create would refuse it.
|
|
fn check(name: &str, meta: serde_json::Value) -> Result<ArchetypeMeta, String> {
|
|
hive_types::Ident::parse(name).map_err(|reason| format!("name: {reason}"))?;
|
|
let declared: DeclaredMeta = serde_json::from_value(meta).map_err(|e| e.to_string())?;
|
|
if declared.description.trim().is_empty() {
|
|
return Err("description is empty".to_owned());
|
|
}
|
|
let permissions = match declared.permissions {
|
|
Some(perms) => {
|
|
agent_permissions::unique(&perms)?;
|
|
Some(
|
|
perms
|
|
.into_iter()
|
|
.map(|p| <&str>::from(p).to_owned())
|
|
.collect(),
|
|
)
|
|
}
|
|
None => None,
|
|
};
|
|
Ok(ArchetypeMeta {
|
|
description: declared.description,
|
|
permissions,
|
|
limits: declared
|
|
.limits
|
|
.map(agent_limits::Caps::into_stored)
|
|
.transpose()
|
|
.map_err(|e| format!("limits: {e}"))?,
|
|
inference: declared
|
|
.inference
|
|
.map(inference::InferenceSettings::into_stored)
|
|
.transpose()
|
|
.map_err(|e| format!("inference: {e}"))?,
|
|
})
|
|
}
|
|
|
|
/// The record to store for one evaluation: every archetype that passes
|
|
/// [`check`], and every one that does not with its reason.
|
|
fn record(source: String, archetypes: serde_json::Map<String, serde_json::Value>) -> SourceRecord {
|
|
let mut out = SourceRecord {
|
|
source,
|
|
..SourceRecord::default()
|
|
};
|
|
for (name, meta) in archetypes {
|
|
match check(&name, meta) {
|
|
Ok(meta) => {
|
|
out.archetypes.insert(name, meta);
|
|
}
|
|
Err(why) => {
|
|
out.refused.insert(name, why);
|
|
}
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Ask the builder once and answer what to store, or why there is nothing to.
|
|
async fn evaluate(client: &async_nats::Client) -> Result<SourceRecord, String> {
|
|
use async_nats::RequestErrorKind;
|
|
|
|
let body = serde_json::to_vec(&EvalArchetypes {}).map_err(|e| e.to_string())?;
|
|
let sent = async_nats::Request::new()
|
|
.payload(body.into())
|
|
.timeout(Some(ARCHETYPES_DEADLINE));
|
|
let msg = client
|
|
.send_request(ARCHETYPES_SUBJECT, sent)
|
|
.await
|
|
.map_err(|e| match e.kind() {
|
|
RequestErrorKind::NoResponders => "no swarm builder is listening: \
|
|
deploy.builder.enable is set on no host, or its service is not connected to \
|
|
the swarm queue"
|
|
.to_owned(),
|
|
_ => format!("asking the swarm builder failed: {e}"),
|
|
})?;
|
|
match serde_json::from_slice(&msg.payload) {
|
|
Ok(ArchetypesReply::Evaluated { source, archetypes }) => Ok(record(source, archetypes)),
|
|
Ok(ArchetypesReply::Failed { error_tail }) => {
|
|
Err(format!("the evaluation failed: {error_tail}"))
|
|
}
|
|
Err(e) => Err(format!("the builder's reply did not decode: {e}")),
|
|
}
|
|
}
|
|
|
|
/// Store `value` under [`BUILTIN_SOURCE`] unless it is already there.
|
|
/// Answers whether it wrote.
|
|
async fn store(client: &async_nats::Client, value: &SourceRecord) -> anyhow::Result<bool> {
|
|
let store = swarm_queue_client::archetypes::open_or_create(client).await?;
|
|
let (current, revision, undecodable) =
|
|
read_or_default::<SourceRecord>(&store, BUILTIN_SOURCE).await?;
|
|
if current == *value && !undecodable {
|
|
return Ok(false);
|
|
}
|
|
// A lost race is another write of the same evaluation; the next pass
|
|
// compares again.
|
|
write(&store, BUILTIN_SOURCE, value, revision).await
|
|
}
|
|
|
|
/// Refresh the record once the queue is connected, then every
|
|
/// [`REFRESH_INTERVAL`].
|
|
pub(crate) fn spawn(client: async_nats::Client) {
|
|
tokio::spawn(async move {
|
|
const CONNECT_POLL: Duration = Duration::from_secs(5);
|
|
loop {
|
|
while swarm_queue_client::ensure_connected(&client).is_err() {
|
|
tokio::time::sleep(CONNECT_POLL).await;
|
|
}
|
|
refresh(&client).await;
|
|
tokio::time::sleep(REFRESH_INTERVAL).await;
|
|
}
|
|
});
|
|
}
|
|
|
|
/// One evaluation, stored when it changed anything.
|
|
async fn refresh(client: &async_nats::Client) {
|
|
let value = match evaluate(client).await {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
tracing::warn!(%error, "archetypes: kept the stored list");
|
|
return;
|
|
}
|
|
};
|
|
for (name, why) in &value.refused {
|
|
tracing::warn!(archetype = %name, %why, "archetypes: refused");
|
|
}
|
|
match store(client, &value).await {
|
|
Ok(true) => tracing::info!(
|
|
count = value.archetypes.len(),
|
|
source = %value.source,
|
|
"archetypes: stored"
|
|
),
|
|
Ok(false) => {}
|
|
Err(e) => tracing::warn!(error = %format!("{e:#}"), "archetypes: storing failed"),
|
|
}
|
|
}
|
|
|
|
/// One archetype the create screen can offer, with the settings it prefills.
|
|
#[derive(Debug, PartialEq, Eq, Serialize, ToSchema)]
|
|
pub(crate) struct ArchetypeView {
|
|
/// The source flake's key, `hyperhive` for the built-ins.
|
|
source: String,
|
|
name: String,
|
|
description: String,
|
|
/// Permission names; absent leaves the default set.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
permissions: Option<Vec<String>>,
|
|
/// As `PUT /api/agents/{name}/limits` takes them.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
#[schema(value_type = Option<Object>)]
|
|
limits: Option<swarm_queue_client::agent_limits::AgentCaps>,
|
|
/// As `PUT /api/agents/{name}/inference` takes them.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
#[schema(value_type = Option<Object>)]
|
|
inference: Option<swarm_wire::agent_inference::AgentInference>,
|
|
}
|
|
|
|
/// One archetype a source declared and this daemon refused.
|
|
#[derive(Debug, PartialEq, Eq, Serialize, ToSchema)]
|
|
pub(crate) struct RefusedView {
|
|
source: String,
|
|
name: String,
|
|
reason: String,
|
|
}
|
|
|
|
#[derive(Debug, Default, PartialEq, Eq, Serialize, ToSchema)]
|
|
pub(crate) struct Archetypes {
|
|
archetypes: Vec<ArchetypeView>,
|
|
refused: Vec<RefusedView>,
|
|
}
|
|
|
|
fn view(key: &str, record: SourceRecord) -> Archetypes {
|
|
let mut out = Archetypes::default();
|
|
for (name, meta) in record.archetypes {
|
|
out.archetypes.push(ArchetypeView {
|
|
source: key.to_owned(),
|
|
name,
|
|
description: meta.description,
|
|
permissions: meta.permissions,
|
|
limits: meta.limits,
|
|
inference: meta.inference,
|
|
});
|
|
}
|
|
for (name, reason) in record.refused {
|
|
out.refused.push(RefusedView {
|
|
source: key.to_owned(),
|
|
name,
|
|
reason,
|
|
});
|
|
}
|
|
out
|
|
}
|
|
|
|
#[utoipa::path(
|
|
get,
|
|
path = "/api/archetypes",
|
|
responses(
|
|
(status = 200, description = "every archetype the swarm builder's evaluation passed, \
|
|
and every one it refused; both empty before the first evaluation", body = Archetypes),
|
|
(status = 503, description = "no swarm queue is configured on this host, or it is \
|
|
not connected (problem+json)", body = String),
|
|
(status = 500, description = "the archetypes could not be read (problem+json)",
|
|
body = String),
|
|
),
|
|
tag = "agents"
|
|
)]
|
|
pub(crate) async fn get_archetypes(
|
|
State(state): State<AppState>,
|
|
) -> Result<Json<Archetypes>, ProblemDetails> {
|
|
let client = queue(&state).map_err(refused)?;
|
|
let failed = |e: anyhow::Error| {
|
|
tracing::warn!(error = %format!("{e:#}"), "archetypes: read failed");
|
|
error_problem(
|
|
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
|
&format!("{e:#}"),
|
|
)
|
|
};
|
|
let Some(store) = swarm_queue_client::archetypes::open(&client)
|
|
.await
|
|
.map_err(|e| failed(e.into()))?
|
|
else {
|
|
return Ok(Json(Archetypes::default()));
|
|
};
|
|
let (record, _, _) = read_or_default::<SourceRecord>(&store, BUILTIN_SOURCE)
|
|
.await
|
|
.map_err(failed)?;
|
|
Ok(Json(view(BUILTIN_SOURCE, record)))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
fn evaluated(archetypes: serde_json::Value) -> SourceRecord {
|
|
let serde_json::Value::Object(map) = archetypes else {
|
|
panic!("an object")
|
|
};
|
|
record("/nix/store/h".to_owned(), map)
|
|
}
|
|
|
|
#[test]
|
|
fn a_valid_archetype_is_stored_with_its_settings() {
|
|
let r = evaluated(json!({
|
|
"reviewer": {
|
|
"description": "reviews",
|
|
"permissions": ["messaging", "meta", "inbox"],
|
|
"limits": { "memory_max": "4G" },
|
|
"inference": { "effort": "high" },
|
|
"ui": { "anything": true }
|
|
}
|
|
}));
|
|
assert!(r.refused.is_empty(), "{:?}", r.refused);
|
|
let meta = &r.archetypes["reviewer"];
|
|
assert_eq!(meta.description, "reviews");
|
|
assert_eq!(
|
|
meta.permissions.as_deref(),
|
|
Some(
|
|
&[
|
|
"messaging".to_owned(),
|
|
"meta".to_owned(),
|
|
"inbox".to_owned()
|
|
][..]
|
|
)
|
|
);
|
|
assert_eq!(
|
|
meta.limits.as_ref().and_then(|l| l.memory_max.as_deref()),
|
|
Some("4G")
|
|
);
|
|
assert!(meta.inference.as_ref().is_some_and(|i| i.effort.is_some()));
|
|
}
|
|
|
|
#[test]
|
|
fn a_bad_archetype_is_refused_by_name_and_the_rest_stay() {
|
|
let r = evaluated(json!({
|
|
"ok": { "description": "fine" },
|
|
"no-description": { "description": " " },
|
|
"bad-permission": { "description": "x", "permissions": ["root"] },
|
|
"twice": { "description": "x", "permissions": ["meta", "meta"] },
|
|
"bad-limits": { "description": "x", "limits": { "cpu_quota": "lots" } },
|
|
"typo-limits": { "description": "x", "limits": { "memory": "4G" } },
|
|
"bad-inference": { "description": "x", "inference": { "effort": "huge" } },
|
|
"Bad_Name": { "description": "x" }
|
|
}));
|
|
assert_eq!(r.archetypes.keys().collect::<Vec<_>>(), ["ok"]);
|
|
let refused: Vec<&str> = r.refused.keys().map(String::as_str).collect();
|
|
assert_eq!(
|
|
refused,
|
|
[
|
|
"Bad_Name",
|
|
"bad-inference",
|
|
"bad-limits",
|
|
"bad-permission",
|
|
"no-description",
|
|
"twice",
|
|
"typo-limits"
|
|
]
|
|
);
|
|
assert!(r.refused["bad-limits"].starts_with("limits: "));
|
|
assert!(r.refused["twice"].contains("listed twice"));
|
|
}
|
|
|
|
/// Against a real server: no builder listening names the option; a
|
|
/// listener's evaluation is checked and stored, and storing the same
|
|
/// record again writes nothing.
|
|
#[tokio::test]
|
|
async fn an_evaluation_is_stored_and_rewritten_only_when_it_changes() {
|
|
use futures_util::StreamExt as _;
|
|
|
|
let nats = crate::test_nats::TestNats::start();
|
|
let client = nats.client().await;
|
|
let unanswered = evaluate(&client).await;
|
|
assert!(
|
|
matches!(&unanswered, Err(e) if e.contains("deploy.builder.enable")),
|
|
"{unanswered:?}"
|
|
);
|
|
|
|
let builder = nats.client().await;
|
|
let mut requests = builder
|
|
.queue_subscribe(
|
|
ARCHETYPES_SUBJECT,
|
|
swarm_wire::builder::QUEUE_GROUP.to_owned(),
|
|
)
|
|
.await
|
|
.expect("subscribe");
|
|
// A message the builder sends itself and receives back proves the
|
|
// server processed the subscription before it.
|
|
let mut probe = builder
|
|
.subscribe("archetypes-test.probe")
|
|
.await
|
|
.expect("subscribe");
|
|
builder
|
|
.publish("archetypes-test.probe", "".into())
|
|
.await
|
|
.expect("publish");
|
|
probe.next().await.expect("the probe comes back");
|
|
tokio::spawn(async move {
|
|
let msg = requests.next().await.expect("a request");
|
|
let serde_json::Value::Object(archetypes) = json!({
|
|
"reviewer": { "description": "reviews", "permissions": ["meta"] },
|
|
"broken": { "description": "" }
|
|
}) else {
|
|
unreachable!()
|
|
};
|
|
let reply = serde_json::to_vec(&ArchetypesReply::Evaluated {
|
|
source: "/nix/store/h".to_owned(),
|
|
archetypes,
|
|
})
|
|
.expect("encodes");
|
|
builder
|
|
.publish(msg.reply.expect("a reply subject"), reply.into())
|
|
.await
|
|
.expect("publish");
|
|
});
|
|
let value = evaluate(&client).await.expect("evaluated");
|
|
assert_eq!(value.archetypes.keys().collect::<Vec<_>>(), ["reviewer"]);
|
|
assert_eq!(value.refused.keys().collect::<Vec<_>>(), ["broken"]);
|
|
|
|
assert!(store(&client, &value).await.expect("stored"));
|
|
assert!(!store(&client, &value).await.expect("compared"));
|
|
let changed = SourceRecord {
|
|
source: "/nix/store/h2".to_owned(),
|
|
..value
|
|
};
|
|
assert!(store(&client, &changed).await.expect("stored"));
|
|
let bucket = swarm_queue_client::archetypes::open(&client)
|
|
.await
|
|
.expect("opens")
|
|
.expect("exists");
|
|
let (stored, _, _) = read_or_default::<SourceRecord>(&bucket, BUILTIN_SOURCE)
|
|
.await
|
|
.expect("reads");
|
|
assert_eq!(stored, changed);
|
|
}
|
|
|
|
#[test]
|
|
fn the_view_lists_each_archetype_under_its_source() {
|
|
let r = evaluated(json!({
|
|
"reviewer": { "description": "reviews" },
|
|
"broken": { "description": "" }
|
|
}));
|
|
let v = view(BUILTIN_SOURCE, r);
|
|
assert_eq!(v.archetypes.len(), 1);
|
|
assert_eq!(v.archetypes[0].source, "hyperhive");
|
|
assert_eq!(v.archetypes[0].name, "reviewer");
|
|
assert_eq!(v.archetypes[0].permissions, None);
|
|
assert_eq!(v.refused.len(), 1);
|
|
assert_eq!(v.refused[0].name, "broken");
|
|
}
|
|
}
|