open-standards/mesh/mesh-rs/crates/mesh-engine/src/lib.rs

//! The serving-engine interface (SPEC §5).
//!
//! A backend implements [`ServingEngine`]: it advertises a fidelity level and a
//! modeled per-token cost, and it serves a [`ServeRequest`] producing tokens and
//! a [`ServeOutcome`] (the counters the receipt's J/accepted-token figure is
//! built from). The mesh meters energy *around* the `serve` call, so an engine
//! does not self-report energy — it just does the work honestly.
//!
//! [`ReferenceEngine`] is the conformance engine: deterministic output, no model
//! weights, but real CPU work proportional to the tokens served so a meter reads
//! a non-trivial span. Production backends (a GGUF/Metal decoder, a vLLM-class
//! server) implement the same trait and drop in unchanged.

#![forbid(unsafe_code)]

use mesh_schema::{BackendId, ContentId, Fidelity, ServeOutcome, ServeRequest};

/// The content address of a model's weights (SPEC §3.6), for a backend that
/// loads them from files: BLAKE3 over the canonical list of the files its
/// forward pass reads, each named by its path relative to `dir` and addressed
/// by the BLAKE3 of its bytes, in path order.
///
/// Naming each file keeps the address honest about structure: the same bytes
/// split across differently named shards address differently, as they load
/// differently. Hashing every byte is the point. A digest over a sample, a
/// header or a file size is cheaper and would look like a binding while
/// leaving the weights free to change underneath it, so this reads them all,
/// once, when the backend loads them.
///
/// `files` are the paths the forward pass reads, relative to `dir`; the order
/// they are listed in does not matter.
pub fn weights_digest(dir: &std::path::Path, files: &[&str]) -> std::io::Result<ContentId> {
    #[derive(serde::Serialize)]
    struct WeightsManifest<'a> {
        files: Vec<(&'a str, ContentId)>,
    }
    let mut named: Vec<&str> = files.to_vec();
    named.sort_unstable();
    named.dedup();
    let mut manifest = WeightsManifest { files: Vec::with_capacity(named.len()) };
    for name in named {
        let mut hasher = blake3::Hasher::new();
        hasher.update_reader(std::fs::File::open(dir.join(name))?)?;
        manifest.files.push((name, ContentId(*hasher.finalize().as_bytes())));
    }
    ContentId::of(&manifest).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))
}

/// What an engine returns from one `serve`: the generated token ids and the
/// serving counters.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Served {
    pub tokens: Vec<u32>,
    pub outcome: ServeOutcome,
}

/// A backend that can serve a metered inference.
pub trait ServingEngine {
    /// This backend's stable identifier.
    fn backend_id(&self) -> BackendId;
    /// The recipe (composed model) this engine serves.
    fn recipe_id(&self) -> String;
    /// The fidelity level this engine is rated to deliver — an advertisement,
    /// which the router admits the backend on (§4).
    ///
    /// This is a promise made *before* the work. What the backend actually
    /// served belongs in [`ServeOutcome::delivered_fidelity`], and a conforming
    /// engine MUST NOT report a delivered level above this rating: the router
    /// admitted it on this basis, so it cannot upgrade the basis afterwards.
    fn fidelity(&self) -> Fidelity;
    /// Modeled microjoules per token (the router's a-priori cost estimate).
    fn est_uj_per_token(&self) -> u64;

