botstadia-interface 0.2.0

//! A bot's bundle: what a builder uploads.
//!
//! One deterministic ZIP holding `component.wasm` (the bot, a component for
//! the `agent` world), the manifest `botstadia.toml`, and optional ONNX models
//! under `artifacts/`. A trained-policy bundle's `component.wasm` is its own
//! copy of the policy shell.
//!
//! ```toml
//! format = 1
//! arena = "dance-off"
//! payload_schema_version = 1
//! kind = "policy"            # policy | rust | c
//! nonce = "second-try"       # optional: the same code as a new bot
//!
//! [[model]]                  # in order: its index is `infer`'s `model`
//! name = "policy"
//! path = "artifacts/policy.onnx"
//! ```
//!
//! The manifest's model order is the `model` index the `agent` world's
//! `infer` takes. The ZIP is built deterministically (entries in name order,
//! fixed times and permissions, one compression), so the same files always
//! give the same bundle digest.

use serde::{Deserialize, Serialize};

/// The manifest's file name inside a bundle.
pub const MANIFEST: &str = "botstadia.toml";
/// The bot's program inside a bundle.
pub const COMPONENT: &str = "component.wasm";
/// Where models sit inside a bundle.
pub const ARTIFACTS_DIR: &str = "artifacts";
/// The manifest format this crate reads and writes.
pub const MANIFEST_FORMAT: u32 = 1;
/// The largest bundle, in bytes.
pub const MAX_BUNDLE_BYTES: u64 = 25_000_000;
/// The most ONNX models one bundle may carry.
pub const MAX_MODELS: usize = 8;

/// How a bot was made.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BotKind {
    /// An ONNX policy run by the generic policy shell.
    Policy,
    Rust,
    C,
}

impl BotKind {
    pub fn name(self) -> &'static str {
        match self {
            Self::Policy => "policy",
            Self::Rust => "rust",
            Self::C => "c",
        }
    }

    /// What the site and the CLI call it.
    pub fn label(self) -> &'static str {
        match self {
            Self::Policy => "Trained policy",
            Self::Rust => "Rust",
            Self::C => "C",
        }
    }
}

/// One ONNX model in a bundle.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ModelEntry {
    pub name: String,
    /// Inside the bundle, under `artifacts/`.
    pub path: String,
}

/// `botstadia.toml`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Manifest {
    pub format: u32,
    /// The arena's slug.
    pub arena: String,
    pub payload_schema_version: u32,
    pub kind: BotKind,
    /// Any text, so the same code can go up again as a new bot.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub nonce: Option<String>,
    /// In order: a model's index here is `infer`'s `model` argument.
    #[serde(default, rename = "model", skip_serializing_if = "Vec::is_empty")]
    pub models: Vec<ModelEntry>,
}

impl Manifest {
    pub fn new(arena: &str, payload_schema_version: u32, kind: BotKind) -> Self {
        Self {
            format: MANIFEST_FORMAT,
            arena: arena.into(),
            payload_schema_version,
            kind,
            nonce: None,
            models: Vec::new(),
        }
    }

    #[cfg(feature = "bundle")]
    pub fn to_toml(&self) -> String {
        toml::to_string(self).expect("a manifest always serialises")
    }

    #[cfg(feature = "bundle")]
    pub fn from_toml(text: &str) -> Result<Self, String> {
        toml::from_str(text).map_err(|error| format!("{MANIFEST}: {error}"))
    }

    /// Everything wrong with the manifest on its own. Empty means it's sound.
    pub fn problems(&self) -> Vec<String> {
        let mut problems = Vec::new();
        if self.format != MANIFEST_FORMAT {
            problems.push(format!(
                "manifest format {} isn't {MANIFEST_FORMAT}",
                self.format
            ));
        }
        if !crate::declaration::is_slug(&self.arena) {
            problems.push(format!("arena `{}` isn't a slug", self.arena));
        }
        if self.models.len() > MAX_MODELS {
            problems.push(format!(
                "{} models, but a bundle carries at most {MAX_MODELS}",
                self.models.len()
            ));
        }
        if self.kind == BotKind::Policy && self.models.is_empty() {
            problems.push("a trained policy needs its model".into());
        }
        let mut seen = std::collections::HashSet::new();
        for model in &self.models {
            if !seen.insert(model.name.as_str()) {
                problems.push(format!("model `{}` is listed twice", model.name));
            }
            if !model_path_ok(&model.path) {
                problems.push(format!(
                    "model `{}`: `{}` must be a plain name ending .onnx under {ARTIFACTS_DIR}/",
                    model.name, model.path
                ));
            }
        }
        if self.nonce.as_ref().is_some_and(|nonce| nonce.len() > 128) {
            problems.push("the nonce is longer than 128 characters".into());
        }
        problems
    }
}

/// `artifacts/<name>.onnx`, with a plain file name: no directories, no `..`.
fn model_path_ok(path: &str) -> bool {
    let Some(name) = path.strip_prefix(&format!("{ARTIFACTS_DIR}/")) else {
        return false;
    };
    name.ends_with(".onnx")
        && name.len() > ".onnx".len()
        && name
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
        && !name.starts_with('.')
}

/// Whether a ZIP entry may be in a bundle: the component, the manifest, or a
/// model the manifest lists.
pub fn entry_allowed(manifest: &Manifest, entry: &str) -> bool {
    entry == COMPONENT || entry == MANIFEST || manifest.models.iter().any(|m| m.path == entry)
}

#[cfg(all(test, feature = "bundle"))]
mod tests {
    use super::*;

    #[test]
    fn round_trips_and_checks_paths() {
        let mut manifest = Manifest::new("dance-off", 1, BotKind::Policy);
        manifest.models.push(ModelEntry {
            name: "policy".into(),
            path: "artifacts/policy.onnx".into(),
        });
        let text = manifest.to_toml();
        assert!(text.contains("[[model]]"));
        assert_eq!(Manifest::from_toml(&text).unwrap(), manifest);
        assert!(manifest.problems().is_empty());
        manifest.models[0].path = "artifacts/../x.onnx".into();
        assert_eq!(manifest.problems().len(), 1);
        assert!(!entry_allowed(&manifest, "extra.txt"));
        assert!(entry_allowed(&manifest, COMPONENT));
    }

    #[test]
    fn a_policy_needs_a_model() {
        let manifest = Manifest::new("boxing", 1, BotKind::Policy);
        assert!(manifest.problems()[0].contains("needs its model"));
        assert!(
            Manifest::new("boxing", 1, BotKind::Rust)
                .problems()
                .is_empty()
        );
    }
}