hyperhive/swarm-controller/src/agents_md.rs

//! `GET` and `PUT /api/agents/{name}/agents-md`: an agent's `AGENTS.md`.
//!
//! **The file in the agent's state dir is the truth, and only its harness
//! writes it.** A read serves the copy the harness publishes into the
//! `agents-md` bucket at start and on every change
//! ([`swarm_queue_client::agents_md`]), so it answers for a stopped agent
//! too. An edit is a git diff sent to the harness as a command on its
//! control subject ([`crate::agent_control`]); the harness applies it to the
//! file and says whether it did, so an edit needs the agent running.
//!
//! No hive is on either path.

use axum::Json;
use axum::extract::{Path, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse as _, Response};
use problem_details::ProblemDetails;
use swarm_controller_api::agents_md::{AgentsMdEdit, AgentsMdSaved, AgentsMdView};
use swarm_queue_client::AgentControl;
use swarm_queue_client::agents_md::{Apply, ApplyReply, Mirror};

use crate::agent_control::{ControlSink, Outcome};
use crate::routes::{AppState, error_problem, wanted_error_status};

/// Reads the `agents-md` bucket, the way [`crate::agent_icon::AgentIconReader`]
/// reads the icon bucket, and creates it: the agents writing it cannot.
pub struct AgentsMdReader {
    client: async_nats::Client,
    store: tokio::sync::OnceCell<async_nats::jetstream::kv::Store>,
}

impl AgentsMdReader {
    #[must_use]
    pub fn new(client: async_nats::Client) -> Self {
        Self {
            client,
            store: tokio::sync::OnceCell::new(),
        }
    }

    async fn store(
        &self,
    ) -> std::result::Result<&async_nats::jetstream::kv::Store, swarm_queue_client::Error> {
        self.store
            .get_or_try_init(|| swarm_queue_client::agents_md::open_or_create(&self.client))
            .await
    }

    /// Open the bucket, creating it, as soon as the queue is connected. Until
    /// then every agent's write is refused and retried.
    pub async fn create_when_connected(&self) {
        const POLL: std::time::Duration = std::time::Duration::from_secs(5);
        loop {
            if swarm_queue_client::ensure_connected(&self.client).is_ok() {
                match self.store().await {
                    Ok(_) => return,
                    Err(e) => tracing::warn!(
                        error = swarm_queue_client::chain(&e),
                        "creating the agents-md bucket failed; retrying"
                    ),
                }
            }
            tokio::time::sleep(POLL).await;
        }
    }

    /// `agent`'s published `AGENTS.md`, or `None` when it has published none.
    async fn get(&self, agent: &str) -> anyhow::Result<Option<Mirror>> {
        use anyhow::Context as _;

        // An unconnected client hangs on a JetStream request rather than
        // failing it.
        swarm_queue_client::ensure_connected(&self.client)?;
        let store = self.store().await?;
        let Some(bytes) = store
            .get(agent)
            .await
            .with_context(|| format!("reading the agents-md entry for {agent}"))?
        else {
            return Ok(None);
        };
        let mirror = serde_json::from_slice(&bytes)
            .with_context(|| format!("the agents-md entry for {agent} is not a mirror"))?;
        Ok(Some(mirror))
    }
}