    /// The content address of the weights this engine runs, when it knows
    /// them before serving. Mesh names this model in the JCP frame a grant is
    /// decided over, so a grant's `ModelIn` caveat can admit or refuse the
    /// backend before any work runs (SPEC §2). `None`, the default, names no
    /// model, and a grant that permits only named models refuses it.
    ///
    /// An engine that reports one here MUST report the same digest in
    /// [`ServeOutcome::weights_digest`]: Mesh refuses to seal a serve whose
    /// engine ran other weights than the ones the grant was decided over.
    fn weights_digest(&self) -> Option<ContentId> {
        None
    }
    /// Whether serving the same request and prompt again produces the same
    /// tokens. A sealed serve from a deterministic engine can be re-executed
    /// and checked against its receipt's output commitment
    /// (`mesh_core::replay`). `false`, the default, for any engine that
    /// samples: re-running it settles nothing, and a replay says so rather
    /// than reporting a divergence that is only the sampler.
    fn deterministic(&self) -> bool {
        false
    }
    /// Serve the request. `prompt` is the input token payload — passed
    /// alongside the request, never sealed into the receipt in the clear (the
    /// receipt records the content-addressed `prompt_hash`, §3.1, and the
    /// payload's digest only inside its output commitment, §6.4).
    /// Implementations MUST do work proportional to the tokens produced so the
    /// surrounding meter reads a real span.
    fn serve(&mut self, req: &ServeRequest, prompt: &[u32]) -> Served;
}

/// Deterministic, weights-free reference engine.
///
/// Output is a pure function of the request's `prompt_hash` and `max_tokens`, so
/// conformance vectors are stable. It models two real serving behaviours:
/// a **prefix/result cache** (a repeated prompt serves cached tokens, recorded
/// as `cache_hit_tokens`) and an optional **speculative-decode** accept ratio
/// (recorded as `spec_proposed` / `spec_accepted`). Greedy decode accepts every
/// token, so by default `tokens_accepted == tokens_out`.
pub struct ReferenceEngine {
    id: BackendId,
    recipe: String,
    fidelity: Fidelity,
    uj_per_token: u64,
    work_per_token: u32,
    /// Permille of proposed speculative tokens accepted (0 = no speculation).
    spec_accept_permille: u32,
    /// Whether to report `delivered_fidelity`. Off by default: the published
    /// `mesh/1` vectors are generated from this engine, and a receipt that
    /// carries the field seals to different bytes, so an engine that reported
    /// it unconditionally would change every published vector and break every
    /// verifier that byte-compares them.
    report_delivered: bool,
    /// What to report as the weights' address and the arithmetic's precision.
    /// Both `None` by default, for the same reason as `report_delivered`: the
    /// published `mesh/1` vectors are generated from this engine, and a receipt
    /// that carries either field seals to different bytes. This engine answers
    /// from a table, so it has no weights of its own to address — the switch
    /// exists to make the fields exercisable end to end.
    weights_digest: Option<ContentId>,
    precision: Option<String>,
    /// Seeds served before — a repeat counts as a cache hit.
    seen: std::collections::HashSet<u64>,
}

impl ReferenceEngine {
    /// Build a reference engine.
    ///
    /// * `work_per_token` — integer-mixing iterations per token; turn this up to
    ///   make the metered span larger and more stable.
    /// * `spec_accept_permille` — set non-zero to model speculative decoding;
    ///   `0` is plain greedy decode.
    pub fn new(
        id: impl Into<String>,
        recipe: impl Into<String>,
        fidelity: Fidelity,
        uj_per_token: u64,
        work_per_token: u32,
    ) -> Self {
        Self {
            id: BackendId::new(id),
            recipe: recipe.into(),
            fidelity,
            uj_per_token,
            work_per_token,
            spec_accept_permille: 0,
            report_delivered: false,
            weights_digest: None,
            precision: None,
            seen: std::collections::HashSet::new(),
        }
    }

    /// Enable modeled speculative decoding at the given accept rate (permille).
    /// Report `delivered_fidelity` on every serve.
    ///
    /// This engine serves at its rating by construction — its output is a pure
    /// function of the request — so what it reports is the rating. The point of
    /// the switch is that it makes the *field* exercisable: a conformance vector
    /// built with it carries the delivered figure, and one built without it
    /// keeps the bytes the pack shipped with. A real engine that can degrade
    /// does not need a switch; it reports what it served.
    pub fn reporting_delivered_fidelity(mut self) -> Self {
        self.report_delivered = true;
        self
    }

