set an agent's initial agents.md at create

swarm-controller/src/agents_md_initial.rs

@@ -0,0 +1,252 @@
+//! 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);
+    }
+}