/// Read an agent's `AGENTS.md` as its harness last published it.
#[utoipa::path(
    get,
    path = "/api/agents/{name}/agents-md",
    params(("name" = String, Path, description = "agent name")),
    responses(
        (status = 200, description = "the file's text and revision, as the agent last \
            published them; answered while the agent is stopped too", body = AgentsMdView),
        (status = 400, description = "not an identifier (problem+json)", body = String),
        (status = 404, description = "this agent has published no AGENTS.md: its harness \
            has not run since it started mirroring the file (problem+json)", body = String),
        (status = 500, description = "the bucket could not be read (problem+json)", body = String),
        (status = 503, description = "no swarm queue is wired up, or it is not connected \
            (problem+json)", body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn get_agents_md(
    State(state): State<AppState>,
    Path(name): Path<String>,
) -> Result<Json<AgentsMdView>, ProblemDetails> {
    let Some(reader) = state.agents_md.as_ref() else {
        return Err(error_problem(
            StatusCode::SERVICE_UNAVAILABLE,
            "no swarm queue is wired up on this host",
        ));
    };
    let agent = hive_types::Ident::parse(&name)
        .map_err(|reason| error_problem(StatusCode::BAD_REQUEST, reason))?
        .into_string();
    let mirror = reader.get(&agent).await.map_err(|e| {
        tracing::warn!(%agent, error = %format!("{e:#}"), "reading the agents-md entry failed");
        error_problem(wanted_error_status(&e), &format!("{e:#}"))
    })?;
    let Some(Mirror { content, revision }) = mirror else {
        return Err(error_problem(
            StatusCode::NOT_FOUND,
            "this agent has published no AGENTS.md",
        ));
    };
    Ok(Json(AgentsMdView { content, revision }))
}

/// The revision a 409 names, as a problem+json extension member.
#[derive(serde::Serialize)]
struct CurrentRevision {
    revision: String,
}

/// Edit an agent's `AGENTS.md` with a git diff, which the agent's harness
/// applies to the file.
#[utoipa::path(
    put,
    path = "/api/agents/{name}/agents-md",
    params(("name" = String, Path, description = "agent whose AGENTS.md to edit")),
    request_body = AgentsMdEdit,
    responses(
        (status = 200, description = "the diff is applied; the body is the file's new \
            revision", body = AgentsMdSaved),
        (status = 400, description = "not an identifier (problem+json)", body = String),
        (status = 401, description = "the request carries no `Remote-User` (problem+json)",
            body = String),
        (status = 409, description = "the file is no longer at `base_revision`, so nothing \
            was applied; `revision` in the problem+json body is the one it is at", body = String),
        (status = 422, description = "`git apply` refused the diff, or it touches something \
            other than AGENTS.md; `detail` is git's message and the file is unchanged \
            (problem+json)", body = String),
        (status = 503, description = "the agent is not running, or its harness is not \
            connected to the swarm queue, or no swarm queue is configured on this host or it \
            could not be reached (problem+json)", body = String),
        (status = 504, description = "the harness held the edit and did not answer in time, \
            so whether it was applied is unknown (problem+json)", body = String),
        (status = 500, description = "the queue refused or failed the request, or the harness \
            could not read or write the file (problem+json)", body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn put_agents_md(
    State(state): State<AppState>,
    Path(agent): Path<String>,
    headers: HeaderMap,
    Json(edit): Json<AgentsMdEdit>,
) -> Result<Json<AgentsMdSaved>, Response> {
    let problem = |s: StatusCode, d: &str| error_problem(s, d).into_response();
    let (agent, from) =
        crate::agent_control::target(&headers, &agent).map_err(|(s, d)| problem(s, d))?;
    let Some(status) = state.status.as_ref() else {
        return Err(problem(
            StatusCode::SERVICE_UNAVAILABLE,
            "no swarm queue is configured on this host",
        ));
    };
    let client = status.queue_client();
    swarm_queue_client::ensure_connected(&client).map_err(|e| {
        problem(
            StatusCode::SERVICE_UNAVAILABLE,
            &swarm_queue_client::chain(&e),
        )
    })?;
    send(&client, &agent, &from, edit).await
}

async fn send(
    sink: &impl ControlSink,
    agent: &str,
    from: &str,
    edit: AgentsMdEdit,
) -> Result<Json<AgentsMdSaved>, Response> {
    let problem = |s: StatusCode, d: &str| error_problem(s, d).into_response();
    let subject = swarm_queue_client::agent_control_subject(agent);
    let command = AgentControl::ApplyAgentsMd(Apply {
        base_revision: edit.base_revision,
        diff: edit.diff,
    });
    let payload = serde_json::to_vec(&command).map_err(|e| {
        problem(
            StatusCode::INTERNAL_SERVER_ERROR,
            &format!("encoding the edit failed: {e}"),
        )
    })?;
    let reply = match sink.request(subject.clone(), payload).await {
        Outcome::Answered(reply) => serde_json::from_slice::<ApplyReply>(&reply).map_err(|e| {
            problem(
                StatusCode::INTERNAL_SERVER_ERROR,
                &format!("{agent}'s answer was unreadable: {e}"),
            )
        })?,
        Outcome::NotRunning => {
            return Err(problem(
                StatusCode::SERVICE_UNAVAILABLE,
                &format!(
                    "{agent} is not running: its harness is not listening for commands, so \
                     AGENTS.md cannot be edited until it is"
                ),
            ));
        }
        Outcome::Unconfirmed => {
            tracing::warn!(%agent, %from, "agents-md edit unanswered");
            return Err(problem(
                StatusCode::GATEWAY_TIMEOUT,
                &format!("{agent} did not answer in time: whether the edit was applied is unknown"),
            ));
        }
        Outcome::Failed(e) => {
            tracing::warn!(%agent, %subject, error = %e, "agents-md edit failed");
            return Err(problem(
                StatusCode::INTERNAL_SERVER_ERROR,
                &format!("sending to {subject} failed: {e}"),
            ));
        }
    };
    match reply {
        ApplyReply::Applied { revision } => {
            tracing::info!(%agent, %from, %revision, "agents-md edit applied");
            Ok(Json(AgentsMdSaved { revision }))
        }
        ApplyReply::Stale { revision } => {
            Err(ProblemDetails::from_status_code(StatusCode::CONFLICT)
                .with_detail(format!(
                    "{agent}'s AGENTS.md changed since the edit was made: it is at {revision}. \
                     Reload it and make the edit again"
                ))
                .with_extensions(CurrentRevision { revision })
                .into_response())
        }
        ApplyReply::Rejected { error } => Err(problem(
            StatusCode::UNPROCESSABLE_ENTITY,
            &format!("the diff does not apply to {agent}'s AGENTS.md: {error}"),
        )),
        ApplyReply::Failed { error } => {
            tracing::warn!(%agent, %error, "agents-md edit failed in the harness");
            Err(problem(
                StatusCode::INTERNAL_SERVER_ERROR,
                &format!("{agent} could not apply the edit: {error}"),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use axum::http::StatusCode;
    use swarm_controller_api::agents_md::AgentsMdEdit;
    use swarm_queue_client::AgentControl;
    use swarm_queue_client::agents_md::{Apply, ApplyReply};

    use super::send;
    use crate::agent_control::{ControlSink, Outcome};

    /// Answers every command with one fixed outcome and records what it got.
    struct FakeSink {
        answer: Box<dyn Fn() -> Outcome + Send + Sync>,
        got: Mutex<Vec<(String, Vec<u8>)>>,
    }

    impl FakeSink {
        fn replying(reply: &ApplyReply) -> Self {
            let bytes = serde_json::to_vec(reply).expect("encodes");
            Self::new(move || Outcome::Answered(bytes.clone()))
        }

        fn new(answer: impl Fn() -> Outcome + Send + Sync + 'static) -> Self {
            Self {
                answer: Box::new(answer),
                got: Mutex::default(),
            }
        }
    }

    impl ControlSink for FakeSink {
        async fn request(&self, subject: String, payload: Vec<u8>) -> Outcome {
            self.got
                .lock()
                .expect("unpoisoned")
                .push((subject, payload));
            (self.answer)()
        }
    }

    fn edit() -> AgentsMdEdit {
        AgentsMdEdit {
            base_revision: "abc".to_owned(),
            diff: "--- a/AGENTS.md\n".to_owned(),
        }
    }

    async fn body(resp: axum::response::Response) -> serde_json::Value {
        let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .expect("body");
        serde_json::from_slice(&bytes).expect("json")
    }

    #[tokio::test]
    async fn an_edit_goes_to_the_agent_s_control_subject_and_its_revision_comes_back() {
        let sink = FakeSink::replying(&ApplyReply::Applied {
            revision: "def".to_owned(),
        });
        let saved = send(&sink, "atlas", "mara", edit()).await.expect("applied");
        assert_eq!(saved.0.revision, "def");
        let got = sink.got.lock().expect("unpoisoned");
        assert_eq!(got[0].0, "$SWARM.control.atlas");
        let sent: AgentControl = serde_json::from_slice(&got[0].1).expect("a command");
        assert_eq!(
            sent,
            AgentControl::ApplyAgentsMd(Apply {
                base_revision: "abc".to_owned(),
                diff: "--- a/AGENTS.md\n".to_owned(),
            })
        );
    }

    #[tokio::test]
    async fn a_stale_base_is_409_naming_the_current_revision() {
        let sink = FakeSink::replying(&ApplyReply::Stale {
            revision: "fff".to_owned(),
        });
        let resp = send(&sink, "atlas", "mara", edit())
            .await
            .expect_err("stale");
        assert_eq!(resp.status(), StatusCode::CONFLICT);
        assert_eq!(body(resp).await["revision"], "fff");
    }

    #[tokio::test]
    async fn a_diff_git_refuses_is_422_with_git_s_message() {
        let sink = FakeSink::replying(&ApplyReply::Rejected {
            error: "error: patch failed: AGENTS.md:1".to_owned(),
        });
        let resp = send(&sink, "atlas", "mara", edit())
            .await
            .expect_err("rejected");
        assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
        let detail = body(resp).await["detail"]
            .as_str()
            .expect("detail")
            .to_owned();
        assert!(detail.contains("patch failed: AGENTS.md:1"), "{detail}");
    }

    #[tokio::test]
    async fn no_subscriber_is_503_and_the_rest_are_504_and_500() {
        let status = |sink: FakeSink| async move {
            send(&sink, "atlas", "mara", edit())
                .await
                .expect_err("not a success")
                .status()
        };
        assert_eq!(
            status(FakeSink::new(|| Outcome::NotRunning)).await,
            StatusCode::SERVICE_UNAVAILABLE
        );
        assert_eq!(
            status(FakeSink::new(|| Outcome::Unconfirmed)).await,
            StatusCode::GATEWAY_TIMEOUT
        );
        assert_eq!(
            status(FakeSink::new(|| Outcome::Failed("refused".to_owned()))).await,
            StatusCode::INTERNAL_SERVER_ERROR
        );
        assert_eq!(
            status(FakeSink::replying(&ApplyReply::Failed {
                error: "disk full".to_owned()
            }))
            .await,
            StatusCode::INTERNAL_SERVER_ERROR
        );
        // A cancel's answer is not an edit's.
        assert_eq!(
            status(FakeSink::new(|| Outcome::Answered(
                br#"{"acted":true}"#.to_vec()
            )))
            .await,
            StatusCode::INTERNAL_SERVER_ERROR
        );
    }

    /// Through a real `nats-server`: nothing is served before the harness
    /// publishes, and its publish, straight to the key's subject as the
    /// harness sends it, is what a read returns.
    #[tokio::test]
    async fn a_read_serves_what_the_harness_published() {
        use swarm_queue_client::agents_md::{Mirror, subject};

        let nats = crate::test_nats::TestNats::start();
        let client = nats.client().await;
        let reader = super::AgentsMdReader::new(client.clone());
        reader.create_when_connected().await;
        assert_eq!(reader.get("atlas").await.expect("read"), None);

        let mirror = Mirror {
            content: "# atlas\n".to_owned(),
            revision: "abc".to_owned(),
        };
        let js = async_nats::jetstream::new(client.clone());
        js.publish(
            subject("atlas"),
            serde_json::to_vec(&mirror).expect("encodes").into(),
        )
        .await
        .expect("sent")
        .await
        .expect("acked");
        assert_eq!(reader.get("atlas").await.expect("read"), Some(mirror));
        assert_eq!(reader.get("iris").await.expect("read"), None);
    }
}