hyperhive/swarm-queue-client/src/agents_md_initial.rs

//! The text an agent's `AGENTS.md` starts with: one key per agent in the
//! `agents-md-initial` KV bucket, set when the agent is created.
//!
//! **swarm-controller is the only writer**: `POST /api/agents` stores the key
//! before it queues the agent's creation, and a destroy deletes it.
//!
//! **The agent's harness reads its own key** when it starts, and writes the
//! text to `AGENTS.md` only when the file does not exist. The file in the
//! agent's state dir stays the truth ([`crate::agents_md`]); this key is not
//! read again once the file exists, and later edits never reach it.
//!
//! Keyed by the agent alone, like [`crate::agent_archetype`]. The agent's
//! credential is granted one direct get on its own key and nothing else of
//! the bucket.

/// The KV bucket the initial texts are stored in, keyed by agent name.
pub const BUCKET: &str = "agents-md-initial";

/// The stream backing [`BUCKET`]. `create_key_value` prefixes `KV_`.
pub const STREAM: &str = "KV_agents-md-initial";

/// The longest initial text, in bytes. The agent mirrors the file into
/// [`crate::agents_md`] whole, so the limit keeps it well under the queue's
/// payload limit.
pub const MAX_BYTES: usize = 64 * 1024;

/// The subject `agent`'s entry is stored under.
#[must_use]
pub fn key_subject(agent: &str) -> String {
    format!("$KV.{BUCKET}.{agent}")
}

/// The request subject that reads `agent`'s current entry.
#[must_use]
pub fn direct_get_subject(agent: &str) -> String {
    format!("$JS.API.DIRECT.GET.{STREAM}.{}", key_subject(agent))
}

/// The value under an agent's key.
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Initial {
    /// The text `AGENTS.md` starts with.
    pub content: String,
}

/// Open the bucket for writing, creating it if nothing has yet.
///
/// **Controller-side only.** `history: 1`: the agent wants the current value.
#[cfg(feature = "kv")]
pub async fn open_or_create(
    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(
        &js,
        BUCKET,
        "agents-md-initial",
        "The text each agent's AGENTS.md starts with, set by the swarm controller".to_owned(),
    )
    .await
}

/// Read `agent`'s own entry with one direct get. `None` when the key is
/// absent or deleted, or when no bucket exists yet.
///
/// # Errors
///
/// A failed request or an undecodable value. A request the queue refused for
/// permissions reaches the client as a timeout, not as a refusal.
#[cfg(feature = "agents-md-initial")]
pub async fn read_own(
    client: &async_nats::Client,
    agent: &str,
) -> Result<Option<Initial>, Box<dyn std::error::Error + Send + Sync>> {
    use async_nats::jetstream::stream::DirectGetErrorKind;

    let js = async_nats::jetstream::new(client.clone());
    let stream = js.get_stream_no_info(STREAM).await?;
    let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
        Ok(message) => message,
        Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => return Ok(None),
        Err(e)
            if matches!(
                e.kind(),
                DirectGetErrorKind::ErrorResponse(status, _)
                    if status == async_nats::StatusCode::NO_RESPONDERS
            ) =>
        {
            return Ok(None);
        }
        Err(e) => return Err(e.into()),
    };
    Ok(decode(
        message
            .headers
            .get("KV-Operation")
            .map(async_nats::HeaderValue::as_str),
        &message.payload,
    )?)
}

/// One stored message as a value: a KV delete or purge is stored as an empty
/// marker carrying the `KV-Operation` header, and reads as no entry.
#[cfg(feature = "agents-md-initial")]
fn decode(operation: Option<&str>, payload: &[u8]) -> serde_json::Result<Option<Initial>> {
    if matches!(operation, Some("DEL" | "PURGE")) {
        return Ok(None);
    }
    serde_json::from_slice(payload).map(Some)
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The names the grant in `swarm-nats-auth` is built from.
    #[test]
    fn names_match_the_queue_grant() {
        assert_eq!(STREAM, format!("KV_{BUCKET}"));
        assert_eq!(key_subject("iris"), "$KV.agents-md-initial.iris");
        assert_eq!(
            direct_get_subject("iris"),
            "$JS.API.DIRECT.GET.KV_agents-md-initial.$KV.agents-md-initial.iris"
        );
    }

    /// A stored text reads back, a delete is no entry.
    #[cfg(feature = "agents-md-initial")]
    #[test]
    fn decoding_reads_the_text_and_a_delete_as_none() {
        assert_eq!(
            decode(None, br##"{"content":"# hi\n","later":1}"##).unwrap(),
            Some(Initial {
                content: "# hi\n".to_owned()
            })
        );
        assert_eq!(decode(Some("DEL"), b"").unwrap(), None);
        assert_eq!(decode(Some("PURGE"), b"").unwrap(), None);
    }
}