    /// Report a weights address and a precision, as a backend serving real
    /// weights would (SPEC §3.6).
    ///
    /// Off by default for the same reason as `reporting_delivered_fidelity`:
    /// the published vectors are built from this engine, and either field
    /// changes the sealed bytes. With it, a vector exercises the binding; a
    /// vector built without it keeps the bytes the pack shipped with. A real
    /// backend needs no switch — it digests its weights once at load and
    /// reports what it ran.
    pub fn reporting_model_identity(mut self, weights_digest: ContentId, precision: &str) -> Self {
        self.weights_digest = Some(weights_digest);
        self.precision = Some(precision.to_string());
        self
    }

    pub fn with_speculation(mut self, accept_permille: u32) -> Self {
        self.spec_accept_permille = accept_permille.min(1000);
        self
    }

    /// Seed the deterministic token stream from the prompt hash.
    fn seed(req: &ServeRequest) -> u64 {
        let b = req.prompt_hash.0;
        u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])
            ^ ((req.max_tokens as u64) << 1)
    }

    /// One step of a splitmix64-style PRNG. Pure; deterministic.
    fn next(state: &mut u64) -> u64 {
        *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = *state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }

    /// Burn `work_per_token` iterations of real integer work so the meter reads
    /// the span. `black_box` keeps the optimiser from eliding it.
    fn burn(&self, mut x: u64) -> u64 {
        for _ in 0..self.work_per_token {
            x = x
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            x = std::hint::black_box(x);
        }
        x
    }
}

impl ServingEngine for ReferenceEngine {
    fn backend_id(&self) -> BackendId {
        self.id.clone()
    }

    fn recipe_id(&self) -> String {
        self.recipe.clone()
    }

    fn fidelity(&self) -> Fidelity {
        self.fidelity
    }

    fn est_uj_per_token(&self) -> u64 {
        self.uj_per_token
    }

    fn weights_digest(&self) -> Option<ContentId> {
        self.weights_digest
    }

    /// Its tokens are a pure function of the request: a SplitMix64 stream
    /// seeded from `prompt_hash` and `max_tokens`. Its cache changes the
    /// cache-hit count, never the tokens.
    fn deterministic(&self) -> bool {
        true
    }

