252 lines
9 KiB
Rust
252 lines
9 KiB
Rust
//! The text an agent's `AGENTS.md` starts with: one key per agent in
|
|
//! `agents-md-initial` (`swarm_queue_client::agents_md_initial`), which this
|
|
//! daemon alone writes. `POST /api/agents` stores it before it queues the
|
|
//! creation; a destroy deletes it. The agent's harness reads its own key at
|
|
//! start and writes the file only when none exists.
|
|
|
|
use problem_details::ProblemDetails;
|
|
use swarm_queue_client::agents_md_initial::{self, Initial, MAX_BYTES};
|
|
|
|
use crate::routes::{AppState, error_problem};
|
|
use crate::schedules::{queue, read_or_default, write_revision};
|
|
|
|
/// How many times a losing writer re-reads and re-applies, as in
|
|
/// `wanted::MAX_ATTEMPTS`.
|
|
const MAX_ATTEMPTS: usize = 5;
|
|
|
|
fn queue_problem(e: &anyhow::Error) -> ProblemDetails {
|
|
tracing::warn!(error = %format!("{e:#}"), "agents-md initial text: queue operation failed");
|
|
error_problem(
|
|
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
|
&format!("{e:#}"),
|
|
)
|
|
}
|
|
|
|
/// Refuse a text longer than [`MAX_BYTES`].
|
|
pub(crate) fn check(content: &str) -> Result<(), String> {
|
|
if content.len() > MAX_BYTES {
|
|
return Err(format!(
|
|
"agents_md is {} bytes; the limit is {MAX_BYTES}",
|
|
content.len()
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// What [`store`] replaced, for [`restore`]; `None` when it wrote nothing.
|
|
pub(crate) type Replaced = Option<crate::schedules::Replaced<Initial>>;
|
|
|
|
/// Record that `agent`'s `AGENTS.md` starts with `content`.
|
|
pub(crate) async fn store(
|
|
client: &async_nats::Client,
|
|
agent: &str,
|
|
content: &str,
|
|
) -> Result<Replaced, ProblemDetails> {
|
|
let store = agents_md_initial::open_or_create(client)
|
|
.await
|
|
.map_err(|e| queue_problem(&e.into()))?;
|
|
let value = Initial {
|
|
content: content.to_owned(),
|
|
};
|
|
for _ in 0..MAX_ATTEMPTS {
|
|
let (current, revision, undecodable) = read_or_default::<Initial>(&store, agent)
|
|
.await
|
|
.map_err(|e| queue_problem(&e))?;
|
|
// An absent key reads as the default, which is not "stored": a first
|
|
// write of an empty text still has to land.
|
|
if current == value && !undecodable && revision.is_some() {
|
|
return Ok(None);
|
|
}
|
|
if let Some(written) = write_revision(&store, agent, &value, revision)
|
|
.await
|
|
.map_err(|e| queue_problem(&e))?
|
|
{
|
|
tracing::info!(%agent, bytes = content.len(), "set agent's initial AGENTS.md at create");
|
|
return Ok(Some(crate::schedules::Replaced::new(
|
|
revision.map(|_| current),
|
|
written,
|
|
)));
|
|
}
|
|
}
|
|
Err(queue_problem(&anyhow::anyhow!(
|
|
"gave up writing the initial AGENTS.md of {agent} after {MAX_ATTEMPTS} conflicting writes"
|
|
)))
|
|
}
|
|
|
|
/// Undo what [`store`] did, unless another write has landed since. Best
|
|
/// effort: a failure is logged, since the caller is already on its way to an
|
|
/// error of its own.
|
|
pub(crate) async fn restore(client: &async_nats::Client, agent: &str, replaced: Replaced) {
|
|
let Some(replaced) = replaced else { return };
|
|
let undone = match agents_md_initial::open_or_create(client).await {
|
|
Ok(store) => replaced.undo(&store, agent).await,
|
|
Err(e) => Err(anyhow::anyhow!("{e:?}")),
|
|
};
|
|
match undone {
|
|
Ok(true) => {}
|
|
Ok(false) => {
|
|
tracing::info!(%agent, "left the agent's initial AGENTS.md: another write landed since");
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(%agent, error = %format!("{e:#}"), "could not put back the agent's initial AGENTS.md");
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Delete a destroyed agent's key, so the next agent of that name does not
|
|
/// start from this one's text. A failure is logged and does not stop the
|
|
/// destroy, as in `agent_archetype::clear_on_destroy`.
|
|
pub(crate) async fn clear_on_destroy(state: &AppState, hive: &str, agent: &str) {
|
|
let cleared = async {
|
|
let client = queue(state).map_err(|(_, why)| anyhow::anyhow!(why))?;
|
|
agents_md_initial::open_or_create(&client)
|
|
.await?
|
|
.delete(agent)
|
|
.await?;
|
|
anyhow::Ok(())
|
|
}
|
|
.await;
|
|
match cleared {
|
|
Ok(()) => tracing::info!(hive, agent, "destroyed agent's initial AGENTS.md cleared"),
|
|
Err(e) => tracing::error!(
|
|
hive,
|
|
agent,
|
|
error = %format!("{e:#}"),
|
|
"clearing this agent's initial AGENTS.md FAILED; it was destroyed anyway, so a new \
|
|
agent of the same name starts from this text. Declaring the agent destroyed again \
|
|
re-runs this."
|
|
),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use axum::extract::{Path, State};
|
|
use axum::http::StatusCode;
|
|
|
|
use super::*;
|
|
use crate::test_nats::TestNats;
|
|
|
|
async fn published(client: &async_nats::Client, agent: &str) -> Option<String> {
|
|
let store = agents_md_initial::open_or_create(client)
|
|
.await
|
|
.expect("bucket");
|
|
store
|
|
.get(agent)
|
|
.await
|
|
.expect("read")
|
|
.map(|v| String::from_utf8(v.to_vec()).expect("utf-8"))
|
|
}
|
|
|
|
fn request(name: &str, agents_md: Option<String>) -> crate::routes::agents::CreateAgentRequest {
|
|
crate::routes::agents::CreateAgentRequest {
|
|
name: name.to_owned(),
|
|
hive: "pr1ma".to_owned(),
|
|
limits: None,
|
|
inference: None,
|
|
permissions: None,
|
|
archetype: None,
|
|
agents_md,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn a_text_over_the_limit_is_refused_and_one_at_it_passes() {
|
|
assert_eq!(check(&"a".repeat(MAX_BYTES)), Ok(()));
|
|
let why = check(&"a".repeat(MAX_BYTES + 1)).unwrap_err();
|
|
assert!(why.contains("the limit is"), "{why}");
|
|
}
|
|
|
|
/// `POST /api/agents` with `agents_md` stores it under the new agent's
|
|
/// key; a text over the limit is refused before anything is stored or
|
|
/// queued.
|
|
#[tokio::test]
|
|
async fn create_stores_the_text_it_is_given_and_refuses_one_too_big() {
|
|
let nats = TestNats::start();
|
|
let client = nats.client().await;
|
|
let state = crate::agent_limits::tests::declared(&client).await;
|
|
|
|
let problem = crate::routes::agents::create_agent(
|
|
State(state.clone()),
|
|
axum::Json(request("iris", Some("a".repeat(MAX_BYTES + 1)))),
|
|
)
|
|
.await
|
|
.expect_err("too big");
|
|
assert_eq!(problem.status, Some(StatusCode::BAD_REQUEST), "{problem:?}");
|
|
assert_eq!(published(&client, "iris").await, None);
|
|
let queued = state
|
|
.jobq
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.graph()
|
|
.nodes()
|
|
.count();
|
|
assert_eq!(queued, 0, "nothing is queued for a refused create");
|
|
|
|
let _ = crate::routes::agents::create_agent(
|
|
State(state.clone()),
|
|
axum::Json(request("iris", Some("# iris\n".to_owned()))),
|
|
)
|
|
.await
|
|
.expect("created");
|
|
assert_eq!(
|
|
published(&client, "iris").await.as_deref(),
|
|
Some(r##"{"content":"# iris\n"}"##)
|
|
);
|
|
|
|
// Control: no `agents_md` writes no key.
|
|
let _ = crate::routes::agents::create_agent(
|
|
State(state.clone()),
|
|
axum::Json(request("damocles", None)),
|
|
)
|
|
.await
|
|
.expect("created");
|
|
assert_eq!(published(&client, "damocles").await, None);
|
|
}
|
|
|
|
/// An empty text is a text: the key is written, so the agent gets an
|
|
/// empty file rather than none.
|
|
#[tokio::test]
|
|
async fn an_empty_text_is_stored() {
|
|
let nats = TestNats::start();
|
|
let client = nats.client().await;
|
|
let state = crate::agent_limits::tests::declared(&client).await;
|
|
let _ = crate::routes::agents::create_agent(
|
|
State(state.clone()),
|
|
axum::Json(request("iris", Some(String::new()))),
|
|
)
|
|
.await
|
|
.expect("created");
|
|
assert_eq!(
|
|
published(&client, "iris").await.as_deref(),
|
|
Some(r#"{"content":""}"#)
|
|
);
|
|
}
|
|
|
|
/// Destroying an agent deletes its key, so a new agent of that name does
|
|
/// not start from the old text.
|
|
#[tokio::test]
|
|
async fn a_recreated_agent_does_not_start_from_the_destroyed_ones_text() {
|
|
let nats = TestNats::start();
|
|
let client = nats.client().await;
|
|
let state = crate::agent_limits::tests::declared(&client).await;
|
|
let _ = crate::routes::agents::create_agent(
|
|
State(state.clone()),
|
|
axum::Json(request("atlas", Some("old".to_owned()))),
|
|
)
|
|
.await
|
|
.expect("created");
|
|
assert!(published(&client, "atlas").await.is_some());
|
|
|
|
let _ = crate::routes::agents::set_agent_state(
|
|
State(state.clone()),
|
|
Path(("pr1ma".to_owned(), "atlas".to_owned())),
|
|
axum::Json(crate::routes::agents::SetAgentStateRequest {
|
|
state: swarm_queue_client::wanted::AgentState::Destroyed,
|
|
}),
|
|
)
|
|
.await
|
|
.expect("destroyed");
|
|
assert_eq!(published(&client, "atlas").await, None);
|
|
}
|
|
}
|