    fn serve(&mut self, req: &ServeRequest, _prompt: &[u32]) -> Served {
        // The reference engine is synthetic: its output is a deterministic
        // function of the request's content id, independent of the prompt
        // payload, so conformance vectors are stable. A real engine consumes
        // `_prompt`.
        let seed = Self::seed(req);
        let cache_hit = self.seen.contains(&seed);
        self.seen.insert(seed);

        let n = req.max_tokens;
        let mut state = seed;
        let mut tokens = Vec::with_capacity(n as usize);
        let mut sink = 0u64;
        for _ in 0..n {
            let r = Self::next(&mut state);
            tokens.push((r % 32_000) as u32); // a plausible vocab id
            sink ^= self.burn(r);
        }
        std::hint::black_box(sink);

        let (spec_proposed, spec_accepted) = if self.spec_accept_permille > 0 {
            let proposed = n;
            let accepted = ((n as u64 * self.spec_accept_permille as u64) / 1000) as u32;
            (proposed, accepted)
        } else {
            (0, 0)
        };

        let outcome = ServeOutcome {
            tokens_in: req.prompt_tokens,
            tokens_out: n,
            tokens_accepted: n, // greedy decode accepts every emitted token
            cache_hit_tokens: if cache_hit { n } else { 0 },
            spec_proposed,
            spec_accepted,
            decision: None,
            delivered_fidelity: self.report_delivered.then_some(self.fidelity),
            // This engine answers from a table: there are no weights to
            // address and no arithmetic whose precision could be reported.
            // A backend serving real weights digests them once at load and
            // reports it here; absent is the honest answer for this one.
            weights_digest: self.weights_digest,
            precision: self.precision.clone(),
        };
        Served { tokens, outcome }
    }
}

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

    /// A scratch directory holding `files`, removed when dropped.
    struct Scratch(std::path::PathBuf);
    impl Scratch {
        fn with(tag: &str, files: &[(&str, &[u8])]) -> Self {
            let dir = std::env::temp_dir().join(format!("mesh-weights-{}-{tag}", std::process::id()));
            let _ = std::fs::remove_dir_all(&dir);
            std::fs::create_dir_all(&dir).unwrap();
            for (name, bytes) in files {
                std::fs::write(dir.join(name), bytes).unwrap();
            }
            Scratch(dir)
        }
    }
    impl Drop for Scratch {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    /// The address is a function of every named file's bytes and name, and of
    /// nothing else: not the order they are listed in.
    #[test]
    fn a_weights_digest_addresses_every_byte_and_every_name() {
        let base = Scratch::with("base", &[("config.json", b"{\"d\":4}"), ("model.safetensors", b"weights")]);
        let files = ["config.json", "model.safetensors"];
        let d = weights_digest(&base.0, &files).unwrap();
        assert_eq!(d, weights_digest(&base.0, &["model.safetensors", "config.json"]).unwrap(), "order-free");
        assert_eq!(d, weights_digest(&base.0, &files).unwrap(), "deterministic");

        let one_byte = Scratch::with("byte", &[("config.json", b"{\"d\":4}"), ("model.safetensors", b"weightz")]);
        assert_ne!(d, weights_digest(&one_byte.0, &files).unwrap(), "one changed byte");
        let renamed = Scratch::with("name", &[("config.json", b"{\"d\":4}"), ("shard.safetensors", b"weights")]);
        assert_ne!(d, weights_digest(&renamed.0, &["config.json", "shard.safetensors"]).unwrap(), "a renamed file");
        assert_ne!(d, weights_digest(&base.0, &["model.safetensors"]).unwrap(), "a file left out");
        assert!(weights_digest(&base.0, &["config.json", "absent.safetensors"]).is_err(), "a missing file is an error");
    }

    /// Each file is addressed exactly as `ContentId::of_bytes` would address
    /// it, though it is streamed rather than read into memory.
    #[test]
    fn a_streamed_file_hash_is_its_content_id() {
        let big: Vec<u8> = (0..300_000u32).map(|i| (i % 251) as u8).collect();
        let dir = Scratch::with("stream", &[("w.bin", &big)]);
        let mut hasher = blake3::Hasher::new();
        hasher.update_reader(std::fs::File::open(dir.0.join("w.bin")).unwrap()).unwrap();
        assert_eq!(ContentId(*hasher.finalize().as_bytes()), ContentId::of_bytes(&big));
    }

    use mesh_schema::ContentId;

    fn req(tag: u8, max_tokens: u32) -> ServeRequest {
        let mut seed = [tag; 9];
        seed[8] = max_tokens as u8;
        ServeRequest {
            prompt_hash: ContentId::of_bytes(&seed),
            prompt_tokens: 8,
            max_tokens,
            fidelity_floor: Fidelity::F0,
            joule_ceiling_uj: u64::MAX,
            decision: None,
        }
    }

    #[test]
    fn deterministic_output() {
        let mut e = ReferenceEngine::new("ref", "tiny", Fidelity::F1, 50_000, 64);
        let a = e.serve(&req(7, 16), &[]);
        let mut e2 = ReferenceEngine::new("ref", "tiny", Fidelity::F1, 50_000, 64);
        let b = e2.serve(&req(7, 16), &[]);
        assert_eq!(a.tokens, b.tokens);
        assert_eq!(a.outcome.tokens_out, 16);
        assert_eq!(a.outcome.tokens_accepted, 16);
    }

    #[test]
    fn repeat_prompt_is_a_cache_hit() {
        let mut e = ReferenceEngine::new("ref", "tiny", Fidelity::F1, 50_000, 8);
        let first = e.serve(&req(3, 8), &[]);
        assert_eq!(first.outcome.cache_hit_tokens, 0);
        let second = e.serve(&req(3, 8), &[]);
        assert_eq!(second.outcome.cache_hit_tokens, 8);
    }

    #[test]
    fn speculation_records_accept_ratio() {
        let mut e =
            ReferenceEngine::new("ref", "tiny", Fidelity::F2, 40_000, 8).with_speculation(750);
        let s = e.serve(&req(5, 100), &[]);
        assert_eq!(s.outcome.spec_proposed, 100);
        assert_eq!(s.outcome.spec_accepted, 75);
    }
}